JackM323
61cb0150f0
new Metrics and SimManager for live training viewing Robot.py usage added to machinelearning
88 lines
3.4 KiB
Python
88 lines
3.4 KiB
Python
"""
|
|
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 |