{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e519d99b",
   "metadata": {},
   "source": [
    "# KOMO-2: Reporting & explaining convergence\n",
    "\n",
    "When designing motion problems using KOMO, it is really important to check feasibility and -- if the point of convergence is infeasible or different to what you expected -- know how to introspect the result. The key here is get information about which constraints were actually active, what their constraint violation is, and \"how much each constraint is pulling\" at the point of convergence. I recently figured a way to quantify the latter (inspecting the gradients the constraints contribute to the underlying Lagrangian), which really \"explains\" the point of convergence as good as it gets.\n",
    "\n",
    "You should have done the [first KOMO tutorial](komo_1_intro.ipynb) first."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14de1955",
   "metadata": {},
   "source": [
    "## Problem specs reporting\n",
    "\n",
    "Let's first define a simple feasible KOMO problem (same as in the 1st tutorial), solve it, and get info on convergence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e07bf36",
   "metadata": {},
   "outputs": [],
   "source": [
    "import robotic as ry\n",
    "import numpy as np\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "059a8ee7",
   "metadata": {},
   "outputs": [],
   "source": [
    "C = ry.Config()\n",
    "C.addFile(ry.raiPath('scenarios/pandaSingle.g'))\n",
    "C.addFrame('box') \\\n",
    "    .setPosition([-.25,.1,1.]) \\\n",
    "    .setShape(ry.ST.ssBox, size=[.06,.06,.06,.005]) \\\n",
    "    .setColor([1,.5,0]) \\\n",
    "    .setContact(True)\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bccb7b55",
   "metadata": {},
   "outputs": [],
   "source": [
    "komo = ry.KOMO(C, 1, 1, 0, False) # minimalisitc IK problem: just 1 slice\n",
    "komo.addControlObjective([], 0, 1e-1) # basic homing (0th order) objective on joint angles q\n",
    "komo.addObjective([], ry.FS.positionDiff, ['l_gripper', 'box'], ry.OT.eq, [1e1]);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69427e67",
   "metadata": {},
   "source": [
    "Even before we run optimization, we can a report containing the komo specs and list of all objectives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97fb23dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "komo.report()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bdb195d",
   "metadata": {},
   "source": [
    "## Constraints error reporting\n",
    "\n",
    "After optimization, the report will contain errors (i.e. constraint violations, or sqr costs) for all objects. In the current case, all constraint violations are very small:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "178e3d42",
   "metadata": {},
   "outputs": [],
   "source": [
    "ret = ry.NLP_Solver(komo.nlp(), verbose=0 ) .solve()\n",
    "q = komo.getPath()\n",
    "C.setJointState(q[0])\n",
    "C.view()\n",
    "print(ret)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5842764",
   "metadata": {},
   "outputs": [],
   "source": [
    "komo.report()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96fe745f",
   "metadata": {},
   "source": [
    "Let's make the problem infeasible by moving the target out of reach and repeating optimization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b727e37c",
   "metadata": {},
   "outputs": [],
   "source": [
    "C.getFrame('box') .setPosition([.8, .8, 1.])\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5adc15f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "komo.updateRootObjects(C)\n",
    "solver = ry.NLP_Solver(komo.nlp(), verbose=0 )\n",
    "ret = solver.solve()\n",
    "q = komo.getPath()\n",
    "C.setJointState(q[0])\n",
    "C.view()\n",
    "print(ret)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b20fc581",
   "metadata": {},
   "outputs": [],
   "source": [
    "komo.report()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46bf2e2b",
   "metadata": {},
   "source": [
    "We clearly see that the `positionDiff` eq-objective is violated (`err` includes the `scale` as factor, so err=4.6 means 0.46 meters error)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a371b003",
   "metadata": {},
   "source": [
    "## Lagrange gradients reporting\n",
    "\n",
    "Seeing only objective errors sometimes does not really explain what is the issues. So instead we can let the solver report what the lagrange gradients w.r.t. each feature are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee59d5f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "solver.reportLagrangeGradients(komo.getFeatureNames())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51c2bca7",
   "metadata": {},
   "source": [
    "Note that this list is always sorted, starting with largest gradients first. For complex problems, this report can be rather insightful. Sometimes one explicitly sees which objectives *fight* against each other in the Lagrangian by both having similar large gradient sizes (typically with opposing directions, which the report does not show; also a matrix with \"angles between gradients\" can be computed)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d67e74c",
   "metadata": {},
   "source": [
    "## Plotting constraint errors over time for paths\n",
    "\n",
    "The above only considers a simple IK problem. Let's look at a path problem (we take the same 4-waypoint problem from the 1st tutorial):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6da9bda",
   "metadata": {},
   "outputs": [],
   "source": [
    "C = ry.Config()\n",
    "C.addFile(ry.raiPath('scenarios/pandaSingle.g'))\n",
    "C.addFrame('way1'). setShape(ry.ST.marker, [.1]) .setPosition([.4, .2, 1.])\n",
    "C.addFrame('way2'). setShape(ry.ST.marker, [.1]) .setPosition([.4, .2, 1.4])\n",
    "C.addFrame('way3'). setShape(ry.ST.marker, [.1]) .setPosition([-.4, .2, 1.])\n",
    "C.addFrame('way4'). setShape(ry.ST.marker, [.1]) .setPosition([-.4, .2, 1.4])\n",
    "q0 = C.getJointState()\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40341e34",
   "metadata": {},
   "outputs": [],
   "source": [
    "C.setJointState(q0)\n",
    "komo = ry.KOMO(C, 4, 10, 2, False)\n",
    "komo.addControlObjective([], 0, 1e-1)\n",
    "komo.addControlObjective([], 2, 1e0)\n",
    "komo.addObjective([1], ry.FS.positionDiff, ['l_gripper', 'way1'], ry.OT.eq, [1e1])\n",
    "komo.addObjective([2], ry.FS.positionDiff, ['l_gripper', 'way2'], ry.OT.eq, [1e1])\n",
    "komo.addObjective([3], ry.FS.positionDiff, ['l_gripper', 'way3'], ry.OT.eq, [1e1])\n",
    "komo.addObjective([4], ry.FS.positionDiff, ['l_gripper', 'way4'], ry.OT.eq, [1e1])\n",
    "komo.addObjective([4], ry.FS.jointState, [], ry.OT.eq, [1e1], [], order=1)\n",
    "\n",
    "solver = ry.NLP_Solver(komo.nlp(), verbose=0 )\n",
    "ret = solver.solve()\n",
    "print(ret)\n",
    "q = komo.getPath()\n",
    "\n",
    "for t in range(q.shape[0]):\n",
    "    C.setJointState(q[t])\n",
    "    C.view(False, f'waypoint {t}')\n",
    "    time.sleep(.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2edd1e16",
   "metadata": {},
   "source": [
    "Using `plotOverTime` in the report method now plots the constraint violations (or sqr costs) over *phase*, using gnuplot. (You need to have `gnuplot` installed on your Ubuntu, using `sudo apt install gnuplot`.)\n",
    "\n",
    "Here, the only interesting (non-zero) signal is from the 2nd-order control costs, which are the acceleration costs. This is a typical acceperation cost profile for optimal paths between waypoints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f9776fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "R = komo.report(False, plotOverTime=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b96a7c9",
   "metadata": {},
   "source": [
    "Finally, the largrange gradients show that the control costs and the final velocity constraints contribute a lot, but each waypoint also contributes a significant gradient that explains the point of convergence. (If you'd increase control cost scaling, the waypoint Lagrange parameters and their reported gradients would equally increase.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6b64321",
   "metadata": {},
   "outputs": [],
   "source": [
    "solver.reportLagrangeGradients(komo.getFeatureNames())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ee8271b",
   "metadata": {},
   "source": [
    "Test the following: Change the timing of the 3rd waypoint to \"[2]\" as well, requiring the robot to be at way2 and way3 at the same time -- which clearly is infeasible. Check the solver return and lagrange gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d08b232",
   "metadata": {},
   "outputs": [],
   "source": [
    "del komo\n",
    "del C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88432b5f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}