standard mpc example is added.
This commit is contained in:
Rooholla-KhorramBakht
parent
d1df87fe26
commit
df649c040a
|
|
@ -261,7 +261,7 @@
|
|||
<inertial pos="-0.0068528 -3.9973e-06 0.039705" quat="0.494949 0.495288 -0.504676 0.504992" mass="0.077892" diaginertia="4.8843e-05 3.8232e-05 1.7707e-05"/>
|
||||
<joint name="Joint6" pos="0 0 0" axis="0 0 -1" range="-2.35 2.35"/>
|
||||
<geom type="mesh" rgba="1 1 1 1" mesh="Link6"/>
|
||||
<body name="EF_dummy" pos="0.0 0 0.0405">
|
||||
<body name="EF_dummy" pos="0.0 0 0.0406">
|
||||
<!-- <inertial pos="0 0 0" mass="0.0001" diaginertia="0.0 0.0 0.0"/> -->
|
||||
<inertial pos="0 0 0" mass="0.04" diaginertia="9.6e-06 9.6e-06 9.6e-06"/>
|
||||
<geom size="0.015" rgba="0 1 0 1"/>
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
|
|
@ -0,0 +1,358 @@
|
|||
from Go2Py.sim.mujoco import Go2Sim
|
||||
import numpy as np
|
||||
import os
|
||||
import sys
|
||||
import mim_solvers
|
||||
import friction_utils
|
||||
import mujoco
|
||||
import pinocchio as pin
|
||||
import crocoddyl
|
||||
import pinocchio
|
||||
import mujoco.viewer
|
||||
import time
|
||||
import mujoco as mj
|
||||
|
||||
|
||||
class Go2MPC:
|
||||
def __init__(self, assets_path, HORIZON=250, friction_mu = 0.2, dt = 0.02):
|
||||
self.assets_path = assets_path
|
||||
self.HORIZON = HORIZON
|
||||
self.max_iterations = 500
|
||||
self.dt = dt
|
||||
self.urdf_path = os.path.join(assets_path, 'urdf/go2_with_arm.urdf')
|
||||
self.xml_path = os.path.join(assets_path, 'mujoco/go2_with_arm.xml')
|
||||
self.pin_robot = pin.RobotWrapper.BuildFromURDF(self.urdf_path, self.assets_path, pin.JointModelFreeFlyer())
|
||||
self.pinRef = pin.LOCAL_WORLD_ALIGNED
|
||||
self.friction_mu = friction_mu
|
||||
self.rmodel = self.pin_robot.model
|
||||
self.rdata = self.pin_robot.data
|
||||
|
||||
self.mpc_to_unitree_name_map = \
|
||||
{'FL_HAA_joint':'FL_hip_joint',
|
||||
'FL_HFE_joint':'FL_thigh_joint',
|
||||
'FL_KFE_joint':'FL_calf_joint',
|
||||
'FR_HAA_joint':'FR_hip_joint',
|
||||
'FR_HFE_joint':'FR_thigh_joint',
|
||||
'FR_KFE_joint':'FR_calf_joint',
|
||||
'HL_HAA_joint':'RL_hip_joint',
|
||||
'HL_HFE_joint':'RL_thigh_joint',
|
||||
'HL_KFE_joint':'RL_calf_joint',
|
||||
'HR_HAA_joint':'RR_hip_joint',
|
||||
'HR_HFE_joint': 'RR_thigh_joint',
|
||||
'HR_KFE_joint': 'RR_calf_joint',
|
||||
'Joint1':'Joint1',
|
||||
'Joint2':'Joint2',
|
||||
'Joint3':'Joint3',
|
||||
'Joint4':'Joint4',
|
||||
'Joint5':'Joint5',
|
||||
'Joint6':'Joint6'
|
||||
}
|
||||
self.unitree_to_mpc_name_map = {val:key for key, val in self.mpc_to_unitree_name_map.items()}
|
||||
self.unitree_joint_order = ['FR', 'FL', 'RR', 'RL']
|
||||
|
||||
mpc_to_unitree_idx_map = {}
|
||||
unitree_id = 0
|
||||
for foot_name in self.unitree_joint_order:
|
||||
for actuator in ['_hip_joint', '_thigh_joint', '_calf_joint']:
|
||||
unitree_actuator_name = foot_name+actuator
|
||||
mpc_joint_name = self.unitree_to_mpc_name_map[unitree_actuator_name]
|
||||
mpc_joint_id = self.rmodel.getJointId(mpc_joint_name)
|
||||
mpc_to_unitree_idx_map[mpc_joint_id-2]=unitree_id
|
||||
unitree_id+=1
|
||||
|
||||
for unitree_actuator_name in ['Joint1', 'Joint2', 'Joint3', 'Joint4', 'Joint5', 'Joint6']:
|
||||
mpc_joint_name = self.unitree_to_mpc_name_map[unitree_actuator_name]
|
||||
mpc_joint_id = self.rmodel.getJointId(mpc_joint_name)
|
||||
mpc_to_unitree_idx_map[mpc_joint_id-2]=unitree_id
|
||||
unitree_id+=1
|
||||
|
||||
self.mpc_to_unitree_idx = np.zeros(18).astype(np.int32) # mpc_state[mpc_to_unitree_idx] -> state/command in unitree order
|
||||
self.unitree_to_mpc_idx = np.zeros(18).astype(np.int32) # unitree_state[unitree_to_mpc] -> state/command in mpc order
|
||||
for mpc_idx, unitree_idx in mpc_to_unitree_idx_map.items():
|
||||
self.mpc_to_unitree_idx[mpc_idx] = unitree_idx
|
||||
self.unitree_to_mpc_idx[unitree_idx] = mpc_idx
|
||||
|
||||
# set contact frame_names and_indices
|
||||
ee_frame_names = ['FL_FOOT', 'FR_FOOT', 'HL_FOOT', 'HR_FOOT', 'Link6']
|
||||
self.lfFootId = self.rmodel.getFrameId(ee_frame_names[0])
|
||||
self.rfFootId = self.rmodel.getFrameId(ee_frame_names[1])
|
||||
self.lhFootId = self.rmodel.getFrameId(ee_frame_names[2])
|
||||
self.rhFootId = self.rmodel.getFrameId(ee_frame_names[3])
|
||||
self.armEEId = self.rmodel.getFrameId(ee_frame_names[4])
|
||||
self.running_models = []
|
||||
self.constraintModels = []
|
||||
|
||||
self.ccdyl_state = crocoddyl.StateMultibody(self.rmodel)
|
||||
self.ccdyl_actuation = crocoddyl.ActuationModelFloatingBase(self.ccdyl_state)
|
||||
self.nu = self.ccdyl_actuation.nu
|
||||
|
||||
def initialize(self, q0=np.array([0.0, 0.0, 0.33, 0.0, 0.0, 0.0, 1.0]
|
||||
+4*[0.0, 0.77832842, -1.56065452] + [0.0, 0.3, -0.3, 0.0, 0.0, 0.0]
|
||||
)):
|
||||
q0[11+2]=0.0
|
||||
self.q0 = q0.copy()
|
||||
self.x0 = np.concatenate([q0, np.zeros(self.rmodel.nv)])
|
||||
pinocchio.forwardKinematics(self.rmodel, self.rdata, q0)
|
||||
pinocchio.updateFramePlacements(self.rmodel, self.rdata)
|
||||
self.rfFootPos0 = self.rdata.oMf[self.rfFootId].translation
|
||||
self.rhFootPos0 = self.rdata.oMf[self.rhFootId].translation
|
||||
self.lfFootPos0 = self.rdata.oMf[self.lfFootId].translation
|
||||
self.lhFootPos0 = self.rdata.oMf[self.lhFootId].translation
|
||||
self.armEEPos0 = self.rdata.oMf[self.armEEId].translation
|
||||
self.armEEOri0 = self.rdata.oMf[self.armEEId].rotation
|
||||
self.supportFeetIds = [self.lfFootId, self.rfFootId, self.lhFootId, self.rhFootId]
|
||||
# self.supportFeePos = [self.lfFootPos0, self.rfFootPos0, self.lhFootPos0, self.rhFootPos0]
|
||||
self.xs = [self.x0]*(self.HORIZON + 1)
|
||||
self.createProblem()
|
||||
self.createSolver()
|
||||
self.us = self.solver.problem.quasiStatic([self.x0]*self.HORIZON)
|
||||
|
||||
def createProblem(self):
|
||||
#First compute the desired state of the robot
|
||||
comRef = (self.rfFootPos0 + self.rhFootPos0 + self.lfFootPos0 + self.lhFootPos0) / 4
|
||||
comRef[2] = pinocchio.centerOfMass(self.rmodel, self.rdata, self.q0)[2].item()
|
||||
eeDes = self.armEEPos0
|
||||
comDes = []
|
||||
dt = self.dt
|
||||
# radius = 0.065
|
||||
radius = 0.0
|
||||
for t in range(self.HORIZON+1):
|
||||
comDes_t = comRef.copy()
|
||||
w = (2 * np.pi) * 0.2 # / T
|
||||
comDes_t[0] += radius * np.sin(w * t * dt)
|
||||
comDes_t[1] += radius * (np.cos(w * t * dt) - 1)
|
||||
comDes += [comDes_t]
|
||||
body_com_ref = comDes
|
||||
arm_eff_ref = [eeDes]*(self.HORIZON+1)
|
||||
# Now define the model
|
||||
for t in range(self.HORIZON+1):
|
||||
self.contactModel = crocoddyl.ContactModelMultiple(self.ccdyl_state, self.nu)
|
||||
costModel = crocoddyl.CostModelSum(self.ccdyl_state, self.nu)
|
||||
|
||||
# Add contacts
|
||||
for frame_idx in self.supportFeetIds:
|
||||
support_contact = crocoddyl.ContactModel3D(self.ccdyl_state, frame_idx, np.array([0., 0., 0.]), self.pinRef, self.nu, np.array([0., 0.]))
|
||||
self.contactModel.addContact(self.rmodel.frames[frame_idx].name + "_contact", support_contact)
|
||||
|
||||
# Contact for the EE
|
||||
support_contact = crocoddyl.ContactModel3D(self.ccdyl_state, self.armEEId, np.array([0., 0., 0.0]), pin.LOCAL_WORLD_ALIGNED, self.nu, np.array([0., 0.]))
|
||||
self.contactModel.addContact(self.rmodel.frames[self.armEEId].name + "_contact", support_contact)
|
||||
# Add state/control regularization costs
|
||||
state_reg_weight, control_reg_weight = 1e-1, 1e-3
|
||||
|
||||
freeFlyerQWeight = [0.]*3 + [500.]*3
|
||||
freeFlyerVWeight = [10.]*6
|
||||
legsQWeight = [0.01]*(self.rmodel.nv - 6)
|
||||
legsWWeights = [1.]*(self.rmodel.nv - 6)
|
||||
stateWeights = np.array(freeFlyerQWeight + legsQWeight + freeFlyerVWeight + legsWWeights)
|
||||
stateResidual = crocoddyl.ResidualModelState(self.ccdyl_state, self.x0, self.nu)
|
||||
|
||||
stateActivation = crocoddyl.ActivationModelWeightedQuad(stateWeights**2)
|
||||
stateReg = crocoddyl.CostModelResidual(self.ccdyl_state, stateActivation, stateResidual)
|
||||
|
||||
if t == self.HORIZON:
|
||||
costModel.addCost("stateReg", stateReg, state_reg_weight*self.dt)
|
||||
else:
|
||||
costModel.addCost("stateReg", stateReg, state_reg_weight)
|
||||
|
||||
if t != self.HORIZON:
|
||||
ctrlResidual = crocoddyl.ResidualModelControl(self.ccdyl_state, self.nu)
|
||||
ctrlReg = crocoddyl.CostModelResidual(self.ccdyl_state, ctrlResidual)
|
||||
costModel.addCost("ctrlReg", ctrlReg, control_reg_weight)
|
||||
|
||||
|
||||
# Body COM Tracking Cost
|
||||
com_residual = crocoddyl.ResidualModelCoMPosition(self.ccdyl_state, body_com_ref[t], self.nu)
|
||||
com_activation = crocoddyl.ActivationModelWeightedQuad(np.array([1., 1., 1.]))
|
||||
com_track = crocoddyl.CostModelResidual(self.ccdyl_state, com_activation, com_residual) # What does it correspond to exactly?
|
||||
# if t == self.HORIZON:
|
||||
# # costModel.addCost("comTrack", com_track, 1e5)
|
||||
# costModel.addCost("comTrack", com_track, 1e5)
|
||||
# else:
|
||||
# costModel.addCost("comTrack", com_track, 1e1)
|
||||
|
||||
# End Effecor Position Tracking Cost
|
||||
# ef_residual = crocoddyl.ResidualModelFrameTranslation(self.ccdyl_state, self.armEEId, arm_eff_ref[t], self.nu) # Check this cost term
|
||||
# ef_activation = crocoddyl.ActivationModelWeightedQuad(np.array([1., 1., 1.]))
|
||||
# ef_track = crocoddyl.CostModelResidual(self.ccdyl_state, ef_activation, ef_residual)
|
||||
|
||||
ef_rotation_residual = crocoddyl.ResidualModelFrameRotation(self.ccdyl_state, self.armEEId, self.armEEOri0, self.nu)
|
||||
ef_rot_activation = crocoddyl.ActivationModelWeightedQuad(np.array([1., 1., 1.]))
|
||||
ef_rot_track = crocoddyl.CostModelResidual(self.ccdyl_state, ef_rot_activation, ef_rotation_residual)
|
||||
if t == self.HORIZON:
|
||||
# costModel.addCost("efTrack", ef_track, 1e5)
|
||||
costModel.addCost("efRotTrack", ef_rot_track, 1e5)
|
||||
|
||||
else:
|
||||
# costModel.addCost("efTrack", ef_track, 1e1)
|
||||
costModel.addCost("efRotTrack", ef_rot_track, 1e1)
|
||||
# Force tracking term
|
||||
self.ef_des_force = pin.Force()
|
||||
self.ef_des_force.linear[2] = 0
|
||||
contact_force_residual = crocoddyl.ResidualModelContactForce(self.ccdyl_state, self.armEEId, self.ef_des_force, 3, self.nu)
|
||||
contact_force_activation = crocoddyl.ActivationModelWeightedQuad(np.array([1., 1., 1.]))
|
||||
contact_force_track = crocoddyl.CostModelResidual(self.ccdyl_state, contact_force_activation, contact_force_residual)
|
||||
if t == self.HORIZON:
|
||||
costModel.addCost("contact_force_track", contact_force_track, 1e5)
|
||||
else:
|
||||
costModel.addCost("contact_force_track", contact_force_track, 1e1)
|
||||
|
||||
|
||||
# Friction Cone Constraints
|
||||
constraintModelManager = crocoddyl.ConstraintModelManager(self.ccdyl_state, self.ccdyl_actuation.nu)
|
||||
if(t != self.HORIZON):
|
||||
for frame_idx in self.supportFeetIds:
|
||||
name = self.rmodel.frames[frame_idx].name + "_contact"
|
||||
residualFriction = friction_utils.ResidualFrictionCone(self.ccdyl_state, name, self.friction_mu, self.ccdyl_actuation.nu)
|
||||
constraintFriction = crocoddyl.ConstraintModelResidual(self.ccdyl_state, residualFriction, np.array([0.]), np.array([np.inf]))
|
||||
constraintModelManager.addConstraint(name + "friction", constraintFriction)
|
||||
#friction model for the arm
|
||||
# name = self.rmodel.frames[self.armEEId].name + "_contact"
|
||||
# residualFriction = friction_utils.ResidualFrictionCone(self.ccdyl_state, name, self.friction_mu, self.ccdyl_actuation.nu)
|
||||
# constraintFriction = crocoddyl.ConstraintModelResidual(self.ccdyl_state, residualFriction, np.array([0.]), np.array([np.inf]))
|
||||
# constraintModelManager.addConstraint(name + "friction", constraintFriction)
|
||||
|
||||
dmodel = crocoddyl.DifferentialActionModelContactFwdDynamics(self.ccdyl_state, self.ccdyl_actuation, self.contactModel, costModel, constraintModelManager, 0., True)
|
||||
model = crocoddyl.IntegratedActionModelEuler(dmodel, self.dt)
|
||||
self.running_models += [model]
|
||||
self.ocp = crocoddyl.ShootingProblem(self.x0, self.running_models[:-1], self.running_models[-1])
|
||||
|
||||
def createSolver(self):
|
||||
solver = mim_solvers.SolverCSQP(self.ocp)
|
||||
solver.max_qp_iters = 25
|
||||
solver.with_callbacks = True
|
||||
solver.use_filter_line_search = False
|
||||
solver.termination_tolerance = 1e-2
|
||||
solver.eps_abs = 1e-2
|
||||
solver.eps_rel = 0.
|
||||
self.solver = solver
|
||||
|
||||
def getSolution(self, k=None):
|
||||
if k is None:
|
||||
x_idx = 1
|
||||
u_idx = 0
|
||||
else:
|
||||
x_idx = k
|
||||
u_idx = k
|
||||
t = self.xs[x_idx][:3]
|
||||
quat = self.xs[x_idx][3:7]
|
||||
qx = quat[0]
|
||||
qy = quat[1]
|
||||
qz = quat[2]
|
||||
qw = quat[3]
|
||||
q = self.xs[x_idx][7:25]
|
||||
eta = self.xs[x_idx][25:25+6]
|
||||
dq = self.xs[x_idx][25+6:]
|
||||
constraint_norm = self.solver.constraint_norm
|
||||
return dict(
|
||||
position=t,
|
||||
orientation=np.array([qw, qx, qy, qz]), #Mujoco and uniree quaternion order
|
||||
velocity = eta[:3],
|
||||
omega = eta[3:],
|
||||
q = q[self.mpc_to_unitree_idx],
|
||||
dq = dq[self.mpc_to_unitree_idx],
|
||||
tau = self.us[u_idx][[self.mpc_to_unitree_idx]],
|
||||
constraint_norm = constraint_norm
|
||||
)
|
||||
|
||||
def updateAndSolve(self, t, quat, q, v, omega, dq):
|
||||
q_ = np.zeros(self.rmodel.nq)
|
||||
dq_ = np.zeros(self.rmodel.nv)
|
||||
qw = quat[0]
|
||||
qx = quat[1]
|
||||
qy = quat[2]
|
||||
qz = quat[3]
|
||||
q_[:3] = t
|
||||
q_[3:7] = np.array([qx, qy, qz, qw])
|
||||
q_[7:] = q[self.unitree_to_mpc_idx]
|
||||
dq_[:3] = v
|
||||
dq_[3:6] = omega
|
||||
dq_[6:] = dq[self.unitree_to_mpc_idx]
|
||||
pin.framesForwardKinematics(self.rmodel, self.rdata, q_)
|
||||
x = np.hstack([q_, dq_])
|
||||
self.solver.problem.x0 = x
|
||||
self.xs = list(self.solver.xs[1:]) + [self.solver.xs[-1]]
|
||||
self.xs[0] = x
|
||||
self.us = list(self.us[1:]) + [self.us[-1]]
|
||||
self.solver.solve(self.xs, self.us, self.max_iterations)
|
||||
self.xs, self.us = self.solver.xs, self.solver.us
|
||||
return self.getSolution()
|
||||
|
||||
def solve(self):
|
||||
self.solver.solve(self.xs, self.us, self.max_iterations)
|
||||
self.xs, self.us = self.solver.xs, self.solver.us
|
||||
return self.getSolution()
|
||||
|
||||
|
||||
def getForceSensor(model, data, site_id):
|
||||
site_id = mj.mj_name2id(model,mj.mjtObj.mjOBJ_SITE, 'EF_force_site')
|
||||
world_R_sensor = data.xmat[site_id].reshape(3,3).T
|
||||
force_in_body = data.sensordata[-3:].reshape(3,1)
|
||||
force_in_world = world_R_sensor@force_in_body
|
||||
return force_in_world
|
||||
|
||||
|
||||
# Instantiate the simulator
|
||||
robot=Go2Sim(with_arm=True)
|
||||
map = np.zeros((1200, 1200))
|
||||
map[:,649:679] = 400
|
||||
robot.updateHeightMap(map)
|
||||
|
||||
# Instantiate the solver
|
||||
assets_path = '/home/Go2py/Go2Py/assets'
|
||||
mpc = Go2MPC(assets_path, HORIZON=20, friction_mu=0.1)
|
||||
mpc.initialize()
|
||||
mpc.solve()
|
||||
m = list(mpc.solver.problem.runningModels) + [mpc.solver.problem.terminalModel]
|
||||
|
||||
# Reset the robot
|
||||
state = mpc.getSolution(0)
|
||||
robot.pos0 = state['position']
|
||||
robot.rot0 = state['orientation']
|
||||
robot.q0 = state['q']
|
||||
robot.reset()
|
||||
|
||||
# Solve for as many iterations as needed for the first step
|
||||
mpc.max_iterations=500
|
||||
|
||||
measured_forces = []
|
||||
forces = np.linspace(0, 25, 500)
|
||||
breakpoint()
|
||||
for i in range(500):
|
||||
# set the force setpoint
|
||||
for action_model in m:
|
||||
action_model.differential.costs.costs['contact_force_track'].cost.residual.reference.linear[:] = np.array([-forces[i], 0, 0])
|
||||
|
||||
state = robot.getJointStates()
|
||||
q = state['q']
|
||||
dq = state['dq']
|
||||
t, quat = robot.getPose()
|
||||
v = robot.data.qvel[:3]
|
||||
omega = robot.data.qvel[3:6]
|
||||
q = np.hstack([q, np.zeros(2)])
|
||||
dq = np.hstack([dq, np.zeros(2)])
|
||||
solution = mpc.updateAndSolve(t, quat, q, v, omega, dq)
|
||||
# Reduce the max iteration count to ensure real-time execution
|
||||
mpc.max_iterations=100
|
||||
q = solution['q']
|
||||
dq = solution['dq']
|
||||
tau = solution['tau'].squeeze()
|
||||
kp = np.ones(18)*0.
|
||||
kv = np.ones(18)*0.
|
||||
# Step the physics
|
||||
|
||||
for j in range(int(mpc.dt//robot.dt)):
|
||||
robot.setCommands(q, dq, kp, kv, tau)
|
||||
robot.step()
|
||||
|
||||
force_site_to_sensor_idx = {'FL_force_site': 0, 'FR_force_site': 1, 'RL_force_site': 2, 'RR_force_site': 3, 'EF_force_site': 4}
|
||||
force_sensor_site = 'EF_force_site'
|
||||
measured_forces.append(getForceSensor(robot.model, robot.data, force_sensor_site).squeeze().copy())
|
||||
measured_forces = np.array(measured_forces)
|
||||
|
||||
# Visualize the measured force against the desired
|
||||
import matplotlib.pyplot as plt
|
||||
plt.plot(measured_forces[:300,0], '*')
|
||||
plt.plot(forces,'k')
|
||||
plt.show()
|
||||
robot.close()
|
||||
Loading…
Reference in New Issue