114 lines
3.9 KiB
Python
114 lines
3.9 KiB
Python
"""
|
|
ml/evaluate.py - Curriculum-aware 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 = True, # Default to True so curriculum commands are sampled
|
|
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 environment with random commands enabled for curriculum progression
|
|
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] = []
|
|
episode_phases: List[str] = []
|
|
|
|
for ep in range(episodes):
|
|
obs, _ = env.reset()
|
|
done = False
|
|
total_reward = 0.0
|
|
steps = 0
|
|
|
|
initial_phase = env.curriculum_phase.name
|
|
print(f"\n--- Starting Evaluation Episode {ep + 1}/{episodes} [Phase: {initial_phase}] ---")
|
|
|
|
while not done:
|
|
action, _ = model.predict(obs, deterministic=True)
|
|
|
|
# Store phase before step to detect live phase transitions
|
|
prev_phase = env.curriculum_phase
|
|
|
|
obs, reward, terminated, truncated, _ = env.step(action)
|
|
|
|
# Done on either physical failure (terminated) or phase step limit (truncated)
|
|
done = terminated or truncated
|
|
total_reward += float(reward)
|
|
steps += 1
|
|
|
|
# Log live phase transition if unlocked during this step
|
|
if env.curriculum_phase != prev_phase:
|
|
print(f" └─ [Eval Milestone] Curriculum advanced to {env.curriculum_phase.name} at episode step {steps}!")
|
|
|
|
if use_gui:
|
|
time.sleep(1.0 / 240.0)
|
|
|
|
status_str = "FAILED (Terminated)" if terminated else "COMPLETED (Truncated)"
|
|
final_phase = env.curriculum_phase.name
|
|
|
|
episode_rewards.append(total_reward)
|
|
episode_lengths.append(steps)
|
|
episode_phases.append(final_phase)
|
|
|
|
print(
|
|
f"Episode {ep + 1} Finished [{status_str}]: "
|
|
f"Phase = {final_phase} | Total Reward = {total_reward:.2f} | Steps = {steps}"
|
|
)
|
|
|
|
env.close()
|
|
|
|
metrics = {
|
|
"model_path": str(model_path),
|
|
"episodes_evaluated": episodes,
|
|
"final_curriculum_phase": env.curriculum_phase.name,
|
|
"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,
|
|
"episode_phases": episode_phases,
|
|
}
|
|
|
|
print("\n" + "=" * 60)
|
|
print(f"EVALUATION COMPLETE ({episodes} Episodes)")
|
|
print(f"Final Reached Phase: {metrics['final_curriculum_phase']}")
|
|
print(f"Mean Reward: {metrics['mean_reward']:.2f} ± {metrics['std_reward']:.2f}")
|
|
print(f"Mean Episode Length: {metrics['mean_episode_length']:.1f} steps")
|
|
print("=" * 60)
|
|
|
|
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 |