updated evaluation code for previous env changes

outdated code from previous changes on env

README.md

@@ -70,12 +70,6 @@ python -m pip install --upgrade pip setuptools wheel
 pip install -r requirements.txt
 ```
 
-> **Linux X11 Headless Note:** If running PyBullet GUI on Linux gives an X11 server connection error (`cannot connect to X server`), ensure your display environment variable is set:
-> ```bash
-> export DISPLAY=:0
-> python main.py
-> ```
-
 ### Windows (PowerShell)
 
 ```powershell
@@ -85,12 +79,6 @@ python -m pip install --upgrade pip setuptools wheel
 pip install -r requirements.txt
 ```
 
-If PowerShell blocks script execution:
-```powershell
-Set-ExecutionPolicy -Scope Process -ExecutionPolicy RemoteSigned
-.\.venv\Scripts\Activate.ps1
-```
-
 ---
 
 ## Usage Guide
@@ -151,10 +139,38 @@ python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 3
 
 ---
 
+## Reward Function & Termination Mechanics (`ml/env.py`)
+
+### 1. Reward Function Formulation
+
+The per-robot step reward ($R_{\text{step}}$) incentivizes tracking directional velocity commands while maintaining body stability and smooth joint actuation. The total environment step reward is the sum of all individual active robot rewards:
+
+$$R_{\text{step}} = R_{\text{alive}} + R_{\text{tracking}} + R_{\text{rotation}} - P_{\text{stability}} - P_{\text{action}}$$
+
+* **Alive Bonus ($R_{\text{alive}} = +0.1$):** A constant positive baseline awarded every timestep the robot remains upright.
+* **Velocity Tracking ($R_{\text{tracking}} = v_{x,\text{cmd}} \cdot v_x + v_{y,\text{cmd}} \cdot v_y$):** Rewards linear movement in the target command direction ($v_x, v_y$).
+* **Yaw Rotation Tracking ($R_{\text{rotation}} = \omega_{\text{cmd}} \cdot \omega_z$):** Rewards turning along the vertical yaw axis according to angular command $\omega_{\text{cmd}}$.
+* **Body Stability Penalty ($P_{\text{stability}} = 0.2 \cdot (|\text{roll}| + |\text{pitch}|)$):** Penalizes tilting away from a level horizontal posture.
+* **Action Energy Penalty ($P_{\text{action}} = 0.01 \cdot \sum a_i^2$):** Penalizes excessive joint delta actions to encourage smooth, energy-efficient leg movements and reduce jitter.
+
+---
+
+### 2. Failure Detection & Termination Logic
+
+In multi-robot vectorized training (`JackBotEnv`), individual robot failures are handled independently to allow maximum simulation efficiency:
+
+* **Individual Failure Masking:** A robot is flagged as failed (`failed_robots_mask[idx] = True`) if either condition is met:
+  * **Severe Tilt:** Base roll or pitch orientation exceeds $0.7\text{ radians}$ ($\approx 40^\circ$).
+  * **Base Collapse:** Base height drops below $0.05\text{ meters}$ above the ground plane.
+* **Visual Failure Feedback:** When a robot fails during GUI execution (`--gui`), its 3D URDF mesh immediately updates to a **semi-transparent dark gray** visual state (`COLOR_FAILED = [0.3, 0.3, 0.3, 0.6]`) to distinguish it from active learners.
+* **Environment Termination (`terminated=True`):** The entire environment step resets when any Robot fails.
+* **Environment Truncation (`truncated=True`):** Occurs when the episode reaches the maximum allowable step budget (`max_episode_steps = 3000`).
+
+---
+
 ## Environment Mechanics & Telemetry (`ml/env.py`)
 
 When training with `--gui`, `JackBotEnv` includes dynamic visual feedback mechanisms:
 
-* **Failure Detection & Graying:** Robots are continuously evaluated for roll/pitch tilt ($> 0.7\text{ rad}$) or base collapse ($< 0.05\text{ m}$ height). When a robot fails, its state mask is flagged and its 3D mesh automatically turns **semi-transparent dark gray**.
-* **Leader Crown ($\text{👑}$):** A floating crown indicator tracks and sits directly above the robot currently achieving the highest cumulative reward in the multi-robot grid.
-* **Termination Threshold:** The environment episode terminates automatically when the percentage of failed robots exceeds the configured threshold (default: $30\%$).
+* **Metrics HUD:** A live on-screen text overlay tracking active episode count, total step rate (FPS), average base height, roll/pitch angles, and cumulative per-robot rewards.
+* **Leader Crown ($\text{👑}$):** A floating crown debug indicator tracks and positions itself directly above the base of whichever robot is achieving the highest cumulative reward in the multi-robot grid.

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()
 
-        self._update_hud()
+        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)
+            
+            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,
+    )

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,8 +33,9 @@ def main():
         num_robots=args.num_robots,
         robot_spacing=args.robot_spacing,
         start_pose=args.start_pose,
+        save_json=args.save_metrics,
     )
 
 
 if __name__ == "__main__":
-    main()
+    main()

ml/run_train.py

@@ -26,7 +26,7 @@ def main():
     parser.add_argument("--device", type=str, default="auto", help="Device to use: 'cpu', 'cuda', or 'auto'")
     parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during training")
     parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the training environment")
-    parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
+    parser.add_argument("--robot-spacing", type=float, default=1.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")
     args = parser.parse_args()