849 lines
32 KiB
Plaintext
849 lines
32 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## OfflineTrajectory Optimization"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 1,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import pinocchio as pin\n",
|
||
"import crocoddyl\n",
|
||
"import pinocchio\n",
|
||
"import numpy as np\n",
|
||
"urdf_root_path = '/home/Go2py/Go2Py/assets'\n",
|
||
"urdf_path = '/home/Go2py/Go2Py/assets/urdf/go2_with_arm.urdf'\n",
|
||
"robot = pin.RobotWrapper.BuildFromURDF(\n",
|
||
"urdf_path, urdf_root_path, pin.JointModelFreeFlyer())"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 2,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import sys\n",
|
||
"import mim_solvers\n",
|
||
"pinRef = pin.LOCAL_WORLD_ALIGNED\n",
|
||
"FRICTION_CSTR = True\n",
|
||
"MU = 0.8 # friction coefficient"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 3,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"ee_frame_names = ['FL_FOOT', 'FR_FOOT', 'HL_FOOT', 'HR_FOOT', 'Link6']\n",
|
||
"rmodel = robot.model\n",
|
||
"rdata = robot.data\n",
|
||
"# # set contact frame_names and_indices\n",
|
||
"lfFootId = rmodel.getFrameId(ee_frame_names[0])\n",
|
||
"rfFootId = rmodel.getFrameId(ee_frame_names[1])\n",
|
||
"lhFootId = rmodel.getFrameId(ee_frame_names[2])\n",
|
||
"rhFootId = rmodel.getFrameId(ee_frame_names[3])\n",
|
||
"efId = rmodel.getFrameId(ee_frame_names[4])"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 4,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"q0 = np.array([0.0, 0.0, 0.3, 0.0, 0.0, 0.0, 1.0] \n",
|
||
" +4*[0.0, 0.77832842, -1.56065452] + 8*[0.0]\n",
|
||
" )\n",
|
||
"x0 = np.concatenate([q0, np.zeros(rmodel.nv)])"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 5,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"q0 = np.array([0.0, 0.0, 0.3, 0.0, 0.0, 0.0, 1.0] \n",
|
||
" +4*[0.0, 0.77832842, -1.56065452] + 8*[0.0]\n",
|
||
" )\n",
|
||
"x0 = np.concatenate([q0, np.zeros(rmodel.nv)])\n",
|
||
"\n",
|
||
"pinocchio.forwardKinematics(rmodel, rdata, q0)\n",
|
||
"pinocchio.updateFramePlacements(rmodel, rdata)\n",
|
||
"rfFootPos0 = rdata.oMf[rfFootId].translation\n",
|
||
"rhFootPos0 = rdata.oMf[rhFootId].translation\n",
|
||
"lfFootPos0 = rdata.oMf[lfFootId].translation\n",
|
||
"lhFootPos0 = rdata.oMf[lhFootId].translation \n",
|
||
"efPos0 = rdata.oMf[efId].translation"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"comRef = (rfFootPos0 + rhFootPos0 + lfFootPos0 + lhFootPos0) / 4\n",
|
||
"comRef[2] = pinocchio.centerOfMass(rmodel, rdata, q0)[2].item() \n",
|
||
"print(f'The desired CoM position is: {comRef}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 7,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"supportFeetIds = [lfFootId, rfFootId, lhFootId, rhFootId]\n",
|
||
"supportFeePos = [lfFootPos0, rfFootPos0, lhFootPos0, rhFootPos0]"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 8,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"state = crocoddyl.StateMultibody(rmodel)\n",
|
||
"actuation = crocoddyl.ActuationModelFloatingBase(state)\n",
|
||
"nu = actuation.nu"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 9,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"comDes = []\n",
|
||
"\n",
|
||
"N_ocp = 250 #100\n",
|
||
"dt = 0.02\n",
|
||
"T = N_ocp * dt\n",
|
||
"radius = 0.065\n",
|
||
"for t in range(N_ocp+1):\n",
|
||
" comDes_t = comRef.copy()\n",
|
||
" w = (2 * np.pi) * 0.2 # / T\n",
|
||
" comDes_t[0] += radius * np.sin(w * t * dt) \n",
|
||
" comDes_t[1] += radius * (np.cos(w * t * dt) - 1)\n",
|
||
" comDes += [comDes_t]"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 10,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import friction_utils\n",
|
||
"running_models = []\n",
|
||
"constraintModels = []\n",
|
||
"\n",
|
||
"for t in range(N_ocp+1):\n",
|
||
" contactModel = crocoddyl.ContactModelMultiple(state, nu)\n",
|
||
" costModel = crocoddyl.CostModelSum(state, nu)\n",
|
||
"\n",
|
||
" # Add contact\n",
|
||
" for frame_idx in supportFeetIds:\n",
|
||
" support_contact = crocoddyl.ContactModel3D(state, frame_idx, np.array([0., 0., 0.]), pinRef, nu, np.array([0., 0.]))\n",
|
||
" # print(\"contact name = \", rmodel.frames[frame_idx].name + \"_contact\")\n",
|
||
" contactModel.addContact(rmodel.frames[frame_idx].name + \"_contact\", support_contact) \n",
|
||
"\n",
|
||
" # Add state/control reg costs\n",
|
||
"\n",
|
||
" state_reg_weight, control_reg_weight = 1e-1, 1e-3\n",
|
||
"\n",
|
||
" freeFlyerQWeight = [0.]*3 + [500.]*3\n",
|
||
" freeFlyerVWeight = [10.]*6\n",
|
||
" legsQWeight = [0.01]*(rmodel.nv - 6)\n",
|
||
" legsWWeights = [1.]*(rmodel.nv - 6)\n",
|
||
" stateWeights = np.array(freeFlyerQWeight + legsQWeight + freeFlyerVWeight + legsWWeights) \n",
|
||
"\n",
|
||
"\n",
|
||
" stateResidual = crocoddyl.ResidualModelState(state, x0, nu)\n",
|
||
" stateActivation = crocoddyl.ActivationModelWeightedQuad(stateWeights**2)\n",
|
||
" stateReg = crocoddyl.CostModelResidual(state, stateActivation, stateResidual)\n",
|
||
"\n",
|
||
" if t == N_ocp:\n",
|
||
" costModel.addCost(\"stateReg\", stateReg, state_reg_weight*dt)\n",
|
||
" else:\n",
|
||
" costModel.addCost(\"stateReg\", stateReg, state_reg_weight)\n",
|
||
"\n",
|
||
" if t != N_ocp:\n",
|
||
" ctrlResidual = crocoddyl.ResidualModelControl(state, nu)\n",
|
||
" ctrlReg = crocoddyl.CostModelResidual(state, ctrlResidual)\n",
|
||
" costModel.addCost(\"ctrlReg\", ctrlReg, control_reg_weight) \n",
|
||
"\n",
|
||
"\n",
|
||
" # Add COM task\n",
|
||
" com_residual = crocoddyl.ResidualModelCoMPosition(state, comDes[t], nu)\n",
|
||
" com_activation = crocoddyl.ActivationModelWeightedQuad(np.array([1., 1., 1.]))\n",
|
||
" com_track = crocoddyl.CostModelResidual(state, com_activation, com_residual) # What does it correspond to exactly?\n",
|
||
" if t == N_ocp:\n",
|
||
" costModel.addCost(\"comTrack\", com_track, 1e5)\n",
|
||
" else:\n",
|
||
" costModel.addCost(\"comTrack\", com_track, 1e5)\n",
|
||
"\n",
|
||
" # End Effecor Position Task\n",
|
||
" ef_residual = crocoddyl.ResidualModelFrameTranslation(state, efId, efPos0, nu)\n",
|
||
" ef_activation = crocoddyl.ActivationModelWeightedQuad(np.array([1., 1., 1.]))\n",
|
||
" ef_track = crocoddyl.CostModelResidual(state, ef_activation, ef_residual)\n",
|
||
" if t == N_ocp:\n",
|
||
" costModel.addCost(\"efTrack\", ef_track, 1e5)\n",
|
||
" else:\n",
|
||
" costModel.addCost(\"efTrack\", ef_track, 1e5)\n",
|
||
" \n",
|
||
"\n",
|
||
" constraintModelManager = crocoddyl.ConstraintModelManager(state, actuation.nu)\n",
|
||
" if(FRICTION_CSTR):\n",
|
||
" if(t != N_ocp):\n",
|
||
" for frame_idx in supportFeetIds:\n",
|
||
" name = rmodel.frames[frame_idx].name + \"_contact\"\n",
|
||
" residualFriction = friction_utils.ResidualFrictionCone(state, name, MU, actuation.nu)\n",
|
||
" constraintFriction = crocoddyl.ConstraintModelResidual(state, residualFriction, np.array([0.]), np.array([np.inf]))\n",
|
||
" constraintModelManager.addConstraint(name + \"friction\", constraintFriction)\n",
|
||
"\n",
|
||
" dmodel = crocoddyl.DifferentialActionModelContactFwdDynamics(state, actuation, contactModel, costModel, constraintModelManager, 0., True)\n",
|
||
" model = crocoddyl.IntegratedActionModelEuler(dmodel, dt)\n",
|
||
"\n",
|
||
" running_models += [model]\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 11,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# Create shooting problem\n",
|
||
"ocp = crocoddyl.ShootingProblem(x0, running_models[:-1], running_models[-1])\n",
|
||
"\n",
|
||
"solver = mim_solvers.SolverCSQP(ocp)\n",
|
||
"solver.max_qp_iters = 1000\n",
|
||
"max_iter = 500\n",
|
||
"solver.with_callbacks = True\n",
|
||
"solver.use_filter_line_search = False\n",
|
||
"solver.termination_tolerance = 1e-4\n",
|
||
"solver.eps_abs = 1e-6\n",
|
||
"solver.eps_rel = 1e-6"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"xs = [x0]*(solver.problem.T + 1)\n",
|
||
"us = solver.problem.quasiStatic([x0]*solver.problem.T) \n",
|
||
"solver.solve(xs, us, max_iter) "
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 77,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"nq, nv, N = rmodel.nq, rmodel.nv, len(xs) \n",
|
||
"jointPos_sol = []\n",
|
||
"jointVel_sol = []\n",
|
||
"jointAcc_sol = []\n",
|
||
"jointTorques_sol = []\n",
|
||
"centroidal_sol = []\n",
|
||
"xs, us = solver.xs, solver.us\n",
|
||
"x = []\n",
|
||
"for time_idx in range (N):\n",
|
||
" q, v = xs[time_idx][:nq], xs[time_idx][nq:]\n",
|
||
" pin.framesForwardKinematics(rmodel, rdata, q)\n",
|
||
" pin.computeCentroidalMomentum(rmodel, rdata, q, v)\n",
|
||
" centroidal_sol += [\n",
|
||
" np.concatenate(\n",
|
||
" [pin.centerOfMass(rmodel, rdata, q, v), np.array(rdata.hg.linear), np.array(rdata.hg.angular)]\n",
|
||
" )\n",
|
||
" ]\n",
|
||
" jointPos_sol += [q]\n",
|
||
" jointVel_sol += [v]\n",
|
||
" x += [xs[time_idx]]\n",
|
||
" if time_idx < N-1:\n",
|
||
" jointAcc_sol += [solver.problem.runningDatas[time_idx].xnext[nq::]] \n",
|
||
" jointTorques_sol += [us[time_idx]]\n",
|
||
"\n",
|
||
"\n",
|
||
"\n",
|
||
"\n",
|
||
"sol = {'x':x, 'centroidal':centroidal_sol, 'jointPos':jointPos_sol, \n",
|
||
" 'jointVel':jointVel_sol, 'jointAcc':jointAcc_sol, \n",
|
||
" 'jointTorques':jointTorques_sol} \n",
|
||
"\n",
|
||
"for frame_idx in supportFeetIds:\n",
|
||
" # print('extract foot id ', frame_idx, \"_name = \", rmodel.frames[frame_idx].name)\n",
|
||
" ct_frame_name = rmodel.frames[frame_idx].name + \"_contact\"\n",
|
||
" datas = [solver.problem.runningDatas[i].differential.multibody.contacts.contacts[ct_frame_name] for i in range(N-1)]\n",
|
||
" ee_forces = [datas[k].f.vector for k in range(N-1)] \n",
|
||
" sol[ct_frame_name] = [ee_forces[i] for i in range(N-1)] \n",
|
||
" "
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import matplotlib.pyplot as plt\n",
|
||
"constrained_sol = sol\n",
|
||
"time_lin = np.linspace(0, T, solver.problem.T)\n",
|
||
"fig, axs = plt.subplots(4, 3, constrained_layout=True)\n",
|
||
"for i, frame_idx in enumerate(supportFeetIds):\n",
|
||
" ct_frame_name = rmodel.frames[frame_idx].name + \"_contact\"\n",
|
||
" forces = np.array(constrained_sol[ct_frame_name])\n",
|
||
" axs[i, 0].plot(time_lin, forces[:, 0], label=\"Fx\")\n",
|
||
" axs[i, 1].plot(time_lin, forces[:, 1], label=\"Fy\")\n",
|
||
" axs[i, 2].plot(time_lin, forces[:, 2], label=\"Fz\")\n",
|
||
" # Add friction cone constraints \n",
|
||
" Fz_lb = (1./MU)*np.sqrt(forces[:, 0]**2 + forces[:, 1]**2)\n",
|
||
" # Fz_ub = np.zeros(time_lin.shape)\n",
|
||
" # axs[i, 2].plot(time_lin, Fz_ub, 'k-.', label='ub')\n",
|
||
" axs[i, 2].plot(time_lin, Fz_lb, 'k-.', label='lb')\n",
|
||
" axs[i, 0].grid()\n",
|
||
" axs[i, 1].grid()\n",
|
||
" axs[i, 2].grid()\n",
|
||
" axs[i, 0].set_ylabel(ct_frame_name)\n",
|
||
"axs[0, 0].legend()\n",
|
||
"axs[0, 1].legend()\n",
|
||
"axs[0, 2].legend()\n",
|
||
"\n",
|
||
"axs[3, 0].set_xlabel(r\"$F_x$\")\n",
|
||
"axs[3, 1].set_xlabel(r\"$F_y$\")\n",
|
||
"axs[3, 2].set_xlabel(r\"$F_z$\")\n",
|
||
"fig.suptitle('Force', fontsize=16)\n",
|
||
"\n",
|
||
"\n",
|
||
"comDes = np.array(comDes)\n",
|
||
"centroidal_sol = np.array(constrained_sol['centroidal'])\n",
|
||
"plt.figure()\n",
|
||
"plt.plot(comDes[:, 0], comDes[:, 1], \"--\", label=\"reference\")\n",
|
||
"plt.plot(centroidal_sol[:, 0], centroidal_sol[:, 1], label=\"solution\")\n",
|
||
"plt.legend()\n",
|
||
"plt.xlabel(\"x\")\n",
|
||
"plt.ylabel(\"y\")\n",
|
||
"plt.title(\"COM trajectory\")\n",
|
||
"plt.show()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"mujoco_joint_names = \\\n",
|
||
"['FL_hip_joint',\n",
|
||
" 'FL_thigh_joint',\n",
|
||
" 'FL_calf_joint',\n",
|
||
" 'FR_hip_joint',\n",
|
||
" 'FR_thigh_joint',\n",
|
||
" 'FR_calf_joint',\n",
|
||
" 'RL_hip_joint',\n",
|
||
" 'RL_thigh_joint',\n",
|
||
" 'RL_calf_joint',\n",
|
||
" 'RR_hip_joint',\n",
|
||
" 'RR_thigh_joint',\n",
|
||
" 'RR_calf_joint',\n",
|
||
" 'Joint1',\n",
|
||
" 'Joint2',\n",
|
||
" 'Joint3',\n",
|
||
" 'Joint4',\n",
|
||
" 'Joint5',\n",
|
||
" 'Joint6',\n",
|
||
" 'Joint7_1',\n",
|
||
" 'Joint7_2']\n",
|
||
"\n",
|
||
"for joint_name in mujoco_joint_names:\n",
|
||
" joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, joint_name)\n",
|
||
" print(f'Joint {joint_name} has id {joint_id}')\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"pinocchio_joint_names = \\\n",
|
||
"['FL_HAA_joint',\n",
|
||
" 'FL_HFE_joint',\n",
|
||
" 'FL_KFE_joint',\n",
|
||
" 'FR_HAA_joint',\n",
|
||
" 'FR_HFE_joint',\n",
|
||
" 'FR_KFE_joint',\n",
|
||
" 'HL_HAA_joint',\n",
|
||
" 'HL_HFE_joint',\n",
|
||
" 'HL_KFE_joint',\n",
|
||
" 'HR_HAA_joint',\n",
|
||
" 'HR_HFE_joint',\n",
|
||
" 'HR_KFE_joint',\n",
|
||
" 'Joint1',\n",
|
||
" 'Joint2',\n",
|
||
" 'Joint3',\n",
|
||
" 'Joint4',\n",
|
||
" 'Joint5',\n",
|
||
" 'Joint6',\n",
|
||
" 'Joint7_1',\n",
|
||
" 'Joint7_2']\n",
|
||
"\n",
|
||
"\n",
|
||
"for joint_name in pinocchio_joint_names:\n",
|
||
" joint_id = rmodel.getJointId(joint_name)\n",
|
||
" print(f'Joint {joint_name} has id {joint_id}')\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 116,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"pin_idx_to_mujoco_idx = []\n",
|
||
"for i in range(len(pinocchio_joint_names)):\n",
|
||
" joint_name = pinocchio_joint_names[i]\n",
|
||
" pin_joint_id = rmodel.getJointId(joint_name)\n",
|
||
" mujoco_joint_name = mujoco_joint_names[i]\n",
|
||
" mujoco_joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, mujoco_joint_name)\n",
|
||
" pin_idx_to_mujoco_idx.append(mujoco_joint_id)\n",
|
||
"\n",
|
||
"pin_idx_to_mujoco_idx = np.array(pin_idx_to_mujoco_idx) - min(pin_idx_to_mujoco_idx)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import mujoco\n",
|
||
"model = mujoco.MjModel.from_xml_path('/home/Go2py/Go2Py/assets/mujoco/go2_with_arm.xml')\n",
|
||
"data = mujoco.MjData(model)\n",
|
||
"viewer = mujoco.viewer.launch_passive(model, data)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 118,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import time\n",
|
||
"for k in range(1000):\n",
|
||
" i = k % len(xs)\n",
|
||
" q0 = xs[i][:27]\n",
|
||
" qx = q0[3]\n",
|
||
" qy = q0[4]\n",
|
||
" qz = q0[5]\n",
|
||
" qw = q0[6]\n",
|
||
" q0[3:7] = [qw, qx, qy, qz]\n",
|
||
" q0[7:] = q0[7:][pin_idx_to_mujoco_idx]\n",
|
||
" data.qpos[:] = q0\n",
|
||
"\n",
|
||
" mujoco.mj_step(model, data)\n",
|
||
" viewer.sync()\n",
|
||
" time.sleep(1/60)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Online MPC"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 1,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import os\n",
|
||
"import sys\n",
|
||
"import mim_solvers\n",
|
||
"import friction_utils\n",
|
||
"import mujoco\n",
|
||
"import pinocchio as pin\n",
|
||
"import crocoddyl\n",
|
||
"import pinocchio\n",
|
||
"import numpy as np\n",
|
||
"import mujoco.viewer\n",
|
||
"\n",
|
||
"class Go2MPC:\n",
|
||
" def __init__(self, assets_path, HORIZON=250, friction_mu = 0.8, dt = 0.02):\n",
|
||
" self.assets_path = assets_path\n",
|
||
" self.HORIZON = HORIZON\n",
|
||
" self.max_iterations = 500\n",
|
||
" self.dt = dt\n",
|
||
" self.urdf_path = os.path.join(assets_path, 'urdf/go2_with_arm.urdf')\n",
|
||
" self.xml_path = os.path.join(assets_path, 'mujoco/go2_with_arm.xml')\n",
|
||
" self.pin_robot = pin.RobotWrapper.BuildFromURDF(self.urdf_path, self.assets_path, pin.JointModelFreeFlyer())\n",
|
||
" self.pinRef = pin.LOCAL_WORLD_ALIGNED\n",
|
||
" self.friction_mu = friction_mu \n",
|
||
" self.rmodel = self.pin_robot.model\n",
|
||
" self.rdata = self.pin_robot.data\n",
|
||
"\n",
|
||
" self.mpc_to_unitree_name_map = \\\n",
|
||
" {'FL_HAA_joint':'FL_hip_joint',\n",
|
||
" 'FL_HFE_joint':'FL_thigh_joint',\n",
|
||
" 'FL_KFE_joint':'FL_calf_joint',\n",
|
||
" 'FR_HAA_joint':'FR_hip_joint',\n",
|
||
" 'FR_HFE_joint':'FR_thigh_joint',\n",
|
||
" 'FR_KFE_joint':'FR_calf_joint',\n",
|
||
" 'HL_HAA_joint':'RL_hip_joint',\n",
|
||
" 'HL_HFE_joint':'RL_thigh_joint',\n",
|
||
" 'HL_KFE_joint':'RL_calf_joint',\n",
|
||
" 'HR_HAA_joint':'RR_hip_joint',\n",
|
||
" 'HR_HFE_joint': 'RR_thigh_joint',\n",
|
||
" 'HR_KFE_joint': 'RR_calf_joint',\n",
|
||
" 'Joint1':'Joint1',\n",
|
||
" 'Joint2':'Joint2',\n",
|
||
" 'Joint3':'Joint3',\n",
|
||
" 'Joint4':'Joint4',\n",
|
||
" 'Joint5':'Joint5',\n",
|
||
" 'Joint6':'Joint6',\n",
|
||
" 'Joint7_1':'Joint7_1',\n",
|
||
" 'Joint7_2':'Joint7_2' }\n",
|
||
" self.unitree_to_mpc_name_map = {val:key for key, val in self.mpc_to_unitree_name_map.items()}\n",
|
||
" self.unitree_joint_order = ['FR', 'FL', 'RR', 'RL']\n",
|
||
"\n",
|
||
" mpc_to_unitree_idx_map = {}\n",
|
||
" unitree_id = 0\n",
|
||
" for foot_name in self.unitree_joint_order:\n",
|
||
" for actuator in ['_hip_joint', '_thigh_joint', '_calf_joint']:\n",
|
||
" unitree_actuator_name = foot_name+actuator\n",
|
||
" mpc_joint_name = self.unitree_to_mpc_name_map[unitree_actuator_name]\n",
|
||
" mpc_joint_id = self.rmodel.getJointId(mpc_joint_name)\n",
|
||
" mpc_to_unitree_idx_map[mpc_joint_id-2]=unitree_id\n",
|
||
" unitree_id+=1\n",
|
||
"\n",
|
||
" for unitree_actuator_name in ['Joint1', 'Joint2', 'Joint3', 'Joint4', 'Joint5', 'Joint6', 'Joint7_1', 'Joint7_2']:\n",
|
||
" mpc_joint_name = self.unitree_to_mpc_name_map[unitree_actuator_name]\n",
|
||
" mpc_joint_id = self.rmodel.getJointId(mpc_joint_name)\n",
|
||
" mpc_to_unitree_idx_map[mpc_joint_id-2]=unitree_id\n",
|
||
" unitree_id+=1\n",
|
||
"\n",
|
||
" self.mpc_to_unitree_idx = np.zeros(20).astype(np.int32) # mpc_state[mpc_to_unitree_idx] -> state/command in unitree order \n",
|
||
" self.unitree_to_mpc_idx = np.zeros(20).astype(np.int32) # unitree_state[unitree_to_mpc] -> state/command in mpc order\n",
|
||
" for mpc_idx, unitree_idx in mpc_to_unitree_idx_map.items():\n",
|
||
" self.mpc_to_unitree_idx[mpc_idx] = unitree_idx\n",
|
||
" self.unitree_to_mpc_idx[unitree_idx] = mpc_idx\n",
|
||
"\n",
|
||
" # set contact frame_names and_indices\n",
|
||
" ee_frame_names = ['FL_FOOT', 'FR_FOOT', 'HL_FOOT', 'HR_FOOT', 'Link6']\n",
|
||
" self.lfFootId = self.rmodel.getFrameId(ee_frame_names[0])\n",
|
||
" self.rfFootId = self.rmodel.getFrameId(ee_frame_names[1])\n",
|
||
" self.lhFootId = self.rmodel.getFrameId(ee_frame_names[2])\n",
|
||
" self.rhFootId = self.rmodel.getFrameId(ee_frame_names[3])\n",
|
||
" self.armEEId = self.rmodel.getFrameId(ee_frame_names[4])\n",
|
||
" self.running_models = []\n",
|
||
" self.constraintModels = []\n",
|
||
" \n",
|
||
" self.ccdyl_state = crocoddyl.StateMultibody(self.rmodel)\n",
|
||
" self.ccdyl_actuation = crocoddyl.ActuationModelFloatingBase(self.ccdyl_state)\n",
|
||
" self.nu = self.ccdyl_actuation.nu\n",
|
||
"\n",
|
||
" def initialize(self, q0=np.array([0.0, 0.0, 0.3, 0.0, 0.0, 0.0, 1.0] \n",
|
||
" +4*[0.0, 0.77832842, -1.56065452] + 8*[0.0]\n",
|
||
" )):\n",
|
||
" self.q0 = q0.copy()\n",
|
||
" self.x0 = np.concatenate([q0, np.zeros(self.rmodel.nv)])\n",
|
||
" pinocchio.forwardKinematics(self.rmodel, self.rdata, q0)\n",
|
||
" pinocchio.updateFramePlacements(self.rmodel, self.rdata)\n",
|
||
" self.rfFootPos0 = self.rdata.oMf[self.rfFootId].translation\n",
|
||
" self.rhFootPos0 = self.rdata.oMf[self.rhFootId].translation\n",
|
||
" self.lfFootPos0 = self.rdata.oMf[self.lfFootId].translation\n",
|
||
" self.lhFootPos0 = self.rdata.oMf[self.lhFootId].translation \n",
|
||
" self.armEEPos0 = self.rdata.oMf[self.armEEId].translation\n",
|
||
" self.supportFeetIds = [self.lfFootId, self.rfFootId, self.lhFootId, self.rhFootId]\n",
|
||
" self.supportFeePos = [self.lfFootPos0, self.rfFootPos0, self.lhFootPos0, self.rhFootPos0]\n",
|
||
" self.xs = [self.x0]*(self.HORIZON + 1)\n",
|
||
" self.createProblem()\n",
|
||
" self.createSolver()\n",
|
||
" self.us = self.solver.problem.quasiStatic([self.x0]*self.HORIZON) \n",
|
||
"\n",
|
||
" def createProblem(self):\n",
|
||
" #First compute the desired state of the robot\n",
|
||
" comRef = (self.rfFootPos0 + self.rhFootPos0 + self.lfFootPos0 + self.lhFootPos0) / 4\n",
|
||
" comRef[2] = pinocchio.centerOfMass(self.rmodel, self.rdata, self.q0)[2].item() \n",
|
||
" eeDes = self.armEEPos0\n",
|
||
" comDes = []\n",
|
||
" dt = self.dt\n",
|
||
" # radius = 0.065\n",
|
||
" radius = 0.0\n",
|
||
" for t in range(self.HORIZON+1):\n",
|
||
" comDes_t = comRef.copy()\n",
|
||
" w = (2 * np.pi) * 0.2 # / T\n",
|
||
" comDes_t[0] += radius * np.sin(w * t * dt) \n",
|
||
" comDes_t[1] += radius * (np.cos(w * t * dt) - 1)\n",
|
||
" comDes += [comDes_t]\n",
|
||
" body_com_ref = comDes\n",
|
||
" arm_eff_ref = [eeDes]*(self.HORIZON+1)\n",
|
||
" # Now define the model\n",
|
||
" for t in range(self.HORIZON+1):\n",
|
||
" self.contactModel = crocoddyl.ContactModelMultiple(self.ccdyl_state, self.nu)\n",
|
||
" costModel = crocoddyl.CostModelSum(self.ccdyl_state, self.nu)\n",
|
||
"\n",
|
||
" # Add contacts\n",
|
||
" for frame_idx in self.supportFeetIds:\n",
|
||
" support_contact = crocoddyl.ContactModel3D(self.ccdyl_state, frame_idx, np.array([0., 0., 0.]), self.pinRef, self.nu, np.array([0., 0.]))\n",
|
||
" self.contactModel.addContact(self.rmodel.frames[frame_idx].name + \"_contact\", support_contact) \n",
|
||
"\n",
|
||
" # Add state/control regularization costs\n",
|
||
" state_reg_weight, control_reg_weight = 1e-1, 1e-3\n",
|
||
"\n",
|
||
" freeFlyerQWeight = [0.]*3 + [500.]*3\n",
|
||
" freeFlyerVWeight = [10.]*6\n",
|
||
" legsQWeight = [0.01]*(self.rmodel.nv - 6)\n",
|
||
" legsWWeights = [1.]*(self.rmodel.nv - 6)\n",
|
||
" stateWeights = np.array(freeFlyerQWeight + legsQWeight + freeFlyerVWeight + legsWWeights) \n",
|
||
" stateResidual = crocoddyl.ResidualModelState(self.ccdyl_state, self.x0, self.nu)\n",
|
||
"\n",
|
||
" stateActivation = crocoddyl.ActivationModelWeightedQuad(stateWeights**2)\n",
|
||
" stateReg = crocoddyl.CostModelResidual(self.ccdyl_state, stateActivation, stateResidual)\n",
|
||
"\n",
|
||
" if t == self.HORIZON:\n",
|
||
" costModel.addCost(\"stateReg\", stateReg, state_reg_weight*self.dt)\n",
|
||
" else:\n",
|
||
" costModel.addCost(\"stateReg\", stateReg, state_reg_weight)\n",
|
||
"\n",
|
||
" if t != self.HORIZON:\n",
|
||
" ctrlResidual = crocoddyl.ResidualModelControl(self.ccdyl_state, self.nu)\n",
|
||
" ctrlReg = crocoddyl.CostModelResidual(self.ccdyl_state, ctrlResidual)\n",
|
||
" costModel.addCost(\"ctrlReg\", ctrlReg, control_reg_weight) \n",
|
||
"\n",
|
||
"\n",
|
||
" # Body COM Tracking Cost\n",
|
||
" com_residual = crocoddyl.ResidualModelCoMPosition(self.ccdyl_state, body_com_ref[t], self.nu)\n",
|
||
" com_activation = crocoddyl.ActivationModelWeightedQuad(np.array([1., 1., 1.]))\n",
|
||
" com_track = crocoddyl.CostModelResidual(self.ccdyl_state, com_activation, com_residual) # What does it correspond to exactly?\n",
|
||
" if t == self.HORIZON:\n",
|
||
" # costModel.addCost(\"comTrack\", com_track, 1e5)\n",
|
||
" costModel.addCost(\"comTrack\", com_track, 1e5)\n",
|
||
" else:\n",
|
||
" costModel.addCost(\"comTrack\", com_track, 1e2)\n",
|
||
"\n",
|
||
" # End Effecor Position Tracking Cost\n",
|
||
" ef_residual = crocoddyl.ResidualModelFrameTranslation(self.ccdyl_state, self.armEEId, arm_eff_ref[t], self.nu) # Check this cost term\n",
|
||
" ef_activation = crocoddyl.ActivationModelWeightedQuad(np.array([1., 1., 1.]))\n",
|
||
" ef_track = crocoddyl.CostModelResidual(self.ccdyl_state, ef_activation, ef_residual)\n",
|
||
" if t == self.HORIZON:\n",
|
||
" costModel.addCost(\"efTrack\", ef_track, 1e5)\n",
|
||
" else:\n",
|
||
" costModel.addCost(\"efTrack\", ef_track, 1e3)\n",
|
||
" \n",
|
||
" # Friction Cone Constraints\n",
|
||
" constraintModelManager = crocoddyl.ConstraintModelManager(self.ccdyl_state, self.ccdyl_actuation.nu)\n",
|
||
" if(t != self.HORIZON):\n",
|
||
" for frame_idx in self.supportFeetIds:\n",
|
||
" name = self.rmodel.frames[frame_idx].name + \"_contact\"\n",
|
||
" residualFriction = friction_utils.ResidualFrictionCone(self.ccdyl_state, name, self.friction_mu, self.ccdyl_actuation.nu)\n",
|
||
" constraintFriction = crocoddyl.ConstraintModelResidual(self.ccdyl_state, residualFriction, np.array([0.]), np.array([np.inf]))\n",
|
||
" constraintModelManager.addConstraint(name + \"friction\", constraintFriction)\n",
|
||
"\n",
|
||
" dmodel = crocoddyl.DifferentialActionModelContactFwdDynamics(self.ccdyl_state, self.ccdyl_actuation, self.contactModel, costModel, constraintModelManager, 0., True)\n",
|
||
" model = crocoddyl.IntegratedActionModelEuler(dmodel, self.dt)\n",
|
||
" self.running_models += [model]\n",
|
||
" self.ocp = crocoddyl.ShootingProblem(self.x0, self.running_models[:-1], self.running_models[-1])\n",
|
||
" \n",
|
||
" def createSolver(self):\n",
|
||
" solver = mim_solvers.SolverCSQP(self.ocp)\n",
|
||
" solver.max_qp_iters = 25\n",
|
||
" solver.with_callbacks = True\n",
|
||
" solver.use_filter_line_search = False\n",
|
||
" solver.termination_tolerance = 1e-4\n",
|
||
" solver.eps_abs = 1e-2\n",
|
||
" solver.eps_rel = 0.\n",
|
||
" self.solver = solver\n",
|
||
"\n",
|
||
" def getSolution(self, k=None):\n",
|
||
" if k is None: \n",
|
||
" x_idx = 1\n",
|
||
" u_idx = 0\n",
|
||
" else:\n",
|
||
" x_idx = k\n",
|
||
" u_idx = k\n",
|
||
" t = self.xs[x_idx][:3]\n",
|
||
" quat = self.xs[x_idx][3:7]\n",
|
||
" qx = quat[0]\n",
|
||
" qy = quat[1]\n",
|
||
" qz = quat[2]\n",
|
||
" qw = quat[3]\n",
|
||
" q = self.xs[x_idx][7:27]\n",
|
||
" eta = self.xs[x_idx][27:27+6]\n",
|
||
" dq = self.xs[x_idx][27+6:]\n",
|
||
" constraint_norm = self.solver.constraint_norm\n",
|
||
" return dict(\n",
|
||
" position=t,\n",
|
||
" orientation=np.array([qw, qx, qy, qz]), #Mujoco and uniree quaternion order\n",
|
||
" velocity = eta[:3],\n",
|
||
" omega = eta[3:],\n",
|
||
" q = q[self.mpc_to_unitree_idx],\n",
|
||
" dq = dq[self.mpc_to_unitree_idx], \n",
|
||
" tau = self.us[u_idx][[self.mpc_to_unitree_idx]],\n",
|
||
" constraint_norm = constraint_norm\n",
|
||
" )\n",
|
||
" \n",
|
||
" def updateAndSolve(self, t, quat, q, v, omega, dq):\n",
|
||
" q_ = np.zeros(self.rmodel.nq)\n",
|
||
" dq_ = np.zeros(self.rmodel.nv)\n",
|
||
" qw = quat[0]\n",
|
||
" qx = quat[1]\n",
|
||
" qy = quat[2]\n",
|
||
" qz = quat[3]\n",
|
||
" q_[:3] = t\n",
|
||
" q_[3:7] = np.array([qx, qy, qz, qw])\n",
|
||
" q_[7:] = q[self.unitree_to_mpc_idx]\n",
|
||
" dq_[:3] = v\n",
|
||
" dq_[3:6] = omega\n",
|
||
" dq_[6:] = dq[self.unitree_to_mpc_idx]\n",
|
||
" pin.framesForwardKinematics(self.rmodel, self.rdata, q_)\n",
|
||
" x = np.hstack([q_, dq_])\n",
|
||
" self.solver.problem.x0 = x\n",
|
||
" self.xs = list(self.solver.xs[1:]) + [self.solver.xs[-1]]\n",
|
||
" self.xs[0] = x\n",
|
||
" self.us = list(self.us[1:]) + [self.us[-1]] \n",
|
||
" self.solver.solve(self.xs, self.us, self.max_iterations)\n",
|
||
" self.xs, self.us = self.solver.xs, self.solver.us\n",
|
||
" return self.getSolution()\n",
|
||
" \n",
|
||
" def solve(self):\n",
|
||
" self.solver.solve(self.xs, self.us, self.max_iterations)\n",
|
||
" self.xs, self.us = self.solver.xs, self.solver.us\n",
|
||
" return self.getSolution()\n",
|
||
" "
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 2,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"assets_path = '/home/Go2py/Go2Py/assets'\n",
|
||
"mpc = Go2MPC(assets_path, HORIZON=20, friction_mu=0.1)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"mpc.initialize()\n",
|
||
"mpc.solve()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 4,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from Go2Py.sim.mujoco import Go2Sim\n",
|
||
"import numpy as np\n",
|
||
"robot=Go2Sim(with_arm=True)\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 5,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"map = np.zeros((1200, 1200))\n",
|
||
"map[:20, :200] = 2000\n",
|
||
"robot.updateHeightMap(map)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 6,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import time\n",
|
||
"for i in range(mpc.HORIZON):\n",
|
||
" q = np.zeros(27)\n",
|
||
" state = mpc.getSolution(i)\n",
|
||
" robot.pos0 = state['position']\n",
|
||
" robot.rot0 = state['orientation']\n",
|
||
" robot.q0 = state['q']\n",
|
||
" robot.reset()\n",
|
||
" # time.sleep(0.5)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 7,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"mpc.max_iterations=10"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"for i in range(1000):\n",
|
||
" \n",
|
||
" state = robot.getJointStates()\n",
|
||
" q = state['q']\n",
|
||
" dq = state['dq']\n",
|
||
" t, quat = robot.getPose()\n",
|
||
" v = robot.data.qvel[:3]\n",
|
||
" omega = robot.data.qvel[3:6]\n",
|
||
" solution = mpc.updateAndSolve(t, quat, q, v, omega, dq)\n",
|
||
" solution=mpc.getSolution()\n",
|
||
" q = solution['q']\n",
|
||
" dq = solution['dq']\n",
|
||
" tau = solution['tau'].squeeze()\n",
|
||
" kp = np.ones(20)*0\n",
|
||
" kv = np.ones(20)*0.0\n",
|
||
" # print(solution['constraint_norm'])\n",
|
||
" for i in range(20):\n",
|
||
" robot.setCommands(q, dq, kp, kv, tau)\n",
|
||
" robot.step()"
|
||
]
|
||
}
|
||
],
|
||
"metadata": {
|
||
"kernelspec": {
|
||
"display_name": "base",
|
||
"language": "python",
|
||
"name": "python3"
|
||
},
|
||
"language_info": {
|
||
"codemirror_mode": {
|
||
"name": "ipython",
|
||
"version": 3
|
||
},
|
||
"file_extension": ".py",
|
||
"mimetype": "text/x-python",
|
||
"name": "python",
|
||
"nbconvert_exporter": "python",
|
||
"pygments_lexer": "ipython3",
|
||
"version": "3.12.7"
|
||
}
|
||
},
|
||
"nbformat": 4,
|
||
"nbformat_minor": 2
|
||
}
|