""" 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, num_robots: int = 1, robot_spacing: float = 0.5, start_pose: str = "init_deg", max_episode_steps: int = 3000, urdf_path: str = cfg.urdf_path, termination_threshold: float = 0.001, # Termination ratio threshold ): super().__init__() self.use_gui = use_gui self.random_command = random_command self.num_robots = num_robots self.robot_spacing = robot_spacing self.start_pose = start_pose self.max_episode_steps = max_episode_steps self.urdf_path = urdf_path self.termination_threshold = termination_threshold self.episode_count = 0 self.step_count = 0 self.total_steps = 0 self.cumulative_reward = 0.0 self.robot_rewards = [0.0 for _ in range(self.num_robots)] self.failed_robots_mask = [False for _ in range(self.num_robots)] 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.num_robots, 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 = self.num_robots * 18 obs_dim = self.num_robots * (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((self.num_robots, 4), dtype=np.float32) self.last_action = np.zeros(action_dim, dtype=np.float32) # 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, num_robots: int = 1, spacing: float = 0.5) -> list[float]: cols = int(math.sqrt(num_robots - 1)) + 1 row = robot_id // cols col = robot_id % cols x = (col - (cols - 1) / 2.0) * spacing y = (row - (cols - 1) / 2.0) * spacing return [x, y, 0.2] def sample_command(self) -> np.ndarray: 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 for _ in range(self.num_robots)] self.failed_robots_mask = [False for _ in range(self.num_robots)] for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)): # Get default spawn position spawn_pos = self._robot_base_position(idx, self.num_robots, self.robot_spacing) spawn_orn = [0, 0, 0, 1] # Teleport base back to start 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 ) # Reset joint angles directly without reloading URDF robot_obj.reset_to_init() # Restore original default visual color (clears failure dark gray) 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) if self.random_command: self.commands = np.stack([self.sample_command() for _ in range(self.num_robots)]) else: self.commands = np.zeros((self.num_robots, 4), dtype=np.float32) for _ in range(15): self.sim_manager.step() if self.use_gui: self.hud.reset() self.leader_crown.reset() self._update_hud() return self._get_obs(), {} 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 step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]: self.step_count += 1 self.total_steps += 1 self.last_action = action.copy() action_per_robot = action.reshape(self.num_robots, 18) for robot, act in zip(self.robots, action_per_robot): robot.apply_rl_action(act) self.sim_manager.step() # Update failure status & gray coloring self._update_robot_failures() obs = self._get_obs() reward, per_robot_step_rewards = self._compute_reward() 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.max_episode_steps self._update_hud() self._update_leader_visuals() return obs, reward, terminated, truncated, {} def _compute_reward(self) -> Tuple[float, list[float]]: rewards = [] for idx, pb_id in enumerate(self.pb_robots): linear_vel, angular_vel = p.getBaseVelocity(pb_id, physicsClientId=self.sim_manager.physics_client) _, orientation = p.getBasePositionAndOrientation(pb_id, physicsClientId=self.sim_manager.physics_client) roll, pitch, _ = p.getEulerFromQuaternion(orientation) command = self.commands[idx] forward_reward = command[0] * linear_vel[0] + command[1] * linear_vel[1] rotation_reward = command[3] * angular_vel[2] stability_penalty = abs(roll) + abs(pitch) action_penalty = float(np.sum(np.square(self.last_action.reshape(self.num_robots, -1)[idx]))) * 0.01 r_step = 0.1 + forward_reward + rotation_reward - 0.2 * stability_penalty - action_penalty rewards.append(r_step) return float(np.sum(rewards)), rewards def _is_done(self) -> bool: """Returns True only when the percentage of failed robots exceeds the threshold.""" failed_count = sum(self.failed_robots_mask) failure_ratio = failed_count / self.num_robots return failure_ratio >= self.termination_threshold 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): """Positions floating crown above top robot without altering active materials.""" if not self.use_gui or self.num_robots <= 1: 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 for each robot and turns failed ones gray.""" for idx, pb_id in enumerate(self.pb_robots): if self.failed_robots_mask[idx]: continue # Already marked failed position, orientation = p.getBasePositionAndOrientation( pb_id, physicsClientId=self.sim_manager.physics_client ) roll, pitch, _ = p.getEulerFromQuaternion(orientation) is_tilted = abs(roll) > 0.7 or abs(pitch) > 0.7 is_collapsed = position[2] < 0.05 if is_tilted or is_collapsed: self.failed_robots_mask[idx] = True if self.use_gui: # Turn failed robot semi-transparent dark gray self._set_robot_color(pb_id, COLOR_FAILED) def close(self): self.sim_manager.disconnect()