[train] 更新新MJCF与第一版完整训练框架
This commit is contained in:
@@ -1,57 +0,0 @@
|
||||
"""Robot constants and control parameters."""
|
||||
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
|
||||
# Paths
|
||||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
SCENE_XML = REPO_ROOT / "mjcf" / "scene.xml"
|
||||
MJCF_PATH = REPO_ROOT / "mjcf" / "wheelleg.xml"
|
||||
|
||||
# Robot geometry
|
||||
WHEEL_RADIUS = 0.10 # m
|
||||
WHEEL_TRACK = 0.32 # m (left-right distance)
|
||||
ROBOT_MASS = 12.3 # kg
|
||||
MAX_TORQUE = 17.0 # Nm per joint
|
||||
MAX_JOINT_VEL = 13.0 # rad/s
|
||||
|
||||
# Leg link lengths (from MJCF)
|
||||
L_THIGH = 0.25 # m
|
||||
L_CALF = 0.20 # m (to wheel center)
|
||||
|
||||
# Leg names and joint ordering
|
||||
LEG_NAMES = ("fl", "fr", "rl", "rr")
|
||||
LEG_JOINTS = ("hip_abduction_joint", "hip_pitch_joint", "knee_joint")
|
||||
WHEEL_JOINT = "wheel_joint"
|
||||
|
||||
# Default standing pose (from go2w_sim2sim: [0, 0.8, -1.5])
|
||||
DEFAULT_JOINT_ANGLES = {
|
||||
"hip_abduction": 0.0,
|
||||
"hip_pitch": 0.93,
|
||||
"knee": -1.65,
|
||||
}
|
||||
|
||||
# Actuator modes (MJCF native):
|
||||
# Leg joints: position PD (kp=120, kd=8), ctrl = target angle
|
||||
# Wheel joints: velocity (gain=0.5), ctrl = target velocity (rad/s)
|
||||
|
||||
# Control rates
|
||||
SIM_DT = 0.002 # 500 Hz (from scene.xml)
|
||||
CTRL_DT = 0.004 # 250 Hz control loop
|
||||
CTRL_DECIMATION = int(CTRL_DT / SIM_DT)
|
||||
|
||||
# Wheel drive
|
||||
WHEEL_VEL_MAX = 10.0 # rad/s max wheel command
|
||||
|
||||
# Body pose control gains (for height/roll/pitch compensation)
|
||||
KP_HEIGHT = 3.0 # rad/m error → joint angle correction
|
||||
KP_ROLL = 0.5 # compensation gain
|
||||
KP_PITCH = 0.5 # compensation gain
|
||||
|
||||
# Gait parameters
|
||||
GAIT_FREQ = 2.5 # Hz
|
||||
GAIT_DUTY = 0.6 # stance fraction
|
||||
SWING_HEIGHT = 0.06 # m
|
||||
|
||||
# Trot phase offsets: FL/RR in phase, FR/RL in phase
|
||||
PHASE_OFFSETS = {"fl": 0.0, "fr": 0.5, "rl": 0.5, "rr": 0.0}
|
||||
@@ -1,287 +0,0 @@
|
||||
"""Main controller: wheel mode + trot mode for wheeled-legged robot.
|
||||
|
||||
Wheel mode: differential drive + leg posture hold (height/roll/pitch compensation)
|
||||
Trot mode: quadruped gait with wheel-assisted propulsion
|
||||
|
||||
Actuator interface:
|
||||
- Leg joints: ctrl = target angle (PD: kp=60, kd=3)
|
||||
- Wheel joints: ctrl = target velocity in rad/s (gain=2.0)
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from robot import Robot, RobotState
|
||||
from dynamics import Dynamics
|
||||
from mpc_controller import MPCController
|
||||
from config import (
|
||||
LEG_NAMES, DEFAULT_JOINT_ANGLES, WHEEL_RADIUS, WHEEL_TRACK,
|
||||
WHEEL_VEL_MAX, KP_ROLL, KP_PITCH,
|
||||
GAIT_FREQ, GAIT_DUTY, SWING_HEIGHT, PHASE_OFFSETS,
|
||||
)
|
||||
|
||||
|
||||
class Controller:
|
||||
"""Wheeled-legged robot controller."""
|
||||
|
||||
def __init__(self, robot: Robot):
|
||||
self.robot = robot
|
||||
self.dynamics = Dynamics()
|
||||
|
||||
# User commands
|
||||
self.vel_x = 0.0 # m/s forward
|
||||
self.vel_y = 0.0 # m/s lateral
|
||||
self.yaw_rate = 0.0 # rad/s
|
||||
self.height = 0.33 # m desired body height
|
||||
|
||||
# Mode: "wheel", "trot", or "mpc"
|
||||
self.mode = "wheel"
|
||||
|
||||
# Prone (lie down) state
|
||||
self.prone = False
|
||||
|
||||
# MPC controller
|
||||
self._mpc_ctrl = MPCController(robot)
|
||||
self._mpc_active = False # track torque mode state
|
||||
|
||||
# Gait state
|
||||
self._gait_phase = 0.0
|
||||
|
||||
# Smoothed commands for trot mode (avoid sudden jumps)
|
||||
self._smooth_vx = 0.0
|
||||
self._smooth_vy = 0.0
|
||||
self._smooth_yaw = 0.0
|
||||
|
||||
# Default leg angles
|
||||
self._default_q = np.array([
|
||||
DEFAULT_JOINT_ANGLES["hip_abduction"],
|
||||
DEFAULT_JOINT_ANGLES["hip_pitch"],
|
||||
DEFAULT_JOINT_ANGLES["knee"],
|
||||
])
|
||||
|
||||
# Swing leg memory
|
||||
self._swing_start_foot = {leg: np.zeros(3) for leg in LEG_NAMES}
|
||||
self._last_contact = {leg: True for leg in LEG_NAMES}
|
||||
|
||||
def compute(self, state: RobotState, dt: float) -> tuple[np.ndarray, np.ndarray]:
|
||||
# Smooth all velocity commands (both modes)
|
||||
alpha = min(dt * 3.0, 1.0) # ~0.33s time constant
|
||||
self._smooth_vx += alpha * (self.vel_x - self._smooth_vx)
|
||||
self._smooth_vy += alpha * (self.vel_y - self._smooth_vy)
|
||||
self._smooth_yaw += alpha * (self.yaw_rate - self._smooth_yaw)
|
||||
|
||||
if self.prone:
|
||||
self._ensure_position_mode()
|
||||
return self._prone_mode()
|
||||
if self.mode == "mpc":
|
||||
return self._mpc_mode(state, dt)
|
||||
if self.mode == "wheel":
|
||||
self._ensure_position_mode()
|
||||
return self._wheel_mode(state, dt)
|
||||
else:
|
||||
self._ensure_position_mode()
|
||||
return self._trot_mode(state, dt)
|
||||
|
||||
def _mpc_mode(self, state: RobotState, dt: float):
|
||||
"""MPC locomotion: MIT motor protocol (PD + MPC feedforward torque)."""
|
||||
# Switch to torque mode if not already
|
||||
if not self._mpc_active:
|
||||
self.robot.enable_torque_mode()
|
||||
self._mpc_active = True
|
||||
|
||||
# Sync commands to MPC controller
|
||||
self._mpc_ctrl.vel_x = self.vel_x
|
||||
self._mpc_ctrl.vel_y = self.vel_y
|
||||
self._mpc_ctrl.yaw_rate = self.yaw_rate
|
||||
self._mpc_ctrl.height = self.height
|
||||
|
||||
# Compute and apply (sets ctrl directly via set_ctrl_mit)
|
||||
self._mpc_ctrl.compute(state, dt)
|
||||
# Return dummy - ctrl already set
|
||||
return np.zeros(12), np.zeros(4)
|
||||
|
||||
def _ensure_position_mode(self):
|
||||
"""Switch back to position PD mode if coming from MPC."""
|
||||
if self._mpc_active:
|
||||
self.robot.enable_position_mode()
|
||||
self._mpc_active = False
|
||||
|
||||
def _prone_mode(self):
|
||||
"""Lie down: actual prone pose from real robot."""
|
||||
leg_targets = np.zeros(12)
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
side = 1.0 if leg[1] == "l" else -1.0
|
||||
leg_targets[i*3] = side * 0.3 # fl/rl: +0.3, fr/rr: -0.3
|
||||
leg_targets[i*3+1] = 1.5 # hip pitch
|
||||
leg_targets[i*3+2] = -2.65 # knee fully folded
|
||||
return leg_targets, np.zeros(4)
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# WHEEL MODE
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
def _wheel_mode(self, state: RobotState, dt: float):
|
||||
"""Wheel drive + leg posture hold.
|
||||
|
||||
vel_y: limited effect in wheel mode (differential drive cannot produce
|
||||
pure lateral motion). Uses hip_abduction lean for small lateral force.
|
||||
For significant lateral motion, use trot mode.
|
||||
"""
|
||||
wheel_targets = self._differential_drive(self._smooth_vx, self._smooth_yaw)
|
||||
leg_targets = self._posture_control(state)
|
||||
return leg_targets, wheel_targets
|
||||
|
||||
def _posture_control(self, state: RobotState) -> np.ndarray:
|
||||
"""Leg joint targets: table-interpolated height control."""
|
||||
leg_targets = np.zeros(12)
|
||||
|
||||
# Calibrated height→angle lookup (measured from simulation)
|
||||
_H = [0.157, 0.248, 0.311, 0.366, 0.411, 0.448]
|
||||
_HIP = [1.5, 1.2, 1.0, 0.8, 0.6, 0.4]
|
||||
_KNEE = [-2.5, -2.1, -1.8, -1.5, -1.2, -0.9]
|
||||
|
||||
h_clamp = np.clip(self.height, _H[0], _H[-1])
|
||||
q_hip_base = float(np.interp(h_clamp, _H, _HIP))
|
||||
q_knee_base = float(np.interp(h_clamp, _H, _KNEE))
|
||||
|
||||
roll_corr = -KP_ROLL * state.rpy[0]
|
||||
pitch_corr = -KP_PITCH * state.rpy[1]
|
||||
lateral_lean = 0.3 * self.vel_y
|
||||
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
side = 1.0 if leg[1] == "l" else -1.0
|
||||
leg_targets[i*3] = np.clip(side * roll_corr + lateral_lean, -0.5, 0.5)
|
||||
leg_targets[i*3+1] = np.clip(q_hip_base + pitch_corr, -1.0, 2.5)
|
||||
leg_targets[i*3+2] = np.clip(q_knee_base, -2.6, -0.3)
|
||||
|
||||
return leg_targets
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# TROT MODE
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
def _trot_mode(self, state: RobotState, dt: float):
|
||||
"""Trot gait with wheel assist."""
|
||||
# Advance gait phase
|
||||
self._gait_phase = (self._gait_phase + dt * GAIT_FREQ) % 1.0
|
||||
|
||||
# Contact state
|
||||
contacts = {}
|
||||
for leg in LEG_NAMES:
|
||||
phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0
|
||||
contacts[leg] = phase < GAIT_DUTY
|
||||
|
||||
# Pinocchio update
|
||||
q_pin, dq_pin = self.robot.get_qpos_qvel_for_pinocchio()
|
||||
self.dynamics.update(q_pin, dq_pin)
|
||||
|
||||
leg_targets = np.zeros(12)
|
||||
wheel_targets = np.zeros(4)
|
||||
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
if contacts[leg]:
|
||||
# Stance: posture hold
|
||||
leg_targets[i*3:(i+1)*3] = self._stance_leg_target(state, leg)
|
||||
self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg)
|
||||
self._last_contact[leg] = True
|
||||
# Wheel: drive with smoothed velocity
|
||||
wheel_targets[i] = self._differential_drive_single(
|
||||
self._smooth_vx, self._smooth_yaw, leg)
|
||||
else:
|
||||
# Swing: IK trajectory
|
||||
swing_phase = self._get_swing_phase(leg)
|
||||
target_foot = self._compute_swing_target(leg, state, swing_phase)
|
||||
q_ik = self.dynamics.inverse_kinematics(leg, target_foot, q_pin)
|
||||
leg_targets[i*3:(i+1)*3] = q_ik
|
||||
self._last_contact[leg] = False
|
||||
# Wheel: zero (free during swing)
|
||||
wheel_targets[i] = 0.0
|
||||
|
||||
return leg_targets, wheel_targets
|
||||
|
||||
def _stance_leg_target(self, state: RobotState, leg: str) -> np.ndarray:
|
||||
"""Stance leg: table-interpolated height + attitude compensation."""
|
||||
_H = [0.157, 0.248, 0.311, 0.366, 0.411, 0.448]
|
||||
_HIP = [1.5, 1.2, 1.0, 0.8, 0.6, 0.4]
|
||||
_KNEE = [-2.5, -2.1, -1.8, -1.5, -1.2, -0.9]
|
||||
|
||||
h_clamp = np.clip(self.height, _H[0], _H[-1])
|
||||
q_hip = float(np.interp(h_clamp, _H, _HIP))
|
||||
q_knee = float(np.interp(h_clamp, _H, _KNEE))
|
||||
|
||||
roll_corr = -KP_ROLL * state.rpy[0]
|
||||
pitch_corr = -KP_PITCH * state.rpy[1]
|
||||
side = 1.0 if leg[1] == "l" else -1.0
|
||||
lateral_lean = 0.3 * self.vel_y
|
||||
|
||||
return np.array([
|
||||
np.clip(side * roll_corr + lateral_lean, -0.5, 0.5),
|
||||
np.clip(q_hip + pitch_corr, -1.0, 2.5),
|
||||
np.clip(q_knee, -2.6, -0.3),
|
||||
])
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# DIFFERENTIAL DRIVE
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
def _differential_drive(self, vel_x: float, yaw_rate: float) -> np.ndarray:
|
||||
"""4 wheel velocities from body commands."""
|
||||
vel_left = (vel_x - 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS
|
||||
vel_right = (vel_x + 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS
|
||||
targets = np.zeros(4)
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
targets[i] = vel_left if leg[1] == "l" else vel_right
|
||||
return np.clip(targets, -WHEEL_VEL_MAX, WHEEL_VEL_MAX)
|
||||
|
||||
def _differential_drive_single(self, vel_x: float, yaw_rate: float, leg: str) -> float:
|
||||
if leg[1] == "l":
|
||||
v = (vel_x - 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS
|
||||
else:
|
||||
v = (vel_x + 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS
|
||||
return np.clip(v, -WHEEL_VEL_MAX, WHEEL_VEL_MAX)
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# SWING TRAJECTORY
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
def _get_swing_phase(self, leg: str) -> float:
|
||||
phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0
|
||||
if phase < GAIT_DUTY:
|
||||
return 0.0
|
||||
return (phase - GAIT_DUTY) / (1.0 - GAIT_DUTY)
|
||||
|
||||
def _compute_swing_target(self, leg: str, state: RobotState,
|
||||
swing_phase: float) -> np.ndarray:
|
||||
"""Swing foot target with Raibert heuristic using COMMANDED velocity."""
|
||||
p_start = self._swing_start_foot[leg]
|
||||
p_end = self._compute_touchdown(leg, state)
|
||||
|
||||
s = swing_phase
|
||||
s_mj = 10*s**3 - 15*s**4 + 6*s**5
|
||||
|
||||
pos = p_start + (p_end - p_start) * s_mj
|
||||
|
||||
# Z lift
|
||||
z_lift = 64.0 * s**3 * (1.0 - s)**3
|
||||
pos[2] = p_start[2] + SWING_HEIGHT * z_lift
|
||||
|
||||
return pos
|
||||
|
||||
def _compute_touchdown(self, leg: str, state: RobotState) -> np.ndarray:
|
||||
"""Raibert heuristic using COMMANDED velocity.
|
||||
|
||||
When commands are zero, foot lands at its takeoff position (no net motion).
|
||||
When commands are nonzero, foot placement is offset by commanded velocity.
|
||||
"""
|
||||
# Base: land where the foot took off (zero net displacement)
|
||||
td = self._swing_start_foot[leg].copy()
|
||||
|
||||
# Add commanded velocity offset (Raibert-style)
|
||||
t_stance = (1.0 / GAIT_FREQ) * GAIT_DUTY
|
||||
yaw = state.rpy[2]
|
||||
c, s = np.cos(yaw), np.sin(yaw)
|
||||
R_z = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])
|
||||
cmd_vel_world = R_z @ np.array([self._smooth_vx, self._smooth_vy, 0.0])
|
||||
|
||||
td[0] += cmd_vel_world[0] * t_stance * 0.5
|
||||
td[1] += cmd_vel_world[1] * t_stance * 0.5
|
||||
td[2] = WHEEL_RADIUS # ground level
|
||||
return td
|
||||
@@ -1,97 +0,0 @@
|
||||
"""Pinocchio dynamics: FK, Jacobian, IK for the wheeled-legged robot."""
|
||||
|
||||
import numpy as np
|
||||
import pinocchio as pin
|
||||
from config import MJCF_PATH, LEG_NAMES
|
||||
|
||||
# Foot frame names in Pinocchio model (wheel link centers)
|
||||
FOOT_FRAMES = {leg: f"{leg}_wheel_Link" for leg in LEG_NAMES}
|
||||
|
||||
# Leg joint names for each leg
|
||||
_LEG_JOINT_NAMES = {
|
||||
leg: [f"{leg}_{jt}" for jt in ("hip_abduction_joint", "hip_pitch_joint", "knee_joint")]
|
||||
for leg in LEG_NAMES
|
||||
}
|
||||
|
||||
|
||||
class Dynamics:
|
||||
"""Pinocchio-based kinematics/dynamics. Deployable on real hardware."""
|
||||
|
||||
def __init__(self):
|
||||
self.model = pin.buildModelFromMJCF(str(MJCF_PATH))
|
||||
self.data = self.model.createData()
|
||||
|
||||
# Cache frame IDs
|
||||
self._foot_fids = {}
|
||||
for leg, fname in FOOT_FRAMES.items():
|
||||
self._foot_fids[leg] = self.model.getFrameId(fname)
|
||||
|
||||
# Cache joint velocity indices for each leg (3 joints)
|
||||
self._leg_v_indices = {}
|
||||
for leg, jnames in _LEG_JOINT_NAMES.items():
|
||||
indices = []
|
||||
for jn in jnames:
|
||||
jid = self.model.getJointId(jn)
|
||||
indices.append(self.model.joints[jid].idx_v)
|
||||
self._leg_v_indices[leg] = indices
|
||||
|
||||
# Cache joint config indices for each leg
|
||||
self._leg_q_indices = {}
|
||||
for leg, jnames in _LEG_JOINT_NAMES.items():
|
||||
indices = []
|
||||
for jn in jnames:
|
||||
jid = self.model.getJointId(jn)
|
||||
indices.append(self.model.joints[jid].idx_q)
|
||||
self._leg_q_indices[leg] = indices
|
||||
|
||||
def update(self, q: np.ndarray, dq: np.ndarray):
|
||||
"""Forward kinematics + Jacobians.
|
||||
|
||||
Args:
|
||||
q: Pinocchio config (nq=23: pos3, quat_xyzw4, joints16)
|
||||
dq: Pinocchio velocity (nv=22: v_body3, w_body3, joints16)
|
||||
"""
|
||||
pin.forwardKinematics(self.model, self.data, q, dq)
|
||||
pin.updateFramePlacements(self.model, self.data)
|
||||
pin.computeJointJacobians(self.model, self.data, q)
|
||||
|
||||
def get_foot_pos(self, leg: str) -> np.ndarray:
|
||||
"""Foot (wheel center) position in world frame (3,)."""
|
||||
return self.data.oMf[self._foot_fids[leg]].translation.copy()
|
||||
|
||||
def get_foot_jacobian_leg(self, leg: str) -> np.ndarray:
|
||||
"""3x3 linear Jacobian of foot w.r.t. 3 leg joints (world frame)."""
|
||||
fid = self._foot_fids[leg]
|
||||
J_full = pin.getFrameJacobian(
|
||||
self.model, self.data, fid, pin.LOCAL_WORLD_ALIGNED)[:3, :]
|
||||
cols = self._leg_v_indices[leg]
|
||||
return J_full[:, cols]
|
||||
|
||||
def inverse_kinematics(self, leg: str, target_pos: np.ndarray,
|
||||
q_current: np.ndarray, max_iter=30, eps=1e-4) -> np.ndarray:
|
||||
"""Numerical IK for one leg. Returns (3,) joint angles.
|
||||
|
||||
Args:
|
||||
leg: Leg name
|
||||
target_pos: Desired foot position in world frame (3,)
|
||||
q_current: Current full Pinocchio config (nq=23)
|
||||
"""
|
||||
q = q_current.copy()
|
||||
fid = self._foot_fids[leg]
|
||||
q_indices = self._leg_q_indices[leg]
|
||||
|
||||
for _ in range(max_iter):
|
||||
pin.forwardKinematics(self.model, self.data, q)
|
||||
pin.updateFramePlacements(self.model, self.data)
|
||||
err = target_pos - self.data.oMf[fid].translation
|
||||
if np.linalg.norm(err) < eps:
|
||||
break
|
||||
pin.computeJointJacobians(self.model, self.data, q)
|
||||
J = pin.getFrameJacobian(
|
||||
self.model, self.data, fid, pin.LOCAL_WORLD_ALIGNED)[:3, :]
|
||||
J_leg = J[:, self._leg_v_indices[leg]]
|
||||
dq = np.linalg.solve(J_leg.T @ J_leg + 1e-6 * np.eye(3), J_leg.T @ err)
|
||||
for i, idx in enumerate(q_indices):
|
||||
q[idx] += dq[i]
|
||||
|
||||
return np.array([q[idx] for idx in q_indices])
|
||||
@@ -1,127 +0,0 @@
|
||||
"""GUI control panel for the wheeled-legged robot."""
|
||||
|
||||
import tkinter as tk
|
||||
from tkinter import ttk
|
||||
|
||||
|
||||
class GUI:
|
||||
"""Tkinter control panel: sliders + gait buttons + status display."""
|
||||
|
||||
def __init__(self, controller):
|
||||
self.ctrl = controller
|
||||
self.root = tk.Tk()
|
||||
self.root.title("WheelLeg Control")
|
||||
self.root.geometry("400x500")
|
||||
self.root.protocol("WM_DELETE_WINDOW", self._on_close)
|
||||
self._closed = False
|
||||
|
||||
self._build()
|
||||
|
||||
def _build(self):
|
||||
# Mode buttons
|
||||
mf = ttk.LabelFrame(self.root, text="Mode")
|
||||
mf.pack(fill="x", padx=8, pady=4)
|
||||
for mode in ("wheel", "trot", "mpc"):
|
||||
ttk.Button(mf, text=mode.upper(),
|
||||
command=lambda m=mode: self._set_mode(m)
|
||||
).pack(side="left", padx=4, expand=True)
|
||||
ttk.Button(mf, text="PRONE/STAND",
|
||||
command=self._toggle_prone).pack(side="left", padx=4, expand=True)
|
||||
|
||||
# Command sliders
|
||||
cf = ttk.LabelFrame(self.root, text="Commands")
|
||||
cf.pack(fill="x", padx=8, pady=4)
|
||||
|
||||
self.vel_x_var = tk.DoubleVar(value=0.0)
|
||||
self.vel_y_var = tk.DoubleVar(value=0.0)
|
||||
self.yaw_var = tk.DoubleVar(value=0.0)
|
||||
self.height_var = tk.DoubleVar(value=self.ctrl.height)
|
||||
|
||||
self._slider(cf, "Vel X", self.vel_x_var, -1.5, 1.5)
|
||||
self._slider(cf, "Vel Y*", self.vel_y_var, -0.5, 0.5)
|
||||
self._slider(cf, "Yaw", self.yaw_var, -2.0, 2.0)
|
||||
self._slider(cf, "Height", self.height_var, 0.16, 0.45)
|
||||
|
||||
ttk.Label(cf, text="* Vel Y: trot mode only (diff-drive can't sidestep)",
|
||||
font=("", 8)).pack(anchor="w", padx=8)
|
||||
|
||||
ttk.Button(cf, text="Reset", command=self._reset).pack(pady=4)
|
||||
|
||||
# Status display
|
||||
sf = ttk.LabelFrame(self.root, text="Status")
|
||||
sf.pack(fill="both", expand=True, padx=8, pady=4)
|
||||
self.status_text = tk.Text(sf, height=12, width=45, font=("Consolas", 9))
|
||||
self.status_text.pack(fill="both", expand=True, padx=4, pady=4)
|
||||
|
||||
def _slider(self, parent, label, var, lo, hi):
|
||||
f = ttk.Frame(parent)
|
||||
f.pack(fill="x", padx=4, pady=2)
|
||||
ttk.Label(f, text=label, width=7).pack(side="left")
|
||||
ttk.Scale(f, from_=lo, to=hi, variable=var,
|
||||
command=lambda *_: self._sync()).pack(side="left", fill="x", expand=True)
|
||||
lbl = ttk.Label(f, text="0.00", width=6)
|
||||
lbl.pack(side="left")
|
||||
var.trace_add("write", lambda *_, v=var, l=lbl: l.config(text=f"{v.get():.2f}"))
|
||||
|
||||
def _set_mode(self, mode):
|
||||
self.ctrl.mode = mode
|
||||
self.ctrl.prone = False
|
||||
|
||||
def _toggle_prone(self):
|
||||
self.ctrl.prone = not self.ctrl.prone
|
||||
|
||||
def _sync(self):
|
||||
self.ctrl.vel_x = self.vel_x_var.get()
|
||||
self.ctrl.vel_y = self.vel_y_var.get()
|
||||
self.ctrl.yaw_rate = self.yaw_var.get()
|
||||
self.ctrl.height = self.height_var.get()
|
||||
|
||||
def _reset(self):
|
||||
self.vel_x_var.set(0.0)
|
||||
self.vel_y_var.set(0.0)
|
||||
self.yaw_var.set(0.0)
|
||||
self._sync()
|
||||
|
||||
def _on_close(self):
|
||||
self._closed = True
|
||||
self.root.destroy()
|
||||
|
||||
@property
|
||||
def closed(self):
|
||||
return self._closed
|
||||
|
||||
def update_status(self, state, step):
|
||||
"""Update status text with current robot state."""
|
||||
txt = (
|
||||
f"Mode: {self.ctrl.mode} Step: {step}\n"
|
||||
f"Pos: x={state.pos[0]:.3f} y={state.pos[1]:.3f} z={state.pos[2]:.3f}\n"
|
||||
f"RPY: r={np.degrees(state.rpy[0]):.1f}° p={np.degrees(state.rpy[1]):.1f}° "
|
||||
f"y={np.degrees(state.rpy[2]):.1f}°\n"
|
||||
f"Vel: vx={state.lin_vel[0]:.3f} vy={state.lin_vel[1]:.3f} vz={state.lin_vel[2]:.3f}\n"
|
||||
f"Cmd: vx={self.ctrl.vel_x:.2f} yaw={self.ctrl.yaw_rate:.2f} h={self.ctrl.height:.3f}\n"
|
||||
f"─────────────────────────────────\n"
|
||||
)
|
||||
# Joint angles (compact)
|
||||
for i, leg in enumerate(("FL", "FR", "RL", "RR")):
|
||||
q = state.joint_pos[i*4:i*4+3]
|
||||
w = state.joint_vel[i*4+3]
|
||||
txt += f"{leg}: [{q[0]:+.2f} {q[1]:+.2f} {q[2]:+.2f}] w={w:+.1f}\n"
|
||||
|
||||
self.status_text.delete("1.0", tk.END)
|
||||
self.status_text.insert(tk.END, txt)
|
||||
|
||||
def tick(self):
|
||||
"""Process GUI events. Returns False if window closed."""
|
||||
if self._closed:
|
||||
return False
|
||||
try:
|
||||
self.root.update_idletasks()
|
||||
self.root.update()
|
||||
return True
|
||||
except tk.TclError:
|
||||
self._closed = True
|
||||
return False
|
||||
|
||||
|
||||
# Need numpy for degrees conversion in update_status
|
||||
import numpy as np
|
||||
@@ -1,242 +0,0 @@
|
||||
"""Convex MPC solver for wheeled-legged robot.
|
||||
|
||||
Centroidal dynamics: single rigid body model with 4 contact forces.
|
||||
State: x = [pos(3), rpy(3), vel(3), omega(3)] = 12
|
||||
Input: u = [f1(3), f2(3), f3(3), f4(3)] = 12
|
||||
Friction pyramid constraints on each foot.
|
||||
|
||||
Reference: MIT Cheetah 3 Convex MPC (Di Carlo et al.)
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from scipy import sparse
|
||||
from scipy.linalg import block_diag
|
||||
import osqp
|
||||
|
||||
from config import ROBOT_MASS, LEG_NAMES
|
||||
|
||||
# MPC parameters
|
||||
MPC_HORIZON = 10 # prediction steps
|
||||
MPC_DT = 0.02 # 50 Hz MPC update
|
||||
MU = 0.6 # friction coefficient
|
||||
FZ_MAX = 200.0 # max vertical force per leg
|
||||
FZ_MIN = 10.0 # min vertical force (stance)
|
||||
NX = 12 # state dim
|
||||
NU = 12 # input dim (4 legs × 3D force)
|
||||
|
||||
# Cost weights: [pos_x, pos_y, pos_z, roll, pitch, yaw, vx, vy, vz, wx, wy, wz]
|
||||
Q_WEIGHTS = np.array([2.0, 2.0, 50.0, 50.0, 50.0, 10.0, 2.0, 2.0, 1.0, 1.0, 1.0, 1.0])
|
||||
R_WEIGHTS = np.array([1e-6] * 12)
|
||||
|
||||
|
||||
def _skew(v):
|
||||
return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
|
||||
|
||||
|
||||
class ConvexMPC:
|
||||
"""Convex MPC: solves QP for optimal ground reaction forces."""
|
||||
|
||||
def __init__(self, mass=ROBOT_MASS, inertia=None):
|
||||
self.mass = mass
|
||||
# Approximate body inertia (diagonal, world-aligned)
|
||||
if inertia is None:
|
||||
self.I_body = np.diag([0.07, 0.26, 0.24])
|
||||
else:
|
||||
self.I_body = np.array(inertia).reshape(3, 3)
|
||||
|
||||
self.N = MPC_HORIZON
|
||||
self.dt = MPC_DT
|
||||
self.Q = np.diag(Q_WEIGHTS)
|
||||
self.R = np.diag(R_WEIGHTS)
|
||||
self.gravity = np.array([0, 0, -9.81])
|
||||
|
||||
self._last_forces = np.zeros(NU)
|
||||
|
||||
def solve(self, x0, x_ref, foot_positions, contact_schedule):
|
||||
"""Solve MPC QP.
|
||||
|
||||
Args:
|
||||
x0: (12,) current state [pos, rpy, vel, omega]
|
||||
x_ref: (12, N) reference trajectory over horizon
|
||||
foot_positions: (4, 3) foot positions in world frame (relative to CoM)
|
||||
contact_schedule: (4, N) binary contact table (1=stance)
|
||||
|
||||
Returns:
|
||||
forces: (12,) optimal forces for current timestep [f1x,f1y,f1z,...,f4x,f4y,f4z]
|
||||
"""
|
||||
N = self.N
|
||||
|
||||
# Build dynamics matrices
|
||||
Ad, Bd_list, gd = self._discretize_dynamics(x0, foot_positions)
|
||||
|
||||
# Build QP: min 0.5 z'Hz + f'z s.t. lb <= Az <= ub, lbx <= z <= ubx
|
||||
# Decision variables: z = [x1,...,xN, u0,...,uN-1]
|
||||
nvars = N * NX + N * NU
|
||||
|
||||
# --- Hessian ---
|
||||
H_diag = np.concatenate([np.tile(2 * Q_WEIGHTS, N), np.tile(2 * R_WEIGHTS, N)])
|
||||
H = sparse.diags(H_diag, format='csc')
|
||||
|
||||
# --- Gradient ---
|
||||
g = np.zeros(nvars)
|
||||
for k in range(N):
|
||||
g[k*NX:(k+1)*NX] = -2 * self.Q @ x_ref[:, k]
|
||||
|
||||
# --- Dynamics equality constraints ---
|
||||
# x_{k+1} = Ad @ x_k + Bd_k @ u_k + gd
|
||||
# Rewrite: x_{k+1} - Ad @ x_k - Bd_k @ u_k = gd (for k>0)
|
||||
# x_1 - Bd_0 @ u_0 = Ad @ x0 + gd (for k=0)
|
||||
n_eq = N * NX
|
||||
A_eq = np.zeros((n_eq, nvars))
|
||||
b_eq = np.zeros(n_eq)
|
||||
|
||||
# k=0: x_1 = Ad @ x0 + Bd_0 @ u_0 + gd
|
||||
A_eq[0:NX, 0:NX] = np.eye(NX) # x_1
|
||||
A_eq[0:NX, N*NX:N*NX+NU] = -Bd_list[0] # -Bd_0 @ u_0
|
||||
b_eq[0:NX] = Ad @ x0 + gd
|
||||
|
||||
for k in range(1, N):
|
||||
row = k * NX
|
||||
# x_{k+1}
|
||||
A_eq[row:row+NX, k*NX:(k+1)*NX] = np.eye(NX)
|
||||
# -Ad @ x_k
|
||||
A_eq[row:row+NX, (k-1)*NX:k*NX] = -Ad
|
||||
# -Bd_k @ u_k
|
||||
A_eq[row:row+NX, N*NX+k*NU:N*NX+(k+1)*NU] = -Bd_list[k]
|
||||
b_eq[row:row+NX] = gd
|
||||
|
||||
# --- Friction pyramid inequality constraints ---
|
||||
# For each stance leg at each timestep: 4 faces
|
||||
# fx - mu*fz <= 0, -fx - mu*fz <= 0, fy - mu*fz <= 0, -fy - mu*fz <= 0
|
||||
n_ineq_max = 4 * 4 * N
|
||||
A_ineq = np.zeros((n_ineq_max, nvars))
|
||||
u_ineq = np.zeros(n_ineq_max)
|
||||
|
||||
row = 0
|
||||
for k in range(N):
|
||||
u_base = N * NX + k * NU
|
||||
for leg in range(4):
|
||||
if contact_schedule[leg, k] == 1:
|
||||
fx_idx = u_base + leg * 3
|
||||
fy_idx = u_base + leg * 3 + 1
|
||||
fz_idx = u_base + leg * 3 + 2
|
||||
|
||||
# Friction pyramid: stance leg
|
||||
A_ineq[row, fx_idx] = 1.0
|
||||
A_ineq[row, fz_idx] = -MU
|
||||
row += 1
|
||||
A_ineq[row, fx_idx] = -1.0
|
||||
A_ineq[row, fz_idx] = -MU
|
||||
row += 1
|
||||
A_ineq[row, fy_idx] = 1.0
|
||||
A_ineq[row, fz_idx] = -MU
|
||||
row += 1
|
||||
A_ineq[row, fy_idx] = -1.0
|
||||
A_ineq[row, fz_idx] = -MU
|
||||
row += 1
|
||||
|
||||
A_ineq = A_ineq[:row]
|
||||
u_ineq = u_ineq[:row]
|
||||
|
||||
# Stack constraints
|
||||
A_full = np.vstack([A_eq, A_ineq])
|
||||
l_full = np.concatenate([b_eq, -np.inf * np.ones(row)])
|
||||
u_full = np.concatenate([b_eq, u_ineq])
|
||||
|
||||
# --- Box constraints on forces (as identity rows in A) ---
|
||||
A_box = np.zeros((N * NU, nvars))
|
||||
l_box = -np.inf * np.ones(N * NU)
|
||||
u_box = np.inf * np.ones(N * NU)
|
||||
|
||||
for k in range(N):
|
||||
u_base = N * NX + k * NU
|
||||
for leg in range(4):
|
||||
idx = u_base + leg * 3
|
||||
box_row = k * NU + leg * 3
|
||||
# Identity rows for fx, fy, fz
|
||||
for j in range(3):
|
||||
A_box[box_row + j, idx + j] = 1.0
|
||||
|
||||
if contact_schedule[leg, k] == 1:
|
||||
# Stance: fz bounded
|
||||
l_box[box_row + 2] = FZ_MIN
|
||||
u_box[box_row + 2] = FZ_MAX
|
||||
else:
|
||||
# Swing: all forces = 0
|
||||
l_box[box_row:box_row+3] = 0.0
|
||||
u_box[box_row:box_row+3] = 0.0
|
||||
|
||||
# Final constraint matrix
|
||||
A_full = np.vstack([A_full, A_box])
|
||||
l_full = np.concatenate([l_full, l_box])
|
||||
u_full = np.concatenate([u_full, u_box])
|
||||
|
||||
# --- Solve with OSQP ---
|
||||
A_sparse = sparse.csc_matrix(A_full)
|
||||
H_sparse = sparse.triu(H, format='csc')
|
||||
|
||||
solver = osqp.OSQP()
|
||||
solver.setup(H_sparse, g, A_sparse, l_full, u_full,
|
||||
eps_abs=1e-4, eps_rel=1e-4,
|
||||
max_iter=500, polish=True, verbose=False,
|
||||
warm_start=True)
|
||||
|
||||
# Warm start with previous solution
|
||||
if self._last_forces is not None:
|
||||
x_warm = np.zeros(nvars)
|
||||
x_warm[N*NX:N*NX+NU] = self._last_forces
|
||||
solver.warm_start(x=x_warm)
|
||||
|
||||
result = solver.solve()
|
||||
|
||||
if result.info.status == 'solved' or result.info.status == 'solved_inaccurate':
|
||||
# Extract first timestep forces
|
||||
forces = result.x[N*NX:N*NX+NU]
|
||||
self._last_forces = forces.copy()
|
||||
else:
|
||||
forces = self._last_forces
|
||||
|
||||
return forces
|
||||
|
||||
def _discretize_dynamics(self, x0, foot_positions):
|
||||
"""Build discrete-time centroidal dynamics.
|
||||
|
||||
State: [pos, rpy, vel, omega] (12)
|
||||
Continuous: dx/dt = Ac @ x + Bc @ u + gc
|
||||
Discrete: x_{k+1} = Ad @ x + Bd @ u + gd
|
||||
"""
|
||||
m = self.mass
|
||||
I_inv = np.linalg.inv(self.I_body)
|
||||
dt = self.dt
|
||||
yaw = x0[5]
|
||||
cy, sy = np.cos(yaw), np.sin(yaw)
|
||||
|
||||
# Rotation for rpy rate ≈ R_z^T @ omega
|
||||
R_zT = np.array([[cy, sy, 0], [-sy, cy, 0], [0, 0, 1]])
|
||||
|
||||
# Ac (12×12)
|
||||
Ac = np.zeros((NX, NX))
|
||||
Ac[0:3, 6:9] = np.eye(3) # pos_dot = vel
|
||||
Ac[3:6, 9:12] = R_zT # rpy_dot ≈ R_z^T @ omega
|
||||
|
||||
# Ad = I + Ac*dt (first-order)
|
||||
Ad = np.eye(NX) + Ac * dt
|
||||
|
||||
# Bc varies per timestep (foot positions change contact point)
|
||||
Bd_list = []
|
||||
for k in range(self.N):
|
||||
Bc = np.zeros((NX, NU))
|
||||
for leg in range(4):
|
||||
r = foot_positions[leg]
|
||||
# vel_dot += f/m
|
||||
Bc[6:9, leg*3:(leg+1)*3] = np.eye(3) / m
|
||||
# omega_dot += I^{-1} @ (r × f)
|
||||
Bc[9:12, leg*3:(leg+1)*3] = I_inv @ _skew(r)
|
||||
Bd = Bc * dt
|
||||
Bd_list.append(Bd)
|
||||
|
||||
# Gravity contribution
|
||||
gd = np.zeros(NX)
|
||||
gd[6:9] = self.gravity * dt # vel += g*dt
|
||||
|
||||
return Ad, Bd_list, gd
|
||||
@@ -1,261 +0,0 @@
|
||||
"""MPC controller integration for wheeled-legged robot.
|
||||
|
||||
Integrates: gait scheduler + reference trajectory + ConvexMPC solver +
|
||||
swing leg control + stance force mapping + wheel drive.
|
||||
|
||||
Architecture (following go2-convex-mpc):
|
||||
- MPC runs at ~50 Hz (every MPC_DECIMATION control steps)
|
||||
- Swing/stance leg controller runs at control rate (250 Hz)
|
||||
- Wheel drive: stance legs use differential drive, swing legs coast
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from robot import Robot, RobotState
|
||||
from dynamics import Dynamics
|
||||
from mpc import ConvexMPC, MPC_DT
|
||||
from config import (
|
||||
LEG_NAMES, WHEEL_RADIUS, WHEEL_TRACK, WHEEL_VEL_MAX,
|
||||
CTRL_DT, GAIT_FREQ, GAIT_DUTY, SWING_HEIGHT, PHASE_OFFSETS,
|
||||
DEFAULT_JOINT_ANGLES, ROBOT_MASS,
|
||||
)
|
||||
|
||||
# MPC update decimation (relative to control loop)
|
||||
MPC_DECIMATION = max(1, int(MPC_DT / CTRL_DT)) # ~5 steps at 250Hz
|
||||
|
||||
|
||||
class MPCController:
|
||||
"""Convex MPC locomotion controller for wheeled-legged robot."""
|
||||
|
||||
def __init__(self, robot: Robot):
|
||||
self.robot = robot
|
||||
self.dynamics = Dynamics()
|
||||
self.mpc = ConvexMPC(mass=ROBOT_MASS)
|
||||
|
||||
# User commands
|
||||
self.vel_x = 0.0
|
||||
self.vel_y = 0.0
|
||||
self.yaw_rate = 0.0
|
||||
self.height = 0.35 # actual standing height with default joint angles
|
||||
|
||||
# Gait state - start at phase 0 with all legs in stance (duty=0.6)
|
||||
self._gait_phase = 0.0
|
||||
self._step_count = 0
|
||||
self._initialized = False
|
||||
|
||||
# MPC solution cache - initialize with gravity compensation
|
||||
self._mpc_forces = np.zeros(12)
|
||||
self._init_gravity_comp()
|
||||
|
||||
# Swing trajectory state
|
||||
self._swing_start_foot = {leg: np.zeros(3) for leg in LEG_NAMES}
|
||||
self._swing_start_time = {leg: 0.0 for leg in LEG_NAMES}
|
||||
self._last_contact = {leg: True for leg in LEG_NAMES}
|
||||
|
||||
# Smoothed commands
|
||||
self._smooth_vx = 0.0
|
||||
self._smooth_vy = 0.0
|
||||
self._smooth_yaw = 0.0
|
||||
|
||||
def _init_gravity_comp(self):
|
||||
"""Pre-fill MPC forces with static gravity compensation."""
|
||||
fz_per_leg = ROBOT_MASS * 9.81 / 4.0
|
||||
for i in range(4):
|
||||
self._mpc_forces[i*3 + 2] = fz_per_leg
|
||||
|
||||
def compute(self, state: RobotState, dt: float):
|
||||
"""Main MPC control loop.
|
||||
|
||||
Uses MIT motor protocol: tau = kp*(q_des-q) + kd*(dq_des-dq) + tau_ff
|
||||
where tau_ff comes from MPC force mapping via Jacobian transpose.
|
||||
|
||||
Returns:
|
||||
tau_legs: (12,) feedforward torques for MIT mode
|
||||
wheel_targets: (4,) wheel velocity targets
|
||||
"""
|
||||
# Smooth commands
|
||||
alpha = min(dt * 3.0, 1.0)
|
||||
self._smooth_vx += alpha * (self.vel_x - self._smooth_vx)
|
||||
self._smooth_vy += alpha * (self.vel_y - self._smooth_vy)
|
||||
self._smooth_yaw += alpha * (self.yaw_rate - self._smooth_yaw)
|
||||
|
||||
# Update Pinocchio
|
||||
q_pin, dq_pin = self.robot.get_qpos_qvel_for_pinocchio()
|
||||
self.dynamics.update(q_pin, dq_pin)
|
||||
|
||||
# Initialize foot positions on first call
|
||||
if not self._initialized:
|
||||
for leg in LEG_NAMES:
|
||||
self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg)
|
||||
self._initialized = True
|
||||
|
||||
# Decide if we should trot or just stand
|
||||
moving = (abs(self._smooth_vx) > 0.02 or
|
||||
abs(self._smooth_vy) > 0.02 or
|
||||
abs(self._smooth_yaw) > 0.05)
|
||||
|
||||
if moving:
|
||||
self._gait_phase = (self._gait_phase + dt * GAIT_FREQ) % 1.0
|
||||
else:
|
||||
self._gait_phase = 0.0 # all legs in stance
|
||||
|
||||
# Contact schedule
|
||||
contacts = {}
|
||||
for leg in LEG_NAMES:
|
||||
phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0
|
||||
contacts[leg] = phase < GAIT_DUTY
|
||||
|
||||
# Get foot positions relative to CoM
|
||||
foot_positions = np.zeros((4, 3))
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
foot_positions[i] = self.dynamics.get_foot_pos(leg) - state.pos
|
||||
|
||||
# --- Run MPC at lower rate ---
|
||||
if self._step_count % MPC_DECIMATION == 0:
|
||||
x0 = self._build_state_vector(state)
|
||||
x_ref = self._build_reference(state)
|
||||
contact_table = self._build_contact_table()
|
||||
self._mpc_forces = self.mpc.solve(x0, x_ref, foot_positions, contact_table)
|
||||
|
||||
self._step_count += 1
|
||||
|
||||
# --- Compute feedforward torques and desired joint positions ---
|
||||
tau_ff = np.zeros(12)
|
||||
q_des = np.zeros(12)
|
||||
dq_des = np.zeros(12)
|
||||
kp = np.zeros(12)
|
||||
kd = np.zeros(12)
|
||||
wheel_targets = np.zeros(4)
|
||||
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
if contacts[leg]:
|
||||
# Stance: MPC force → feedforward torque, PD holds posture
|
||||
f_leg = self._mpc_forces[i*3:(i+1)*3]
|
||||
J = self.dynamics.get_foot_jacobian_leg(leg)
|
||||
tau_ff[i*3:(i+1)*3] = J.T @ (-f_leg)
|
||||
|
||||
# PD target: default standing angles (posture hold)
|
||||
q_des[i*3] = DEFAULT_JOINT_ANGLES["hip_abduction"]
|
||||
q_des[i*3+1] = DEFAULT_JOINT_ANGLES["hip_pitch"]
|
||||
q_des[i*3+2] = DEFAULT_JOINT_ANGLES["knee"]
|
||||
kp[i*3:(i+1)*3] = [40.0, 40.0, 40.0]
|
||||
kd[i*3:(i+1)*3] = [3.0, 3.0, 3.0]
|
||||
|
||||
# Record foot position
|
||||
self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg)
|
||||
self._last_contact[leg] = True
|
||||
|
||||
# Wheel drive
|
||||
wheel_targets[i] = self._wheel_cmd(leg)
|
||||
else:
|
||||
# Swing: IK target position, strong PD, no feedforward
|
||||
if self._last_contact[leg]:
|
||||
self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg)
|
||||
self._swing_start_time[leg] = state.time
|
||||
self._last_contact[leg] = False
|
||||
|
||||
q_ik = self._swing_leg_ik(leg, state, q_pin)
|
||||
q_des[i*3:(i+1)*3] = q_ik
|
||||
kp[i*3:(i+1)*3] = [60.0, 60.0, 60.0] # strong PD for swing
|
||||
kd[i*3:(i+1)*3] = [3.0, 3.0, 3.0]
|
||||
# tau_ff stays 0 for swing
|
||||
|
||||
wheel_targets[i] = 0.0
|
||||
|
||||
# Use MIT protocol via robot interface
|
||||
self.robot.set_ctrl_mit(q_des, dq_des, kp, kd, tau_ff, wheel_targets)
|
||||
# Return dummy (actual ctrl is set directly above)
|
||||
return None, None
|
||||
|
||||
def _build_state_vector(self, state: RobotState):
|
||||
"""Build MPC state: [pos, rpy, vel, omega]."""
|
||||
return np.concatenate([state.pos, state.rpy, state.lin_vel, state.ang_vel])
|
||||
|
||||
def _build_reference(self, state: RobotState):
|
||||
"""Build reference trajectory over MPC horizon."""
|
||||
N = self.mpc.N
|
||||
x_ref = np.zeros((12, N))
|
||||
|
||||
yaw = state.rpy[2]
|
||||
cy, sy = np.cos(yaw), np.sin(yaw)
|
||||
R_z = np.array([[cy, -sy, 0], [sy, cy, 0], [0, 0, 1]])
|
||||
vel_world = R_z @ np.array([self._smooth_vx, self._smooth_vy, 0.0])
|
||||
|
||||
for k in range(N):
|
||||
t = (k + 1) * self.mpc.dt
|
||||
# Position: integrate from current
|
||||
x_ref[0, k] = state.pos[0] + vel_world[0] * t
|
||||
x_ref[1, k] = state.pos[1] + vel_world[1] * t
|
||||
x_ref[2, k] = self.height
|
||||
# RPY: keep roll/pitch zero, integrate yaw
|
||||
x_ref[3, k] = 0.0
|
||||
x_ref[4, k] = 0.0
|
||||
x_ref[5, k] = yaw + self._smooth_yaw * t
|
||||
# Velocity
|
||||
x_ref[6, k] = vel_world[0]
|
||||
x_ref[7, k] = vel_world[1]
|
||||
x_ref[8, k] = 0.0
|
||||
# Angular velocity
|
||||
x_ref[9, k] = 0.0
|
||||
x_ref[10, k] = 0.0
|
||||
x_ref[11, k] = self._smooth_yaw
|
||||
|
||||
return x_ref
|
||||
|
||||
def _build_contact_table(self):
|
||||
"""Build contact schedule over MPC horizon."""
|
||||
N = self.mpc.N
|
||||
table = np.zeros((4, N), dtype=int)
|
||||
for k in range(N):
|
||||
future_phase = (self._gait_phase + (k + 1) * self.mpc.dt * GAIT_FREQ) % 1.0
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
leg_phase = (future_phase + PHASE_OFFSETS[leg]) % 1.0
|
||||
table[i, k] = 1 if leg_phase < GAIT_DUTY else 0
|
||||
return table
|
||||
|
||||
def _swing_leg_ik(self, leg: str, state: RobotState, q_pin: np.ndarray):
|
||||
"""Swing leg: compute IK target joint angles for trajectory."""
|
||||
swing_phase = self._get_swing_phase(leg)
|
||||
|
||||
p_start = self._swing_start_foot[leg]
|
||||
p_end = self._compute_touchdown(leg, state)
|
||||
|
||||
s = swing_phase
|
||||
s_mj = 10*s**3 - 15*s**4 + 6*s**5
|
||||
|
||||
pos_des = p_start + (p_end - p_start) * s_mj
|
||||
# Z lift
|
||||
z_lift = 64.0 * s**3 * (1.0 - s)**3
|
||||
pos_des[2] = p_start[2] + SWING_HEIGHT * z_lift
|
||||
|
||||
# IK to get joint angles
|
||||
q_ik = self.dynamics.inverse_kinematics(leg, pos_des, q_pin)
|
||||
return q_ik
|
||||
|
||||
def _get_swing_phase(self, leg: str) -> float:
|
||||
phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0
|
||||
if phase < GAIT_DUTY:
|
||||
return 0.0
|
||||
return (phase - GAIT_DUTY) / (1.0 - GAIT_DUTY)
|
||||
|
||||
def _compute_touchdown(self, leg: str, state: RobotState) -> np.ndarray:
|
||||
"""Raibert heuristic for touchdown position."""
|
||||
td = self._swing_start_foot[leg].copy()
|
||||
t_stance = GAIT_DUTY / GAIT_FREQ
|
||||
|
||||
yaw = state.rpy[2]
|
||||
cy, sy = np.cos(yaw), np.sin(yaw)
|
||||
R_z = np.array([[cy, -sy, 0], [sy, cy, 0], [0, 0, 1]])
|
||||
cmd_vel_world = R_z @ np.array([self._smooth_vx, self._smooth_vy, 0.0])
|
||||
|
||||
td[0] += cmd_vel_world[0] * t_stance * 0.5
|
||||
td[1] += cmd_vel_world[1] * t_stance * 0.5
|
||||
td[2] = WHEEL_RADIUS
|
||||
return td
|
||||
|
||||
def _wheel_cmd(self, leg: str) -> float:
|
||||
"""Differential drive for a single wheel."""
|
||||
if leg[1] == "l":
|
||||
v = (self._smooth_vx - 0.5 * WHEEL_TRACK * self._smooth_yaw) / WHEEL_RADIUS
|
||||
else:
|
||||
v = (self._smooth_vx + 0.5 * WHEEL_TRACK * self._smooth_yaw) / WHEEL_RADIUS
|
||||
return np.clip(v, -WHEEL_VEL_MAX, WHEEL_VEL_MAX)
|
||||
@@ -1,229 +0,0 @@
|
||||
"""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 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)
|
||||
|
||||
|
||||
@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 2R FK
|
||||
# leg_length = sqrt(L1^2 + L2^2 - 2*L1*L2*cos(pi + knee))
|
||||
import math
|
||||
L1, L2 = 0.25, 0.20
|
||||
knee = DEFAULT_JOINT_ANGLES["knee"]
|
||||
leg_length = math.sqrt(L1**2 + L2**2 - 2*L1*L2*math.cos(math.pi + knee))
|
||||
# base_z = wheel_radius + leg_length - hip_z_offset
|
||||
base_z = 0.10 + leg_length - 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
|
||||
@@ -1,111 +0,0 @@
|
||||
"""Main entry point: wheeled-legged robot simulation.
|
||||
|
||||
Controls:
|
||||
Mode: wheel (default) - differential drive + posture hold
|
||||
trot - quadruped gait with wheel assist
|
||||
mpc - convex MPC locomotion (torque control)
|
||||
|
||||
Keyboard (in MuJoCo viewer):
|
||||
W/S: vel_x ±0.1
|
||||
A/D: yaw_rate ±0.2
|
||||
Q/E: height ±0.02
|
||||
1: wheel mode
|
||||
2: trot mode
|
||||
3: MPC mode
|
||||
4: prone toggle
|
||||
Z: reset commands
|
||||
"""
|
||||
|
||||
import time
|
||||
import numpy as np
|
||||
import mujoco.viewer as mjv
|
||||
|
||||
from robot import Robot
|
||||
from controller import Controller
|
||||
from gui import GUI
|
||||
from config import CTRL_DECIMATION
|
||||
|
||||
|
||||
def main():
|
||||
robot = Robot()
|
||||
robot.reset()
|
||||
ctrl = Controller(robot)
|
||||
gui = GUI(ctrl)
|
||||
|
||||
step = 0
|
||||
sim_steps_per_ctrl = CTRL_DECIMATION
|
||||
|
||||
def key_callback(keycode):
|
||||
"""Called from MuJoCo render thread - only modify ctrl directly, not tkinter."""
|
||||
try:
|
||||
c = chr(keycode).lower()
|
||||
except (ValueError, OverflowError):
|
||||
return
|
||||
if c == 'w':
|
||||
ctrl.vel_x = min(ctrl.vel_x + 0.1, 1.5)
|
||||
elif c == 's':
|
||||
ctrl.vel_x = max(ctrl.vel_x - 0.1, -1.5)
|
||||
elif c == 'a':
|
||||
ctrl.yaw_rate = min(ctrl.yaw_rate + 0.2, 2.0)
|
||||
elif c == 'd':
|
||||
ctrl.yaw_rate = max(ctrl.yaw_rate - 0.2, -2.0)
|
||||
elif c == 'q':
|
||||
ctrl.height = min(ctrl.height + 0.02, 0.45)
|
||||
elif c == 'e':
|
||||
ctrl.height = max(ctrl.height - 0.02, 0.15)
|
||||
elif c == '1':
|
||||
ctrl.mode = "wheel"; ctrl.prone = False
|
||||
elif c == '2':
|
||||
ctrl.mode = "trot"; ctrl.prone = False
|
||||
elif c == '3':
|
||||
ctrl.mode = "mpc"; ctrl.prone = False
|
||||
elif c == '4':
|
||||
ctrl.prone = not ctrl.prone
|
||||
elif c == 'z':
|
||||
ctrl.vel_x = 0.0; ctrl.vel_y = 0.0; ctrl.yaw_rate = 0.0
|
||||
|
||||
with mjv.launch_passive(robot.model, robot.data, key_callback=key_callback) as viewer:
|
||||
viewer.cam.distance = 2.5
|
||||
viewer.cam.elevation = -20
|
||||
viewer.cam.azimuth = 135
|
||||
|
||||
last_time = robot.data.time
|
||||
|
||||
while viewer.is_running() and not gui.closed:
|
||||
t_start = time.perf_counter()
|
||||
|
||||
# Detect viewer reset (Backspace) - time jumps back to 0
|
||||
if robot.data.time < last_time:
|
||||
robot.reset()
|
||||
last_time = robot.data.time
|
||||
|
||||
# Get state and compute control
|
||||
state = robot.get_state()
|
||||
leg_targets, wheel_targets = ctrl.compute(state, robot.dt * sim_steps_per_ctrl)
|
||||
|
||||
# Apply control and step simulation
|
||||
# MPC mode sets ctrl directly via set_ctrl_mit, skip set_ctrl
|
||||
if ctrl.mode != "mpc":
|
||||
robot.set_ctrl(leg_targets, wheel_targets)
|
||||
for _ in range(sim_steps_per_ctrl):
|
||||
robot.step()
|
||||
|
||||
viewer.sync()
|
||||
step += 1
|
||||
|
||||
# Update GUI every 25 steps (~10 Hz)
|
||||
if step % 25 == 0:
|
||||
state = robot.get_state()
|
||||
gui.update_status(state, step)
|
||||
if not gui.tick():
|
||||
break
|
||||
|
||||
# Real-time sync
|
||||
elapsed = time.perf_counter() - t_start
|
||||
target_dt = robot.dt * sim_steps_per_ctrl
|
||||
if elapsed < target_dt:
|
||||
time.sleep(target_dt - elapsed)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user