""" 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 config import cfg from Robot import Robot, PyBulletBackend from ml.SimManager import SimManager from ml.MetricsOverlay import MetricsHUD, LeaderCrown # Color Palette RGBA for Terminated/Failed Robots COLOR_FAILED = [0.3, 0.3, 0.3, 0.6] # Collapsed / Tilted Robot (Dark Semi-Transparent Gray) 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 JackBot hexapods with floating text & crown overlay.""" 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, 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.episode_count = 0 self.step_count = 0 self.total_steps = 0 self.cumulative_reward = 0.0 self.robot_rewards = [0.0] self.failed_robots_mask = [False] self.episode_height_sum = 0.0 self._first_reset = True # 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 ) # 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_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.last_action = np.zeros(action_dim, dtype=np.float32) self.target_height = 0.14 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.max_survival_steps = 0 # 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 "max_displacement": 0.25, # Must remain within 0.25m radius "min_avg_height_ratio": 0.90, # Average height >= 90% of target "stability_roll_pitch": 0.18, # Max ~10 degrees tilt }, 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_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, }, } # 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 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 vx = np.random.uniform(0.5, 1.0) vy = 0.0 vz = 0.0 omega = 0.0 elif phase == CurriculumPhase.TURN_AND_DIRECTION: # Phase 2: Forward/Backward + Turning vx = np.random.uniform(-1.0, 1.0) vy = 0.0 vz = 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) vy = np.random.uniform(-0.5, 0.5) vz = 0.0 omega = np.random.uniform(-1.0, 1.0) 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 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_rewards = [0.0] self.failed_robots_mask = [False] self.episode_height_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] self.sim_manager.reset_robot_base(pb_id, spawn_pos, spawn_orn) robot_obj.reset_to_init() if self.use_gui: self.sim_manager.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) 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) 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])] # 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 ) 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 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) 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) 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)]) action_per_robot = action.reshape(1, 18) for robot, act in zip(self.robots, action_per_robot): robot.apply_rl_action(act) if self.step_count % 60 == 0: random_force = np.random.uniform(-2.0, 2.0, size=2) # X and Y push (Newtons) self.sim_manager.apply_external_force( body_id=self.pb_robots[0], force=[random_force[0], random_force[1], 0.0] ) self.sim_manager.step() self._update_robot_failures() self.get_robot_distance_metrics() obs = self._get_obs() reward, per_robot_step_rewards = 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 terminated = self._is_done() truncated = self.step_count >= self.curriculum_episode_limit self._update_hud() self._update_leader_visuals() return obs, reward, terminated, truncated, {} def _phase_progress_ready(self, next_phase: CurriculumPhase) -> bool: if next_phase not in self.curriculum_stage_requirements or not self.pb_robots: return False req = self.curriculum_stage_requirements[next_phase] # 1. Survival Step 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"] # 3. Average Height Ratio Check Across the Episode 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 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"] 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 def _update_curriculum(self): self._curriculum_advanced = False if self.curriculum_phase < CurriculumPhase.FORWARD and self._phase_progress_ready(CurriculumPhase.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}") 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) 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 = command[0] cmd_vy = command[1] cmd_yaw = command[3] # 1. LINEAR VECTOR SPEED MAXIMIZATION cmd_dir = np.array([cmd_vx, cmd_vy], dtype=np.float32) cmd_norm = np.linalg.norm(cmd_dir) if cmd_norm > 0.05: unit_cmd_dir = cmd_dir / cmd_norm actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32) aligned_speed = float(np.dot(actual_vel_2d, unit_cmd_dir)) linear_speed_reward = 2.5 * aligned_speed perp_vel = actual_vel_2d - aligned_speed * unit_cmd_dir drift_penalty = 0.35 * float(np.dot(perp_vel, perp_vel)) else: linear_speed_reward = 0.0 drift_penalty = 1.0 * (linear_vel[0]**2 + linear_vel[1]**2) # 2. TURNING SPEED MAXIMIZATION actual_yaw_rate = angular_vel[2] if abs(cmd_yaw) > 0.05: turning_reward = 1.5 * (actual_yaw_rate * cmd_yaw) else: turning_reward = -0.6 * (actual_yaw_rate ** 2) # 3. SMALL ALIVE BONUS + AGE-RAMPED STILLNESS PENALTY alive_reward = 0.5 age_ratio = min(1.0, self.step_count / max(1, self.curriculum_episode_limit)) still_penalty = 0.0 if (cmd_norm > 0.1 or abs(cmd_yaw) > 0.1) and (abs(linear_vel[0]) < 0.02 and abs(actual_yaw_rate) < 0.05): still_penalty = 0.35 + 0.85 * age_ratio # 4. POSTURE & STABILITY PENALTIES height_penalty : float min_valid_h = self.target_height * 0.90 # 90% threshold if pos[2] < min_valid_h: # Normalized drop below the 90% mark drop = (min_valid_h - pos[2]) / self.target_height # Linear + quadratic penalty that rapidly outweighs the +0.50 alive bonus height_penalty = 4.0 * drop + 20.0 * (drop ** 2) else: # Zero penalty inside the valid 90% - 110% zone! height_penalty = 0.0 stability_penalty = 1.5 * (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 ) rewards.append(r_step) 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 def _update_hud(self): if not self.use_gui or not self.pb_robots: 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))) self.hud.update( episode=self.episode_count, step=self.total_steps, robot_rewards=self.robot_rewards, cmd_vel=self.commands[0], fps=fps, avg_height=avg_height, roll_pitch=avg_roll_pitch ) def _update_leader_visuals(self): if not self.use_gui: 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) 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(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) def close(self): self.sim_manager.disconnect()