From 747e43ec7abcd030ced8eabc77b23d5b28509951 Mon Sep 17 00:00:00 2001 From: Rooholla-KhorramBakht Date: Fri, 6 Dec 2024 14:42:24 -0500 Subject: [PATCH] mpc example cleaned up --- Go2Py/sim/mujoco.py | 2 +- examples/09-crododdyl.ipynb | 427 +++--------------------------------- 2 files changed, 34 insertions(+), 395 deletions(-) diff --git a/Go2Py/sim/mujoco.py b/Go2Py/sim/mujoco.py index a9a21ee..23d8f80 100644 --- a/Go2Py/sim/mujoco.py +++ b/Go2Py/sim/mujoco.py @@ -5,7 +5,7 @@ import numpy as np from Go2Py import ASSETS_PATH import os from scipy.spatial.transform import Rotation -# import cv2 +import cv2 pnt = np.array([-0.2, 0, 0.05]) lidar_angles = np.linspace(0.0, 2 * np.pi, 1024).reshape(-1, 1) diff --git a/examples/09-crododdyl.ipynb b/examples/09-crododdyl.ipynb index 93e338c..b36ce19 100644 --- a/examples/09-crododdyl.ipynb +++ b/examples/09-crododdyl.ipynb @@ -4,383 +4,7 @@ "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" + "## OfflineTrajectory Optimization" ] }, { @@ -839,12 +463,12 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### MPC" + "## Online MPC" ] }, { "cell_type": "code", - "execution_count": 124, + "execution_count": 1, "metadata": {}, "outputs": [], "source": [ @@ -1020,7 +644,7 @@ " if t == self.HORIZON:\n", " costModel.addCost(\"efTrack\", ef_track, 1e5)\n", " else:\n", - " costModel.addCost(\"efTrack\", ef_track, 1e1)\n", + " costModel.addCost(\"efTrack\", ef_track, 1e3)\n", " \n", " # Friction Cone Constraints\n", " constraintModelManager = crocoddyl.ConstraintModelManager(self.ccdyl_state, self.ccdyl_actuation.nu)\n", @@ -1038,12 +662,12 @@ " \n", " def createSolver(self):\n", " solver = mim_solvers.SolverCSQP(self.ocp)\n", - " solver.max_qp_iters = 20\n", + " solver.max_qp_iters = 25\n", " solver.with_callbacks = True\n", " solver.use_filter_line_search = False\n", - " solver.termination_tolerance = 1e-2\n", + " solver.termination_tolerance = 1e-4\n", " solver.eps_abs = 1e-2\n", - " solver.eps_rel = 1e-2\n", + " solver.eps_rel = 0.\n", " self.solver = solver\n", "\n", " def getSolution(self, k=None):\n", @@ -1062,6 +686,7 @@ " 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", @@ -1069,7 +694,8 @@ " 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", + " 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", @@ -1104,17 +730,17 @@ }, { "cell_type": "code", - "execution_count": 125, + "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "assets_path = '/home/Go2py/Go2Py/assets'\n", - "mpc = Go2MPC(assets_path, HORIZON=20, friction_mu=0.4)" + "mpc = Go2MPC(assets_path, HORIZON=20, friction_mu=0.1)" ] }, { "cell_type": "code", - "execution_count": 126, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -1124,17 +750,29 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "from Go2Py.sim.mujoco import Go2Sim\n", - "robot=Go2Sim(with_arm=True)" + "import numpy as np\n", + "robot=Go2Sim(with_arm=True)\n" ] }, { "cell_type": "code", - "execution_count": 128, + "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": [ @@ -1151,11 +789,11 @@ }, { "cell_type": "code", - "execution_count": 129, + "execution_count": 7, "metadata": {}, "outputs": [], "source": [ - "mpc.max_iterations=1" + "mpc.max_iterations=10" ] }, { @@ -1177,8 +815,9 @@ " 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", + " 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()"