""" ml/env.py - Gymnasium Environment for JackBot Hexapod RL Training """ import time import math from typing import Optional, Tuple, Dict, Any, List import gymnasium as gym from gymnasium import spaces import numpy as np import pybullet as p 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 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._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 self.start_positions = [[0.0, 0.0, 0.0]] self.max_distance_from_start = [0.0] self.max_survival_steps = 0 self.curriculum_phase = 0 self.curriculum_episode_limit = 150 self.curriculum_stage_requirements = { 1: {"survival_steps": 200, "distance": 0.00, "stability_roll_pitch": 0.35}, 2: {"survival_steps": 350, "distance": 0.50, "stability_roll_pitch": 0.30}, 3: {"survival_steps": 550, "distance": 0.80, "stability_roll_pitch": 0.25}, } # 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 _set_robot_color(self, pb_id: int, rgba: List[float]): """Helper to change the visual color of a robot body and all its links.""" num_joints = p.getNumJoints(pb_id, physicsClientId=self.sim_manager.physics_client) p.changeVisualShape(pb_id, -1, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client) for j in range(num_joints): p.changeVisualShape(pb_id, j, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client) 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 == 0: # Phase 1: Forward Walking Focus vx = np.random.uniform(0.5, 1.0) vy = 0.0 vz = 0.0 omega = 0.0 elif phase == 1: # 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) else: # 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) 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] 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, 1] p.resetBasePositionAndOrientation( pb_id, spawn_pos, spawn_orn, physicsClientId=self.sim_manager.physics_client ) p.resetBaseVelocity( pb_id, linearVelocity=[0, 0, 0], angularVelocity=[0, 0, 0], physicsClientId=self.sim_manager.physics_client ) robot_obj.reset_to_init() if self.use_gui: self._set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0]) self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32) self.max_distance_from_start = [0.0] self.max_survival_steps = 0 if self._first_reset: self.curriculum_phase = 0 self.curriculum_episode_limit = min(self.max_episode_steps, 150) 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, robot_obj) in enumerate(zip(self.pb_robots, self.robots)): pos, _ = p.getBasePositionAndOrientation( pb_id, physicsClientId=self.sim_manager.physics_client ) self.start_positions[idx] = [float(pos[0]), float(pos[1]), float(pos[2])] for _ in range(200): self.sim_manager.step() self.target_height = self._measure_settled_height() if self.target_height <= 0.0: self.target_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 = p.getBaseVelocity(pb_id, physicsClientId=self.sim_manager.physics_client) 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, _ = p.getBasePositionAndOrientation( pb_id, physicsClientId=self.sim_manager.physics_client ) 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, _ = p.getBasePositionAndOrientation( pb_id, physicsClientId=self.sim_manager.physics_client ) linear_vel, angular_vel = p.getBaseVelocity( pb_id, physicsClientId=self.sim_manager.physics_client ) 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_states = p.getJointStates( pb_id, joint_indices, physicsClientId=self.sim_manager.physics_client ) joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32) robot_obs = np.concatenate([joint_angles, self.commands[idx]]) obs_list.append(robot_obs) return np.concatenate(obs_list).astype(np.float32) def _measure_settled_height(self) -> float: heights = [] for pb_id in self.pb_robots: pos, _ = p.getBasePositionAndOrientation( pb_id, physicsClientId=self.sim_manager.physics_client ) heights.append(pos[2]) return float(np.mean(heights)) if heights else 0.14 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) 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, phase: int) -> bool: if phase not in self.curriculum_stage_requirements: return False if not self.pb_robots or not self.max_distance_from_start: return False req = self.curriculum_stage_requirements[phase] survival_ok = self.max_survival_steps >= req["survival_steps"] distance_ok = self.max_distance_from_start[0] >= req["distance"] position, orientation = p.getBasePositionAndOrientation( self.pb_robots[0], physicsClientId=self.sim_manager.physics_client ) roll, pitch, _ = p.getEulerFromQuaternion(orientation) stability_ok = abs(roll) <= req["stability_roll_pitch"] and abs(pitch) <= req["stability_roll_pitch"] height_ok = position[2] >= max(0.09, self.target_height * 0.85) return survival_ok and distance_ok and stability_ok and height_ok def _update_curriculum(self): self._curriculum_advanced = False phase_labels = { 0: "stand-and-forward", 1: "turn-and-direction", 2: "omni-direction", 3: "full-command", } if self.curriculum_phase < 1 and self._phase_progress_ready(1): self.curriculum_phase = 1 self.curriculum_episode_limit = min(self.max_episode_steps, 400) self._curriculum_advanced = True print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}") elif self.curriculum_phase < 2 and self._phase_progress_ready(2): self.curriculum_phase = 2 self.curriculum_episode_limit = min(self.max_episode_steps, 700) self._curriculum_advanced = True print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}") elif self.curriculum_phase < 3 and self._phase_progress_ready(3): self.curriculum_phase = 3 self.curriculum_episode_limit = min(self.max_episode_steps, self.max_episode_steps) self._curriculum_advanced = True print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) 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, orientation = p.getBasePositionAndOrientation( pb_id, physicsClientId=self.sim_manager.physics_client ) linear_vel, angular_vel = p.getBaseVelocity( pb_id, physicsClientId=self.sim_manager.physics_client ) roll, pitch, _ = p.getEulerFromQuaternion(orientation) command = self.commands[idx] cmd_vx = command[0] cmd_vy = command[1] cmd_yaw = command[3] # ------------------------------------------------------------- # 1. LINEAR VECTOR SPEED MAXIMIZATION (Magnitude + Direction) # ------------------------------------------------------------- cmd_dir = np.array([cmd_vx, cmd_vy], dtype=np.float32) cmd_norm = np.linalg.norm(cmd_dir) if cmd_norm > 0.05: # Normalize target direction vector unit_cmd_dir = cmd_dir / cmd_norm actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32) # Speed aligned with target direction (m/s) aligned_speed = float(np.dot(actual_vel_2d, unit_cmd_dir)) # Moderate movement reward to encourage directional motion linear_speed_reward = 2.5 * aligned_speed # Penalize sideways drift (perpendicular velocity to commanded direction) perp_vel = actual_vel_2d - aligned_speed * unit_cmd_dir drift_penalty = 0.35 * float(np.dot(perp_vel, perp_vel)) else: # If no linear command given, penalize all horizontal movement 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] # rad/s in PyBullet Z-axis if abs(cmd_yaw) > 0.05: # Reward turning in the commanded direction, but less aggressively turning_reward = 1.5 * (actual_yaw_rate * cmd_yaw) else: # Penalize unwanted rotation when joystick turn is centered turning_reward = -0.6 * (actual_yaw_rate ** 2) # ------------------------------------------------------------- # 3. SMALL ALIVE BONUS + AGE-RAMPED STILLNESS PENALTY # ------------------------------------------------------------- alive_reward = 0.02 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 = 4.0 * ((self.target_height - pos[2]) ** 2) if pos[2] < self.target_height else 0.0 stability_penalty = 1.5 * (roll**2 + pitch**2) # Penalize only large control jumps; allow continuous low-amplitude motion 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 # Combined Step Reward 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, orient = p.getBasePositionAndOrientation(pb_id, physicsClientId=self.sim_manager.physics_client) roll, pitch, _ = p.getEulerFromQuaternion(orient) 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, _ = p.getBasePositionAndOrientation( leader_pb_id, physicsClientId=self.sim_manager.physics_client ) 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, orientation = p.getBasePositionAndOrientation( pb_id, physicsClientId=self.sim_manager.physics_client ) roll, pitch, _ = p.getEulerFromQuaternion(orientation) 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._set_robot_color(pb_id, COLOR_FAILED) def close(self): self.sim_manager.disconnect()