""" ml/env.py - Gymnasium Environment for JackBot Hexapod RL Training """ import time import math from enum import IntEnum 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 # Color Palette RGBA for Terminated/Failed Robots COLOR_FAILED = [0.3, 0.3, 0.3, 0.6] class CurriculumPhase(IntEnum): STAND_ONLY = 0 FORWARD = 1 TURN_AND_DIRECTION = 2 OMNI_DIRECTION = 3 FULL_COMMAND = 4 class JackBotEnv(gym.Env): """Gymnasium environment wrapping a single JackBot hexapod.""" def __init__( self, use_gui: bool = True, random_command: bool = True, max_episode_steps: int = 3000, urdf_path: str = cfg.urdf_path, ): super().__init__() self.use_gui = use_gui self.random_command = random_command self.max_episode_steps = max_episode_steps self.urdf_path = urdf_path self.max_robot_speed = 0.6 self.step_count = 0 self.total_steps = 0 self.cumulative_reward = 0.0 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 & 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 ) # 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.command = np.zeros(4, dtype=np.float32) self.last_action = np.zeros(action_dim, dtype=np.float32) self.target_height = 0.122 self.collapse_height_fraction = 0.55 self.tilt_failure_rad = 0.9 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_stage_requirements = { CurriculumPhase.FORWARD: { "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, "max_lateral_drift": 0.8, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.25, }, CurriculumPhase.OMNI_DIRECTION: { "survival_steps": 600, "min_distance": 5.0, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.25, }, CurriculumPhase.FULL_COMMAND: { "survival_steps": 750, "min_distance": 8.0, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.20, }, } self.hud = MetricsHUD(physics_client_id=self.sim_manager.physics_client) self.last_time = time.time() 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: phase = self.curriculum_phase 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: 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: vx = np.random.uniform(-0.8, 0.8) vy = np.random.uniform(-0.5, 0.5) vz = 0.0 omega = np.random.uniform(-0.8, 0.8) else: vx = np.random.uniform(-1.0, 1.0) vy = np.random.uniform(-1.0, 1.0) vz = 0.0 omega = np.random.uniform(-1.0, 1.0) return np.array([vx, vy, vz, omega], dtype=np.float32) def reset(self, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None): super().reset(seed=seed) self.episode_count += 1 self.step_count = 0 self.cumulative_reward = 0.0 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 spawn_pos = self._robot_base_position(0) spawn_orn = [0.0, 0.0, 0.0, 1.0] 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(self.pb_robot, [1.0, 1.0, 1.0, 1.0]) 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_survival_steps = 0 if self._first_reset: self.curriculum_phase = CurriculumPhase.STAND_ONLY self._first_reset = False self.command = np.zeros(4, dtype=np.float32) self.next_cmd_resample_step = self.initial_stand_steps pos, _ = self.sim_manager.get_robot_pose(self.pb_robot) self.start_position = [float(pos[0]), float(pos[1]), float(pos[2])] self.target_height = self.sim_manager.settle_and_measure_height( [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._update_hud() return self._get_obs(), {} def get_current_robot_metrics(self) -> list: """Returns metric summary for the callback.""" if self.is_failed: return [] 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: 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() 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 self.robot.apply_rl_action(action) if self.step_count % 60 == 0: random_force = np.random.uniform(-2.0, 2.0, size=2) self.sim_manager.apply_external_force( 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_failure() self._update_distance_metrics() self._update_curriculum() obs = self._get_obs() reward = self._compute_reward(action, previous_action) self.cumulative_reward += reward self.robot_reward += reward terminated = self.is_failed truncated = self.step_count >= self.max_episode_steps 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: return False req = self.curriculum_stage_requirements[next_phase] # 1. Survival Check survival_ok = self.max_survival_steps >= req["survival_steps"] # 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 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 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 if "min_forward_distance" in req: distance_ok = dx >= req["min_forward_distance"] elif "min_distance" in req: distance_ok = dist_2d >= req["min_distance"] 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 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) or forced_forward): self.curriculum_phase = CurriculumPhase.FORWARD self._curriculum_advanced = True 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_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_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_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) -> 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 ) cmd_vx, cmd_vy, _, cmd_yaw = self.command cmd_norm = math.hypot(cmd_vx, cmd_vy) # 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 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 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] # 3. Posture & Stability Sub-Rewards height_error = pos[2] - self.target_height r_height = math.exp(-150.0 * (height_error ** 2)) orientation_error = roll**2 + pitch**2 r_stability = math.exp(-25.0 * orientation_error) joint_error = np.mean(np.square(current_joints - self.default_joint_angles)) r_pose = math.exp(-2.0 * joint_error) action_delta = np.mean(np.square(action - previous_action)) r_smoothness = math.exp(-0.1 * action_delta) # 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 total_reward = ( (w_height * r_height) + (w_stability * r_stability) + (w_pose * r_pose) + (w_smoothness * r_smoothness) ) else: # WALKING / TURNING MODE is_moving = (filtered_speed > 0.0) or (abs(filtered_yaw_rate) > 0.0) # 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: return now = time.time() fps = 1.0 / max(now - self.last_time, 1e-5) self.last_time = now 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_reward], cmd_vel=self.command, fps=fps, height=pos[2], roll_pitch=(math.degrees(roll), math.degrees(pitch)) ) def _update_robot_failure(self): if self.is_failed: return if self.step_count < 15: return 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 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() 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