243 lines
8.2 KiB
Python
243 lines
8.2 KiB
Python
"""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
|