From c93c524a10b7cfc68c0e4214ab47339231633178 Mon Sep 17 00:00:00 2001 From: JackM323 Date: Wed, 5 Aug 2026 22:44:05 +0200 Subject: [PATCH] code reduction for training and eval reward and everything else changed again wont be the last time --- .gitignore | 2 +- ml/MetricsOverlay.py | 80 ++---- ml/SimManager.py | 40 +++ ml/__init__.py | 6 +- ml/env.py | 664 ++++++++++++++++++++++--------------------- ml/evaluate.py | 114 -------- ml/run_eval.py | 113 ++++++-- ml/run_train.py | 141 +++++++-- ml/train.py | 257 ----------------- requirements.txt | 3 +- 10 files changed, 606 insertions(+), 814 deletions(-) delete mode 100644 ml/evaluate.py delete mode 100644 ml/train.py diff --git a/.gitignore b/.gitignore index 3098c3f..415b53e 100644 --- a/.gitignore +++ b/.gitignore @@ -4,7 +4,7 @@ node_modules/ .vscode/ __pycache__/ ml/checkpoints/ -ml/tensorboard/ +ml/logs/ # Ignore environment files with private passwords/keys .env diff --git a/ml/MetricsOverlay.py b/ml/MetricsOverlay.py index 73263dc..e7ba4ae 100644 --- a/ml/MetricsOverlay.py +++ b/ml/MetricsOverlay.py @@ -44,7 +44,7 @@ class MetricsHUD: robot_rewards: List[float], cmd_vel: np.ndarray, fps: float = 0.0, - avg_height: float = 0.0, + height: float = 0.0, roll_pitch: Tuple[float, float] = (0.0, 0.0) ) -> None: """Updates floating black text block in 3D space with billboarding.""" @@ -65,7 +65,7 @@ class MetricsHUD: f"----------------------\n" f"Top Rewards: [{top1}, {top2}, {top3}]\n" f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]\n" - f"Height: {avg_height:.3f} m\n" + f"Height: {height:.3f} m\n" f"Roll/Pitch: {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°" ) @@ -76,69 +76,25 @@ class MetricsHUD: # Calculate dynamic orientation to align text flat against camera plane text_orientation = self._get_camera_facing_orientation() - if self._text_id is None: - self._text_id = p.addUserDebugText( - text=hud_text, - textPosition=text_position, - textColorRGB=text_color, - textSize=0.1, - textOrientation=text_orientation, - physicsClientId=self.client_id - ) - else: - self._text_id = p.addUserDebugText( - text=hud_text, - textPosition=text_position, - textColorRGB=text_color, - textSize=0.1, - textOrientation=text_orientation, - replaceItemUniqueId=self._text_id, - physicsClientId=self.client_id - ) + # Safely remove the old text to prevent PyBullet ghosting/overlapping + if self._text_id is not None: + try: + p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id) + except Exception: + pass + + # Draw fresh text + self._text_id = p.addUserDebugText( + text=hud_text, + textPosition=text_position, + textColorRGB=text_color, + textSize=0.1, + textOrientation=text_orientation, + physicsClientId=self.client_id + ) def reset(self) -> None: """Removes the active debug text item so a new episode starts with a clean overlay.""" - if self._text_id is not None: - try: - p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id) - except Exception: - pass - self._text_id = None - - -class LeaderCrown: - """Renders a floating crown or star emoji above the leading robot in PyBullet.""" - - def __init__(self, physics_client_id: int = 0): - self.client_id = physics_client_id - self._text_id = None - - def update(self, leader_pos: list[float]): - """Positions a floating crown ~0.35m directly above the lead robot's base.""" - crown_pos = [leader_pos[0], leader_pos[1], leader_pos[2] + 0.35] - - # You can use "👑 CROWN", "⭐ LEADER", or "★ TOP1" - crown_text = "👑" - - if self._text_id is None: - self._text_id = p.addUserDebugText( - text=crown_text, - textPosition=crown_pos, - textColorRGB=[1.0, 0.84, 0.0], - textSize=2.0, - physicsClientId=self.client_id - ) - else: - self._text_id = p.addUserDebugText( - text=crown_text, - textPosition=crown_pos, - textColorRGB=[1.0, 0.84, 0.0], - textSize=2.0, - replaceItemUniqueId=self._text_id, - physicsClientId=self.client_id - ) - - def reset(self): if self._text_id is not None: try: p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id) diff --git a/ml/SimManager.py b/ml/SimManager.py index 3f01a0e..fbd2ed9 100644 --- a/ml/SimManager.py +++ b/ml/SimManager.py @@ -82,11 +82,40 @@ class SimManager: def step(self): p.stepSimulation(physicsClientId=self.physics_client) + def set_rendering(self, enabled: bool) -> None: + """Toggles PyBullet 3D rendering to speed up simulation.""" + if self.physics_client is not None and p.isConnected(self.physics_client): + p.configureDebugVisualizer( + p.COV_ENABLE_RENDERING, + 1 if enabled else 0, + physicsClientId=self.physics_client + ) + def disconnect(self): if self.physics_client is not None and p.isConnected(self.physics_client): p.disconnect(self.physics_client) self.physics_client = None + def get_contact_points( + self, + bodyA: int = -1, + bodyB: int = -1, + linkIndexA: int = -1, + linkIndexB: int = -1, + ): + """Wrapper around pybullet.getContactPoints bound to this simulation client.""" + kwargs = {"physicsClientId": self.physics_client} + if bodyA != -1: + kwargs["bodyA"] = bodyA + if bodyB != -1: + kwargs["bodyB"] = bodyB + if linkIndexA != -1: + kwargs["linkIndexA"] = linkIndexA + if linkIndexB != -1: + kwargs["linkIndexB"] = linkIndexB + + return p.getContactPoints(**kwargs) + # --- ROBOT GETTERS AND SETTERS --- def reset_robot_base( @@ -124,6 +153,17 @@ class SimManager: roll, pitch, yaw = p.getEulerFromQuaternion(orn) return float(roll), float(pitch), float(yaw) + def get_foot_link_indices(self, body_id: int) -> list[int]: + """Inspects URDF joint structure to extract link IDs for leg tips and tibias.""" + foot_indices = [] + num_joints = p.getNumJoints(body_id, physicsClientId=self.physics_client) + for j_idx in range(num_joints): + info = p.getJointInfo(body_id, j_idx, physicsClientId=self.physics_client) + link_name = info[12].decode("utf-8") + if "tip" in link_name or "tibia" in link_name: + foot_indices.append(j_idx) + return foot_indices + def get_robot_pose_and_rpy(self, body_id: int) -> Tuple[List[float], Tuple[float, float, float]]: """Returns base position and (roll, pitch, yaw) tuple in radians.""" pos, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client) diff --git a/ml/__init__.py b/ml/__init__.py index 952554b..8f6caba 100644 --- a/ml/__init__.py +++ b/ml/__init__.py @@ -1,6 +1,4 @@ -from .env import JackBotEnv +from .env import JackBotEnv, CurriculumCallback from .model import ActorCritic -from .train import train -from .evaluate import evaluate -__all__ = ["JackBotEnv", "ActorCritic", "train", "evaluate"] +__all__ = ["JackBotEnv", "CurriculumCallback", "ActorCritic"] diff --git a/ml/env.py b/ml/env.py index a1713e2..73c5d3b 100644 --- a/ml/env.py +++ b/ml/env.py @@ -9,14 +9,15 @@ from typing import Optional, Tuple, Dict, Any, List import gymnasium as gym from gymnasium import spaces import numpy as np +from stable_baselines3.common.callbacks import BaseCallback from config import cfg from Robot import Robot, PyBulletBackend from ml.SimManager import SimManager -from ml.MetricsOverlay import MetricsHUD, LeaderCrown +from ml.MetricsOverlay import MetricsHUD # Color Palette RGBA for Terminated/Failed Robots -COLOR_FAILED = [0.3, 0.3, 0.3, 0.6] # Collapsed / Tilted Robot (Dark Semi-Transparent Gray) +COLOR_FAILED = [0.3, 0.3, 0.3, 0.6] class CurriculumPhase(IntEnum): @@ -28,93 +29,92 @@ class CurriculumPhase(IntEnum): class JackBotEnv(gym.Env): - """Gymnasium environment wrapping JackBot hexapods with floating text & crown overlay.""" + """Gymnasium environment wrapping a single JackBot hexapod.""" def __init__( self, use_gui: bool = True, random_command: bool = True, - robot_spacing: float = 0.5, - start_pose: str = "init_deg", - max_episode_steps: int = 5000, + max_episode_steps: int = 3000, urdf_path: str = cfg.urdf_path, ): super().__init__() self.use_gui = use_gui self.random_command = random_command - self.robot_spacing = robot_spacing - self.start_pose = start_pose self.max_episode_steps = max_episode_steps self.urdf_path = urdf_path - self.max_robot_speed = 0.8 + self.max_robot_speed = 0.6 - self.episode_count = 0 self.step_count = 0 self.total_steps = 0 self.cumulative_reward = 0.0 - self.robot_rewards = [0.0] - self.consecutive_still_steps = [0] - self.failed_robots_mask = [False] + self.robot_reward = 0.0 + self.consecutive_still_steps = 0 + self.is_failed = False + + self.episode_count = 0 self.episode_height_sum = 0.0 + self.episode_roll_sum = 0.0 + self.episode_pitch_sum = 0.0 + self._curriculum_advanced = False self._first_reset = True + # Dynamic Command Resampling Timing (60 Hz control loop) + self.control_freq = 60 + self.min_cmd_hold_steps = int(2.0 * self.control_freq) # 120 steps (2s) + self.max_cmd_hold_steps = int(6.0 * self.control_freq) # 360 steps (6s) + self.next_cmd_resample_step = 0 + self.initial_stand_steps = 120 # Mandatory 2s standing window at episode reset + # Initialize Simulation Manager self.sim_manager = SimManager(use_gui=self.use_gui) self.sim_manager.connect() - # Connect physics world - self.plane, self.pb_robots, self.robot_joint_indices = self.sim_manager.load_scene( - self.urdf_path, self.robot_spacing, self._robot_base_position + # Connect physics world & load single robot + self.plane, pb_robots, robot_joint_indices = self.sim_manager.load_scene( + self.urdf_path, 0.0, self._robot_base_position + ) + self.pb_robot = pb_robots[0] + self.joint_indices = robot_joint_indices[0] + + # Instantiate Robot Python wrapper (start_pose is managed inside Robot.py) + self.robot = Robot( + backend_type=PyBulletBackend(self.sim_manager, body_id=self.pb_robot), + urdf_path=self.urdf_path ) - # Instantiate Robot Python wrappers per PyBullet body ID - self.robots = [ - Robot( - backend_type=PyBulletBackend(self.sim_manager, body_id=pb_id), - start_pose=self.start_pose, - urdf_path=self.urdf_path - ) - for pb_id in self.pb_robots - ] - - # Action (18 joint deltas per robot) & Observation (18 angles + 4 command dims per robot) + # Action (18 joint deltas) & Observation (18 angles + 4 command dims) action_dim = 18 obs_dim = 18 + 4 self.action_space = spaces.Box(-1.0, 1.0, shape=(action_dim,), dtype=np.float32) self.observation_space = spaces.Box(-np.inf, np.inf, shape=(obs_dim,), dtype=np.float32) - self.commands = np.zeros((1, 4), dtype=np.float32) + self.command = np.zeros(4, dtype=np.float32) self.last_action = np.zeros(action_dim, dtype=np.float32) - self.target_height = 0.14 + self.target_height = 0.122 self.collapse_height_fraction = 0.55 self.tilt_failure_rad = 0.9 - # Small random exploration pulse settings - self.exploration_bonus_prob = 0.03 - self.exploration_bonus_interval = 120 - self.exploration_bonus_scale = 0.08 - self.exploration_bonus_active = False - - self.start_positions = [[0.0, 0.0, 0.0]] - self.max_distance_from_start = [0.0] + self.foot_link_indices = self._find_foot_link_indices() + self.start_position = [0.0, 0.0, 0.0] + self.max_distance_from_start = 0.0 self.max_survival_steps = 0 + self.default_joint_angles = np.zeros(18, dtype=np.float32) # Curriculum Initialization via Enum self.curriculum_phase = CurriculumPhase.STAND_ONLY - self.curriculum_episode_limit = 300 # 300 steps limit gives headroom for 400-step requirement - # Gates required to unlock each target phase self.curriculum_stage_requirements = { CurriculumPhase.FORWARD: { - "survival_steps": 400, # Must survive ~8 seconds - "min_avg_height_ratio": 0.90, # Average height >= 90% of target - "stability_roll_pitch": 0.18, # Max ~10 degrees tilt + "survival_steps": 300, + "min_avg_height_ratio": 0.88, + "max_avg_roll_pitch": 0.18, # ~10 degrees average }, CurriculumPhase.TURN_AND_DIRECTION: { "survival_steps": 500, - "min_forward_distance": 2.5, # Must walk +2.5m forward (+X) - "max_lateral_drift": 0.8, # Max 0.8m drift on Y axis + "min_forward_distance": 2.5, + "max_lateral_drift": 0.8, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.25, }, @@ -132,51 +132,42 @@ class JackBotEnv(gym.Env): }, } - # Floating HUD & Leader Crown Visualizers self.hud = MetricsHUD(physics_client_id=self.sim_manager.physics_client) - self.leader_crown = LeaderCrown(physics_client_id=self.sim_manager.physics_client) self.last_time = time.time() - def _robot_base_position(self, robot_id: int, spacing: float = 0.5) -> list[float]: - return [0.0, 0.0, 0.14] + def _robot_base_position(self, robot_id: int, spacing: float = 0.0) -> list[float]: + return [0.0, 0.0, 0.13] + + def _find_foot_link_indices(self) -> list: + return self.sim_manager.get_foot_link_indices(self.pb_robot) def sample_command(self) -> np.ndarray: - """Curriculum command sampler with survival-gated difficulty progression.""" phase = self.curriculum_phase - if phase == CurriculumPhase.STAND_ONLY: - # Phase 0: Pure Standing (Zero commands) - vx = 0.0 - vy = 0.0 - vz = 0.0 - omega = 0.0 - elif phase == CurriculumPhase.FORWARD: - # Phase 1: Straight Forward Walking + Standing Checks - if np.random.random() < 0.30: - # 30% chance: Stand Still Command (vx = 0) - vx = 0.0 - else: - # 70% chance: Forward Walk Command - vx = np.random.uniform(0.15, 0.50) - - vy, vz, omega = 0.0, 0.0, 0.0 + stand_probabilities = { + CurriculumPhase.STAND_ONLY: 1.0, + CurriculumPhase.FORWARD: 0.25, + CurriculumPhase.TURN_AND_DIRECTION: 0.20, + CurriculumPhase.OMNI_DIRECTION: 0.15, + CurriculumPhase.FULL_COMMAND: 0.15, + } + if np.random.random() < stand_probabilities.get(phase, 0.15): + return np.zeros(4, dtype=np.float32) + + if phase == CurriculumPhase.FORWARD: + vx = np.random.uniform(0.15, 0.50) + vy, vz, omega = 0.0, 0.0, 0.0 elif phase == CurriculumPhase.TURN_AND_DIRECTION: - # Phase 2: Forward/Backward + Turning + Standing - if np.random.random() < 0.20: - vx, omega = 0.0, 0.0 - else: - vx = np.random.uniform(-1.0, 1.0) - omega = np.random.uniform(-0.8, 0.8) + vx = np.random.uniform(-0.8, 0.8) vy, vz = 0.0, 0.0 + omega = np.random.uniform(-0.8, 0.8) elif phase == CurriculumPhase.OMNI_DIRECTION: - # Phase 3: Full Omnidirectional Movement - vx = np.random.uniform(-1.0, 1.0) + vx = np.random.uniform(-0.8, 0.8) vy = np.random.uniform(-0.5, 0.5) vz = 0.0 - omega = np.random.uniform(-1.0, 1.0) + omega = np.random.uniform(-0.8, 0.8) else: - # Phase 4: Full Unconstrained Commands vx = np.random.uniform(-1.0, 1.0) vy = np.random.uniform(-1.0, 1.0) vz = 0.0 @@ -189,177 +180,158 @@ class JackBotEnv(gym.Env): self.episode_count += 1 self.step_count = 0 self.cumulative_reward = 0.0 - self.robot_rewards = [0.0] - self.failed_robots_mask = [False] + self.robot_reward = 0.0 + self.is_failed = False + self.episode_height_sum = 0.0 + self.episode_roll_sum = 0.0 + self.episode_pitch_sum = 0.0 - for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)): - spawn_pos = self._robot_base_position(idx, self.robot_spacing) - spawn_orn = [0.0, 0.0, 0.0, 1.0] + spawn_pos = self._robot_base_position(0) + spawn_orn = [0.0, 0.0, 0.0, 1.0] - self.sim_manager.reset_robot_base(pb_id, spawn_pos, spawn_orn) - robot_obj.reset_to_init() + self.sim_manager.reset_robot_base(self.pb_robot, spawn_pos, spawn_orn) + self.robot.reset_to_init() - if self.use_gui: - self.sim_manager.set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0]) + if self.use_gui: + self.sim_manager.set_robot_color(self.pb_robot, [1.0, 1.0, 1.0, 1.0]) - self.consecutive_still_steps = [0 for _ in range(len(self.pb_robots))] + self.consecutive_still_steps = 0 self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32) - self.max_distance_from_start = [0.0] + self.max_distance_from_start = 0.0 self.max_survival_steps = 0 - self.exploration_bonus_active = False if self._first_reset: self.curriculum_phase = CurriculumPhase.STAND_ONLY - self.curriculum_episode_limit = min(self.max_episode_steps, 500) self._first_reset = False - if self.random_command: - self.commands = np.stack([self.sample_command() for _ in range(1)]) - else: - self.commands = np.zeros((1, 4), dtype=np.float32) + self.command = np.zeros(4, dtype=np.float32) + self.next_cmd_resample_step = self.initial_stand_steps - for idx, pb_id in enumerate(self.pb_robots): - pos, _ = self.sim_manager.get_robot_pose(pb_id) - self.start_positions[idx] = [float(pos[0]), float(pos[1]), float(pos[2])] + pos, _ = self.sim_manager.get_robot_pose(self.pb_robot) + self.start_position = [float(pos[0]), float(pos[1]), float(pos[2])] - # Settle for 200 steps after dropping in a standing position, then measure target height self.target_height = self.sim_manager.settle_and_measure_height( - self.pb_robots, steps=200, fallback_height=0.14 + [self.pb_robot], steps=200, fallback_height=0.122 + ) + + self.default_joint_angles = np.array( + self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices), + dtype=np.float32 ) if self.use_gui: self.hud.reset() - self.leader_crown.reset() self._update_hud() return self._get_obs(), {} - def get_robot_velocities(self) -> list: - """Exposes velocities for the SB3 metrics callback.""" - vels = [] - for pb_id in self.pb_robots: - lin_v, ang_v = self.sim_manager.get_robot_velocity(pb_id) - vels.append((lin_v, ang_v)) - return vels - - def get_robot_distance_metrics(self) -> list: - """Exposes distance-from-start metrics for logging without affecting reward.""" - metrics = [] - for idx, pb_id in enumerate(self.pb_robots): - pos, _ = self.sim_manager.get_robot_pose(pb_id) - if idx == 0: - self.episode_height_sum += float(pos[2]) - - start_x, start_y, _ = self.start_positions[idx] - dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32))) - self.max_distance_from_start[idx] = max(self.max_distance_from_start[idx], dist) - metrics.append((dist, self.max_distance_from_start[idx])) - - self.max_survival_steps = max(self.max_survival_steps, self.step_count) - return metrics - def get_current_robot_metrics(self) -> list: - """Returns current per-robot reward, distance, and velocity summaries for alive robots only.""" - metrics = [] - for idx, pb_id in enumerate(self.pb_robots): - if self.failed_robots_mask[idx]: - continue + """Returns metric summary for the callback.""" + if self.is_failed: + return [] - pos, _ = self.sim_manager.get_robot_pose(pb_id) - linear_vel, angular_vel = self.sim_manager.get_robot_velocity(pb_id) - start_x, start_y, _ = self.start_positions[idx] - dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32))) - speed = float(np.linalg.norm(np.array([linear_vel[0], linear_vel[1]], dtype=np.float32))) - yaw_rate = float(abs(angular_vel[2])) - metrics.append({ - "reward": float(self.robot_rewards[idx]), - "distance_from_start": dist, - "speed": speed, - "yaw_rate": yaw_rate, - "alive": True, - "survival_steps": int(self.step_count), - }) - - return metrics - - def get_survival_steps(self) -> int: - """Returns the current survival length for the environment's current episode.""" - return int(self.step_count) + pos, _ = self.sim_manager.get_robot_pose(self.pb_robot) + linear_vel, angular_vel = self.sim_manager.get_robot_velocity(self.pb_robot) + start_x, start_y, _ = self.start_position + dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32))) + speed = float(np.linalg.norm(np.array([linear_vel[0], linear_vel[1]], dtype=np.float32))) + yaw_rate = float(abs(angular_vel[2])) + + return [{ + "reward": float(self.robot_reward), + "distance_from_start": dist, + "speed": speed, + "yaw_rate": yaw_rate, + "alive": True, + "survival_steps": int(self.step_count), + "phase_name": self.curriculum_phase.name, + }] def _get_obs(self) -> np.ndarray: - obs_list = [] - for idx, (pb_id, joint_indices) in enumerate(zip(self.pb_robots, self.robot_joint_indices)): - joint_angles = self.sim_manager.get_robot_joint_angles(pb_id, joint_indices) - robot_obs = np.concatenate([joint_angles, self.commands[idx]]) - obs_list.append(robot_obs) - - return np.concatenate(obs_list).astype(np.float32) + joint_angles = self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices) + return np.concatenate([joint_angles, self.command]).astype(np.float32) def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]: self.step_count += 1 self.total_steps += 1 previous_action = self.last_action.copy() self.last_action = action.copy() - self._update_curriculum() - # Resample commands every 300 steps during long episodes - if self.random_command and (self.step_count % 300 == 0 or self._curriculum_advanced): - self.commands = np.stack([self.sample_command() for _ in range(1)]) + if self.random_command and (self.step_count >= self.next_cmd_resample_step or self._curriculum_advanced): + self.command = self.sample_command() + random_interval = np.random.randint(self.min_cmd_hold_steps, self.max_cmd_hold_steps + 1) + self.next_cmd_resample_step = self.step_count + random_interval - action_per_robot = action.reshape(1, 18) - - for robot, act in zip(self.robots, action_per_robot): - robot.apply_rl_action(act) + self.robot.apply_rl_action(action) if self.step_count % 60 == 0: - random_force = np.random.uniform(-2.0, 2.0, size=2) # X and Y push (Newtons) + random_force = np.random.uniform(-2.0, 2.0, size=2) self.sim_manager.apply_external_force( - body_id=self.pb_robots[0], + body_id=self.pb_robot, force=[random_force[0], random_force[1], 0.0] ) + render_freq = 10 # Only draw 1 in every 10 frames + + if self.use_gui and self.step_count % render_freq != 0: + self.sim_manager.set_rendering(False) + self.sim_manager.step() + + if self.use_gui and self.step_count % render_freq == 0: + self.sim_manager.set_rendering(True) - self._update_robot_failures() - self.get_robot_distance_metrics() + self._update_robot_failure() + self._update_distance_metrics() + self._update_curriculum() obs = self._get_obs() - reward, per_robot_step_rewards = self._compute_reward(action, previous_action) + reward = self._compute_reward(action, previous_action) self.cumulative_reward += reward - for idx, r_step in enumerate(per_robot_step_rewards): - self.robot_rewards[idx] += r_step + self.robot_reward += reward - terminated = self._is_done() - truncated = self.step_count >= self.curriculum_episode_limit + terminated = self.is_failed + truncated = self.step_count >= self.max_episode_steps - self._update_hud() - self._update_leader_visuals() + if self.step_count % 120 == 0 and self.use_gui: + self._update_hud() return obs, reward, terminated, truncated, {} + def _update_distance_metrics(self): + pos, _ = self.sim_manager.get_robot_pose(self.pb_robot) + self.episode_height_sum += float(pos[2]) + start_x, start_y, _ = self.start_position + dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32))) + self.max_distance_from_start = max(self.max_distance_from_start, dist) + self.max_survival_steps = max(self.max_survival_steps, self.step_count) + def _phase_progress_ready(self, next_phase: CurriculumPhase) -> bool: - if next_phase not in self.curriculum_stage_requirements or not self.pb_robots: + if next_phase not in self.curriculum_stage_requirements: return False req = self.curriculum_stage_requirements[next_phase] - # 1. Survival Step Check + # 1. Survival Check survival_ok = self.max_survival_steps >= req["survival_steps"] - # 2. Body Posture & Tilt Check - position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robots[0]) - stability_ok = abs(roll) <= req["stability_roll_pitch"] and abs(pitch) <= req["stability_roll_pitch"] + # 2. Smooth Average Stability Checks (Prevents 1-frame spikes from failing curriculum) + avg_roll = self.episode_roll_sum / max(1, self.step_count) + avg_pitch = self.episode_pitch_sum / max(1, self.step_count) + max_allowed_angle = req.get("max_avg_roll_pitch", 0.20) + stability_ok = (avg_roll <= max_allowed_angle) and (avg_pitch <= max_allowed_angle) - # 3. Average Height Ratio Check Across the Episode + # 3. Average Height Check avg_height = self.episode_height_sum / max(1, self.step_count) required_min_avg_height = self.target_height * req.get("min_avg_height_ratio", 0.85) height_ok = avg_height >= required_min_avg_height - # 4. Distance and Displacement Checks - start_x, start_y, _ = self.start_positions[0] - dx = position[0] - start_x - dy = position[1] - start_y + # 4. Distance and Drift Checks + pos, _ = self.sim_manager.get_robot_pose(self.pb_robot) + start_x, start_y, _ = self.start_position + dx = pos[0] - start_x + dy = pos[1] - start_y dist_2d = math.hypot(dx, dy) distance_ok = True @@ -368,203 +340,243 @@ class JackBotEnv(gym.Env): elif "min_distance" in req: distance_ok = dist_2d >= req["min_distance"] - displacement_ok = True - if "max_displacement" in req: - displacement_ok = dist_2d <= req["max_displacement"] - drift_ok = True if "max_lateral_drift" in req: drift_ok = abs(dy) <= req["max_lateral_drift"] - return survival_ok and height_ok and stability_ok and displacement_ok and distance_ok and drift_ok - + return survival_ok and height_ok and stability_ok and distance_ok and drift_ok + def _update_curriculum(self): self._curriculum_advanced = False + forced_forward = (self.curriculum_phase < CurriculumPhase.FORWARD) and (self.step_count >= 2500) - if self.curriculum_phase < CurriculumPhase.FORWARD and self._phase_progress_ready(CurriculumPhase.FORWARD): + if self.curriculum_phase < CurriculumPhase.FORWARD and (self._phase_progress_ready(CurriculumPhase.FORWARD) or forced_forward): self.curriculum_phase = CurriculumPhase.FORWARD - self.curriculum_episode_limit = min(self.max_episode_steps, 600) self._curriculum_advanced = True - print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}") + reason = "FORCED (2500 steps)" if forced_forward else "MET" + print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked [{reason}] at step {self.step_count}") elif self.curriculum_phase < CurriculumPhase.TURN_AND_DIRECTION and self._phase_progress_ready(CurriculumPhase.TURN_AND_DIRECTION): self.curriculum_phase = CurriculumPhase.TURN_AND_DIRECTION - self.curriculum_episode_limit = min(self.max_episode_steps, 800) self._curriculum_advanced = True print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}") elif self.curriculum_phase < CurriculumPhase.OMNI_DIRECTION and self._phase_progress_ready(CurriculumPhase.OMNI_DIRECTION): self.curriculum_phase = CurriculumPhase.OMNI_DIRECTION - self.curriculum_episode_limit = min(self.max_episode_steps, 1000) self._curriculum_advanced = True print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}") elif self.curriculum_phase < CurriculumPhase.FULL_COMMAND and self._phase_progress_ready(CurriculumPhase.FULL_COMMAND): self.curriculum_phase = CurriculumPhase.FULL_COMMAND - self.curriculum_episode_limit = min(self.max_episode_steps, self.max_episode_steps) self._curriculum_advanced = True print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}") - def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> Tuple[float, list[float]]: - rewards = [] - current_actions = action.reshape(1, 18) - previous_actions = previous_action.reshape(1, 18) + def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> float: + """ + Calculates task rewards using normalized Exponential Kernels. + Includes a deadband filter for jittering and zero-reward gating when stationary. + """ + # 1. Fetch Robot State + pos, (roll, pitch, yaw) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot) + linear_vel, angular_vel = self.sim_manager.get_robot_velocity(self.pb_robot) + current_joints = np.array( + self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices), + dtype=np.float32 + ) - for idx, pb_id in enumerate(self.pb_robots): - pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id) - linear_vel, angular_vel = self.sim_manager.get_robot_velocity(pb_id) - command = self.commands[idx] + cmd_vx, cmd_vy, _, cmd_yaw = self.command + cmd_norm = math.hypot(cmd_vx, cmd_vy) - cmd_vx = command[0] # Normalized -1.0 to 1.0 - cmd_vy = command[1] # Normalized -1.0 to 1.0 - cmd_yaw = command[3] + # 2. Velocity Deadband Filtering (Ignores jittering & micro-movements) + VEL_DEADBAND = 0.04 # 4 cm/s threshold + YAW_DEADBAND = 0.05 # 0.05 rad/s threshold - cmd_dir = np.array([cmd_vx, cmd_vy], dtype=np.float32) - cmd_norm = np.linalg.norm(cmd_dir) + raw_speed = math.hypot(linear_vel[0], linear_vel[1]) + if raw_speed < VEL_DEADBAND: + filtered_vx, filtered_vy = 0.0, 0.0 + filtered_speed = 0.0 + else: + filtered_vx, filtered_vy = linear_vel[0], linear_vel[1] + filtered_speed = raw_speed - actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32) + raw_yaw_rate = abs(angular_vel[2]) + if raw_yaw_rate < YAW_DEADBAND: + filtered_yaw_rate = 0.0 + else: + filtered_yaw_rate = angular_vel[2] - # ----------------------------------------------------------------- - # 1. NORMALIZED SPEED TRACKING (1.0 = v_max) - # ----------------------------------------------------------------- - if cmd_norm > 0.05: - unit_cmd_dir = cmd_dir / cmd_norm - aligned_speed = float(np.dot(actual_vel_2d, unit_cmd_dir)) + # 3. Posture & Stability Sub-Rewards + height_error = pos[2] - self.target_height + r_height = math.exp(-150.0 * (height_error ** 2)) - target_speed = cmd_norm * self.max_robot_speed # e.g., 0.35 * 0.8 = 0.28 m/s + orientation_error = roll**2 + pitch**2 + r_stability = math.exp(-25.0 * orientation_error) - if aligned_speed <= 0.0: - # Zero or backward movement gets ZERO reward - linear_speed_reward = 0.0 - elif aligned_speed <= target_speed: - # Smoothly scales from 0.0 to +5.0 as speed approaches target - linear_speed_reward = 5.0 * (aligned_speed / target_speed) - else: - # Gently penalize overspeeding beyond target - overspeed_ratio = (aligned_speed - target_speed) / target_speed - linear_speed_reward = max(0.0, 5.0 - 2.5 * overspeed_ratio) + joint_error = np.mean(np.square(current_joints - self.default_joint_angles)) + r_pose = math.exp(-2.0 * joint_error) - # Drift penalty for sideways sliding - perp_vel = actual_vel_2d - aligned_speed * unit_cmd_dir - drift_penalty = 0.35 * float(np.dot(perp_vel, perp_vel)) - else: - aligned_speed = 0.0 - linear_speed_reward = 0.0 - drift_penalty = 1.0 * (linear_vel[0]**2 + linear_vel[1]**2) + action_delta = np.mean(np.square(action - previous_action)) + r_smoothness = math.exp(-0.1 * action_delta) - # ----------------------------------------------------------------- - # 2. TIME-ACCUMULATING STILLNESS PENALTY - # ----------------------------------------------------------------- - actual_yaw_rate = angular_vel[2] - is_moving = abs(aligned_speed) > 0.03 or abs(actual_yaw_rate) > 0.08 + # 4. Mode Logic + if cmd_norm < 0.05 and abs(cmd_yaw) < 0.05: + # STANDING MODE: Reward clean posture, height, and stability + w_height = 0.35 + w_stability = 0.35 + w_pose = 0.20 + w_smoothness = 0.10 - if (cmd_norm > 0.05 or abs(cmd_yaw) > 0.05) and not is_moving: - # Increment counter every step the robot stalls under command - self.consecutive_still_steps[idx] += 1 - else: - # Reset counter as soon as robot makes a valid move! - self.consecutive_still_steps[idx] = 0 - - # Penalty grows by 0.05 every step frozen, capped at -3.5 per step - still_penalty = min(3.5, 0.05 * self.consecutive_still_steps[idx]) - - # ----------------------------------------------------------------- - # 3. TURNING & POSTURE PENALTIES - # ----------------------------------------------------------------- - if abs(cmd_yaw) > 0.05: - turning_reward = 1.5 * (actual_yaw_rate * cmd_yaw) - else: - turning_reward = -0.6 * (actual_yaw_rate ** 2) - - alive_reward = 0.20 - - # Height drop penalty - min_valid_h = self.target_height * 0.90 - if pos[2] < min_valid_h: - drop = (min_valid_h - pos[2]) / self.target_height - height_penalty = 4.0 * drop + 20.0 * (drop ** 2) - else: - height_penalty = 0.0 - - stability_penalty = 8 * (roll**2 + pitch**2) - - control_delta = np.abs(current_actions[idx] - previous_actions[idx]) - large_delta_mask = control_delta > 0.12 - large_delta_penalty = 0.002 * float(np.sum(np.square(control_delta[large_delta_mask]))) if np.any(large_delta_mask) else 0.0 - - r_step = ( - alive_reward - + linear_speed_reward - + turning_reward - - drift_penalty - - still_penalty - - height_penalty - - stability_penalty - - large_delta_penalty + total_reward = ( + (w_height * r_height) + + (w_stability * r_stability) + + (w_pose * r_pose) + + (w_smoothness * r_smoothness) ) - rewards.append(r_step) + else: + # WALKING / TURNING MODE + is_moving = (filtered_speed > 0.0) or (abs(filtered_yaw_rate) > 0.0) - return float(np.sum(rewards)), rewards - - def _is_done(self) -> bool: - """Returns True when the robot has entered a failed state.""" - return bool(self.failed_robots_mask[0]) if self.failed_robots_mask else False + # HARD GATE: If commanded to move but standing still/jittering, reward is strictly 0.0 + if not is_moving: + return 0.0 + + target_vx = cmd_vx * self.max_robot_speed + target_vy = cmd_vy * self.max_robot_speed + + lin_vel_error = (filtered_vx - target_vx)**2 + (filtered_vy - target_vy)**2 + r_lin_vel = math.exp(-25.0 * lin_vel_error) + + ang_vel_error = (filtered_yaw_rate - cmd_yaw)**2 + r_ang_vel = math.exp(-15.0 * ang_vel_error) + + # Stillness Check: Commanded to move, but staying virtually still + stillness_penalty = 0.0 + if math.hypot(target_vx, target_vy) > 0.08 and math.hypot(linear_vel[0], linear_vel[1]) < 0.03: + r_lin_vel = 0.0 # Strip velocity credit completely + stillness_penalty = -0.25 + + w_lin_vel = 0.55 + w_ang_vel = 0.15 + w_height = 0.10 + w_stability = 0.12 + w_smoothness = 0.08 + + total_reward = ( + (w_lin_vel * r_lin_vel) + + (w_ang_vel * r_ang_vel) + + (w_height * r_height) + + (w_stability * r_stability) + + (w_smoothness * r_smoothness) + + stillness_penalty + ) + + # Scaled reward for policy stability + return float(total_reward / 10.0) def _update_hud(self): - if not self.use_gui or not self.pb_robots: + if not self.use_gui: return now = time.time() fps = 1.0 / max(now - self.last_time, 1e-5) self.last_time = now - heights = [] - rolls = [] - pitches = [] - for pb_id in self.pb_robots: - pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id) - heights.append(pos[2]) - rolls.append(math.degrees(roll)) - pitches.append(math.degrees(pitch)) - - avg_height = float(np.mean(heights)) - avg_roll_pitch = (float(np.mean(rolls)), float(np.mean(pitches))) - + pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot) + self.hud.update( episode=self.episode_count, step=self.total_steps, - robot_rewards=self.robot_rewards, - cmd_vel=self.commands[0], + robot_rewards=[self.robot_reward], + cmd_vel=self.command, fps=fps, - avg_height=avg_height, - roll_pitch=avg_roll_pitch + height=pos[2], + roll_pitch=(math.degrees(roll), math.degrees(pitch)) ) - def _update_leader_visuals(self): - if not self.use_gui: + def _update_robot_failure(self): + if self.is_failed: return - best_idx = int(np.argmax(self.robot_rewards)) - leader_pb_id = self.pb_robots[best_idx] - leader_pos, _ = self.sim_manager.get_robot_pose(leader_pb_id) - self.leader_crown.update(leader_pos) + if self.step_count < 15: + return - def _update_robot_failures(self): - """Checks failure condition and colors failed robots dark gray.""" - for idx, pb_id in enumerate(self.pb_robots): - if self.failed_robots_mask[idx]: - continue + position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot) + collapse_threshold = max(0.04, self.collapse_height_fraction * self.target_height) + is_tilted = abs(roll) > self.tilt_failure_rad or abs(pitch) > self.tilt_failure_rad + is_collapsed = position[2] < collapse_threshold - position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id) - - collapse_threshold = max(0.06, self.collapse_height_fraction * self.target_height) - is_tilted = abs(roll) > self.tilt_failure_rad or abs(pitch) > self.tilt_failure_rad - is_collapsed = position[2] < collapse_threshold - - if is_tilted or is_collapsed: - self.failed_robots_mask[idx] = True - if self.use_gui: - self.sim_manager.set_robot_color(pb_id, COLOR_FAILED) + if is_tilted or is_collapsed: + self.is_failed = True + if self.use_gui: + self.sim_manager.set_robot_color(self.pb_robot, COLOR_FAILED) def close(self): - self.sim_manager.disconnect() \ No newline at end of file + self.sim_manager.disconnect() + + +class CurriculumCallback(BaseCallback): + """Logs curriculum phase breakdown and best performance metrics to TensorBoard.""" + + def __init__(self, verbose=0): + super().__init__(verbose) + self.best_speed = 0.0 + self.best_yaw_rate = 0.0 + self.best_distance = 0.0 + self.best_survival_steps = 0.0 + self.best_reward = -float('inf') + + def _on_step(self) -> bool: + return True + + def _on_rollout_end(self) -> bool: + try: + vec_env = self.training_env + alive_metrics = vec_env.env_method("get_current_robot_metrics") + + self.best_speed = 0.0 + self.best_yaw_rate = 0.0 + self.best_distance = 0.0 + self.best_survival_steps = 0.0 + self.best_reward = -float('inf') + + phase_counts = { + "stand_only": 0, + "forward": 0, + "turn_and_direction": 0, + "omni_direction": 0, + "full_command": 0, + } + + for worker_res in alive_metrics: + for metrics in worker_res: + if not metrics.get("alive", False): + continue + + phase_key = metrics.get("phase_name", "STAND_ONLY").lower() + if phase_key in phase_counts: + phase_counts[phase_key] += 1 + + if metrics["reward"] > self.best_reward: + self.best_reward = float(metrics["reward"]) + if metrics["speed"] > self.best_speed: + self.best_speed = float(metrics["speed"]) + if metrics["yaw_rate"] > self.best_yaw_rate: + self.best_yaw_rate = float(metrics["yaw_rate"]) + if metrics["distance_from_start"] > self.best_distance: + self.best_distance = float(metrics["distance_from_start"]) + if metrics["survival_steps"] > self.best_survival_steps: + self.best_survival_steps = float(metrics["survival_steps"]) + + for phase_name, count in phase_counts.items(): + self.logger.record(f"phase/{phase_name}", count) + + self.logger.record("custom/best_reward", float(self.best_reward) if np.isfinite(self.best_reward) else 0.0) + self.logger.record("custom/best_survival_steps", float(self.best_survival_steps)) + self.logger.record("custom/best_distance_from_start_m", float(self.best_distance)) + self.logger.record("custom/best_speed_mps", float(self.best_speed)) + self.logger.record("custom/best_yaw_rate_rads", float(self.best_yaw_rate)) + + except Exception: + pass + + return True \ No newline at end of file diff --git a/ml/evaluate.py b/ml/evaluate.py deleted file mode 100644 index c3e413f..0000000 --- a/ml/evaluate.py +++ /dev/null @@ -1,114 +0,0 @@ -""" -ml/evaluate.py - Curriculum-aware 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, - robot_spacing: float = 0.5, - start_pose: str = "init_deg", - random_command: bool = True, # Default to True so curriculum commands are sampled - 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}") - # Force device="cpu" to prevent AMD ROCm/hipBLASLt matrix multiplication crashes - model = PPO.load(model_path, device="cpu") - - # Initialize environment with random commands enabled for curriculum progression - env = JackBotEnv( - use_gui=use_gui, - random_command=random_command, - robot_spacing=robot_spacing, - start_pose=start_pose, - ) - - episode_rewards: List[float] = [] - episode_lengths: List[int] = [] - episode_phases: List[str] = [] - - for ep in range(episodes): - obs, _ = env.reset() - done = False - total_reward = 0.0 - steps = 0 - - initial_phase = env.curriculum_phase.name - print(f"\n--- Starting Evaluation Episode {ep + 1}/{episodes} [Phase: {initial_phase}] ---") - - while not done: - action, _ = model.predict(obs, deterministic=True) - - # Store phase before step to detect live phase transitions - prev_phase = env.curriculum_phase - - obs, reward, terminated, truncated, _ = env.step(action) - - # Done on either physical failure (terminated) or phase step limit (truncated) - done = terminated or truncated - total_reward += float(reward) - steps += 1 - - # Log live phase transition if unlocked during this step - if env.curriculum_phase != prev_phase: - print(f" └─ [Eval Milestone] Curriculum advanced to {env.curriculum_phase.name} at episode step {steps}!") - - if use_gui: - time.sleep(1.0 / 240.0) - - status_str = "FAILED (Terminated)" if terminated else "COMPLETED (Truncated)" - final_phase = env.curriculum_phase.name - - episode_rewards.append(total_reward) - episode_lengths.append(steps) - episode_phases.append(final_phase) - - print( - f"Episode {ep + 1} Finished [{status_str}]: " - f"Phase = {final_phase} | Total Reward = {total_reward:.2f} | Steps = {steps}" - ) - - env.close() - - metrics = { - "model_path": str(model_path), - "episodes_evaluated": episodes, - "final_curriculum_phase": env.curriculum_phase.name, - "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, - "episode_phases": episode_phases, - } - - print("\n" + "=" * 60) - print(f"EVALUATION COMPLETE ({episodes} Episodes)") - print(f"Final Reached Phase: {metrics['final_curriculum_phase']}") - print(f"Mean Reward: {metrics['mean_reward']:.2f} ± {metrics['std_reward']:.2f}") - print(f"Mean Episode Length: {metrics['mean_episode_length']:.1f} steps") - print("=" * 60) - - 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 \ No newline at end of file diff --git a/ml/run_eval.py b/ml/run_eval.py index 5bd9f74..7e140e6 100644 --- a/ml/run_eval.py +++ b/ml/run_eval.py @@ -5,38 +5,109 @@ Usage: """ import argparse +import json +import time import sys from pathlib import Path - -ROOT_DIR = Path(__file__).resolve().parent.parent -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) - -from ml.evaluate import evaluate - +import numpy as np +from stable_baselines3 import PPO +sys.path.append(str(Path(__file__).resolve().parent.parent)) +from ml.env import JackBotEnv def main(): - parser = argparse.ArgumentParser(description="JackBot Curriculum 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=5, help="Number of evaluation episodes to run.") - parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during evaluation") - 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("--no-random-command", dest="random_command", action="store_false", help="Disable command sampling and lock to zero commands") + parser = argparse.ArgumentParser(description="JackBot Policy Evaluator") + parser.add_argument("--model", type=str, required=True, help="Path to trained model checkpoint (.zip)") + parser.add_argument("--episodes", type=int, default=5, help="Number of evaluation episodes to run") + parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering") + parser.add_argument("--no-random-command", dest="random_command", action="store_false", help="Lock commands to zero (disable random command sampling)") parser.set_defaults(random_command=True) - parser.add_argument("--save-metrics", type=str, default=None, help="Optional path to save JSON metrics report") + parser.add_argument("--save-metrics", type=str, default=None, help="Optional JSON path to output evaluation summary") args = parser.parse_args() - evaluate( - model_path=args.model, - episodes=args.episodes, + print(f"[Eval] Loading policy model from: {args.model}") + model = PPO.load(args.model, device="cpu") + + # Instantiate single evaluation environment matching train setup + env = JackBotEnv( use_gui=args.gui, - robot_spacing=args.robot_spacing, - start_pose=args.start_pose, random_command=args.random_command, - save_json=args.save_metrics, ) + episode_rewards = [] + episode_lengths = [] + episode_phases = [] + + try: + for ep in range(args.episodes): + obs, _ = env.reset() + done = False + total_reward = 0.0 + steps = 0 + initial_phase = env.curriculum_phase.name + + print(f"\n--- Starting Evaluation Episode {ep + 1}/{args.episodes} [Phase: {initial_phase}] ---") + + while not done: + # Deterministic prediction matches evaluation standards + action, _ = model.predict(obs, deterministic=True) + + prev_phase = env.curriculum_phase + obs, reward, terminated, truncated, _ = env.step(action) + done = terminated or truncated + + total_reward += float(reward) + steps += 1 + + if env.curriculum_phase != prev_phase: + print(f" └─ [Eval Milestone] Reached {env.curriculum_phase.name} at step {steps}!") + + if args.gui: + time.sleep(1.0 / 240.0) + + status_str = "FAILED (Terminated)" if terminated else "COMPLETED (Truncated)" + final_phase = env.curriculum_phase.name + + episode_rewards.append(total_reward) + episode_lengths.append(steps) + episode_phases.append(final_phase) + + print( + f"Episode {ep + 1} Finished [{status_str}]: " + f"Phase = {final_phase} | Total Reward = {total_reward:.2f} | Steps = {steps}" + ) + + finally: + env.close() + + # Calculate metrics report + mean_reward = float(np.mean(episode_rewards)) + std_reward = float(np.std(episode_rewards)) + mean_length = float(np.mean(episode_lengths)) + + print("\n" + "=" * 60) + print(f"EVALUATION COMPLETE ({args.episodes} Episodes)") + print(f"Final Reached Phase: {env.curriculum_phase.name}") + print(f"Mean Reward: {mean_reward:.2f} ± {std_reward:.2f}") + print(f"Mean Episode Length: {mean_length:.1f} steps") + print("=" * 60) + + if args.save_metrics: + metrics = { + "model_path": str(args.model), + "episodes_evaluated": args.episodes, + "final_curriculum_phase": env.curriculum_phase.name, + "mean_reward": mean_reward, + "std_reward": std_reward, + "mean_episode_length": mean_length, + "raw_rewards": episode_rewards, + "episode_phases": episode_phases, + } + out_path = Path(args.save_metrics) + 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 report to: {out_path.resolve()}") + if __name__ == "__main__": main() \ No newline at end of file diff --git a/ml/run_train.py b/ml/run_train.py index c325097..66e5155 100644 --- a/ml/run_train.py +++ b/ml/run_train.py @@ -1,47 +1,132 @@ """Minimal training launcher for quick experiments. Usage: - python ml/run_train.py --timesteps 500000 --model ml/checkpoints/ppo_joint_command - -This is a convenience wrapper around `ml.train.train` with friendly defaults -for interactive experimentation. + python ml/run_train.py --total-timesteps 1500000 --gui """ import argparse +import os +import re import sys +import warnings from pathlib import Path +from stable_baselines3 import PPO +from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv +from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback -ROOT_DIR = Path(__file__).resolve().parent.parent -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) +sys.path.append(str(Path(__file__).resolve().parent.parent)) +from ml.env import JackBotEnv, CurriculumCallback -from ml.train import train +# Silence SB3's UserWarning about SubprocVecEnv vs DummyVecEnv +warnings.filterwarnings("ignore", category=UserWarning, module="stable_baselines3") + + +def get_next_run_number(save_dir: str) -> int: + """Scans the save directory for existing ppo patterns and returns the next integer.""" + if not os.path.exists(save_dir): + return 1 + + existing_numbers = [] + for item in os.listdir(save_dir): + # Match pattern ppo followed by numbers (e.g., ppo1, ppo_1, jackbot_ppo12) + matches = re.findall(r"ppo_?(\d+)", item, re.IGNORECASE) + for m in matches: + existing_numbers.append(int(m)) + + return max(existing_numbers, default=0) + 1 + + +def make_env(rank: int, use_gui: bool = False, seed: int = 0): + """Utility helper to instantiate parallel JackBot environments.""" + def _init(): + env = JackBotEnv( + use_gui=use_gui, + random_command=True, + ) + env.reset(seed=seed + rank) + return env + return _init def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--timesteps", type=int, default=500000, help="Total number of 'practice steps'.") - parser.add_argument("--model", type=str, default="ml/checkpoints/ppo_joint_command", help="Path to save the trained model") - parser.add_argument("--seed", type=int, default=0, help="Random seed for reproducibility") - 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-workers", type=int, default=1, help="Number of training environment workers") - 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") + parser = argparse.ArgumentParser(description="JackBot PPO Curriculum Trainer") + parser.add_argument("--num-workers", type=int, default=16, help="Number of parallel sub-process environments") + parser.add_argument("--total-timesteps", type=int, default=1_500_000, help="Total training timesteps") + parser.add_argument("--log-dir", type=str, default="ml/logs", help="Directory for TensorBoard logs") + parser.add_argument("--save-dir", type=str, default="ml/checkpoints", help="Directory for model checkpoints") + parser.add_argument("--save-freq", type=int, default=50_000, help="Checkpoint save frequency (steps)") + parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering") args = parser.parse_args() - Path(args.model).parent.mkdir(parents=True, exist_ok=True) - train( - total_timesteps=args.timesteps, - model_path=args.model, - seed=args.seed, - device=args.device, - use_gui=args.gui, - num_workers=args.num_workers, - robot_spacing=args.robot_spacing, - start_pose=args.start_pose, + os.makedirs(args.log_dir, exist_ok=True) + os.makedirs(args.save_dir, exist_ok=True) + + # Automatically determine next run number (e.g. ppo1, ppo2, ppo3...) + run_num = get_next_run_number(args.save_dir) + ppo_name = f"ppo{run_num}" + + print(f"[Train] Initializing Run #{run_num} ('{ppo_name}') with {args.num_workers} parallel workers...") + + env_fns = [ + make_env(rank=i, use_gui=(args.gui if i == 0 else False)) + for i in range(args.num_workers) + ] + vec_env = SubprocVecEnv(env_fns) + + # Initialize PPO Policy Hyperparameters + model = PPO( + policy="MlpPolicy", + env=vec_env, + learning_rate=3e-4, + n_steps=256, + batch_size=256, + n_epochs=10, + gamma=0.99, + gae_lambda=0.95, + clip_range=0.2, + ent_coef=0.03, + target_kl=0.05, + vf_coef=0.5, + max_grad_norm=0.5, + verbose=1, + tensorboard_log=args.log_dir, + device="cpu", ) + # Setup Callbacks with ppo naming + checkpoint_callback = CheckpointCallback( + save_freq=max(1, args.save_freq // args.num_workers), + save_path=args.save_dir, + name_prefix=f"jackbot_{ppo_name}", + ) + curriculum_callback = CurriculumCallback() + + eval_env = DummyVecEnv([lambda: JackBotEnv(use_gui=False, random_command=True)]) + best_model_path = os.path.join(args.save_dir, f"best_model_{ppo_name}") + + eval_callback = EvalCallback( + eval_env, + best_model_save_path=best_model_path, + log_path="ml/logs/results", + eval_freq=max(1, 10_000 // args.num_workers), + deterministic=True, + render=False, + ) + + print(f"[Train] Starting training for {args.total_timesteps} timesteps...") + try: + model.learn( + total_timesteps=args.total_timesteps, + callback=[checkpoint_callback, curriculum_callback, eval_callback], + progress_bar=True, + ) + final_model_path = os.path.join(args.save_dir, f"jackbot_{ppo_name}_final.zip") + model.save(final_model_path) + print(f"[Train] Training complete! Saved final model to {final_model_path}") + finally: + vec_env.close() + eval_env.close() + if __name__ == "__main__": - main() + main() \ No newline at end of file diff --git a/ml/train.py b/ml/train.py deleted file mode 100644 index 7361589..0000000 --- a/ml/train.py +++ /dev/null @@ -1,257 +0,0 @@ -import os -import argparse -from pathlib import Path - -import numpy as np -from stable_baselines3.common.callbacks import BaseCallback -from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv -from .env import JackBotEnv - - -class MilestoneCheckpointCallback(BaseCallback): - """ - Saves a model checkpoint the FIRST time total_timesteps - crosses every multiple of step_interval (e.g., 100,000). - Saves inside the matching PPO_X subfolder as created by TensorBoard. - """ - def __init__(self, save_path: str, name_prefix: str = "ppo_jackbot", step_interval: int = 100_000, verbose: int = 1): - super().__init__(verbose) - self.base_save_path = save_path - self.run_save_path = save_path - self.name_prefix = name_prefix - self.step_interval = step_interval - self.last_milestone = 0 - - def _on_training_start(self) -> None: - """Executed right before training loop starts. Resolves TensorBoard's run folder name (e.g. PPO_1).""" - if self.logger and self.logger.dir: - run_folder_name = Path(self.logger.dir).name # Extracts "PPO_1", "PPO_2", etc. - self.run_save_path = os.path.join(self.base_save_path, run_folder_name) - else: - self.run_save_path = self.base_save_path - - os.makedirs(self.run_save_path, exist_ok=True) - - def _on_step(self) -> bool: - current_milestone = self.num_timesteps // self.step_interval - - if current_milestone > self.last_milestone: - self.last_milestone = current_milestone - milestone_step = current_milestone * self.step_interval - - save_file = os.path.join( - self.run_save_path, - f"{self.name_prefix}_{milestone_step}_steps.zip" - ) - self.model.save(save_file) - - if self.verbose > 0: - print(f"\n[Checkpoint] Saved milestone model at {self.num_timesteps} steps -> {save_file}\n") - - return True - - -class JackBotMetricsCallback(BaseCallback): - """ - Tracks the best current alive-robot performance for the most recent rollout, - instead of logging lifetime maxima from the entire training run. - """ - def __init__(self, verbose=0): - super().__init__(verbose) - self.best_alive_speed = 0.0 - self.best_alive_yaw_rate = 0.0 - self.best_alive_distance = 0.0 - self.best_alive_survival_steps = 0.0 - self.best_alive_reward = -float('inf') - - def _on_step(self) -> bool: - return True - - def _on_rollout_end(self) -> bool: - try: - vec_env = self.training_env - alive_metrics = vec_env.env_method("get_current_robot_metrics") - - self.best_alive_speed = 0.0 - self.best_alive_yaw_rate = 0.0 - self.best_alive_distance = 0.0 - self.best_alive_survival_steps = 0.0 - self.best_alive_reward = -float('inf') - - for worker_res in alive_metrics: - for metrics in worker_res: - if not metrics.get("alive", False): - continue - - if metrics["reward"] > self.best_alive_reward: - self.best_alive_reward = float(metrics["reward"]) - if metrics["speed"] > self.best_alive_speed: - self.best_alive_speed = float(metrics["speed"]) - if metrics["yaw_rate"] > self.best_alive_yaw_rate: - self.best_alive_yaw_rate = float(metrics["yaw_rate"]) - if metrics["distance_from_start"] > self.best_alive_distance: - self.best_alive_distance = float(metrics["distance_from_start"]) - if metrics["survival_steps"] > self.best_alive_survival_steps: - self.best_alive_survival_steps = float(metrics["survival_steps"]) - - self.logger.record("custom/best_alive_speed_mps", float(self.best_alive_speed)) - self.logger.record("custom/best_alive_yaw_rate_rads", float(self.best_alive_yaw_rate)) - self.logger.record("custom/best_alive_distance_from_start_m", float(self.best_alive_distance)) - self.logger.record("custom/best_alive_survival_steps", float(self.best_alive_survival_steps)) - self.logger.record("custom/best_alive_reward", float(self.best_alive_reward) if np.isfinite(self.best_alive_reward) else 0.0) - - self.logger.record("custom/max_speed_mps", float(self.best_alive_speed)) - self.logger.record("custom/max_yaw_rate_rads", float(self.best_alive_yaw_rate)) - self.logger.record("custom/max_distance_from_start_m", float(self.best_alive_distance)) - self.logger.record("custom/max_survival_steps", float(self.best_alive_survival_steps)) - self.logger.record("custom/best_episode_reward", float(self.best_alive_reward) if np.isfinite(self.best_alive_reward) else 0.0) - except Exception: - pass - - return True - - -def parse_args(): - parser = argparse.ArgumentParser(description="Train a joint-command policy for JackBot.") - parser.add_argument("--timesteps", type=int, default=500_000, help="Total training timesteps") - parser.add_argument("--model-path", type=str, default="ml/checkpoints/ppo_joint_command", help="Where to save the trained model") - parser.add_argument("--seed", type=int, default=0, help="Random seed") - parser.add_argument("--device", type=str, default="auto", help="Device to use: 'cpu', 'cuda', or 'auto' to autodetect") - parser.add_argument("--gui", "--use-gui", dest="use_gui", action="store_true", help="Enable PyBullet GUI during training") - parser.add_argument("--num-workers", type=int, default=8, help="Number of parallel CPU worker processes") - parser.add_argument("--robot-spacing", type=float, default=3.0, 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") - return parser.parse_args() - - -def make_env(robot_spacing, start_pose, use_gui, rank, seed=0): - def _init(): - env = JackBotEnv( - use_gui=use_gui if rank == 0 else False, - random_command=True, - robot_spacing=robot_spacing, - start_pose=start_pose, - ) - env.reset(seed=seed + rank) - return env - return _init - - -def train( - total_timesteps: int, - model_path: str, - seed: int = 0, - device: str = "auto", - use_gui: bool = False, - num_workers: int = 8, - robot_spacing: float = 0.5, - start_pose: str = "init_deg", -): - try: - from stable_baselines3 import PPO - except ImportError as exc: - raise ImportError("stable-baselines3 is required. Install with: pip install stable-baselines3") from exc - - if num_workers > 1: - env_fns = [ - make_env(robot_spacing, start_pose, use_gui, rank=i, seed=seed) - for i in range(num_workers) - ] - env = SubprocVecEnv(env_fns) - else: - env = DummyVecEnv([ - make_env(robot_spacing, start_pose, use_gui, rank=0, seed=seed) - ]) - - def resolve_device(requested_device: str) -> str: - try: - import torch - except ImportError: - if requested_device != "cpu": - raise RuntimeError("PyTorch is not installed.") - return "cpu" - - hip_supported = getattr(torch.version, "hip", None) is not None - cuda_available = torch.cuda.is_available() - hip_available = hip_supported and getattr(torch.backends, "hip", None) is not None and torch.backends.hip.is_available() - - if requested_device == "auto": - return "cuda" if (hip_available or cuda_available) else "cpu" - - if requested_device in {"cuda", "gpu", "hip"}: - if hip_available or cuda_available: - return "cuda" - raise RuntimeError(f"GPU requested ({requested_device}) but not available.") - - return "cpu" - - device = resolve_device(device) - - policy_kwargs = dict( - log_std_init=-1.5, - net_arch=dict(pi=[256, 256], vf=[256, 256]) - ) - - model = PPO( - "MlpPolicy", - env, - verbose=1, - seed=seed, - learning_rate=1.5e-4, - n_steps=256, - batch_size=256, - n_epochs=10, - gamma=0.99, - gae_lambda=0.95, - clip_range=0.2, - target_kl=0.03, - ent_coef=0.03, - vf_coef=0.5, - max_grad_norm=0.5, - device=device, - policy_kwargs=policy_kwargs, - tensorboard_log=str(Path(__file__).resolve().parent / "tensorboard"), - ) - - # Base folder where model runs will be stored - save_dir = str(Path(model_path).parent) - model_prefix = Path(model_path).stem - - milestone_cb = MilestoneCheckpointCallback( - save_path=save_dir, - name_prefix=model_prefix, - step_interval=100_000 - ) - - metrics_callback = JackBotMetricsCallback() - - model.learn(total_timesteps=total_timesteps, callback=[milestone_cb, metrics_callback]) - - # Save the final model inside the matching PPO_X directory as well - if model.logger and model.logger.dir: - run_folder_name = Path(model.logger.dir).name - final_dir = Path(model_path).parent / run_folder_name - else: - final_dir = Path(model_path).parent - - final_dir.mkdir(parents=True, exist_ok=True) - final_save_path = final_dir / f"{model_prefix}_final.zip" - model.save(str(final_save_path)) - if model.verbose > 0: - print(f"[Training Complete] Saved final model to -> {final_save_path}") - - env.close() - - -if __name__ == "__main__": - args = parse_args() - train( - args.timesteps, - args.model_path, - seed=args.seed, - device=args.device, - use_gui=args.use_gui, - num_workers=args.num_workers, - robot_spacing=args.robot_spacing, - start_pose=args.start_pose, - ) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index f3bed07..2c9bb85 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,6 +8,7 @@ matplotlib # Deep learning / RL torch stable-baselines3 -gymnasium[box2d] +stable-baselines3[extra] +gymnasium shimmy tensorboard