""" simulation.py - PyBullet Simulation Interface & Physics Engine Consolidates scene management, physics queries, motor control, and rendering. """ import time import math from typing import List, Tuple, Optional, Union import numpy as np import pybullet as p import pybullet_data from config import cfg import DataTypes as dt class Simulation: """Single PyBullet simulation manager providing getters/setters for JackBot.""" def __init__(self, urdf_path: str = cfg.urdf_path, use_gui: bool = True): self.urdf_path = urdf_path self.use_gui = use_gui self.physics_client: Optional[int] = None self.plane_id: Optional[int] = None self.robot_id: Optional[int] = None self.revolute_joints: List[int] = [] self.connect() def connect(self) -> None: """Establishes connection to PyBullet GUI or DIRECT mode.""" 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) p.resetDebugVisualizerCamera( cameraDistance=1.0, cameraYaw=50, cameraPitch=-35, cameraTargetPosition=[0, 0, 0], physicsClientId=self.physics_client ) def load_scene(self, spawn_pos: Optional[List[float]] = None) -> Tuple[int, int, List[int]]: """Loads plane and robot URDF, discovering revolute joint indices dynamically.""" if spawn_pos is None: spawn_pos = [0.0, 0.0, 0.20] self.plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client) self.robot_id = p.loadURDF(self.urdf_path, spawn_pos, physicsClientId=self.physics_client) self.revolute_joints = [] 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.revolute_joints.append(j) return self.plane_id, self.robot_id, self.revolute_joints # --- ACTUATION (SETTERS) --- def set_robot_joint_angles( self, target_angles: Union[np.ndarray, List[float], dt.RadArray] ) -> None: """Applies motor torque to pull joints toward target position angles.""" if isinstance(target_angles, dt.RadArray): radflat = target_angles.data.flatten() elif isinstance(target_angles, np.ndarray): radflat = target_angles.flatten() else: radflat = target_angles for joint_index, target_angle in zip(self.revolute_joints, radflat): p.setJointMotorControl2( bodyIndex=self.robot_id, jointIndex=joint_index, controlMode=p.POSITION_CONTROL, targetPosition=float(target_angle), force=500, physicsClientId=self.physics_client ) def hard_reset_joint_angles(self, target_angles: Union[np.ndarray, dt.RadArray]) -> None: """Instantly teleports joint angles to target positions, clearing velocity state.""" radflat = target_angles.data.flatten() if isinstance(target_angles, dt.RadArray) else target_angles.flatten() for joint_index, target_angle in zip(self.revolute_joints, radflat): p.resetJointState( bodyUniqueId=self.robot_id, jointIndex=joint_index, targetValue=float(target_angle), targetVelocity=0.0, physicsClientId=self.physics_client ) def reset_robot_base( self, pos: Optional[List[float]] = None, orn: Optional[List[float]] = None, linear_velocity: Optional[List[float]] = None, angular_velocity: Optional[List[float]] = None ) -> None: """Resets root torso position, orientation quaternion, and clears base velocities.""" if pos is None: pos = [0.0, 0.0, 0.20] if orn is None: orn = [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(self.robot_id, pos, orn, physicsClientId=self.physics_client) p.resetBaseVelocity(self.robot_id, linearVelocity=lin_v, angularVelocity=ang_v, physicsClientId=self.physics_client) # --- TELEMETRY (GETTERS) --- def get_robot_pose(self) -> Tuple[List[float], List[float]]: pos, orn = p.getBasePositionAndOrientation(self.robot_id, physicsClientId=self.physics_client) return list(pos), list(orn) def get_robot_rpy(self) -> Tuple[float, float, float]: _, orn = p.getBasePositionAndOrientation(self.robot_id, physicsClientId=self.physics_client) roll, pitch, yaw = p.getEulerFromQuaternion(orn) return float(roll), float(pitch), float(yaw) def get_robot_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]: pos, orn = p.getBasePositionAndOrientation(self.robot_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) -> Tuple[List[float], List[float]]: lin_v, ang_v = p.getBaseVelocity(self.robot_id, physicsClientId=self.physics_client) return list(lin_v), list(ang_v) def get_robot_joint_angles(self) -> np.ndarray: joint_states = p.getJointStates(self.robot_id, self.revolute_joints, physicsClientId=self.physics_client) return np.array([state[0] for state in joint_states], dtype=np.float32) # --- SIMULATION LIFECYCLE CONTROLS --- def step(self) -> None: """Advances physics simulation by 1 time step.""" p.stepSimulation(physicsClientId=self.physics_client) def set_robot_color(self, rgba: List[float]) -> None: num_joints = p.getNumJoints(self.robot_id, physicsClientId=self.physics_client) p.changeVisualShape(self.robot_id, -1, rgbaColor=rgba, physicsClientId=self.physics_client) for j in range(num_joints): p.changeVisualShape(self.robot_id, j, rgbaColor=rgba, physicsClientId=self.physics_client) def settle_and_measure_height(self, steps: int = 200, fallback_height: float = 0.122) -> float: for _ in range(steps): self.step() pos, _ = self.get_robot_pose() return pos[2] if pos[2] > 0.0 else fallback_height def apply_external_force( self, force: Union[List[float], np.ndarray], link_index: int = -1, position: Tuple[float, float, float] = (0.0, 0.0, 0.0) ) -> None: p.applyExternalForce( objectUniqueId=self.robot_id, linkIndex=link_index, forceObj=list(force), posObj=list(position), flags=p.WORLD_FRAME, 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 close(self) -> None: self.disconnect() if __name__ == "__main__": from Robot import Robot, PyBulletBackend sim_instance = Simulation(use_gui=True) sim_instance.load_scene() backend = PyBulletBackend(sim_instance) robot = Robot(backend_type=backend, mode="kinematics") robot.reset_to_init() # Forward velocity command robot.vector_dirmov = [0.3, 0.0, 0.0] try: while True: robot.tick() time.sleep(1.0 / 60.0) except KeyboardInterrupt: sim_instance.disconnect()