""" 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 from collections import defaultdict import gymnasium as gym from gymnasium import spaces import numpy as np from config import cfg from simulation import Simulation from Robot import Robot, PyBulletBackend from ml.MetricsOverlay import MetricsHUD 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 JackBot hexapod simulation.""" def __init__( self, use_gui: bool = True, random_command: bool = True, max_episode_steps: int = 3000, urdf_path: str = cfg.urdf_path, robot_mode: str = "direct", # "direct", "residual", or "kinematics" ): super().__init__() self.robot_mode = robot_mode 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.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 self.last_reward_components: Dict[str, float] = {} self.episode_reward_components_sum: Dict[str, float] = defaultdict(float) self.control_freq = 60 self.min_cmd_hold_steps = int(2.0 * self.control_freq) self.max_cmd_hold_steps = int(6.0 * self.control_freq) self.next_cmd_resample_step = 0 self.initial_stand_steps = 120 # --- INITIALIZE PYBULLET SIMULATION ENGINE --- self.sim = Simulation(urdf_path=self.urdf_path, use_gui=self.use_gui) self.plane_id, self.pb_robot, self.revolute_joints = self.sim.load_scene() # --- INITIALIZE ROBOT WITH BACKEND --- self.backend = PyBulletBackend(self.sim) self.robot = Robot(backend_type=self.backend, urdf_path=self.urdf_path, mode=self.robot_mode) action_dim = 18 obs_dim = 18 + 3 # 18 Joint Angles + 3 Command Inputs [vx, vy, omega] 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(3, dtype=np.float32) # [vx, vy, omega] 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.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) 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}, 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.physics_client) self.last_time = time.time() def sample_command(self) -> np.ndarray: """Samples a command vector [vx, vy, omega] based on active curriculum phase.""" phase = self.curriculum_phase stand_probs = { 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_probs.get(phase, 0.15): return np.zeros(3, dtype=np.float32) if phase == CurriculumPhase.FORWARD: vx, vy, omega = np.random.uniform(0.15, 0.50), 0.0, 0.0 elif phase == CurriculumPhase.TURN_AND_DIRECTION: vx, vy, omega = np.random.uniform(-0.8, 0.8), 0.0, np.random.uniform(-0.8, 0.8) elif phase == CurriculumPhase.OMNI_DIRECTION: vx, vy, omega = np.random.uniform(-0.8, 0.8), np.random.uniform(-0.5, 0.5), np.random.uniform(-0.8, 0.8) else: vx, vy, omega = np.random.uniform(-1.0, 1.0), np.random.uniform(-1.0, 1.0), np.random.uniform(-1.0, 1.0) return np.array([vx, vy, 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 self.last_reward_components = {} self.episode_reward_components_sum = defaultdict(float) spawn_pos = [0.0, 0.0, 0.20] spawn_orn = [0.0, 0.0, 0.0, 1.0] # 1. Reset base pose and velocities self.sim.reset_robot_base(spawn_pos, spawn_orn) # 2. Reset internal kinematics & hard reset joints in PyBullet self.robot.reset_to_init() if self.use_gui: self.sim.set_robot_color([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 = CurriculumPhase.STAND_ONLY self._first_reset = False self.command = np.zeros(3, dtype=np.float32) self.next_cmd_resample_step = self.initial_stand_steps pos, _ = self.sim.get_robot_pose() self.start_position = list(pos) # Drop settlement self.target_height = self.sim.settle_and_measure_height(steps=200, fallback_height=0.122) self.default_joint_angles = self.sim.get_robot_joint_angles() if self.use_gui: self.hud.reset() self._update_hud() return self._get_obs(), {} def _get_obs(self) -> np.ndarray: return self.robot.get_observation(command=self.command) 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() # Command resampling 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 # Extract [vx, vy, omega] cmd_vx, cmd_vy, cmd_omega = self.command # Mirror main.py input resolution logic: update robot_state and vector_dirmov directly self.robot.robot_state = "walking" if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01) else "idle" self.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)] # Delegate execution tick to Robot instance self.robot.tick(action=action) if self.robot_mode != "kinematics" and self.step_count % 60 == 0: random_force = np.random.uniform(-2.0, 2.0, size=2) self.sim.apply_external_force(force=[random_force[0], random_force[1], 0.0]) 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 info = {"reward_components": self.last_reward_components.copy()} if self.step_count % 120 == 0 and self.use_gui: self._update_hud() return obs, reward, terminated, truncated, info def _update_distance_metrics(self): pos, _ = self.sim.get_robot_pose() 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] survival_ok = self.max_survival_steps >= req["survival_steps"] 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) 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 pos, _ = self.sim.get_robot_pose() dx, dy = pos[0] - self.start_position[0], pos[1] - self.start_position[1] 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 = abs(dy) <= req["max_lateral_drift"] if "max_lateral_drift" in req else True 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 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 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 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 def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> float: pos, (roll, pitch, yaw) = self.sim.get_robot_pose_and_rpy() linear_vel, angular_vel = self.sim.get_robot_velocity() current_joints = self.sim.get_robot_joint_angles() cmd_vx, cmd_vy, cmd_yaw = self.command cmd_norm = math.hypot(cmd_vx, cmd_vy) raw_speed = math.hypot(linear_vel[0], linear_vel[1]) filtered_vx, filtered_vy = (linear_vel[0], linear_vel[1]) if raw_speed >= 0.04 else (0.0, 0.0) filtered_speed = raw_speed if raw_speed >= 0.04 else 0.0 raw_yaw_rate = abs(angular_vel[2]) filtered_yaw_rate = angular_vel[2] if raw_yaw_rate >= 0.05 else 0.0 height_error = pos[2] - self.target_height r_height = math.exp(-150.0 * (height_error ** 2)) r_stability = math.exp(-25.0 * (roll**2 + pitch**2)) r_pose = math.exp(-2.0 * np.mean(np.square(current_joints - self.default_joint_angles))) r_smoothness = math.exp(-0.1 * np.mean(np.square(action - previous_action))) r_lin_vel, r_ang_vel, stillness_penalty = 0.0, 0.0, 0.0 if cmd_norm < 0.05 and abs(cmd_yaw) < 0.05: w_height, w_stability, w_pose, w_smoothness = 0.35, 0.35, 0.20, 0.10 total_reward = (w_height * r_height) + (w_stability * r_stability) + (w_pose * r_pose) + (w_smoothness * r_smoothness) else: is_moving = (filtered_speed > 0.0) or (abs(filtered_yaw_rate) > 0.0) target_vx, target_vy = cmd_vx * self.max_robot_speed, cmd_vy * self.max_robot_speed target_speed = math.hypot(target_vx, target_vy) if not is_moving: total_reward = 0.0 else: lin_vel_error = (filtered_vx - target_vx)**2 + (filtered_vy - target_vy)**2 r_lin_vel = math.exp(-25.0 * lin_vel_error) r_ang_vel = math.exp(-15.0 * ((filtered_yaw_rate - cmd_yaw)**2)) if target_speed > 0.08 and raw_speed < 0.03: r_lin_vel = 0.0 stillness_penalty = -0.25 w_lin_vel, w_ang_vel, w_height, w_stability, w_smoothness = 0.55, 0.15, 0.10, 0.12, 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 ) final_reward = float(total_reward / 10.0) self.last_reward_components = { "height": float(r_height), "stability": float(r_stability), "pose": float(r_pose), "smoothness": float(r_smoothness), "lin_vel": float(r_lin_vel), "ang_vel": float(r_ang_vel), "total": final_reward, } for k, v in self.last_reward_components.items(): self.episode_reward_components_sum[k] += v return final_reward def get_reward_component_averages(self) -> Dict[str, float]: steps = max(1, self.step_count) return {k: v / steps for k, v in self.episode_reward_components_sum.items()} def get_current_robot_metrics(self) -> List[Dict[str, Any]]: linear_vel, angular_vel = self.sim.get_robot_velocity() speed = float(math.hypot(linear_vel[0], linear_vel[1])) yaw_rate = float(abs(angular_vel[2])) metrics = { "alive": not self.is_failed, "phase_name": self.curriculum_phase.name, "reward": float(self.cumulative_reward), "speed": speed, "yaw_rate": yaw_rate, "distance_from_start": float(self.max_distance_from_start), "survival_steps": int(self.max_survival_steps), } return [metrics] 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.get_robot_pose_and_rpy() 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 or self.step_count < 15: return position, (roll, pitch, _) = self.sim.get_robot_pose_and_rpy() 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.set_robot_color(COLOR_FAILED) def close(self): self.sim.disconnect()