optimization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bbd92319",
   "metadata": {},
   "source": [
    "## Example of Optimization\n",
    "\n",
    "### This code snippets provide the examples in Section 5.4.2 of the SLAM Handbook.\n",
    "\n",
    "Consider a pose estimation problem, where the goal is to find the robot's pose ${y}=\\{{R}, {t}\\}$ that minimizes the error between the observed measurements of landmarks and the predicted measurements based on the current robot pose. This can be formulated as an optimization problem with a cost function:\n",
    "    \\begin{equation}\n",
    "        \\min_{\\theta} \\sum_i \\left(R {p}_i + {t} - {q}_i\\right)^T  \\left({R} {p}_i + {t} - {q}_i\\right),\n",
    "    \\end{equation}\n",
    "where ${R} \\in SO(3)$ represents the rotation of the robot's body and ${t} \\in \\mathbb{R}^3$ represents the translation. As illustrated below, the term ${p}_i$ is the position of the $i$-th landmark in the robot's local frame, and ${q}_i$ is the observed position of the $i$-th landmark.\n",
    "\n",
    "![title](figs/diff-opt.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b9e41219",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "import torch, pypose as pp\n",
    "from pypose.optim import LM\n",
    "from pypose.optim.strategy import Constant\n",
    "from pypose.optim.scheduler import StopOnPlateau"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e3e73f9",
   "metadata": {},
   "source": [
    "### Model preparation\n",
    "\n",
    "Define the regression model ${h}({ y}, {p}_i) = {R} {p}_i + {t}$ that computes the predicted landmarks based on the robot's pose."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cf7283df",
   "metadata": {},
   "outputs": [],
   "source": [
    "class InvNet(nn.Module):\n",
    "    def __init__(self, *dim):\n",
    "        super().__init__()\n",
    "        init = pp.randn_SE3(*dim)\n",
    "        self.pose = pp.Parameter(init)\n",
    "\n",
    "    def forward(self, input):\n",
    "        error = (self.pose @ input).Log()\n",
    "        return error.tensor() # Returns Tensor as Error"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8770fa39",
   "metadata": {},
   "source": [
    "### Data preparation\n",
    "\n",
    "Initialize the parameters with an initial guess for the pose parameters ${ y} = \\{{R_0}, {t_0}\\}$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8f8a912c",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cpu\") # can also be cuda\n",
    "input = pp.randn_SE3(2, 2, device=device)\n",
    "invnet = InvNet(2, 2).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1710cc9",
   "metadata": {},
   "source": [
    "### Optimizer preparation\n",
    "\n",
    "Set up the LM optimizer, including robust kernel functions and correctors to handle outliers. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c68fdb4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "strategy = Constant(damping=1e-4)\n",
    "optimizer = LM(invnet, strategy=strategy)\n",
    "scheduler = StopOnPlateau(optimizer, steps=10, patience=3, decreasing=1e-3, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd916136",
   "metadata": {},
   "source": [
    "### Optimization\n",
    "\n",
    "Perform the optimization by iteratively updating the parameters. \n",
    "The following sample code illustrates the use of PyPose to solve this problem by providing two options for calling the optimizer, assuming the target pose is the identity pose for simplicity.\n",
    "\n",
    "The first one is full optimization, suitable for application.\n",
    "\n",
    "The second one is step optimization, suitable for visulizing intermedium output.\n",
    "\n",
    "Users only need to select one of them in application."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b1b8ae08",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "StopOnPlateau on step 0 Loss 2.333283e+01 --> Loss 9.887070e-07 (reduction/loss: 1.0000e+00).\n",
      "StopOnPlateau on step 1 Loss 9.887070e-07 --> Loss 3.617341e-14 (reduction/loss: 1.0000e+00).\n",
      "StopOnPlateau on step 2 Loss 3.617341e-14 --> Loss 1.781908e-14 (reduction/loss: 5.0740e-01).\n",
      "StopOnPlateau on step 3 Loss 1.781908e-14 --> Loss 5.551115e-17 (reduction/loss: 9.9688e-01).\n",
      "StopOnPlateau: Maximum patience steps reached, Quiting..\n"
     ]
    }
   ],
   "source": [
    "# 1st option, full optimization\n",
    "scheduler.optimize(input=input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6e97e384",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2nd option, step optimization\n",
    "while scheduler.continual():\n",
    "    loss = optimizer.step(input)\n",
    "    scheduler.step(loss)"
   ]
  }
 ],
 "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}