120 lines
4.1 KiB
Python
120 lines
4.1 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 = 1,
|
|
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}")
|
|
model = PPO.load(model_path)
|
|
|
|
# Initialize environment with active command sampling so robot actually walks
|
|
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:
|
|
# Deterministic evaluation (no exploration noise)
|
|
action, _ = model.predict(obs, deterministic=True)
|
|
obs, reward, terminated, truncated, _ = env.step(action)
|
|
|
|
done = terminated or truncated
|
|
total_reward += float(reward)
|
|
steps += 1
|
|
|
|
# Give PyBullet GUI frame pacing if running visually
|
|
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()
|
|
|
|
# Calculate summary statistics
|
|
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)
|
|
|
|
# Optional JSON metrics export
|
|
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
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Evaluate a trained JackBot policy.")
|
|
parser.add_argument("--model-path", type=str, required=True, help="Path to trained PPO model zip")
|
|
parser.add_argument("--episodes", type=int, default=5, help="Number of evaluation rounds")
|
|
parser.add_argument("--gui", action="store_true", help="Render GUI simulation")
|
|
parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in sim")
|
|
parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
|
|
parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg")
|
|
parser.add_argument("--save-json", type=str, default=None, help="Optional output JSON path for metrics")
|
|
args = parser.parse_args()
|
|
|
|
evaluate(
|
|
model_path=args.model_path,
|
|
episodes=args.episodes,
|
|
use_gui=args.gui,
|
|
num_robots=args.num_robots,
|
|
robot_spacing=args.robot_spacing,
|
|
start_pose=args.start_pose,
|
|
save_json=args.save_json,
|
|
) |