JackM323
c5ca79a354
Training environment to make a walk model for the hexapod generated code that will be checked
105 lines
3.8 KiB
Python
105 lines
3.8 KiB
Python
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() |