""" ml/SimManager.py - PyBullet Simulation & Multi-Body Manager """ from typing import Dict, List, Tuple, Optional import pybullet as p import pybullet_data import numpy as np 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 set_rendering(self, enabled: bool) -> None: """Toggles PyBullet 3D rendering to speed up simulation.""" if self.physics_client is not None and p.isConnected(self.physics_client): p.configureDebugVisualizer( p.COV_ENABLE_RENDERING, 1 if enabled else 0, 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 def get_contact_points( self, bodyA: int = -1, bodyB: int = -1, linkIndexA: int = -1, linkIndexB: int = -1, ): """Wrapper around pybullet.getContactPoints bound to this simulation client.""" kwargs = {"physicsClientId": self.physics_client} if bodyA != -1: kwargs["bodyA"] = bodyA if bodyB != -1: kwargs["bodyB"] = bodyB if linkIndexA != -1: kwargs["linkIndexA"] = linkIndexA if linkIndexB != -1: kwargs["linkIndexB"] = linkIndexB return p.getContactPoints(**kwargs) # --- ROBOT GETTERS AND SETTERS --- def reset_robot_base( self, body_id: int, position: List[float], orientation: Optional[List[float]] = None, linear_velocity: Optional[List[float]] = None, angular_velocity: Optional[List[float]] = None ) -> None: """Resets a robot body's base position, orientation, and velocities.""" if orientation is None: orientation = [0.0, 0.0, 0.0, 1.0] if linear_velocity is None: linear_velocity = [0.0, 0.0, 0.0] if angular_velocity is None: angular_velocity = [0.0, 0.0, 0.0] p.resetBasePositionAndOrientation( body_id, position, orientation, physicsClientId=self.physics_client ) p.resetBaseVelocity( body_id, linearVelocity=linear_velocity, angularVelocity=angular_velocity, physicsClientId=self.physics_client ) def get_robot_pose(self, body_id: int) -> Tuple[List[float], List[float]]: """Returns base position (x, y, z) and orientation quaternion (x, y, z, w).""" pos, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client) return list(pos), list(orn) def get_robot_rpy(self, body_id: int) -> Tuple[float, float, float]: """Returns roll, pitch, yaw angles in radians for the given robot body.""" _, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client) roll, pitch, yaw = p.getEulerFromQuaternion(orn) return float(roll), float(pitch), float(yaw) def get_foot_link_indices(self, body_id: int) -> list[int]: """Inspects URDF joint structure to extract link IDs for leg tips and tibias.""" foot_indices = [] num_joints = p.getNumJoints(body_id, physicsClientId=self.physics_client) for j_idx in range(num_joints): info = p.getJointInfo(body_id, j_idx, physicsClientId=self.physics_client) link_name = info[12].decode("utf-8") if "tip" in link_name or "tibia" in link_name: foot_indices.append(j_idx) return foot_indices def get_robot_pose_and_rpy(self, body_id: int) -> Tuple[List[float], Tuple[float, float, float]]: """Returns base position and (roll, pitch, yaw) tuple in radians.""" pos, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client) roll, pitch, yaw = p.getEulerFromQuaternion(orn) return list(pos), (float(roll), float(pitch), float(yaw)) def get_robot_velocity(self, body_id: int) -> Tuple[List[float], List[float]]: """Returns linear velocity (vx, vy, vz) and angular velocity (wx, wy, wz).""" lin_v, ang_v = p.getBaseVelocity(body_id, physicsClientId=self.physics_client) return list(lin_v), list(ang_v) def get_robot_joint_angles(self, body_id: int, joint_indices: Optional[List[int]] = None) -> np.ndarray: """Returns joint angles as a 1D numpy array float32 for specified or registered joint indices.""" if joint_indices is None: joint_indices = self.robot_joints.get(body_id, list(range(18))) joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=self.physics_client) return np.array([state[0] for state in joint_states], dtype=np.float32) def set_robot_color(self, body_id: int, rgba: List[float]) -> None: """Changes visual color RGBA of base link and all joints of the specified robot body.""" num_joints = p.getNumJoints(body_id, physicsClientId=self.physics_client) p.changeVisualShape(body_id, -1, rgbaColor=rgba, physicsClientId=self.physics_client) for j in range(num_joints): p.changeVisualShape(body_id, j, rgbaColor=rgba, physicsClientId=self.physics_client) def measure_robot_heights(self, robot_ids: List[int]) -> List[float]: """Gets current Z height for all specified robot body IDs.""" heights = [] for body_id in robot_ids: pos, _ = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client) heights.append(pos[2]) return heights def settle_and_measure_height(self, robot_ids: List[int], steps: int = 200, fallback_height: float = 0.14) -> float: """Steps simulation for designated steps so robot settles, then calculates target standing height.""" for _ in range(steps): self.step() heights = self.measure_robot_heights(robot_ids) mean_height = float(np.mean(heights)) if heights else fallback_height return mean_height if mean_height > 0.0 else fallback_height def apply_external_force( self, body_id: int, force: list[float] | np.ndarray, link_index: int = -1, position: list[float] | np.ndarray = (0.0, 0.0, 0.0), frame: int = p.WORLD_FRAME, ): """ Applies a 3D force vector (in Newtons) to a robot link. :param body_id: PyBullet body ID. :param force: [fx, fy, fz] force vector in Newtons. :param link_index: Target link index (-1 refers to the base/torso). :param position: Offset [x, y, z] relative to link center where force is applied. :param frame: p.WORLD_FRAME (global axes) or p.LINK_FRAME (robot's body axes). """ p.applyExternalForce( objectUniqueId=body_id, linkIndex=link_index, forceObj=list(force), posObj=list(position), flags=frame, physicsClientId=self.physics_client, )