JackM323
acb3d671be
new reward/penalty system learning phases with curriculum learning new training parameters cleanup of old code better logging while training multiple environments instead of robots (they could bumb into each other)
87 lines
3.3 KiB
Python
87 lines
3.3 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, robot_spacing: float, base_pos_fn
|
|
) -> Tuple[int, List[int], List[List[int]]]:
|
|
"""Loads the plane and a single hexapod body into the simulation scene."""
|
|
plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client)
|
|
robots = []
|
|
robot_joint_indices = []
|
|
self.robot_joints.clear()
|
|
|
|
base_pos = base_pos_fn(0, 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 |