{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TKvYiJgnYExi"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://underactuated.csail.mit.edu/pend.html).  I recommend having both windows open, side-by-side!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A4QOaw_zYLfI"
   },
   "outputs": [],
   "source": [
    "from copy import copy\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import mpld3\n",
    "import numpy as np\n",
    "from IPython.display import display\n",
    "from pydrake.all import (DiagramBuilder, Linearize, LinearQuadraticRegulator,\n",
    "                         MeshcatVisualizer, RigidTransform, RotationMatrix,\n",
    "                         Saturation, SceneGraph, Simulator, StartMeshcat,\n",
    "                         VectorLogSink, VectorSystem, wrap_to)\n",
    "from pydrake.examples.pendulum import (PendulumGeometry, PendulumParams,\n",
    "                                       PendulumPlant)\n",
    "\n",
    "from underactuated import plot_2d_phase_portrait, running_as_notebook\n",
    "from underactuated.meshcat_cpp_utils import MeshcatSliders, interact\n",
    "\n",
    "if running_as_notebook:\n",
    "    mpld3.enable_notebook()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer (run this cell only once, each instance consumes a port)\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C1sdq2R88C16"
   },
   "source": [
    "# Dynamics of the Simple Pendulum\n",
    "\n",
    "I find it extremely useful to use simulation to get physical intuition about these systems.  Let's make sure we understand how the simple pendulum moves when it is exposed to a torque.\n",
    "\n",
    "For the duration of this notebook, we'll use the equations of motion $$ml^2 \\ddot\\theta + b\\dot\\theta + mgl \\sin\\theta = u,$$ where $u$ is our torque input.\n",
    "\n",
    "## Basic simulation\n",
    "\n",
    "The pendulum is a core example in Drake.  We could certainly load it from a .urdf file, but Drake offers a Pendulum implementation that makes it convenient to manipulate the parameters (and visualize the system with different parameters)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8wj7ZlyEw_AB"
   },
   "outputs": [],
   "source": [
    "def pendulum_simulation():\n",
    "    builder = DiagramBuilder()\n",
    "    pendulum = builder.AddSystem(PendulumPlant())\n",
    "\n",
    "    # Setup visualization\n",
    "    scene_graph = builder.AddSystem(SceneGraph())\n",
    "    PendulumGeometry.AddToBuilder(builder, pendulum.get_state_output_port(),\n",
    "                                  scene_graph)\n",
    "    MeshcatVisualizer.AddToBuilder(builder, scene_graph,\n",
    "                                                   meshcat)\n",
    "    meshcat.Delete()\n",
    "    meshcat.Set2dRenderMode(X_WC=RigidTransform(RotationMatrix.MakeZRotation(np.pi), [0, 1, 0]))\n",
    "\n",
    "    # Setup slider input\n",
    "    meshcat.AddSlider('u', min=-5, max=5, step=.1, value=0.0)\n",
    "    torque_system = builder.AddSystem(MeshcatSliders(meshcat,['u']))\n",
    "    builder.Connect(torque_system.get_output_port(), pendulum.get_input_port())\n",
    "\n",
    "    diagram = builder.Build()\n",
    "\n",
    "    # Set up a simulator to run this diagram\n",
    "    simulator = Simulator(diagram)\n",
    "    context = simulator.get_mutable_context()\n",
    "\n",
    "    meshcat.AddButton('Stop Simulation')\n",
    "\n",
    "    # Set the initial conditions\n",
    "    context.SetContinuousState([0.5, 0]) # theta, thetadot\n",
    "\n",
    "    if running_as_notebook:  # Then we're not just running as a test on CI.\n",
    "        simulator.set_target_realtime_rate(1.0)\n",
    "\n",
    "        print('Use the slider in the MeshCat controls to apply elbow torque.')\n",
    "        print(\"Press 'Stop Simulation' in MeshCat to continue.\")\n",
    "        while meshcat.GetButtonClicks('Stop Simulation') < 1:\n",
    "            simulator.AdvanceTo(simulator.get_context().get_time() + 1.0)\n",
    "    else:\n",
    "        simulator.AdvanceTo(0.1)\n",
    "\n",
    "    meshcat.DeleteAddedControls()\n",
    "\n",
    "pendulum_simulation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aNHdvSloKUSn"
   },
   "source": [
    "After running the cell above, you should use the \"Open Controls\" panel in meshcat to see the interactive slider that will control the torque.  Press the \"Stop Simulation\" button to regain control and continue in the notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zKmlDpfeKUSo"
   },
   "source": [
    "# Autapse\n",
    "\n",
    "The simplest recurrent neural network model.  $$\\dot{x} + x = \\tanh(wx + u)$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ftbxk7ddKUSo"
   },
   "outputs": [],
   "source": [
    "\n",
    "def autapse(x, w=1, u=0):\n",
    "    \"\"\"Args:\n",
    "    w is feedback weight\n",
    "    u is input\n",
    "    \"\"\"\n",
    "    return -x + np.tanh(w * x + u)\n",
    "\n",
    "Autapse = np.vectorize(autapse)\n",
    "xmax = 2.\n",
    "ymax = 1.\n",
    "x = np.arange(-xmax, xmax, 0.01)\n",
    "\n",
    "meshcat.Delete()\n",
    "meshcat.Set2dRenderMode(xmax=xmax, xmin=-xmax, ymin=-ymax, ymax=ymax)\n",
    "meshcat.SetProperty('/Grid', 'visible', True)\n",
    "meshcat.SetProperty('/Axes', 'visible', True)\n",
    "\n",
    "def update(w=1, u=0):\n",
    "    # TODO(russt): Visualize fixed points here, too.\n",
    "    meshcat.SetLine(\"autapse\", np.vstack([x,0*x, Autapse(x, w=w, u=u)]), line_width=4.0)\n",
    "\n",
    "interact(meshcat, update, w=(0,3,0.1), u=(-1.5,1.5,0.1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CvoRbIK1KUSq"
   },
   "source": [
    "# Recurrent neural network units: LSTM and JANET\n",
    "\n",
    "JANET (Just Another NETwork) is a simplied version of the famous Long Short-Term Memory (LSTM) model as described in https://arxiv.org/abs/1804.04849 .  Relative to the autapse model, it adds a multiplicative \"forget gate\":\n",
    "$$\\dot{x} + x = f (1-\\alpha) x + (1-f)\\tanh(wx + u),\\\\ f = \\sigma(w_f x + u_f).$$\n",
    "I've written it here in continuous time and also added a small \"leak term\", $\\alpha>0$.  $\\sigma()$ is the sigmoid function, with range (0,1).  When $f=0$ we have the autapse model, where \"memory\" is possible.  When $f=1$ we have $\\dot{x} = -\\alpha x$, which \"forgets\" by implementing a dynamics with pushing the activations back towards zero.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meshcat.DeleteAddedControls()\n",
    "\n",
    "def sigma(x):\n",
    "  return 1./(1+np.exp(-x))\n",
    "\n",
    "def janet(x, w, u, wf, uf, alpha):\n",
    "    f = sigma(wf*x + uf)\n",
    "    return -x + f*(1-alpha)*x + (1-f)*np.tanh(w * x + u)\n",
    "\n",
    "Janet = np.vectorize(janet)\n",
    "xmax = 2.\n",
    "ymax = 1.\n",
    "x = np.arange(-xmax, xmax, 0.01)\n",
    "\n",
    "meshcat.Delete()\n",
    "meshcat.Set2dRenderMode(xmax=xmax, xmin=-xmax, ymin=-ymax, ymax=ymax)\n",
    "meshcat.SetProperty('/Grid', 'visible', True)\n",
    "meshcat.SetProperty('/Axes', 'visible', True)\n",
    "\n",
    "def update(w=1, u=0, wf=0, uf=0, alpha=0.1):\n",
    "    # TODO(russt): Visualize fixed points here, too.\n",
    "    meshcat.SetLine(\"janet\", np.vstack([x,0*x, Janet(x, w=w, u=u, wf=wf, uf=uf, alpha=alpha)]), line_width=4.0)\n",
    "\n",
    "interact(meshcat, update, w=(0,3,0.1), u=(-1.5,1.5,0.1), wf=(-2, 2, 0.1), uf=(-5,5,0.1), alpha=(0,1, 0.02))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Energy Shaping Control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class EnergyShapingController(VectorSystem):\n",
    "\n",
    "    def __init__(self, pendulum):\n",
    "        VectorSystem.__init__(self, 2, 1)\n",
    "        self.pendulum = pendulum\n",
    "        self.pendulum_context = pendulum.CreateDefaultContext()\n",
    "        self.SetPendulumParams(PendulumParams())\n",
    "\n",
    "    def SetPendulumParams(self, params):\n",
    "        self.pendulum_context.get_mutable_numeric_parameter(0).SetFromVector(\n",
    "            params.CopyToVector())\n",
    "        self.pendulum_context.SetContinuousState([np.pi, 0])\n",
    "        self.desired_energy = self.pendulum.EvalPotentialEnergy(\n",
    "            self.pendulum_context)\n",
    "\n",
    "    def DoCalcVectorOutput(self, context, pendulum_state, unused, output):\n",
    "        self.pendulum_context.SetContinuousState(pendulum_state)\n",
    "        params = self.pendulum_context.get_numeric_parameter(0)\n",
    "        theta = pendulum_state[0]\n",
    "        thetadot = pendulum_state[1]\n",
    "        total_energy = self.pendulum.EvalPotentialEnergy(\n",
    "            self.pendulum_context) + self.pendulum.EvalKineticEnergy(\n",
    "                self.pendulum_context)\n",
    "        output[:] = (params.damping() * thetadot - .1 * thetadot *\n",
    "                     (total_energy - self.desired_energy))\n",
    "\n",
    "\n",
    "def PhasePlot(pendulum):\n",
    "    phase_plot = plt.figure()\n",
    "    ax = phase_plot.gca()\n",
    "    theta_lim = [-np.pi, 3. * np.pi]\n",
    "    ax.set_xlim(theta_lim)\n",
    "    ax.set_ylim(-10., 10.)\n",
    "\n",
    "    theta = np.linspace(theta_lim[0], theta_lim[1], 601)  # 4*k + 1\n",
    "    thetadot = np.zeros(theta.shape)\n",
    "    context = pendulum.CreateDefaultContext()\n",
    "    params = context.get_numeric_parameter(0)\n",
    "    context.SetContinuousState([np.pi, 0])\n",
    "    E_upright = pendulum.EvalPotentialEnergy(context)\n",
    "    E = [E_upright, .1 * E_upright, 1.5 * E_upright]\n",
    "    for e in E:\n",
    "        for i in range(theta.size):\n",
    "            v = ((e + params.mass() * params.gravity() * params.length() *\n",
    "                  np.cos(theta[i])) /\n",
    "                 (.5 * params.mass() * params.length() * params.length()))\n",
    "            if (v >= 0):\n",
    "                thetadot[i] = np.sqrt(v)\n",
    "            else:\n",
    "                thetadot[i] = float(\"nan\")\n",
    "        ax.plot(theta, thetadot, color=[.6, .6, .6])\n",
    "        ax.plot(theta, -thetadot, color=[.6, .6, .6])\n",
    "\n",
    "    return ax\n",
    "\n",
    "\n",
    "def energy_shaping_demo():\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    pendulum = builder.AddSystem(PendulumPlant())\n",
    "    ax = PhasePlot(pendulum)\n",
    "    saturation = builder.AddSystem(Saturation(min_value=[-3], max_value=[3]))\n",
    "    builder.Connect(saturation.get_output_port(0), pendulum.get_input_port(0))\n",
    "    controller = builder.AddSystem(EnergyShapingController(pendulum))\n",
    "    builder.Connect(pendulum.get_output_port(0), controller.get_input_port(0))\n",
    "    builder.Connect(controller.get_output_port(0), saturation.get_input_port(0))\n",
    "\n",
    "    logger = builder.AddSystem(VectorLogSink(2))\n",
    "    builder.Connect(pendulum.get_output_port(0), logger.get_input_port(0))\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    simulator = Simulator(diagram)\n",
    "    context = simulator.get_mutable_context()\n",
    "\n",
    "    for i in range(5):\n",
    "        context.SetTime(0.)\n",
    "        context.SetContinuousState(np.random.randn(2,))\n",
    "        simulator.Initialize()\n",
    "        simulator.AdvanceTo(4)\n",
    "        log = logger.FindLog(context)\n",
    "        ax.plot(log.data()[0, :], log.data()[1, :])\n",
    "        log.Clear()\n",
    "\n",
    "    display(mpld3.display())  \n",
    "\n",
    "energy_shaping_demo()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Swing-up and balance\n",
    "\n",
    "Now we will combine our simple energy shaping controller with a linear controller that stabilizes the upright fixed point once we get close enough.  We'll read more about this approach in the Acrobot and Cart-Pole notes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def BalancingLQR(pendulum):\n",
    "    context = pendulum.CreateDefaultContext()\n",
    "\n",
    "    pendulum.get_input_port(0).FixValue(context, [0])\n",
    "    context.SetContinuousState([np.pi, 0])\n",
    "\n",
    "    Q = np.diag((10., 1.))\n",
    "    R = [1]\n",
    "\n",
    "    linearized_pendulum = Linearize(pendulum, context)\n",
    "    (K, S) = LinearQuadraticRegulator(linearized_pendulum.A(),\n",
    "                                      linearized_pendulum.B(), Q, R)\n",
    "    return (K, S)\n",
    "\n",
    "\n",
    "class SwingUpAndBalanceController(VectorSystem):\n",
    "\n",
    "    def __init__(self, pendulum):\n",
    "        VectorSystem.__init__(self, 2, 1)\n",
    "        (self.K, self.S) = BalancingLQR(pendulum)\n",
    "        self.energy_shaping = EnergyShapingController(pendulum)\n",
    "        self.energy_shaping_context = self.energy_shaping.CreateDefaultContext()\n",
    "\n",
    "        # TODO(russt): Add a witness function to tell the simulator about the\n",
    "        # discontinuity when switching to LQR.\n",
    "\n",
    "    def DoCalcVectorOutput(self, context, pendulum_state, unused, output):\n",
    "        xbar = copy(pendulum_state)\n",
    "        xbar[0] = wrap_to(xbar[0], 0, 2. * np.pi) - np.pi\n",
    "\n",
    "        # If x'Sx <= 2, then use the LQR controller\n",
    "        if (xbar.dot(self.S.dot(xbar)) < 2.):\n",
    "            output[:] = -self.K.dot(xbar)\n",
    "        else:\n",
    "            self.energy_shaping.get_input_port(0).FixValue(self.energy_shaping_context, pendulum_state)\n",
    "            output[:] = self.energy_shaping.get_output_port(0).Eval(self.energy_shaping_context)\n",
    "\n",
    "def swing_up_and_balance_demo():\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    pendulum = builder.AddSystem(PendulumPlant())\n",
    "    ax = PhasePlot(pendulum)\n",
    "    saturation = builder.AddSystem(Saturation(min_value=[-3], max_value=[3]))\n",
    "    builder.Connect(saturation.get_output_port(0), pendulum.get_input_port(0))\n",
    "    controller = builder.AddSystem(SwingUpAndBalanceController(pendulum))\n",
    "    builder.Connect(pendulum.get_output_port(0), controller.get_input_port(0))\n",
    "    builder.Connect(controller.get_output_port(0), saturation.get_input_port(0))\n",
    "\n",
    "    logger = builder.AddSystem(VectorLogSink(2))\n",
    "    builder.Connect(pendulum.get_output_port(0), logger.get_input_port(0))\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    simulator = Simulator(diagram)\n",
    "    context = simulator.get_mutable_context()\n",
    "\n",
    "    for i in range(5):\n",
    "        context.SetTime(0.)\n",
    "        context.SetContinuousState(np.random.randn(2,))\n",
    "        simulator.Initialize()\n",
    "        simulator.AdvanceTo(4)\n",
    "        log = logger.FindLog(context)\n",
    "        ax.plot(log.data()[0, :], log.data()[1, :])\n",
    "        log.Clear()\n",
    "\n",
    "    ax.set_xlim(np.pi - 3., np.pi + 3.)\n",
    "    ax.set_ylim(-5., 5.)\n",
    "    display(mpld3.display())\n",
    "\n",
    "swing_up_and_balance_demo()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Unstable equilibrium point that attracts all trajectories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dynamics(x):\n",
    "    r = np.sqrt(x[0]**2 + x[1]**2)\n",
    "    theta = np.arctan2(x[1], x[0])\n",
    "    rdot = r*(1-r)\n",
    "    thetadot = np.sin(theta/2)**2\n",
    "    return rdot*np.cos(theta) - r*thetadot*np.sin(theta), rdot*np.sin(theta) + r*thetadot*np.cos(theta)\n",
    "\n",
    "plot_2d_phase_portrait(dynamics, x1lim=(-1.5, 1.5), x2lim=(-1.5, 1.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Underactuated Robotics - The Simple Pendulum.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}