Files
JackBot/ml/evaluate.py
T
JackM323 101004ead1 reward ajustments
robot starts balancing without changing phase and farms alive bonus
fix ->external force and bigger height penalty
2026-08-04 22:17:33 +02:00

91 lines
2.8 KiB
Python

"""
ml/evaluate.py - Evaluation routine for trained JackBot PPO policies.
"""
import argparse
import json
import time
from pathlib import Path
from typing import Optional, Dict, List, Any
import numpy as np
def evaluate(
model_path: str,
episodes: int = 5,
use_gui: bool = True,
robot_spacing: float = 0.5,
start_pose: str = "init_deg",
random_command: bool = False,
save_json: Optional[str] = None,
) -> Dict[str, Any]:
try:
from stable_baselines3 import PPO
except ImportError as exc:
raise ImportError(
"stable-baselines3 is required for evaluation. Install with: pip install stable-baselines3"
) from exc
from .env import JackBotEnv
print(f"[Eval] Loading policy model from: {model_path}")
# Force device="cpu" to prevent AMD ROCm/hipBLASLt matrix multiplication crashes
model = PPO.load(model_path, device="cpu")
# Initialize standard environment (returns single array of shape (352,))
env = JackBotEnv(
use_gui=use_gui,
random_command=random_command,
robot_spacing=robot_spacing,
start_pose=start_pose,
)
episode_rewards: List[float] = []
episode_lengths: List[int] = []
for ep in range(episodes):
obs, _ = env.reset()
terminated = False
total_reward = 0.0
steps = 0
print(f"\n--- Starting Evaluation Episode {ep + 1}/{episodes} ---")
while not terminated:
action, _ = model.predict(obs, deterministic=True)
obs, reward, terminated, truncated, _ = env.step(action)
total_reward += float(reward)
steps += 1
if use_gui:
time.sleep(1.0 / 240.0)
episode_rewards.append(total_reward)
episode_lengths.append(steps)
print(f"Episode {ep + 1} Finished: Total Reward = {total_reward:.2f} | Steps = {steps}")
env.close()
metrics = {
"model_path": str(model_path),
"episodes_evaluated": episodes,
"mean_reward": float(np.mean(episode_rewards)),
"std_reward": float(np.std(episode_rewards)),
"mean_episode_length": float(np.mean(episode_lengths)),
"raw_rewards": episode_rewards,
}
print("\n" + "=" * 50)
print(f"EVALUATION COMPLETE ({episodes} Episodes)")
print(f"Mean Reward: {metrics['mean_reward']:.2f} ± {metrics['std_reward']:.2f}")
print(f"Mean Episode Length: {metrics['mean_episode_length']:.1f} steps")
print("=" * 50)
if save_json:
out_path = Path(save_json)
out_path.parent.mkdir(parents=True, exist_ok=True)
with open(out_path, "w") as f:
json.dump(metrics, f, indent=4)
print(f"[Eval] Saved evaluation metrics to: {out_path.resolve()}")
return metrics