95 lines
2.9 KiB
Python
95 lines
2.9 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,
|
|
num_robots: int = 16, # Default to 16 to match your trained (352,) observation space
|
|
robot_spacing: float = 0.5,
|
|
start_pose: str = "init_deg",
|
|
random_command: bool = True,
|
|
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,
|
|
num_robots=num_robots,
|
|
robot_spacing=robot_spacing,
|
|
start_pose=start_pose,
|
|
)
|
|
|
|
episode_rewards: List[float] = []
|
|
episode_lengths: List[int] = []
|
|
|
|
for ep in range(episodes):
|
|
obs, _ = env.reset()
|
|
done = False
|
|
total_reward = 0.0
|
|
steps = 0
|
|
|
|
print(f"\n--- Starting Evaluation Episode {ep + 1}/{episodes} ---")
|
|
|
|
while not done:
|
|
action, _ = model.predict(obs, deterministic=True)
|
|
obs, reward, terminated, truncated, _ = env.step(action)
|
|
|
|
done = terminated or truncated
|
|
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,
|
|
"num_robots": num_robots,
|
|
"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 |