Training uses new Robot.py

new Metrics and SimManager for live training viewing
Robot.py usage added to machinelearning

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,11 +50,16 @@ 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:
+            # 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:
@@ -61,8 +67,8 @@ class PyBulletBackend:
             self.sim.step()
 
     def cleanup(self) -> None:
-        if self.sim and hasattr(self.sim, 'close'):
+        if self.sim and hasattr(self.sim, 'disconnect'):
-            self.sim.close()
+            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()

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

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

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):
+    """Gymnasium environment wrapping JackBot hexapods with floating text & crown overlay."""
-        self.sim_manager = SimManager(use_gui=use_gui)
+
+    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
+        # Connect physics world
-        self.plane, self.pb_robots, self.joint_indices = self.sim_manager.load_scene(
+        self.plane, self.pb_robots, self.robot_joint_indices = self.sim_manager.load_scene(
-            cfg.urdf_path, num_robots=num_robots
+            self.urdf_path, self.num_robots, self.robot_spacing, self._robot_base_position
         )
 
-        # Instantiate dedicated Robot Python object for EACH spawned robot
+        # Instantiate Robot Python wrappers per PyBullet body ID
         self.robots = [
-            Robot(backend=PyBulletBackend(self.sim_manager, body_id=pb_id))
+            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 & Observation Spaces
+        # Action (18 joint deltas per robot) & Observation (18 angles + 4 command dims per robot)
-        action_dim = num_robots * 18
+        action_dim = self.num_robots * 18
-        obs_dim = num_robots * (18 + 4)
+        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]:
-            """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]
+        return [x, y, 0.2]
-
-    def _connect_sim(self):
-        self.sim_manager.connect()
-        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
-        )
-
-        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
-            )
-        else:
-            self._first_reset = False
-
-        self.last_action = np.zeros(self.num_robots * len(self.robot_joint_indices[0]), dtype=np.float32)
-
-        if seed is not None:
-            self.seed(seed)
-
-        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._reset_pose()
-        self._update_gui_hud(reward=0.0)
-
-        return self._get_obs(), {}
-
-    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 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):
+        return np.concatenate(obs_list).astype(np.float32)
-        action_per_robot = action.reshape(len(self.robots), 18)
+
+    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
+        # Update failure status & gray coloring
-        obs = np.concatenate([
+        self._update_robot_failures()
-            robot.get_observation(command=np.zeros(4)) 
-            for robot in self.robots
-        ])
 
-        reward = self._compute_reward()
+        obs = self._get_obs()
-        done = False
+        reward, per_robot_step_rewards = self._compute_reward()
-        return obs, reward, done, False, {}
         
-    def _update_gui_hud(self, reward: float):
+        self.cumulative_reward += reward
-        """Passes current state metrics to the SimManager HUD renderer."""
+        for idx, r_step in enumerate(per_robot_step_rewards):
-        linear_vel, _ = p.getBaseVelocity(self.robots[0])
+            self.robot_rewards[idx] += r_step
-        cmd = self.commands[0]
 
-        stats = {
+        terminated = self._is_done()
-            "Episode": self.episode_count,
+        truncated = self.step_count >= self.max_episode_steps
-            "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)
 
-    def _compute_reward(self) -> float:
+        self._update_hud()
+        self._update_leader_visuals()
+        return obs, reward, terminated, truncated, {}
+
+    def _compute_reward(self) -> Tuple[float, list[float]]:
         rewards = []
-        for robot_id, robot in enumerate(self.robots):
+        for idx, pb_id in enumerate(self.pb_robots):
-            linear_vel, angular_vel = p.getBaseVelocity(robot)
+            linear_vel, angular_vel = p.getBaseVelocity(pb_id, physicsClientId=self.sim_manager.physics_client)
-            _, orientation = p.getBasePositionAndOrientation(robot)
+            _, 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:
+        """Returns True only when the percentage of failed robots exceeds the threshold."""
-            _, orientation = p.getBasePositionAndOrientation(robot)
+        failed_count = sum(self.failed_robots_mask)
-            roll, pitch, _ = p.getEulerFromQuaternion(orientation)
+        failure_ratio = failed_count / self.num_robots
-            if abs(roll) > 0.7 or abs(pitch) > 0.7:
+        return failure_ratio >= self.termination_threshold
-                return True
-            position, _ = p.getBasePositionAndOrientation(robot)
-            if position[2] < 0.02:
-                return True
 
-        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()

ml/sim_manager.py

@@ -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()