166 lines
7.4 KiB
Python
166 lines
7.4 KiB
Python
"""
|
|
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 disconnect(self):
|
|
if self.physics_client is not None and p.isConnected(self.physics_client):
|
|
p.disconnect(self.physics_client)
|
|
self.physics_client = None
|
|
|
|
# --- 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_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 |