""" 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.min_joint_limits, self.max_joint_limits = self.sim._get_urdf_joint_limits() self.command = np.zeros(3, dtype=np.float32) # [vx, vy, omega] self.last_action = np.zeros(action_dim, dtype=np.float32) self.last_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.15] 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]) action_dim = self.action_space.shape[0] self.last_action = np.zeros(action_dim, dtype=np.float32) self.last_last_action = np.zeros(action_dim, 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( target_angles=self.robot.current_rad, steps=300, fallback_height=0.122 ) self.default_joint_angles = self.sim.get_robot_joint_angles() self.default_action = np.clip(self.default_joint_angles / (np.pi / 2.0), -1.0, 1.0) 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]]: previous_action = self.last_action.copy() self.step_count += 1 self.total_steps += 1 self.last_last_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 target_angles = np.where( action < 0.0, self.default_joint_angles + action * (self.default_joint_angles - self.min_joint_limits), self.default_joint_angles + action * (self.max_joint_limits - self.default_joint_angles) ) self.robot.tick(action=target_angles) 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 if self.curriculum_phase < CurriculumPhase.FORWARD and (self._phase_progress_ready(CurriculumPhase.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]) # Forgiving zones: ignore noise below 0.04 m/s and 0.05 rad/s 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 = raw_yaw_rate if raw_yaw_rate >= 0.05 else 0.0 # --- 1. FIXED JITTER PENALTY --- # Scaled way down (0.005) and capped so it can never dominate the reward action_accel = action - 2.0 * self.last_action + self.last_last_action raw_jitter = float(np.mean(np.square(action_accel))) jitter_penalty = -0.005 * min(raw_jitter, 10.0) # --- Base Components --- 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 stand_penalty = 0.0 # --- 2. COMMAND IS ZERO: STANDING MODE --- 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 base_reward = (w_height * r_height) + (w_stability * r_stability) + (w_pose * r_pose) + (w_smoothness * r_smoothness) # Soft quadratic penalty on filtered speed (gives a forgiving gap near 0) stand_penalty = -0.5 * filtered_speed - 0.1 * filtered_yaw_rate total_reward = base_reward + stand_penalty # --- 3. COMMAND IS NON-ZERO: WALKING MODE --- else: is_moving = (filtered_speed > 0.0) or (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.1 # Softened from -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 ) step_reward = float(total_reward / 10.0) alive_bonus = 0.01 final_reward = max(0.0, step_reward + jitter_penalty + alive_bonus) 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), "jitter_penalty": float(jitter_penalty), "stand_penalty": float(stand_penalty), "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: 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()