1_invgeom.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inverse geometry\n",
    "\n",
    "Inverse geometry is the problem of computing a configuration $q$ of a robot model that satisfies a set of 'tasks', which we formalize as error functions $e(q)$ that are zero when the task is achieved.\n",
    "\n",
    "- $q \\in \\mathcal{C}$: configuration of the robot.\n",
    "- $e_i(q) \\in \\mathbb{R}$: error function for task $i$. In this tutorial our error function is real valued.\n",
    "\n",
    "Typical tasks for limbed robots include putting a foot/end-effector on a surface, moving the center of mass (CoM) to a target location, etc. For example, an error function for a foot position task could be:\n",
    "\n",
    "$$\n",
    "e(q) = \\|p_\\mathrm{foot}^* - p_\\mathrm{foot}(q)\\|\n",
    "$$\n",
    "\n",
    "In this tutorial, we will define a task function $e(q)$ for the end-effector pose of a manipulator and solve the inverse geometry problem\n",
    "\n",
    "$$\n",
    "\\underset{q}{\\mathrm{minimize}}\\ e(q)\n",
    "$$\n",
    "\n",
    "using an off-the-shelf nonlinear optimizer that evaluates gradients by finite differences (slower), then a more efficient formulation using CasADi (faster)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Before we start...\n",
    "\n",
    "%load is a \"magic\" command of Jupyter. When we will invite you to think to a solution by yourselves, we will indicate it by the %do_not_load command. Feel free to transform it to a plain %load if you want the answer. This is define in the following import."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import magic_donotload"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's import everything we will need in this tutorial:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load tp1/generated/invgeom_imports\n",
    "import time\n",
    "import unittest\n",
    "import casadi\n",
    "import tp1.meshcat_shapes as meshcat_shapes\n",
    "import example_robot_data as robex\n",
    "import meshcat\n",
    "import numpy as np\n",
    "import pinocchio as pin\n",
    "from numpy.linalg import norm\n",
    "from scipy.optimize import fmin_bfgs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up\n",
    "\n",
    "Let us start by loading the UR5 robot model from `example-robot-data` and define its neutral configuration as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load tp1/generated/invgeom_robot\n",
    "robot = robex.load(\"ur5\")\n",
    "robot.q0 = np.array([0, -np.pi / 2, 0, 0, 0, 0])\n",
    "model = robot.model\n",
    "data = robot.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will also use the MeshCat visualizer to watch the optimization unfold. Let's set up a visualizer for our robot model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load tp1/generated/invgeom_visualizer\n",
    "viz = pin.visualize.MeshcatVisualizer(\n",
    "    robot.model, robot.collision_model, robot.visual_model\n",
    ")\n",
    "robot.setVisualizer(viz, init=False)\n",
    "viz.initViewer(open=False)\n",
    "viz.loadViewerModel()\n",
    "viz.display(robot.q0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viz.viewer.jupyter_cell()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: BFGS with finite difference\n",
    "\n",
    "In this first part, we will define an inverse geometry problem for the UR5 and use the BFGS optimizer from [`scipy.optimize.fmin_bfgs`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.fmin_bfgs.html#scipy-optimize-fmin-bfgs) to solve it.\n",
    "\n",
    "### Task target\n",
    "\n",
    "Our task will be to move the end effector *frame* of the UR5 to a desired pose (pose := position + orientation). For this purpose, we define:\n",
    "\n",
    "- The name of the corresponding frame, which is ``\"tool0\"`` for the UR5.\n",
    "- The target pose of this frame, which we define by an affine transform ${}_{world} M_{target} \\in SE(3)$ (also known as a \"placement\").\n",
    "\n",
    "The robot will start from an arbitrary initial configuration (hint: to make the solver's life easier, avoid the singular $[0\\ 0\\ \\ldots\\ 0]$ configuration) $q_0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load tp1/generated/invgeom_task_params\n",
    "tool_id = model.getFrameId(\"tool0\")\n",
    "\n",
    "transform_target_to_world = pin.SE3(\n",
    "    pin.utils.rotate(\"x\", np.pi / 4),\n",
    "    np.array([-0.5, 0.1, 0.2]),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task error\n",
    "\n",
    "For a given robot configuration $q$, the end effector will be located at ${}_{world} M_{tool}(q)$, which is a function of $q$ we can compute with Pinocchio's forward kinematics function. The goal of our task will be to make ${}_{world} M_{tool}(q)$ coincide with ${}_{world} M_{target}$, which we can define mathematically as: minimize the (Lie-group) difference ${}_{world} M_{target} \\ominus {}_{world} M_{tool}(q)$, or equivalently:\n",
    "\n",
    "$$\n",
    "\\underset{q}{\\mathrm{minimize}} \\ e(q) := \\| {}_{world} M_{target} \\ominus {}_{world} M_{tool}(q) \\|\n",
    "$$\n",
    "\n",
    "Remember the definition of the $\\ominus$ (left) Lie-group difference? No? Let's write it in Python anyway!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load tp1/generated/invgeom_error_function\n",
    "\n",
    "\n",
    "def error(q: np.ndarray) -> float:\n",
    "    \"\"\"Function that the final configuration should minimize.\"\"\"\n",
    "    pin.framesForwardKinematics(model, data, q)\n",
    "    transform_tool_to_world = data.oMf[tool_id]\n",
    "    return norm(\n",
    "        pin.log(transform_tool_to_world.inverse() * transform_target_to_world).vector\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sidetrack: visualization\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load tp1/generated/invgeom_visualizer_callback\n",
    "viewer = viz.viewer\n",
    "meshcat_shapes.frame(viewer[\"target\"], opacity=1.0)\n",
    "meshcat_shapes.frame(viewer[\"current\"], opacity=0.5)\n",
    "\n",
    "\n",
    "def callback(q: np.ndarray):\n",
    "    pin.framesForwardKinematics(model, data, q)\n",
    "    transform_frame_to_world = data.oMf[tool_id]\n",
    "    viewer[\"target\"].set_transform(transform_target_to_world.np)\n",
    "    viewer[\"current\"].set_transform(transform_frame_to_world.np)\n",
    "    viz.display(q)\n",
    "    time.sleep(1e-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now ready to optimize for real!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### BFGS optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load tp1/generated/invgeom_fmin_bfgs\n",
    "qguess = np.array([0.12, -2.2, -1.45, 1.82, -0.95, 0.17])\n",
    "qopt = fmin_bfgs(error, qguess, callback=callback)\n",
    "\n",
    "print(\n",
    "    \"The robot finally reached effector placement at\\n\",\n",
    "    robot.placement(qopt, 6),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If everything went well, our robot converged to a configuration $q$ that achieves the desired end effector placement, *i.e.* $e(q) \\approx 0 \\Leftrightarrow {}_{world} M_{target} \\approx {{}_{world} M_{tool}}(q)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: IPOPT with automatic differentiation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We wil now define our problem in [CasADi](https://web.casadi.org), a framework that allows us to formulate problems as computation graphs from which it can compute derivatives (notably first-order derivatives, *a.k.a.* gradients, and second-order derivatives, *a.k.a.* Hessians) automatically. This feature is called *automatic differentiation*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Modelling with CasADi\n",
    "\n",
    "Pinocchio, in the version 3 we are using in today's tutorial, provides a `pinocchio.casadi` interface to represent configurations (such as $q \\in \\mathcal{C}$ or $M \\in SE(3)$) in CasADi computation graphs, and differentiate them appropriately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load tp1/generated/invgeom_casadi_imports\n",
    "import casadi\n",
    "from pinocchio import casadi as cpin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "CasADi being designed for Euclidean variables, differentiating properly over manifolds is not as straightforward as representing everything as matrices and vectors in $\\mathbb{R}^n$ and taking their usual derivatives. Pinocchio allows us to take derivatives over manifolds properly, without worrying about what happens below the surface. To be able to do that, we first need to redefine our model as a CasADi one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load tp1/generated/invgeom_casadi_model\n",
    "cmodel = cpin.Model(model)\n",
    "cdata = cmodel.createData()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our configuration variable $q$ is now a CasADi ``SX`` symbol of size $n_q$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load tp1/generated/invgeom_cq\n",
    "cq = casadi.SX.sym(\"q\", model.nq, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We complete the computation graph from ``cmodel`` to ``cq`` by calling the corresponding forward-kinematics function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load tp1/generated/invgeom_casadi_fk\n",
    "cpin.framesForwardKinematics(cmodel, cdata, cq)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task error\n",
    "\n",
    "We define the same task error as above:\n",
    "\n",
    "$$\n",
    "\\underset{q}{\\mathrm{minimize}} \\ e(q) := \\| {}_{world} M_{target} \\ominus {}_{world} M_{tool}(q) \\|\n",
    "$$\n",
    "\n",
    "But this time as a CasADi function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load tp1/generated/invgeom_casadi_error\n",
    "error_tool = casadi.Function(\n",
    "    \"etool\",\n",
    "    [cq],\n",
    "    [\n",
    "        cpin.log6(\n",
    "            cdata.oMf[tool_id].inverse() * cpin.SE3(transform_target_to_world)\n",
    "        ).vector\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we define a computation graph from our optimization variable to our error function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load tp1/generated/invgeom_casadi_computation_graph\n",
    "opti = casadi.Opti()\n",
    "var_q = opti.variable(model.nq)\n",
    "totalcost = casadi.sumsqr(error_tool(var_q))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### IPOPT optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load tp1/generated/invgeom_ipopt\n",
    "opti.minimize(totalcost)\n",
    "opti.solver(\"ipopt\")  # select the backend solver\n",
    "opti.callback(lambda i: callback(opti.debug.value(var_q)))\n",
    "\n",
    "# Caution: in case the solver does not converge, we are picking the candidate values\n",
    "# at the last iteration in opti.debug, and they are NO guarantee of what they mean.\n",
    "try:\n",
    "    sol = opti.solve_limited()\n",
    "    sol_q = opti.value(var_q)\n",
    "except:\n",
    "    print(\"ERROR in convergence, plotting debug info.\")\n",
    "    sol_q = opti.debug.value(var_q)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can check that the solver converged to the same solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load tp1/generated/invgeom_check_final_placement\n",
    "print(\n",
    "    \"The robot finally reached effector placement at\\n\",\n",
    "    robot.placement(sol_q, 6),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To conclude, count the number of iterations and number of function/gradient evaluations. This is automatic differentiation 🤠"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}