{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "## Without ARM" ] }, { "cell_type": "code", "execution_count": null, "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_reordered.urdf'\n", "robot = pin.RobotWrapper.BuildFromURDF(\n", "urdf_path, urdf_root_path, pin.JointModelFreeFlyer())" ] }, { "cell_type": "code", "execution_count": 62, "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": 28, "metadata": {}, "outputs": [], "source": [ "ee_frame_names = ['FL_foot', 'FR_foot', 'RL_foot', 'RR_foot']\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])" ] }, { "cell_type": "code", "execution_count": 29, "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]\n", " )\n", "x0 = np.concatenate([q0, np.zeros(rmodel.nv)])" ] }, { "cell_type": "code", "execution_count": 30, "metadata": {}, "outputs": [], "source": [ "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 " ] }, { "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": 35, "metadata": {}, "outputs": [], "source": [ "supportFeetIds = [lfFootId, rfFootId, lhFootId, rhFootId]\n", "supportFeePos = [lfFootPos0, rfFootPos0, lhFootPos0, rhFootPos0]" ] }, { "cell_type": "code", "execution_count": 36, "metadata": {}, "outputs": [], "source": [ "state = crocoddyl.StateMultibody(rmodel)\n", "actuation = crocoddyl.ActuationModelFloatingBase(state)\n", "nu = actuation.nu" ] }, { "cell_type": "code", "execution_count": 37, "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": 38, "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)\n", " if t == N_ocp:\n", " costModel.addCost(\"comTrack\", com_track, 1e5)\n", " else:\n", " costModel.addCost(\"comTrack\", com_track, 1e5)\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": 39, "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": 41, "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": "markdown", "metadata": {}, "source": [ "### Animate The Motion" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "from Go2Py.sim.mujoco import Go2Sim\n", "robot_sim = Go2Sim()" ] }, { "cell_type": "code", "execution_count": 70, "metadata": {}, "outputs": [], "source": [ "import mujoco\n", "import time\n", "for i in range(len(xs)):\n", " q0 = xs[i][:19]\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", " robot_sim.reset(q0)\n", " time.sleep(0.01)" ] }, { "cell_type": "code", "execution_count": 71, "metadata": {}, "outputs": [], "source": [ "from Go2Py.robot.model import Go2Model\n", "model = Go2Model()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## With Arm" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Trajectory 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": [ "### MPC" ] }, { "cell_type": "code", "execution_count": 124, "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, 1e1)\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 = 20\n", " solver.with_callbacks = True\n", " solver.use_filter_line_search = False\n", " solver.termination_tolerance = 1e-2\n", " solver.eps_abs = 1e-2\n", " solver.eps_rel = 1e-2\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", " 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", " )\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": 125, "metadata": {}, "outputs": [], "source": [ "assets_path = '/home/Go2py/Go2Py/assets'\n", "mpc = Go2MPC(assets_path, HORIZON=20, friction_mu=0.4)" ] }, { "cell_type": "code", "execution_count": 126, "metadata": {}, "outputs": [], "source": [ "mpc.initialize()\n", "mpc.solve()" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [], "source": [ "from Go2Py.sim.mujoco import Go2Sim\n", "robot=Go2Sim(with_arm=True)" ] }, { "cell_type": "code", "execution_count": 128, "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": 129, "metadata": {}, "outputs": [], "source": [ "mpc.max_iterations=1" ] }, { "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)*10\n", " kv = np.ones(20)*0.3\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 }