"""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