""" ml/env.py - Gymnasium environment for JackBot RL training and evaluation. This file defines JackBotEnv, the main training/evaluation environment used by PPO. It wraps the PyBullet simulation and Robot interfaces into a Gymnasium-compatible step/reset loop, manages command sampling, curriculum progression, and reward calculation, and exposes metrics that the training callbacks can log. """ 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 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() 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: # Read raw joint angles from backend raw_angles = np.asarray(self.robot.backend.get_joint_angles(), dtype=np.float32).flatten() min_lim = self.min_joint_limits.flatten() max_lim = self.max_joint_limits.flatten() # Map raw joint radians [min, max] -> normalized [-1, 1] normalized_joints = 2.0 * (raw_angles - min_lim) / (max_lim - min_lim) - 1.0 normalized_joints = np.clip(normalized_joints, -1.0, 1.0) # Concatenate normalized joints with active command vector obs = np.concatenate([normalized_joints, self.command]).astype(np.float32) return obs 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 ONLY if random_command is True 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 active [vx, vy, omega] cmd_vx, cmd_vy, cmd_omega = self.command # Mirror input resolution logic to keep robot state synchronized 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)] # Direct Mode: Target joint scaling action_flat = np.asarray(action, dtype=np.float32).flatten() action_clipped = np.clip(action_flat, -1.0, 1.0) if self.robot_mode == "direct": # Map [-1, 1] linearly to physical joint limits [min, max] min_lim = self.min_joint_limits.flatten() max_lim = self.max_joint_limits.flatten() target_angles = min_lim + (action_clipped + 1.0) * 0.5 * (max_lim - min_lim) else: # Residual mode mapping logic target_angles = self.default_joint_angles.flatten() + action_clipped * 0.20 # Apply target joint angles to physics engine self.robot.tick(action=target_angles, physics_substeps=4) 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() # Build next observation preserving active command obs = self._get_obs() reward = self._compute_reward(action_flat, 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() if self.use_gui: time.sleep(1.0 / self.control_freq) 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: # If the command says move but the robot stays effectively still, # give a real penalty instead of a neutral reward. stillness_penalty = -0.10 total_reward = stillness_penalty 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.05 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 = 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 def close(self): self.sim.disconnect()