""" ml/SimManager.py - PyBullet Simulation Manager (Single Robot Dedicated) """ from typing import List, Tuple, Optional, Union import pybullet as p import pybullet_data import numpy as np import DataTypes as dt class SimManager: """Manages PyBullet simulation lifecycle for a single JackBot hexapod.""" def __init__(self, use_gui: bool = True): self.use_gui = use_gui self.physics_client = None self.plane: Optional[int] = None self.joint_indices: List[int] = [] self.foot_indices: List[int] = [] def connect(self) -> None: """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: 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, spacing: float = 0.0, position_func=None ) -> Tuple[int, List[int], List[List[int]]]: """Loads ground plane and the single robot URDF with high ground friction and force settings.""" # 1. Load ground plane and set explicit friction self.plane = p.loadURDF("plane.urdf", physicsClientId=self.physics_client) # 2. Determine starting position (matching simulation.py starting height of 0.20m) spawn_pos = position_func(0) if position_func is not None else [0.0, 0.0, 0.20] # 3. Spawn single robot body self.robot_id = p.loadURDF(urdf_path, spawn_pos, physicsClientId=self.physics_client) # 4. Retrieve revolute joint indices self.joint_indices = [] for j in range(p.getNumJoints(self.robot_id, physicsClientId=self.physics_client)): info = p.getJointInfo(self.robot_id, j, physicsClientId=self.physics_client) if info[2] == p.JOINT_REVOLUTE: self.joint_indices.append(j) # Return format maintains 100% compatibility with JackBotEnv unpack sequence return self.plane, [self.robot_id], [self.joint_indices] def _resolve_body_id(self, body_id: Optional[int] = None) -> int: """Internal helper to return the single active robot ID.""" if body_id is not None: return body_id if self.robot_id is not None: return self.robot_id raise RuntimeError("No robot loaded in SimManager. Call load_scene() first.") def set_robot_joint_angles( self, target_angles: Union[np.ndarray, List[float]], joint_indices: Optional[List[int]] = None, body_id: Optional[int] = None, ) -> None: """Applies position control with force=500 matching simulation.py logic.""" bid = self._resolve_body_id(body_id) j_indices = joint_indices if joint_indices is not None else self.joint_indices radflat = target_angles.flatten() if isinstance(target_angles, np.ndarray) else target_angles for joint_index, target_angle in zip(j_indices, radflat): p.setJointMotorControl2( bodyIndex=bid, jointIndex=joint_index, controlMode=p.POSITION_CONTROL, targetPosition=float(target_angle), force=500, physicsClientId=self.physics_client ) def updatePosForBody( self, body_id_or_rad: Union[int, dt.RadArray], current_rad: Optional[dt.RadArray] = None ) -> None: """Updates robot joint positions using RadArray input.""" if isinstance(body_id_or_rad, dt.RadArray): rad_data = body_id_or_rad bid = self.robot_id else: bid = self._resolve_body_id(body_id_or_rad) rad_data = current_rad if rad_data is not None: self.set_robot_joint_angles(rad_data.data, body_id=bid) def step(self) -> None: p.stepSimulation(physicsClientId=self.physics_client) def set_rendering(self, enabled: bool) -> None: """Toggles PyBullet 3D rendering visualizer.""" 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) -> None: 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) # Add to ml/SimManager.py def hard_reset_joint_angles(self, target_angles: np.ndarray, body_id: Optional[int] = None) -> None: bid = self._resolve_body_id(body_id) radflat = target_angles.flatten() for joint_index, target_angle in zip(self.joint_indices, radflat): p.resetJointState( bodyUniqueId=bid, jointIndex=joint_index, targetValue=float(target_angle), targetVelocity=0.0, physicsClientId=self.physics_client ) # --- SINGLE ROBOT GETTERS & SETTERS --- def reset_robot_base( self, body_id_or_pos: Union[int, List[float]], position_or_orn: Optional[List[float]] = None, orientation: Optional[List[float]] = None, linear_velocity: Optional[List[float]] = None, angular_velocity: Optional[List[float]] = None ) -> None: """Resets the robot base pose and clears linear/angular velocities.""" if isinstance(body_id_or_pos, int): bid = body_id_or_pos pos = position_or_orn if position_or_orn is not None else [0.0, 0.0, 0.20] orn = orientation if orientation is not None else [0.0, 0.0, 0.0, 1.0] else: bid = self.robot_id pos = body_id_or_pos orn = position_or_orn if position_or_orn is not None else [0.0, 0.0, 0.0, 1.0] lin_v = linear_velocity if linear_velocity is not None else [0.0, 0.0, 0.0] ang_v = angular_velocity if angular_velocity is not None else [0.0, 0.0, 0.0] p.resetBasePositionAndOrientation( bid, pos, orn, physicsClientId=self.physics_client ) p.resetBaseVelocity( bid, linearVelocity=lin_v, angularVelocity=ang_v, physicsClientId=self.physics_client ) def get_robot_pose(self, body_id: Optional[int] = None) -> Tuple[List[float], List[float]]: """Returns base position (x, y, z) and orientation quaternion (x, y, z, w).""" bid = self._resolve_body_id(body_id) pos, orn = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client) return list(pos), list(orn) def get_robot_rpy(self, body_id: Optional[int] = None) -> Tuple[float, float, float]: """Returns roll, pitch, yaw angles in radians for the robot.""" bid = self._resolve_body_id(body_id) _, orn = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client) roll, pitch, yaw = p.getEulerFromQuaternion(orn) return float(roll), float(pitch), float(yaw) def get_robot_pose_and_rpy(self, body_id: Optional[int] = None) -> Tuple[List[float], Tuple[float, float, float]]: """Returns base position and (roll, pitch, yaw) tuple in radians.""" bid = self._resolve_body_id(body_id) pos, orn = p.getBasePositionAndOrientation(bid, 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: Optional[int] = None) -> Tuple[List[float], List[float]]: """Returns linear velocity (vx, vy, vz) and angular velocity (wx, wy, wz).""" bid = self._resolve_body_id(body_id) lin_v, ang_v = p.getBaseVelocity(bid, physicsClientId=self.physics_client) return list(lin_v), list(ang_v) def get_robot_joint_angles( self, body_id: Optional[int] = None, joint_indices: Optional[List[int]] = None ) -> np.ndarray: """Returns joint angles as a 1D numpy float32 array.""" bid = self._resolve_body_id(body_id) j_indices = joint_indices if joint_indices is not None else self.joint_indices joint_states = p.getJointStates(bid, j_indices, physicsClientId=self.physics_client) return np.array([state[0] for state in joint_states], dtype=np.float32) def get_foot_link_indices(self, body_id: Optional[int] = None) -> List[int]: """Inspects URDF joint structure to extract link IDs for leg tips and tibias.""" bid = self._resolve_body_id(body_id) foot_indices = [] num_joints = p.getNumJoints(bid, physicsClientId=self.physics_client) for j_idx in range(num_joints): info = p.getJointInfo(bid, 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 set_robot_color( self, body_id_or_rgba: Union[int, List[float]], rgba: Optional[List[float]] = None ) -> None: """Changes visual color RGBA of base link and all joints.""" if isinstance(body_id_or_rgba, int): bid = body_id_or_rgba color = rgba else: bid = self.robot_id color = body_id_or_rgba if color is None: color = [1.0, 1.0, 1.0, 1.0] num_joints = p.getNumJoints(bid, physicsClientId=self.physics_client) p.changeVisualShape(bid, -1, rgbaColor=color, physicsClientId=self.physics_client) for j in range(num_joints): p.changeVisualShape(bid, j, rgbaColor=color, physicsClientId=self.physics_client) def measure_robot_height(self, body_id: Optional[int] = None) -> float: """Gets current Z height of the robot base.""" bid = self._resolve_body_id(body_id) pos, _ = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client) return pos[2] def settle_and_measure_height( self, robot_ids_or_steps: Union[List[int], int] = 200, steps: int = 200, fallback_height: float = 0.122 ) -> float: """Steps simulation so robot settles onto floor, then calculates target standing height.""" # Intelligently resolve whether the first argument was passed as robot_ids list or step count actual_steps = robot_ids_or_steps if isinstance(robot_ids_or_steps, int) else steps for _ in range(actual_steps): self.step() height = self.measure_robot_height() return height if height > 0.0 else fallback_height def apply_external_force( self, force: Union[List[float], np.ndarray], body_id: Optional[int] = None, link_index: int = -1, position: Union[List[float], np.ndarray] = (0.0, 0.0, 0.0), frame: int = p.WORLD_FRAME, ) -> None: """Applies external force vector to target link (defaults to base link).""" bid = self._resolve_body_id(body_id) p.applyExternalForce( objectUniqueId=bid, linkIndex=link_index, forceObj=list(force), posObj=list(position), flags=frame, physicsClientId=self.physics_client, )