[sim] 添加后期MuJoCo姿态与MPC工具集
This commit is contained in:
@@ -0,0 +1,231 @@
|
||||
"""MuJoCo interface for the wheeled-legged robot.
|
||||
|
||||
Configures actuators as proper PD controllers at runtime:
|
||||
- Leg joints: force = kp*(ctrl - qpos) - kd*qvel, ctrl = target angle
|
||||
- Wheel joints: force = gain*(ctrl - qvel), ctrl = target velocity (rad/s)
|
||||
"""
|
||||
|
||||
import math
|
||||
import numpy as np
|
||||
import mujoco
|
||||
from dataclasses import dataclass
|
||||
from config import (SCENE_XML, LEG_NAMES, LEG_JOINTS, WHEEL_JOINT,
|
||||
DEFAULT_JOINT_ANGLES, WHEEL_RADIUS, WHEEL_TRACK,
|
||||
L_THIGH, L_CALF)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RobotState:
|
||||
"""Robot state from MuJoCo."""
|
||||
pos: np.ndarray # (3,) world position
|
||||
quat: np.ndarray # (4,) quaternion (w,x,y,z) MuJoCo convention
|
||||
rot: np.ndarray # (3,3) body→world rotation
|
||||
rpy: np.ndarray # (3,) roll, pitch, yaw
|
||||
lin_vel: np.ndarray # (3,) world frame linear velocity
|
||||
ang_vel: np.ndarray # (3,) body frame angular velocity
|
||||
joint_pos: np.ndarray # (16,) all joint positions [fl3+wheel, fr3+wheel, rl3+wheel, rr3+wheel]
|
||||
joint_vel: np.ndarray # (16,) all joint velocities
|
||||
time: float
|
||||
|
||||
|
||||
class Robot:
|
||||
"""MuJoCo simulation interface with proper PD actuator configuration."""
|
||||
|
||||
# Leg PD gains (tuned for 12.3kg robot)
|
||||
LEG_KP = 60.0
|
||||
LEG_KD = 3.0
|
||||
# Wheel velocity gain
|
||||
WHEEL_KP = 2.0
|
||||
|
||||
def __init__(self, xml_path=None):
|
||||
self.model = mujoco.MjModel.from_xml_path(str(xml_path or SCENE_XML))
|
||||
self.data = mujoco.MjData(self.model)
|
||||
|
||||
# Cache IDs
|
||||
self._base_bid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "base_link")
|
||||
self._actuator_ids = {} # name → actuator index
|
||||
self._joint_qpos_adr = {} # name → qpos address
|
||||
self._joint_qvel_adr = {} # name → qvel address
|
||||
|
||||
# Build joint/actuator maps
|
||||
self._ctrl_order = []
|
||||
for leg in LEG_NAMES:
|
||||
for jt in (*LEG_JOINTS, WHEEL_JOINT):
|
||||
name = f"{leg}_{jt}"
|
||||
aid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, name)
|
||||
jid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, name)
|
||||
self._actuator_ids[name] = aid
|
||||
self._joint_qpos_adr[name] = self.model.jnt_qposadr[jid]
|
||||
self._joint_qvel_adr[name] = self.model.jnt_dofadr[jid]
|
||||
self._ctrl_order.append(name)
|
||||
|
||||
# Configure actuators as proper PD controllers
|
||||
self._configure_actuators()
|
||||
|
||||
def _configure_actuators(self):
|
||||
"""Set actuators to proper PD mode.
|
||||
|
||||
Leg joints: force = kp*(ctrl - qpos) - kd*qvel
|
||||
Wheels: force = gain*(ctrl - qvel) (velocity tracking)
|
||||
"""
|
||||
for i in range(self.model.nu):
|
||||
name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, i)
|
||||
self.model.actuator_biastype[i] = 1 # affine bias
|
||||
self.model.actuator_gaintype[i] = 0 # fixed gain
|
||||
self.model.actuator_forcelimited[i] = 0 # no force clamp (17Nm is in actuatorfrcrange)
|
||||
|
||||
if 'wheel' not in name:
|
||||
self.model.actuator_gainprm[i, 0] = self.LEG_KP
|
||||
self.model.actuator_biasprm[i, 0] = 0.0
|
||||
self.model.actuator_biasprm[i, 1] = -self.LEG_KP
|
||||
self.model.actuator_biasprm[i, 2] = -self.LEG_KD
|
||||
self.model.actuator_ctrlrange[i] = [-3.14, 3.14]
|
||||
else:
|
||||
self.model.actuator_gainprm[i, 0] = self.WHEEL_KP
|
||||
self.model.actuator_biasprm[i, 0] = 0.0
|
||||
self.model.actuator_biasprm[i, 1] = 0.0
|
||||
self.model.actuator_biasprm[i, 2] = -self.WHEEL_KP
|
||||
self.model.actuator_ctrlrange[i] = [-20.0, 20.0]
|
||||
|
||||
@property
|
||||
def dt(self):
|
||||
return self.model.opt.timestep
|
||||
|
||||
def reset(self):
|
||||
"""Reset to standing pose at correct height for default joint angles."""
|
||||
mujoco.mj_resetData(self.model, self.data)
|
||||
|
||||
# Set default leg angles
|
||||
for leg in LEG_NAMES:
|
||||
for jt, key in zip(LEG_JOINTS, ("hip_abduction", "hip_pitch", "knee")):
|
||||
name = f"{leg}_{jt}"
|
||||
adr = self._joint_qpos_adr[name]
|
||||
self.data.qpos[adr] = DEFAULT_JOINT_ANGLES[key]
|
||||
|
||||
# Compute correct base height from default angles using exact FK
|
||||
# z_base = wheel_radius + L1*cos(hip) + L2*cos(hip+knee) - hip_z_offset
|
||||
hip = DEFAULT_JOINT_ANGLES["hip_pitch"]
|
||||
knee = DEFAULT_JOINT_ANGLES["knee"]
|
||||
base_z = (WHEEL_RADIUS
|
||||
+ L_THIGH * math.cos(hip)
|
||||
+ L_CALF * math.cos(hip + knee)
|
||||
- 0.054)
|
||||
self.data.qpos[2] = base_z
|
||||
self.data.qpos[3] = 1.0 # quat w
|
||||
|
||||
mujoco.mj_forward(self.model, self.data)
|
||||
|
||||
# Set ctrl to match initial pose (so PD doesn't jerk)
|
||||
for leg in LEG_NAMES:
|
||||
for jt, key in zip(LEG_JOINTS, ("hip_abduction", "hip_pitch", "knee")):
|
||||
name = f"{leg}_{jt}"
|
||||
self.data.ctrl[self._actuator_ids[name]] = DEFAULT_JOINT_ANGLES[key]
|
||||
# Wheels: zero velocity
|
||||
self.data.ctrl[self._actuator_ids[f"{leg}_{WHEEL_JOINT}"]] = 0.0
|
||||
|
||||
def get_state(self) -> RobotState:
|
||||
"""Extract robot state."""
|
||||
pos = self.data.xpos[self._base_bid].copy()
|
||||
quat = self.data.xquat[self._base_bid].copy() # (w,x,y,z)
|
||||
rot = self.data.xmat[self._base_bid].reshape(3, 3).copy()
|
||||
|
||||
rpy = np.array([
|
||||
np.arctan2(rot[2, 1], rot[2, 2]),
|
||||
np.arctan2(-rot[2, 0], np.sqrt(rot[2, 1]**2 + rot[2, 2]**2)),
|
||||
np.arctan2(rot[1, 0], rot[0, 0]),
|
||||
])
|
||||
|
||||
# Base velocity (world frame)
|
||||
lin_vel = self.data.qvel[0:3].copy()
|
||||
ang_vel = self.data.qvel[3:6].copy()
|
||||
|
||||
# Joint states (16 joints: 4 legs × 4 joints each)
|
||||
joint_pos = np.zeros(16)
|
||||
joint_vel = np.zeros(16)
|
||||
for i, name in enumerate(self._ctrl_order):
|
||||
joint_pos[i] = self.data.qpos[self._joint_qpos_adr[name]]
|
||||
joint_vel[i] = self.data.qvel[self._joint_qvel_adr[name]]
|
||||
|
||||
return RobotState(
|
||||
pos=pos, quat=quat, rot=rot, rpy=rpy,
|
||||
lin_vel=lin_vel, ang_vel=ang_vel,
|
||||
joint_pos=joint_pos, joint_vel=joint_vel,
|
||||
time=self.data.time,
|
||||
)
|
||||
|
||||
def set_ctrl(self, leg_targets: np.ndarray, wheel_targets: np.ndarray):
|
||||
"""Set actuator commands (position PD mode).
|
||||
|
||||
Args:
|
||||
leg_targets: (12,) target joint angles for legs [fl3, fr3, rl3, rr3]
|
||||
wheel_targets: (4,) target wheel velocities [fl, fr, rl, rr] in rad/s
|
||||
"""
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
for j, jt in enumerate(LEG_JOINTS):
|
||||
name = f"{leg}_{jt}"
|
||||
self.data.ctrl[self._actuator_ids[name]] = leg_targets[i * 3 + j]
|
||||
name = f"{leg}_{WHEEL_JOINT}"
|
||||
self.data.ctrl[self._actuator_ids[name]] = wheel_targets[i]
|
||||
|
||||
def set_ctrl_mit(self, q_des: np.ndarray, dq_des: np.ndarray,
|
||||
kp: np.ndarray, kd: np.ndarray, tau_ff: np.ndarray,
|
||||
wheel_targets: np.ndarray):
|
||||
"""MIT motor protocol: tau = kp*(q_des-q) + kd*(dq_des-dq) + tau_ff.
|
||||
|
||||
Computes torque in software, sends to actuators in torque mode.
|
||||
Call enable_torque_mode() first.
|
||||
|
||||
Args:
|
||||
q_des: (12,) desired joint angles
|
||||
dq_des: (12,) desired joint velocities
|
||||
kp: (12,) position gains (0 for pure torque)
|
||||
kd: (12,) velocity gains
|
||||
tau_ff: (12,) feedforward torques
|
||||
wheel_targets: (4,) wheel velocity targets
|
||||
"""
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
for j, jt in enumerate(LEG_JOINTS):
|
||||
name = f"{leg}_{jt}"
|
||||
aid = self._actuator_ids[name]
|
||||
idx = i * 3 + j
|
||||
q = self.data.qpos[self._joint_qpos_adr[name]]
|
||||
dq = self.data.qvel[self._joint_qvel_adr[name]]
|
||||
tau = (kp[idx] * (q_des[idx] - q)
|
||||
+ kd[idx] * (dq_des[idx] - dq)
|
||||
+ tau_ff[idx])
|
||||
self.data.ctrl[aid] = np.clip(tau, -17.0, 17.0)
|
||||
name = f"{leg}_{WHEEL_JOINT}"
|
||||
self.data.ctrl[self._actuator_ids[name]] = wheel_targets[i]
|
||||
|
||||
def enable_torque_mode(self):
|
||||
"""Switch leg actuators to direct torque mode (for MPC/MIT)."""
|
||||
for i in range(self.model.nu):
|
||||
name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, i)
|
||||
if 'wheel' not in name:
|
||||
self.model.actuator_gainprm[i, 0] = 1.0
|
||||
self.model.actuator_biasprm[i, :3] = [0, 0, 0]
|
||||
self.model.actuator_biastype[i] = 0
|
||||
self.model.actuator_ctrlrange[i] = [-17.0, 17.0]
|
||||
|
||||
def enable_position_mode(self):
|
||||
"""Switch leg actuators back to position PD mode."""
|
||||
self._configure_actuators()
|
||||
|
||||
def step(self):
|
||||
"""Advance one simulation timestep."""
|
||||
mujoco.mj_step(self.model, self.data)
|
||||
|
||||
def get_qpos_qvel_for_pinocchio(self):
|
||||
"""Get full qpos/qvel for Pinocchio (reorder quaternion)."""
|
||||
qpos = self.data.qpos.copy()
|
||||
qvel = self.data.qvel.copy()
|
||||
# MuJoCo quat: (w,x,y,z) → Pinocchio: (x,y,z,w)
|
||||
w, x, y, z = qpos[3], qpos[4], qpos[5], qpos[6]
|
||||
q_pin = np.concatenate([qpos[0:3], [x, y, z, w], qpos[7:]])
|
||||
# MuJoCo vel is already [lin_world(3), ang_body(3), joints(16)]
|
||||
# Pinocchio wants [lin_body(3), ang_body(3), joints(16)]
|
||||
from scipy.spatial.transform import Rotation
|
||||
R = Rotation.from_quat([x, y, z, w]).as_matrix()
|
||||
v_body = R.T @ qvel[0:3]
|
||||
dq_pin = np.concatenate([v_body, qvel[3:]])
|
||||
return q_pin, dq_pin
|
||||
Reference in New Issue
Block a user