298 lines
12 KiB
Python
298 lines
12 KiB
Python
"""
|
|
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,
|
|
) |