diff --git a/Robot.py b/Robot.py index f00ce6c..d06f164 100644 --- a/Robot.py +++ b/Robot.py @@ -6,6 +6,7 @@ Handles state, kinematics, backends (Hardware/Simulation), and motion execution. from typing import Protocol, Optional, Union import numpy as np import math +import pybullet as p from states import STATE_REGISTRY from states.State import State @@ -49,20 +50,25 @@ class HardwareBackend: class PyBulletBackend: """Backend for PyBullet simulation execution.""" - def __init__(self, sim_instance): + def __init__(self, sim_instance, body_id: Optional[int] = None): self.sim = sim_instance + self.body_id = body_id def send_angles(self, rad_array: dt.RadArray) -> None: if self.sim: - self.sim.updatePos(rad_array) + # If body_id is set, target that specific robot body + if self.body_id is not None and hasattr(self.sim, 'updatePosForBody'): + self.sim.updatePosForBody(self.body_id, rad_array) + else: + self.sim.updatePos(rad_array) def step_simulation(self) -> None: if self.sim: self.sim.step() def cleanup(self) -> None: - if self.sim and hasattr(self.sim, 'close'): - self.sim.close() + if self.sim and hasattr(self.sim, 'disconnect'): + self.sim.disconnect() class Robot: @@ -102,6 +108,9 @@ class Robot: ri.get_center_points() if hasattr(ri, "get_center_points") else ri.center_points ) + # RL configuration + self.action_scale = 0.1 # Joint delta step size (radians) + # Gait / motion variables self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"]) self.robot_state = "idle" @@ -158,6 +167,51 @@ class Robot: self.transition_to(next_state_key) self.step_sim() + # --- RL METHODS --- + + def apply_rl_action(self, action: np.ndarray) -> None: + """ + Applies continuous RL action deltas [-1, 1] to current joint angles. + """ + action = np.asarray(action, dtype=np.float32) + scaled_action = np.clip(action, -1.0, 1.0) * self.action_scale + + current_flat = self.current_rad.data.flatten() + updated_flat = np.clip( + current_flat + scaled_action, + -np.pi / 2, + np.pi / 2 + ) + + new_rad = dt.RadArray(data=updated_flat.reshape(self.current_rad.data.shape)) + self.set_joint_angles(new_rad) + + def get_observation(self, command: Optional[np.ndarray] = None) -> np.ndarray: + """ + Returns observation vector [18 joint angles] + [optional 4 command dimensions]. + Queries PyBullet if backend is PyBulletBackend; otherwise falls back to internal state. + """ + if isinstance(self.backend, PyBulletBackend) and self.backend.sim and hasattr(self.backend.sim, 'physics_client'): + physics_client = self.backend.sim.physics_client + body_id = self.backend.body_id if self.backend.body_id is not None else 0 + + # Retrieve joint mapping from SimManager/Simulation if available + if hasattr(self.backend.sim, 'robot_joints') and body_id in self.backend.sim.robot_joints: + joint_indices = self.backend.sim.robot_joints[body_id] + else: + joint_indices = list(range(18)) + + joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=physics_client) + joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32) + else: + joint_angles = self.current_rad.data.flatten().astype(np.float32) + + if command is not None: + cmd = np.asarray(command, dtype=np.float32).flatten() + return np.concatenate([joint_angles, cmd]).astype(np.float32) + + return joint_angles.astype(np.float32) + def cleanup(self) -> None: if self.backend and hasattr(self.backend, 'cleanup'): self.backend.cleanup() \ No newline at end of file diff --git a/ml/MetricsOverlay.py b/ml/MetricsOverlay.py new file mode 100644 index 0000000..31afc47 --- /dev/null +++ b/ml/MetricsOverlay.py @@ -0,0 +1,137 @@ +""" +ml/MetricsOverlay.py - Camera-Facing (Billboard) 3D Floating Text Overlay +""" +from typing import List, Tuple, Optional +import numpy as np +import pybullet as p + + +class MetricsHUD: + """Renders real-time telemetry as black text floating in 3D, always facing the active camera.""" + + def __init__(self, physics_client_id: int = 0): + self.client_id = physics_client_id + self._text_id: Optional[int] = None + + def _get_camera_facing_orientation(self) -> List[float]: + """Calculates a quaternion that points the text towards the current GUI camera.""" + try: + cam_info = p.getDebugVisualizerCamera(physicsClientId=self.client_id) + # cam_info index 8: yaw, index 9: pitch + yaw = cam_info[8] + pitch = cam_info[9] + + # Orient the text normal toward the camera view direction + # PyBullet text default faces local +Z/-Y depending on roll, + # converting visualizer yaw/pitch to Euler angles (roll, pitch, yaw in radians) + roll_rad = 0.0 + pitch_rad = np.radians(pitch + 90.0) + yaw_rad = np.radians(yaw) + + text_orientation = p.getQuaternionFromEuler( + [pitch_rad, roll_rad, yaw_rad], + physicsClientId=self.client_id + ) + return text_orientation + except Exception: + # Fallback default orientation if camera info call fails + return [0.0, 0.0, 0.0, 1.0] + + def update( + self, + episode: int, + step: int, + robot_rewards: List[float], + cmd_vel: np.ndarray, + fps: float = 0.0, + avg_height: float = 0.0, + roll_pitch: Tuple[float, float] = (0.0, 0.0) + ) -> None: + """Updates floating black text block in 3D space with billboarding.""" + sorted_rewards = sorted(robot_rewards, reverse=True) + top1 = f"{sorted_rewards[0]:+.2f}" if len(sorted_rewards) > 0 else "0.00" + top2 = f"{sorted_rewards[1]:+.2f}" if len(sorted_rewards) > 1 else "0.00" + top3 = f"{sorted_rewards[2]:+.2f}" if len(sorted_rewards) > 2 else "0.00" + + vx = cmd_vel[0] if len(cmd_vel) > 0 else 0.0 + vy = cmd_vel[1] if len(cmd_vel) > 1 else 0.0 + omega = cmd_vel[3] if len(cmd_vel) > 3 else 0.0 + + hud_text = ( + f"=== JACKBOT METRICS ===\n" + f"Episode: {episode}\n" + f"Global Step: {step}\n" + f"FPS: {fps:.1f}\n" + f"----------------------\n" + f"Top Rewards: [{top1}, {top2}, {top3}]\n" + f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]\n" + f"Height: {avg_height:.3f} m\n" + f"Roll/Pitch: {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°" + ) + + # Position above origin in simulation world + text_position = [-0.8, -0.8, 1.2] + text_color = [0, 0, 0] # Pure black + + # Calculate dynamic orientation to align text flat against camera plane + text_orientation = self._get_camera_facing_orientation() + + if self._text_id is None: + self._text_id = p.addUserDebugText( + text=hud_text, + textPosition=text_position, + textColorRGB=text_color, + textSize=0.1, + textOrientation=text_orientation, + physicsClientId=self.client_id + ) + else: + self._text_id = p.addUserDebugText( + text=hud_text, + textPosition=text_position, + textColorRGB=text_color, + textSize=0.1, + textOrientation=text_orientation, + replaceItemUniqueId=self._text_id, + physicsClientId=self.client_id + ) + + def reset(self) -> None: + """Clears text reference on environment reset.""" + self._text_id = None + + +class LeaderCrown: + """Renders a floating crown or star emoji above the leading robot in PyBullet.""" + + def __init__(self, physics_client_id: int = 0): + self.client_id = physics_client_id + self._text_id = None + + def update(self, leader_pos: list[float]): + """Positions a floating crown ~0.35m directly above the lead robot's base.""" + crown_pos = [leader_pos[0], leader_pos[1], leader_pos[2] + 0.35] + + # You can use "👑 CROWN", "⭐ LEADER", or "★ TOP1" + crown_text = "👑" + + if self._text_id is None: + self._text_id = p.addUserDebugText( + text=crown_text, + textPosition=crown_pos, + textColorRGB=[1.0, 0.84, 0.0], + textSize=2.0, + physicsClientId=self.client_id + ) + else: + self._text_id = p.addUserDebugText( + text=crown_text, + textPosition=crown_pos, + textColorRGB=[1.0, 0.84, 0.0], + textSize=2.0, + replaceItemUniqueId=self._text_id, + physicsClientId=self.client_id + ) + + def reset(self): + self._text_id = None \ No newline at end of file diff --git a/ml/SimManager.py b/ml/SimManager.py new file mode 100644 index 0000000..d712ce7 --- /dev/null +++ b/ml/SimManager.py @@ -0,0 +1,88 @@ +""" +ml/SimManager.py - PyBullet Simulation & Multi-Body Manager +""" +from typing import Dict, List, Tuple +import pybullet as p +import pybullet_data +import DataTypes as dt + + +class SimManager: + """Manages PyBullet simulation lifecycle and multi-robot physics.""" + + def __init__(self, use_gui: bool = True): + self.use_gui = use_gui + self.physics_client = None + self.robot_joints: Dict[int, List[int]] = {} + + def connect(self): + """Connects to PyBullet and hides side GUI panels.""" + if self.physics_client is not None and p.isConnected(self.physics_client): + return + + flags = p.GUI if self.use_gui else p.DIRECT + self.physics_client = p.connect(flags) + + p.setAdditionalSearchPath(pybullet_data.getDataPath()) + p.setGravity(0, 0, -9.81, physicsClientId=self.physics_client) + + if self.use_gui: + # Disable PyBullet side panel and preview windows + p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.physics_client) + p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0, physicsClientId=self.physics_client) + p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client) + p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client) + + def load_scene( + self, urdf_path: str, num_robots: int, robot_spacing: float, base_pos_fn + ) -> Tuple[int, List[int], List[List[int]]]: + """Loads plane and hexapod bodies into the simulation scene.""" + plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client) + robots = [] + robot_joint_indices = [] + self.robot_joints.clear() + + for r_id in range(num_robots): + base_pos = base_pos_fn(r_id, num_robots, robot_spacing) + robot = p.loadURDF( + urdf_path, + basePosition=base_pos, + useFixedBase=False, + physicsClientId=self.physics_client + ) + robots.append(robot) + + joint_indices = [ + i for i in range(p.getNumJoints(robot, physicsClientId=self.physics_client)) + if p.getJointInfo(robot, i, physicsClientId=self.physics_client)[2] == p.JOINT_REVOLUTE + ] + robot_joint_indices.append(joint_indices) + self.robot_joints[robot] = joint_indices + + return plane_id, robots, robot_joint_indices + + def updatePosForBody(self, body_id: int, current_rad: dt.RadArray): + """Sets joint motor position targets on individual robot bodies.""" + if body_id not in self.robot_joints: + return + + joint_indices = self.robot_joints[body_id] + radflat = current_rad.data.flatten() + + for joint_index, target_angle in zip(joint_indices, radflat): + p.setJointMotorControl2( + bodyIndex=body_id, + jointIndex=joint_index, + controlMode=p.POSITION_CONTROL, + targetPosition=float(target_angle), + force=250, + physicsClientId=self.physics_client + ) + + def step(self): + p.stepSimulation(physicsClientId=self.physics_client) + + def disconnect(self): + if self.physics_client is not None and p.isConnected(self.physics_client): + p.disconnect(self.physics_client) + self.physics_client = None \ No newline at end of file diff --git a/ml/env.py b/ml/env.py index 897dac8..4df724d 100644 --- a/ml/env.py +++ b/ml/env.py @@ -1,130 +1,104 @@ -# ml/env.py +""" +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 math +import pybullet as p +from config import cfg from Robot import Robot, PyBulletBackend -from ml.sim_manager import SimManager +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): - def __init__(self, use_gui: bool = False, num_robots: int = 1): - self.sim_manager = SimManager(use_gui=use_gui) + """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() - - # Load simulation bodies - self.plane, self.pb_robots, self.joint_indices = self.sim_manager.load_scene( - cfg.urdf_path, num_robots=num_robots - ) - # Instantiate dedicated Robot Python object for EACH spawned robot - self.robots = [ - Robot(backend=PyBulletBackend(self.sim_manager, body_id=pb_id)) - for pb_id in self.pb_robots - ] - - # Action & Observation Spaces - action_dim = num_robots * 18 - obs_dim = 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) - def _robot_base_position(self, robot_id: int, num_robots: int = 1, spacing: float = 0.5) -> list[float]: - """Calculates grid coordinates for spawning multiple robots in PyBullet.""" - 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.1] - - def _connect_sim(self): - self.sim_manager.connect() - self.plane, self.robots, self.robot_joint_indices = self.sim_manager.load_scene( + # 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 ) - joint_count = len(self.robot_joint_indices[0]) - if any(len(indices) != joint_count for indices in self.robot_joint_indices): - raise ValueError("All robots must have the same number of revolute joints.") - - lower_limits, upper_limits = [], [] - for joint_index in self.robot_joint_indices[0]: - info = p.getJointInfo(self.robots[0], joint_index) - lower, upper = info[8], info[9] - if lower >= upper: - lower, upper = -math.pi, math.pi - lower_limits.append(lower) - upper_limits.append(upper) - - limits = np.array([lower_limits, upper_limits], dtype=np.float32) - self.joint_lower = np.tile(limits[0], (self.num_robots, 1)) - self.joint_upper = np.tile(limits[1], (self.num_robots, 1)) - - pose = ri.init_deg if self.start_pose == "init_deg" else ri.init90_deg - initial_pose = pose.to_rad().data.flatten() - self.initial_angles = np.tile(initial_pose, (self.num_robots, 1)).astype(np.float32) - self.initial_angles = np.clip(self.initial_angles, self.joint_lower, self.joint_upper) - - self.last_action = np.zeros(self.num_robots * joint_count, dtype=np.float32) - self.commands = np.zeros((self.num_robots, 4), dtype=np.float32) - - def reset(self, command: np.ndarray | None = None, seed: int | None = None, options: dict | None = None): - self.episode_count += 1 - self.step_count = 0 - self.cumulative_reward = 0.0 - - if not self._first_reset: - # Reset simulation bodies without reconnecting PyBullet - p.resetSimulation() - p.setGravity(0, 0, -9.81) - self.sim_manager.reset_hud() - self.plane, self.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 ) - else: - self._first_reset = False + for pb_id in self.pb_robots + ] - self.last_action = np.zeros(self.num_robots * len(self.robot_joint_indices[0]), dtype=np.float32) + # 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) - if seed is not None: - self.seed(seed) + 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) - if command is not None: - self.commands = np.array(command, dtype=np.float32).reshape(self.num_robots, 4) - elif 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) + self.commands = np.zeros((self.num_robots, 4), dtype=np.float32) + self.last_action = np.zeros(action_dim, dtype=np.float32) - self._reset_pose() - self._update_gui_hud(reward=0.0) + # 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() - return self._get_obs(), {} + 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 seed(self, seed: int | None = None): - if seed is None: - seed = int(np.random.randint(0, 2**31 - 1)) - random.seed(seed) - np.random.seed(seed) - self._seed = seed - return [seed] - - def _reset_pose(self): - for robot, joint_indices, init_angles in zip(self.robots, self.robot_joint_indices, self.initial_angles): - for joint_index, target_angle in zip(joint_indices, init_angles): - p.resetJointState(robot, joint_index, target_angle) - p.setJointMotorControl2( - bodyIndex=robot, - jointIndex=joint_index, - controlMode=p.POSITION_CONTROL, - targetPosition=target_angle, - force=250, - ) - - for _ in range(400): - p.stepSimulation() - - self.joint_angles = self.initial_angles.copy() + 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) @@ -133,72 +107,173 @@ class JackBotEnv(gym.Env): 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)] + + if not self._first_reset: + p.resetSimulation(physicsClientId=self.sim_manager.physics_client) + p.setGravity(0, 0, -9.81, physicsClientId=self.sim_manager.physics_client) + p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.sim_manager.physics_client) + self.hud.reset() + self.leader_crown.reset() + 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 + ) + for robot_obj, pb_id in zip(self.robots, self.pb_robots): + robot_obj.backend = PyBulletBackend(self.sim_manager, body_id=pb_id) + else: + self._first_reset = False + + 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 robot in self.robots: + robot.reset_to_init() + + for _ in range(100): + self.sim_manager.step() + + self._update_hud() + return self._get_obs(), {} + def _get_obs(self) -> np.ndarray: - return np.concatenate([self.joint_angles.flatten(), self.commands.flatten()]).astype(np.float32) + 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) - def step(self, action: np.ndarray): - action_per_robot = action.reshape(len(self.robots), 18) + 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) - # Apply RL actions independently to each Robot object instance for robot, act in zip(self.robots, action_per_robot): robot.apply_rl_action(act) - # Step PyBullet physics engine once self.sim_manager.step() - # Gather observations across all robot objects - obs = np.concatenate([ - robot.get_observation(command=np.zeros(4)) - for robot in self.robots - ]) + # Update failure status & gray coloring + self._update_robot_failures() - reward = self._compute_reward() - done = False - return obs, reward, done, False, {} + 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 - def _update_gui_hud(self, reward: float): - """Passes current state metrics to the SimManager HUD renderer.""" - linear_vel, _ = p.getBaseVelocity(self.robots[0]) - cmd = self.commands[0] + terminated = self._is_done() + truncated = self.step_count >= self.max_episode_steps - stats = { - "Episode": self.episode_count, - "Step": f"{self.step_count} / {self.max_episode_steps}", - "Step Reward": reward, - "Total Reward": self.cumulative_reward, - "Target Cmd (vx,vy,w)": [cmd[0], cmd[1], cmd[3]], - "Actual Vel (vx,vy)": [linear_vel[0], linear_vel[1]], - } - self.sim_manager.update_hud(stats) + self._update_hud() + self._update_leader_visuals() + return obs, reward, terminated, truncated, {} - def _compute_reward(self) -> float: + def _compute_reward(self) -> Tuple[float, list[float]]: rewards = [] - for robot_id, robot in enumerate(self.robots): - linear_vel, angular_vel = p.getBaseVelocity(robot) - _, orientation = p.getBasePositionAndOrientation(robot) + 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[robot_id] + 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)[robot_id]))) * 0.01 + action_penalty = float(np.sum(np.square(self.last_action.reshape(self.num_robots, -1)[idx]))) * 0.01 - rewards.append(0.1 + forward_reward + rotation_reward - 0.2 * stability_penalty - action_penalty) + r_step = 0.1 + forward_reward + rotation_reward - 0.2 * stability_penalty - action_penalty + rewards.append(r_step) - return float(np.sum(rewards)) + return float(np.sum(rewards)), rewards def _is_done(self) -> bool: - for robot in self.robots: - _, orientation = p.getBasePositionAndOrientation(robot) - roll, pitch, _ = p.getEulerFromQuaternion(orientation) - if abs(roll) > 0.7 or abs(pitch) > 0.7: - return True - position, _ = p.getBasePositionAndOrientation(robot) - if position[2] < 0.02: - return True + """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 - return self.step_count >= self.max_episode_steps + 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() \ No newline at end of file diff --git a/ml/sim_manager.py b/ml/sim_manager.py deleted file mode 100644 index f056647..0000000 --- a/ml/sim_manager.py +++ /dev/null @@ -1,105 +0,0 @@ -from typing import Dict, Any, List, Tuple -import pybullet as p -import pybullet_data -import numpy as np - - -class SimManager: - """ - Manages the PyBullet simulation lifecycle and live HUD overlays. - """ - - def __init__(self, use_gui: bool = True): - self.use_gui = use_gui - self.physics_client = None - self.debug_text_ids: Dict[str, int] = {} - - def connect(self): - """Connects to PyBullet and sets up the basic physics world.""" - if self.physics_client is not None: - p.disconnect(self.physics_client) - - flags = p.GUI if self.use_gui else p.DIRECT - self.physics_client = p.connect(flags) - - p.setAdditionalSearchPath(pybullet_data.getDataPath()) - p.setGravity(0, 0, -9.81) - - if self.use_gui: - # Disable unnecessary side panels for a clean UI - p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0) - p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0) - p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0) - p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0) - - def load_scene( - self, urdf_path: str, num_robots: int, robot_spacing: float, base_pos_fn - ) -> Tuple[int, List[int], List[List[int]]]: - """Loads plane and robots into the active PyBullet simulation.""" - plane_id = p.loadURDF("plane.urdf") - robots = [] - robot_joint_indices = [] - - for robot_id in range(num_robots): - base_pos = base_pos_fn(robot_id, num_robots, robot_spacing) - robot = p.loadURDF(urdf_path, basePosition=base_pos, useFixedBase=False) - robots.append(robot) - - joint_indices = [ - i - for i in range(p.getNumJoints(robot)) - if p.getJointInfo(robot, i)[2] == p.JOINT_REVOLUTE - ] - robot_joint_indices.append(joint_indices) - - return plane_id, robots, robot_joint_indices - - def disconnect(self): - """Safely disconnects from the simulation.""" - if self.physics_client is not None: - p.disconnect(self.physics_client) - self.physics_client = None - self.debug_text_ids.clear() - - def update_hud(self, stats: Dict[str, Any]): - """ - Renders live training metrics onto the PyBullet 3D viewport. - Uses replaceItemUniqueId to prevent flicker. - """ - if not self.use_gui or self.physics_client is None: - return - - # Fixed position near the top-left of the origin in 3D world coordinates - x_pos, y_pos, z_start = -1.2, -1.2, 1.6 - line_height = 0.10 - - for i, (label, value) in enumerate(stats.items()): - if isinstance(value, float): - display_text = f"{label}: {value:.3f}" - elif isinstance(value, (list, np.ndarray)): - formatted_vals = ", ".join(f"{v:.2f}" for v in np.atleast_1d(value)) - display_text = f"{label}: [{formatted_vals}]" - else: - display_text = f"{label}: {value}" - - color = [0, 0, 0] # Black text for clear visibility against the light plane - - if label in self.debug_text_ids: - p.addUserDebugText( - display_text, - [x_pos, y_pos, z_start - i * line_height], - textColorRGB=color, - textSize=1.1, - replaceItemUniqueId=self.debug_text_ids[label], - ) - else: - self.debug_text_ids[label] = p.addUserDebugText( - display_text, - [x_pos, y_pos, z_start - i * line_height], - textColorRGB=color, - textSize=1.1, - ) - - def reset_hud(self): - """Clears debug text tracking.""" - self.debug_text_ids.clear() \ No newline at end of file diff --git a/ml/train.py b/ml/train.py index 64435b7..65cdfc1 100644 --- a/ml/train.py +++ b/ml/train.py @@ -110,4 +110,4 @@ if __name__ == "__main__": num_robots=args.num_robots, robot_spacing=args.robot_spacing, start_pose=args.start_pose, - ) + ) \ No newline at end of file