""" 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 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 # 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 self.curriculum_phase = 0 self.curriculum_episode_limit = 150 self.curriculum_stage_requirements = { 1: {"survival_steps": 1000, "distance": 0.00, "stability_roll_pitch": 0.35}, 2: {"survival_steps": 350, "distance": 10.0, "stability_roll_pitch": 0.30}, 3: {"survival_steps": 550, "distance": 15.00, "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 _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.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 = 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 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) 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) 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, (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"] 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, (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 = 10.0 * ((self.target_height - pos[2]) ** 2) if pos[2] < self.target_height else 0.0 stability_penalty = 1.5 * (roll**2 + pitch**2) print(-height_penalty) 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()