diff --git a/ml/env.py b/ml/env.py index 4df724d..5353e66 100644 --- a/ml/env.py +++ b/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: diff --git a/ml/evaluate.py b/ml/evaluate.py index d67ebd3..22f5d3f 100644 --- a/ml/evaluate.py +++ b/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, + ) \ No newline at end of file diff --git a/ml/run_eval.py b/ml/run_eval.py index 20aafd9..b5dc09f 100644 --- a/ml/run_eval.py +++ b/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() \ No newline at end of file