217 lines
9.2 KiB
Python
217 lines
9.2 KiB
Python
"""
|
|
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=30,
|
|
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)
|
|
|
|
def _get_urdf_joint_limits(self) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""Dynamically reads lower and upper limits for all revolute joints from PyBullet."""
|
|
lower_limits = []
|
|
upper_limits = []
|
|
|
|
# Iterate through joints in PyBullet
|
|
for j_idx in range(p.getNumJoints(self.robot_id, physicsClientId=self.physics_client)):
|
|
info = p.getJointInfo(self.robot_id, j_idx, physicsClientId=self.physics_client)
|
|
joint_type = info[2]
|
|
|
|
# Only collect limits for revolute joints
|
|
if joint_type == p.JOINT_REVOLUTE:
|
|
lower_limits.append(info[8]) # Index 8 = jointLowerLimit
|
|
upper_limits.append(info[9]) # Index 9 = jointUpperLimit
|
|
|
|
return np.array(lower_limits, dtype=np.float32), np.array(upper_limits, dtype=np.float32)
|
|
|
|
# --- SIMULATION LIFECYCLE CONTROLS ---
|
|
|
|
def step(self) -> None:
|
|
"""Advances physics simulation by 1 time step."""
|
|
p.stepSimulation(physicsClientId=self.physics_client)
|
|
|
|
def settle_and_measure_height(
|
|
self, target_angles: Optional[np.ndarray] = None, steps: int = 100, fallback_height: float = 0.122
|
|
) -> float:
|
|
"""Settles the robot into the ground while actively holding target joint angles."""
|
|
for _ in range(steps):
|
|
if target_angles is not None:
|
|
self.set_robot_joint_angles(target_angles)
|
|
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() |