import time
import numpy as np
import pinocchio as pin
import pybullet as pyb
from . import WB_MPC_Wrapper
class Result:
"""
Object to store the result of the control loop
It contains what is sent to the robot (gains, desired positions and velocities,
feedforward torques)
"""
def __init__(self, params):
self.P = np.array(params.Kp_main.tolist() * 4)
self.D = np.array(params.Kd_main.tolist() * 4)
self.FF = params.Kff_main * np.ones(12)
self.q_des = np.zeros(12)
self.v_des = np.zeros(12)
self.tau_ff = np.zeros(12)
class Interpolator:
def __init__(self, params):
self.dt = params.dt_mpc
self.type = params.interpolation_type
self.q0 = np.zeros(3)
self.q1 = np.zeros(3)
self.v0 = np.zeros(3)
self.v1 = np.zeros(3)
if self.type == 3:
self.ts = np.repeat(np.linspace(0, 2 * self.dt, 3), 2)
else:
self.alpha = np.zeros(3)
self.beta = np.zeros(3)
self.gamma = np.zeros(3)
self.delta = 1.0
def update(self, x0, x1, x2=None):
self.q0 = x0[:3]
self.q1 = x1[:3]
self.v0 = x0[3:]
self.v1 = x1[3:]
if self.type == 0: # Linear
self.alpha = 0.0
self.beta = self.v1
self.gamma = self.q0
elif self.type == 1: # Quadratic fixed velocity
self.alpha = 2 * (self.q1 - self.q0 - self.v0 * self.dt) / self.dt**2
self.beta = self.v0
self.gamma = self.q0
elif self.type == 2: # Quadratic time variable
for i in range(3):
q0 = self.q0[i]
v0 = self.v0[i]
q1 = self.q1[i]
v1 = self.v1[i]
if (q1 == q0) or (v1 == -v0):
self.alpha[i] = 0.0
self.beta[i] = 0.0
self.gamma[i] = q1
self.delta = 1.0
else:
self.alpha[i] = (v1**2 - v0**2) / (2 * (q1 - q0))
self.beta[i] = v0
self.gamma[i] = q0
self.delta = 2 * (q1 - q0) / (v1 + v0) / self.dt
elif self.type == 3: # Spline interpolation
from scipy.interpolate import KroghInterpolator
if x2 is not None:
self.q2 = x2[:3]
self.v2 = x2[3:]
self.y = [self.q0, self.v0, self.q1, self.v1, self.q2, self.v2]
self.krog = KroghInterpolator(self.ts, np.array(self.y))
else:
self.y = [self.q0, self.v0, self.q1, self.v1]
self.krog = KroghInterpolator(self.ts[:4], np.array(self.y))
def interpolate(self, t):
if self.type == 3:
q = self.krog(t)
v = self.krog.derivative(t)
return q, v
if self.type == 2:
t *= self.delta
q = 1 / 2 * self.alpha * t**2 + self.beta * t + self.gamma
v = self.v1 if self.type == 2 else self.alpha * t + self.beta
return q, v
def plot(self, n, dt):
import matplotlib.pyplot as plt
ts = np.linspace(0.0, 2 * self.dt, 2 * n + 1)
plt.style.use("seaborn")
for i in range(3):
plt.subplot(3, 2, (i * 2) + 1)
plt.title("Position interpolation")
plt.plot(ts, [self.interpolate(t)[0][i] for t in ts])
plt.scatter(y=self.q0[i], x=0.0, color="violet", marker="+")
plt.scatter(y=self.q1[i], x=self.dt, color="violet", marker="+")
if self.type == 3 and self.q2 is not None:
plt.scatter(y=self.q2[i], x=2 * self.dt, color="violet", marker="+")
plt.subplot(3, 2, (i * 2) + 2)
plt.title("Velocity interpolation")
plt.plot(ts, [self.interpolate(t)[1][i] for t in ts])
plt.scatter(y=self.v0[i], x=0.0, color="violet", marker="+")
plt.scatter(y=self.v1[i], x=self.dt, color="violet", marker="+")
if self.type == 3 and self.v2 is not None:
plt.scatter(y=self.v2[i], x=2 * self.dt, color="violet", marker="+")
plt.show()
class DummyDevice:
def __init__(self):
self.imu = self.IMU()
self.joints = self.Joints()
self.base_position = np.zeros(3)
self.base_position[2] = 0.1944
self.b_base_velocity = np.zeros(3)
self.baseState = tuple()
self.baseVel = tuple()
class IMU:
def __init__(self):
self.linear_acceleration = np.zeros(3)
self.gyroscope = np.zeros(3)
self.attitude_euler = np.zeros(3)
self.attitude_quaternion = np.zeros(4)
class Joints:
def __init__(self):
self.positions = np.zeros(12)
self.velocities = np.zeros(12)
class Controller:
def __init__(self, pd, target, params, q_init, t):
"""
Function that computes the reference control (tau, q_des, v_des and gains)
Args:
params (Params object): store parameters
q_init (array): initial position of actuators
t (float): time of the simulation
"""
self.q_security = np.array([1.2, 2.1, 3.14] * 4)
self.pd = pd
self.target = target
self.params = params
self.q_init = pd.q0
self.k = 0
self.error = False
self.initialized = False
self.result = Result(params)
self.result.q_des = self.pd.q0[7:].copy()
self.result.v_des = self.pd.v0[6:].copy()
self.mpc = WB_MPC_Wrapper.MPC_Wrapper(pd, target, params)
self.mpc_solved = False
self.k_result = 0
self.k_solve = 0
if self.params.interpolate_mpc:
self.interpolator = Interpolator(params)
try:
file = np.load("/tmp/init_guess.npy", allow_pickle=True).item()
self.xs_init = list(file["xs"])
self.us_init = list(file["us"])
print("Initial guess loaded \n")
except:
self.xs_init = None
self.us_init = None
print("No initial guess\n")
device = DummyDevice()
device.joints.positions = q_init
self.compute(device)
def compute(self, device, qc=None):
"""
Run one iteration of the main control loop
Args:
device (object): Interface with the masterboard or the simulation
"""
t_start = time.time()
m = self.read_state(device)
t_measures = time.time()
self.t_measures = t_measures - t_start
if self.k % self.pd.mpc_wbc_ratio == 0:
if self.mpc_solved:
self.k_solve = self.k
self.mpc_solved = False
self.target.update(self.k // self.pd.mpc_wbc_ratio)
self.target.shift_gait()
try:
self.mpc.solve(self.k, m["x_m"], self.xs_init, self.us_init)
except ValueError:
self.error = True
print("MPC Problem")
t_mpc = time.time()
self.t_mpc = t_mpc - t_measures
if not self.error:
self.mpc_result = self.mpc.get_latest_result()
xs = self.mpc_result.xs
if self.mpc_result.new_result:
self.mpc_solved = True
self.k_new = self.k
print(f"MPC solved in {self.k - self.k_solve} iterations")
if not self.initialized and self.params.save_guess:
self.save_guess()
self.result.FF = self.params.Kff_main * np.ones(12)
self.result.tau_ff[3:6] = self.compute_torque(m)[:]
if self.params.interpolate_mpc:
if self.mpc_result.new_result:
if self.params.interpolation_type == 3:
self.interpolator.update(xs[0], xs[1], xs[2])
# self.interpolator.plot(self.pd.mpc_wbc_ratio, self.pd.dt_wbc)
t = (self.k - self.k_solve + 1) * self.pd.dt_wbc
q, v = self.interpolator.interpolate(t)
else:
q, v = self.integrate_x(m)
self.result.q_des[3:6] = q[:]
self.result.v_des[3:6] = v[:]
self.xs_init = self.mpc_result.xs[1:] + [self.mpc_result.xs[-1]]
self.us_init = self.mpc_result.us[1:] + [self.mpc_result.us[-1]]
t_send = time.time()
self.t_send = t_send - t_mpc
# self.clamp_result(device)
# self.security_check(m)
if self.error:
self.set_null_control()
# self.pyb_camera(device)
self.t_loop = time.time() - t_start
self.k += 1
self.initialized = True
return self.error
def pyb_camera(self, device):
"""
Update position of PyBullet camera on the robot position to do as if it was
attached to the robot
"""
if self.k > 10 and self.params.enable_pyb_GUI:
pyb.resetDebugVisualizerCamera(
cameraDistance=0.6,
cameraYaw=45,
cameraPitch=-39.9,
cameraTargetPosition=[device.height[0], device.height[1], 0.0],
)
def security_check(self, m):
"""
Check if the command is fine and set the command to zero in case of error
"""
if not self.error:
if (np.abs(m["qj_m"]) > self.q_security).any():
print("-- POSITION LIMIT ERROR --")
print(m["qj_m"])
print(np.abs(m["qj_m"]) > self.q_security)
self.error = True
elif (np.abs(m["vj_m"]) > 1000 * np.pi / 180).any():
print("-- VELOCITY TOO HIGH ERROR --")
print(m["vj_m"])
print(np.abs(m["vj_m"]) > 1000 * np.pi / 180)
self.error = True
elif (np.abs(self.result.FF) > 3.2).any():
print("-- FEEDFORWARD TORQUES TOO HIGH ERROR --")
print(self.result.FF)
print(np.abs(self.result.FF) > 3.2)
self.error = True
def clamp(self, num, min_value=None, max_value=None):
clamped = False
if min_value is not None and num <= min_value:
num = min_value
clamped = True
if max_value is not None and num >= max_value:
num = max_value
clamped = True
return clamped
def clamp_result(self, device, set_error=False):
"""
Clamp the result
"""
hip_max = 120.0 * np.pi / 180.0
knee_min = 5.0 * np.pi / 180.0
for i in range(4):
if self.clamp(self.result.q_des[3 * i + 1], -hip_max, hip_max):
print("Clamping hip n " + str(i))
self.error = set_error
if self.q_init[7 + 3 * i + 2] >= 0.0 and self.clamp(
self.result.q_des[3 * i + 2], knee_min
):
print("Clamping knee n " + str(i))
self.error = set_error
elif self.q_init[7 + 3 * i + 2] <= 0.0 and self.clamp(
self.result.q_des[3 * i + 2], max_value=-knee_min
):
print("Clamping knee n " + str(i))
self.error = set_error
for i in range(12):
if self.clamp(
self.result.q_des[i],
device.joints.positions[i] - 4.0,
device.joints.positions[i] + 4.0,
):
print("Clamping position difference of motor n " + str(i))
self.error = set_error
if self.clamp(
self.result.v_des[i],
device.joints.velocities[i] - 100.0,
device.joints.velocities[i] + 100.0,
):
print("Clamping velocity of motor n " + str(i))
self.error = set_error
if self.clamp(self.result.tau_ff[i], -3.2, 3.2):
print("Clamping torque of motor n " + str(i))
self.error = set_error
def set_null_control(self):
"""
Send default null values to the robot
"""
self.result.P = np.zeros(12)
self.result.D = 0.1 * np.ones(12)
self.result.FF = np.zeros(12)
self.result.q_des[:] = np.zeros(12)
self.result.v_des[:] = np.zeros(12)
self.result.tau_ff[:] = np.zeros(12)
def save_guess(self):
"""
Save the result of the MPC in a file called /tmp/init_guess.npy
"""
np.save(
open("/tmp/init_guess.npy", "wb"),
{"xs": self.mpc_result.xs, "us": self.mpc_result.us},
)
print("Initial guess saved")
def read_state(self, device):
qj_m = device.joints.positions
vj_m = device.joints.velocities
x_m = np.concatenate([qj_m[3:6], vj_m[3:6]])
return {"qj_m": qj_m, "vj_m": vj_m, "x_m": x_m}
def compute_torque(self, m):
"""
Compute the feedforward torque using ricatti gains
"""
# x_diff = np.concatenate(
# [
# pin.difference(
# self.pd.model,
# m["x_m"][: self.pd.nq],
# self.mpc_result.xs[0][: self.pd.nq],
# ),
# m["x_m"][self.pd.nq :] - self.mpc_result.xs[0][self.pd.nq :],
# ]
# )
x_diff = m["x_m"] - self.mpc_result.xs[0]
tau = self.mpc_result.us[0] + np.dot(self.mpc_result.K[0], x_diff)
return tau
def integrate_x(self, m):
"""
Integrate the position and velocity using the acceleration computed from the
feedforward torque
"""
q0 = m["qj_m"][3:6].copy()
v0 = m["vj_m"][3:6].copy()
tau = self.result.tau_ff[3:6].copy()
a = pin.aba(self.pd.model, self.pd.rdata, q0, v0, tau)
v = v0 + a * self.params.dt_wbc
q = q0 + v * self.params.dt_wbc
return q, v