dp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TKvYiJgnYExi"
   },
   "source": [
    "Welcome!  If you are new to Google Colab/Jupyter notebooks, you might take a look at [this notebook](https://colab.research.google.com/notebooks/basic_features_overview.ipynb) first.\n",
    "\n",
    "**I recommend you run the first code cell of this notebook immediately, to start provisioning drake on the cloud machine, then you can leave this window open as you [read the textbook](http://underactuated.csail.mit.edu/dp.html).**\n",
    "\n",
    "# Notebook Setup\n",
    "\n",
    "The following cell will:\n",
    "- on Colab (only), install Drake to `/opt/drake`, install Drake's prerequisites via `apt`, and add pydrake to `sys.path`.  This will take approximately two minutes on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  If you navigate between notebooks using Colab's \"File->Open\" menu, then you can avoid provisioning a separate machine for each notebook.\n",
    "- import packages used throughout the notebook.\n",
    "\n",
    "You will need to rerun this cell if you restart the kernel, but it should be fast (even on Colab) because the machine will already have drake installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A4QOaw_zYLfI"
   },
   "outputs": [],
   "source": [
    "import importlib\n",
    "import sys\n",
    "from urllib.request import urlretrieve\n",
    "\n",
    "# Install drake (and underactuated).\n",
    "if 'google.colab' in sys.modules and importlib.util.find_spec('underactuated') is None:\n",
    "    urlretrieve(f\"http://underactuated.csail.mit.edu/scripts/setup/setup_underactuated_colab.py\",\n",
    "                \"setup_underactuated_colab.py\")\n",
    "    from setup_underactuated_colab import setup_underactuated\n",
    "    setup_underactuated(underactuated_sha='560c2adace05eb20ebd78377582015d5b2d3859a', drake_version='0.25.0', drake_build='releases')\n",
    "\n",
    "server_args = []\n",
    "if 'google.colab' in sys.modules:\n",
    "  server_args = ['--ngrok_http_tunnel']\n",
    "# Start a single meshcat server instance to use for the remainder of this notebook.\n",
    "from meshcat.servers.zmqserver import start_zmq_server_as_subprocess\n",
    "proc, zmq_url, web_url = start_zmq_server_as_subprocess(server_args=server_args)\n",
    "\n",
    "\n",
    "# Imports.\n",
    "from time import sleep\n",
    "from functools import partial\n",
    "from ipywidgets import interact\n",
    "import matplotlib.animation as animation\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "import meshcat\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import numpy as np\n",
    "import pydrake.all\n",
    "from IPython.display import display, HTML\n",
    "from pydrake.all import DiagramBuilder, LinearSystem, Simulator, WrapToSystem\n",
    "from pydrake.systems.controllers import (DynamicProgrammingOptions,\n",
    "                                         FittedValueIteration, PeriodicBoundaryCondition)\n",
    "\n",
    "import underactuated\n",
    "from underactuated.jupyter import AdvanceToAndVisualize, SetupMatplotlibBackend, running_as_notebook\n",
    "from underactuated.meshcat_utils import plot_surface\n",
    "\n",
    "plt.rcParams.update({\"savefig.transparent\": True})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Grid World\n",
    "\n",
    "The setup here is *almost* identical as the simplest version described in the notes.  The only difference is that this agent is allowed to move diagonally in a single step; this is slightly easier to code since I can have two actions (one for left/right, and another for up/down), and write the dynamics as the trivial linear system ${\\bf x}[n+1] = {\\bf u}[n].$  Only the value iteration code needs to know that the states and actions are actually restricted to the integers.\n",
    "\n",
    "The obstacle (pit of despair) is provided by the method below.  Play around with it!  The rest of the code is mostly to support visualization.\n",
    "\n",
    "TODO(russt): Pull a few more of the visualization frills from my (very) old [matlab code](https://github.com/RobotLocomotion/drake/blob/last_sha_with_original_matlab/drake/examples/GridWorld.m).  At very least, I want to draw the vector field of the resulting policy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def grid_world_example():\n",
    "    time_step = 1\n",
    "    # TODO(russt): Support discrete-time systems in the dynamic programming code, and use this properly.\n",
    "    #plant = LinearSystem(A=np.eye(2), B=np.eye(2), C=np.eye(2), D=np.zeros((2,2)), time_period=time_step)\n",
    "    # for now, just cheat because I know how to make the discrete system as a continuous that will be discretized.\n",
    "    plant = LinearSystem(A=np.zeros((2,2)), B=np.eye(2), C=np.eye(2), D=np.zeros((2,2)))\n",
    "    simulator = Simulator(plant)\n",
    "    options = DynamicProgrammingOptions()\n",
    "\n",
    "    xbins = range(0, 21)\n",
    "    ybins = range(0, 21)\n",
    "    state_grid = [set(xbins), set(ybins)]\n",
    "\n",
    "    input_grid = [set([-1, 0, 1]), set([-1, 0, 1])]\n",
    "\n",
    "    goal = [2, 8]\n",
    "\n",
    "    def obstacle(x):\n",
    "        return x[0]>=6 and x[0]<=8 and x[1]>=4 and x[1]<=7\n",
    "\n",
    "    [X, Y] = np.meshgrid(xbins, ybins)\n",
    "\n",
    "    frames=[]\n",
    "    def draw(iteration, mesh, cost_to_go, policy):\n",
    "        J = np.reshape(cost_to_go, X.shape)\n",
    "        artists = [ax.imshow(J, cmap=cm.jet)]\n",
    "        frames.append(artists)\n",
    "\n",
    "    if running_as_notebook:\n",
    "        options.visualization_callback = draw\n",
    "\n",
    "    def min_time_cost(context):\n",
    "        x = context.get_continuous_state_vector().CopyToVector()\n",
    "        x = np.round(x)\n",
    "        if obstacle(x):\n",
    "            return 10\n",
    "        if np.array_equal(x, goal):\n",
    "            return 0\n",
    "        return 1\n",
    "        \n",
    "    cost_function = min_time_cost\n",
    "    options.convergence_tol = .1;\n",
    "\n",
    "    (fig,ax) = plt.subplots()\n",
    "    ax.set_xlabel(\"x\")\n",
    "    ax.set_ylabel(\"y\")\n",
    "    ax.set_title(\"Cost-to-Go\")\n",
    "\n",
    "    policy, cost_to_go = FittedValueIteration(simulator, cost_function, state_grid,\n",
    "                                            input_grid, time_step, options)\n",
    "\n",
    "    J = np.reshape(cost_to_go, X.shape)\n",
    "    artists = [ax.imshow(J, cmap=cm.jet)]\n",
    "    frames.append(artists)\n",
    "\n",
    "    ax.invert_yaxis()\n",
    "    plt.colorbar(artists[0])\n",
    "\n",
    "    # create animation using the animate() function\n",
    "    ani = animation.ArtistAnimation(fig, frames, interval=200, blit=True, repeat=False)\n",
    "    plt.close('all')\n",
    "\n",
    "    display(HTML(ani.to_jshtml()))\n",
    "\n",
    "grid_world_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your turn.  Change the cost.  Change the obstacles."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Value Iteration for the Double Integrator\n",
    "\n",
    "Note that I've inserted a sleep command in the draw method to intentionally slow down the algorithm, so that you can watch the convergence in the visualizer.  If you take out the pause, it's quite fast!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from underactuated.double_integrator import DoubleIntegratorVisualizer\n",
    "\n",
    "def DoubleIntegrator():\n",
    "    return LinearSystem(A=np.mat('0 1; 0 0'),\n",
    "                        B=np.mat('0; 1'),\n",
    "                        C=np.eye(2),\n",
    "                        D=np.zeros((2,1)))\n",
    "plant = DoubleIntegrator()\n",
    "\n",
    "def double_integrator_example(cost_function, convergence_tol, animate=True, plot=True, draw_iterations=True):\n",
    "    simulator = Simulator(plant)\n",
    "    options = DynamicProgrammingOptions()\n",
    "\n",
    "    qbins = np.linspace(-3., 3., 31)\n",
    "    qdotbins = np.linspace(-3., 3., 51)\n",
    "    state_grid = [set(qbins), set(qdotbins)]\n",
    "\n",
    "    input_limit = 1.\n",
    "    input_grid = [set(np.linspace(-input_limit, input_limit, 9))]\n",
    "    timestep = 0.01\n",
    "\n",
    "    [Q, Qdot] = np.meshgrid(qbins, qdotbins)\n",
    "\n",
    "    vis = meshcat.Visualizer(zmq_url=zmq_url, server_args=server_args)\n",
    "    vis['/Background'].set_property(\"visible\", False)\n",
    "\n",
    "    def draw(iteration, mesh, cost_to_go, policy):\n",
    "        # Don't draw every frame.\n",
    "        if iteration % 20 != 0:\n",
    "            return\n",
    "\n",
    "        plot_surface(vis['Cost-to-go'], Q, Qdot, np.reshape(cost_to_go, Q.shape), color=cm.jet, wireframe=True)\n",
    "        plot_surface(vis['Policy'], Q, Qdot, np.reshape(policy, Q.shape), color=0x9999dd)\n",
    "\n",
    "        # Slow down the algorithm so we can visualize the convergence.\n",
    "        sleep(0.1) \n",
    "\n",
    "    def simulate(policy):\n",
    "        # Animate the resulting policy.\n",
    "        SetupMatplotlibBackend()\n",
    "\n",
    "        builder = DiagramBuilder()\n",
    "        plant = builder.AddSystem(DoubleIntegrator())\n",
    "\n",
    "        vi_policy = builder.AddSystem(policy)\n",
    "        builder.Connect(plant.get_output_port(0), vi_policy.get_input_port(0))\n",
    "        builder.Connect(vi_policy.get_output_port(0), plant.get_input_port(0))\n",
    "\n",
    "        visualizer = builder.AddSystem(DoubleIntegratorVisualizer())\n",
    "        builder.Connect(plant.get_output_port(0), visualizer.get_input_port(0))\n",
    "\n",
    "        diagram = builder.Build()\n",
    "        simulator = Simulator(diagram)\n",
    "\n",
    "        simulator.get_mutable_context().SetContinuousState([-10.0, 0.0])\n",
    "\n",
    "        AdvanceToAndVisualize(simulator, visualizer, 10.)\n",
    "\n",
    "    if running_as_notebook and draw_iterations:\n",
    "        options.visualization_callback = draw\n",
    "    options.convergence_tol = convergence_tol\n",
    "\n",
    "    policy, cost_to_go = FittedValueIteration(simulator, cost_function, state_grid,\n",
    "                                            input_grid, timestep, options)\n",
    "\n",
    "    J = np.reshape(cost_to_go, Q.shape)\n",
    "    \n",
    "    plot_surface(vis['Cost-to-go'], Q, Qdot, J, cm.jet)\n",
    "\n",
    "    if animate:\n",
    "        print('Simulating...')\n",
    "        simulate(policy)\n",
    "        \n",
    "    if plot:\n",
    "        fig = plt.figure(1, figsize=(9, 4))\n",
    "        ax1, ax2 = fig.subplots(1, 2, subplot_kw=dict(projection='3d'))\n",
    "        ax1.set_xlabel(\"q\")\n",
    "        ax1.set_ylabel(\"qdot\")\n",
    "        ax1.set_title(\"Cost-to-Go\")\n",
    "        ax2.set_xlabel(\"q\")\n",
    "        ax2.set_ylabel(\"qdot\")\n",
    "        ax2.set_title(\"Policy\")\n",
    "        surf = ax1.plot_surface(Q, Qdot, J, rstride=1, cstride=1, cmap=cm.jet)\n",
    "        Pi = np.reshape(policy.get_output_values(), Q.shape)\n",
    "        surf = ax2.plot_surface(Q, Qdot, Pi, rstride=1, cstride=1, cmap=cm.jet)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def min_time_cost(context):\n",
    "    x = context.get_continuous_state_vector().CopyToVector()\n",
    "    if x.dot(x) < .05:\n",
    "        return 0.\n",
    "    return 1.\n",
    "\n",
    "double_integrator_example(cost_function=min_time_cost, convergence_tol=0.001, animate=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def quadratic_regulator_cost(context):\n",
    "    x = context.get_continuous_state_vector().CopyToVector()\n",
    "    u = plant.EvalVectorInput(context, 0).CopyToVector()\n",
    "    return x.dot(x) + u.dot(u)\n",
    "\n",
    "double_integrator_example(cost_function=quadratic_regulator_cost, convergence_tol=0.1, animate=True)"
   ]
  },
  {
   "source": [
    "# Change the cost function"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def quadratic_regulator_cost(context, q0, qdot0):\n",
    "    x = context.get_continuous_state_vector().CopyToVector()\n",
    "    x = x - np.array([q0, qdot0])\n",
    "    u = plant.EvalVectorInput(context, 0).CopyToVector()\n",
    "    return x.dot(x) + u.dot(u)\n",
    "\n",
    "def update(q0, qdot0):\n",
    "    double_integrator_example(\n",
    "        cost_function=partial(quadratic_regulator_cost, q0=q0, qdot0=qdot0),\n",
    "        convergence_tol=0.1,  \n",
    "        animate=False, plot=False, draw_iterations=False)        \n",
    "\n",
    "interact(update, q0=(-2,2,0.1), qdot0=(-10,10,0.1));\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Value Iteration for the Simple Pendulum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from pydrake.examples.pendulum import PendulumPlant\n",
    "from underactuated.pendulum import PendulumVisualizer\n",
    "\n",
    "def pendulum_swingup_example(min_time=True, animate=True):\n",
    "    plant = PendulumPlant()\n",
    "    simulator = Simulator(plant)\n",
    "    options = DynamicProgrammingOptions()\n",
    "\n",
    "    qbins = np.linspace(0., 2. * np.pi, 51)\n",
    "    qdotbins = np.linspace(-10., 10., 51)\n",
    "    state_grid = [set(qbins), set(qdotbins)]\n",
    "    options.periodic_boundary_conditions = [\n",
    "        PeriodicBoundaryCondition(0, 0., 2. * np.pi),\n",
    "    ]\n",
    "    input_limit = 3.\n",
    "    input_grid = [set(np.linspace(-input_limit, input_limit, 9))]\n",
    "    timestep = 0.01\n",
    "\n",
    "    [Q, Qdot] = np.meshgrid(qbins, qdotbins)\n",
    "\n",
    "    vis = meshcat.Visualizer(zmq_url=zmq_url, server_args=server_args)\n",
    "    vis['/Background'].set_property(\"visible\", False)\n",
    "\n",
    "    def draw(iteration, mesh, cost_to_go, policy):\n",
    "        # Don't draw every frame.\n",
    "        if iteration % 20 != 0:\n",
    "            return\n",
    "\n",
    "        plot_surface(vis['Cost-to-go'], Q, Qdot, np.reshape(cost_to_go, Q.shape), color=cm.jet, wireframe=True)\n",
    "        plot_surface(vis['Policy'], Q, Qdot, np.reshape(policy, Q.shape), color=0x9999dd)\n",
    "\n",
    "        # Slow down the algorithm so we can visualize the convergence.\n",
    "        sleep(0.1) \n",
    "\n",
    "    def simulate(policy):\n",
    "        # Animate the resulting policy.\n",
    "        plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "        builder = DiagramBuilder()\n",
    "        pendulum = builder.AddSystem(PendulumPlant())\n",
    "\n",
    "        wrap = builder.AddSystem(WrapToSystem(2))\n",
    "        wrap.set_interval(0, 0, 2*np.pi)\n",
    "        builder.Connect(pendulum.get_output_port(0), wrap.get_input_port(0))\n",
    "        vi_policy = builder.AddSystem(policy)\n",
    "        builder.Connect(wrap.get_output_port(0), vi_policy.get_input_port(0))\n",
    "        builder.Connect(vi_policy.get_output_port(0), pendulum.get_input_port(0))\n",
    "\n",
    "        visualizer = builder.AddSystem(PendulumVisualizer(show=plt_is_interactive))\n",
    "        builder.Connect(pendulum.get_output_port(0), visualizer.get_input_port(0))\n",
    "\n",
    "        diagram = builder.Build()\n",
    "        simulator = Simulator(diagram)\n",
    "        simulator.get_mutable_context().SetContinuousState([0.1, 0.0])\n",
    "\n",
    "        AdvanceToAndVisualize(simulator, visualizer, 8.)\n",
    "\n",
    "    if running_as_notebook:\n",
    "        options.visualization_callback = draw\n",
    "\n",
    "    def min_time_cost(context):\n",
    "        x = context.get_continuous_state_vector().CopyToVector()\n",
    "        x[0] = x[0] - np.pi\n",
    "        if x.dot(x) < .05:\n",
    "            return 0.\n",
    "        return 1.\n",
    "\n",
    "    def quadratic_regulator_cost(context):\n",
    "        x = context.get_continuous_state_vector().CopyToVector()\n",
    "        x[0] = x[0] - np.pi\n",
    "        u = plant.EvalVectorInput(context, 0).CopyToVector()\n",
    "        return 2 * x.dot(x) + u.dot(u)\n",
    "    \n",
    "    if min_time:\n",
    "        cost_function = min_time_cost\n",
    "        options.convergence_tol = 0.001\n",
    "    else:\n",
    "        cost_function = quadratic_regulator_cost\n",
    "        options.convergence_tol = 0.1\n",
    "\n",
    "    policy, cost_to_go = FittedValueIteration(simulator, cost_function, state_grid,\n",
    "                                            input_grid, timestep, options)\n",
    "\n",
    "    J = np.reshape(cost_to_go, Q.shape)\n",
    "    \n",
    "    plot_surface(vis['Cost-to-go'], Q, Qdot, J, cm.jet)\n",
    "\n",
    "    if animate:\n",
    "        print('Simulating...')\n",
    "        simulate(policy)\n",
    "\n",
    "    fig = plt.figure(figsize=(9, 4))\n",
    "    ax1, ax2 = fig.subplots(1, 2, subplot_kw=dict(projection='3d'))\n",
    "    ax1.set_xlabel(\"q\")\n",
    "    ax1.set_ylabel(\"qdot\")\n",
    "    ax1.set_title(\"Cost-to-Go\")\n",
    "    ax2.set_xlabel(\"q\")\n",
    "    ax2.set_ylabel(\"qdot\")\n",
    "    ax2.set_title(\"Policy\")\n",
    "    surf = ax1.plot_surface(Q, Qdot, J, rstride=1, cstride=1, cmap=cm.jet)\n",
    "    Pi = np.reshape(policy.get_output_values(), Q.shape)\n",
    "    surf = ax2.plot_surface(Q, Qdot, Pi, rstride=1, cstride=1, cmap=cm.jet)\n",
    "    \n",
    "\n",
    "        \n",
    "pendulum_swingup_example(min_time=True, animate=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Underactuated Robotics - The Simple Pendulum.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3.6.9 64-bit",
   "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.6.9-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}