updated evaluation code for previous env changes

outdated code from previous changes on env

ml/evaluate.py

@@ -1,17 +1,24 @@
+"""
+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
 
-from .env import JackBotEnv
-
 
 def evaluate(
     model_path: str,
     episodes: int = 5,
-    use_gui: bool = False,
+    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:
@@ -19,60 +26,95 @@ def evaluate(
             "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=False,
+        random_command=random_command,
         num_robots=num_robots,
         robot_spacing=robot_spacing,
         start_pose=start_pose,
     )
-    model = PPO.load(model_path)
 
-    for episode in range(episodes):
-        reset_res = env.reset()
-        # handle Gym / Gymnasium compatibility: reset may return (obs, info)
-        if isinstance(reset_res, tuple) and len(reset_res) == 2:
-            obs, _ = reset_res
-        else:
-            obs = reset_res
+    episode_rewards: List[float] = []
+    episode_lengths: List[int] = []
 
+    for ep in range(episodes):
+        obs, _ = env.reset()
         done = False
-        episode_reward = 0.0
+        total_reward = 0.0
+        steps = 0
+
+        print(f"\n--- Starting Evaluation Episode {ep + 1}/{episodes} ---")
 
         while not done:
-            # pass only the observation to the policy
+            # 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
 
-            step_res = env.step(action)
-            # Gymnasium-style: (obs, reward, terminated, truncated, info)
-            if isinstance(step_res, tuple) and len(step_res) == 5:
-                obs, reward, terminated, truncated, info = step_res
-                done = bool(terminated or truncated)
-            else:
-                # legacy Gym: (obs, reward, done, info)
-                obs, reward, done, info = step_res
+            # Give PyBullet GUI frame pacing if running visually
+            if use_gui:
+                time.sleep(1.0 / 240.0)
 
-            episode_reward += float(reward)
-
-        print(f"Episode {episode + 1}: reward={episode_reward:.2f}")
+        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)
-    parser.add_argument("--episodes", type=int, default=5)
-    parser.add_argument("--gui", action="store_true")
-    parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the environment")
+    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", help="Initial robot pose at reset")
+    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(
-        args.model_path,
+        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,
+    )