288 lines
12 KiB
Python
288 lines
12 KiB
Python
"""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
|