Machine Learning Trainer

Training environment to make a walk model for the hexapod
generated code that will be checked

ml/sim_manager.py

@@ -0,0 +1,105 @@
+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_OUTPUT_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()