updated evaluation code for previous env changes

outdated code from previous changes on env

ml/env.py

@@ -115,19 +115,26 @@ class JackBotEnv(gym.Env):
         self.robot_rewards = [0.0 for _ in range(self.num_robots)]
         self.failed_robots_mask = [False for _ in range(self.num_robots)]
 
-        if not self._first_reset:
-            p.resetSimulation(physicsClientId=self.sim_manager.physics_client)
-            p.setGravity(0, 0, -9.81, physicsClientId=self.sim_manager.physics_client)
-            p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.sim_manager.physics_client)
-            self.hud.reset()
-            self.leader_crown.reset()
-            self.plane, self.pb_robots, self.robot_joint_indices = self.sim_manager.load_scene(
-                self.urdf_path, self.num_robots, self.robot_spacing, self._robot_base_position
+        for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
+            # Get default spawn position
+            spawn_pos = self._robot_base_position(idx, self.num_robots, self.robot_spacing)
+            spawn_orn = [0, 0, 0, 1]
+
+            # Teleport base back to start
+            p.resetBasePositionAndOrientation(
+                pb_id, spawn_pos, spawn_orn, physicsClientId=self.sim_manager.physics_client
             )
-            for robot_obj, pb_id in zip(self.robots, self.pb_robots):
-                robot_obj.backend = PyBulletBackend(self.sim_manager, body_id=pb_id)
-        else:
-            self._first_reset = False
+            p.resetBaseVelocity(
+                pb_id, linearVelocity=[0, 0, 0], angularVelocity=[0, 0, 0],
+                physicsClientId=self.sim_manager.physics_client
+            )
+
+            # Reset joint angles directly without reloading URDF
+            robot_obj.reset_to_init()
+
+            # Restore original default visual color (clears failure dark gray)
+            if self.use_gui:
+                self._set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0])
 
         self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32)
 
@@ -136,13 +143,14 @@ class JackBotEnv(gym.Env):
         else:
             self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
 
-        for robot in self.robots:
-            robot.reset_to_init()
-
-        for _ in range(100):
+        for _ in range(15):
             self.sim_manager.step()
 
+        if self.use_gui:
+            self.hud.reset()
+            self.leader_crown.reset()
             self._update_hud()
+
         return self._get_obs(), {}
 
     def _get_obs(self) -> np.ndarray:

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)
             
-            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
+            done = terminated or truncated
+            total_reward += float(reward)
+            steps += 1
 
-            episode_reward += float(reward)
+            # Give PyBullet GUI frame pacing if running visually
+            if use_gui:
+                time.sleep(1.0 / 240.0)
 
-        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,
     )

ml/run_eval.py

@@ -16,13 +16,14 @@ from ml.evaluate import evaluate
 
 
 def main():
-    parser = argparse.ArgumentParser()
+    parser = argparse.ArgumentParser(description="JackBot Policy Evaluator Wrapper")
     parser.add_argument("--model", type=str, required=True, help="Path to the trained model file (.zip)")
-    parser.add_argument("--episodes", type=int, default=3, help="Number of 'rounds' to run. One episode lasts from reset until the robot falls over or the time limit is reached.")
+    parser.add_argument("--episodes", type=int, default=3, help="Number of evaluation episodes to run.")
     parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during evaluation")
     parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the evaluation environment")
     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("--save-metrics", type=str, default=None, help="Optional path to save JSON metrics report")
     args = parser.parse_args()
 
     evaluate(
@@ -32,6 +33,7 @@ def main():
         num_robots=args.num_robots,
         robot_spacing=args.robot_spacing,
         start_pose=args.start_pose,
+        save_json=args.save_metrics,
     )