{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "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.7.6"
    },
    "colab": {
      "name": "vibrating_pendulum.ipynb",
      "provenance": []
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FjpF2FAWF-8s"
      },
      "source": [
        "# Pendulum with Vibrating Base"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U6J1YLh3F-85"
      },
      "source": [
        "## Notebook Setup \n",
        "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
        "- On Google's Colaboratory, 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.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
        "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
        "\n",
        "More details are available [here](http://underactuated.mit.edu/drake.html)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2M2F4dQ-F-87",
        "outputId": "7a1da59b-efbc-4be1-d786-02aba42681d0",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "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='0e0ecc5c032ced0261f716e4a7930fc86bf0f80a', drake_version='0.25.0', drake_build='releases')\n",
        "\n",
        "# Setup matplotlib.  \n",
        "from IPython import get_ipython\n",
        "if get_ipython() is not None: get_ipython().run_line_magic(\"matplotlib\", \"inline\")"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into '/opt/underactuated'...\n",
            "\n",
            "HEAD is now at 0e0ecc5 update requirements for colab.  hoping to clean up first cell.\n",
            "\n",
            "\n",
            "WARNING: apt does not have a stable CLI interface. Use with caution in scripts.\n",
            "\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2XoUBihWF-8-"
      },
      "source": [
        "# python libraries\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from IPython.display import HTML\n",
        "\n",
        "# pydrake imports\n",
        "from pydrake.all import (AddMultibodyPlantSceneGraph, DiagramBuilder, Parser,\n",
        "                         PlanarSceneGraphVisualizer, Simulator, VectorSystem,\n",
        "                         Multiplexer, MatrixGain, LogOutput, plot_system_graphviz)\n",
        "\n",
        "# underactuated imports\n",
        "from underactuated import FindResource, ManipulatorDynamics"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OrYm44UPF-9A"
      },
      "source": [
        "## Problem Description\n",
        "In this problem you will write the controller to make the pendulum with vibrating base spin at constant velocity.\n",
        "At the end of the notebook, you will be able to check your work in a simulation environment we set up for you.\n",
        "\n",
        "**These are the main steps of the exercise:**\n",
        "1. Construct the physical model of the vibrating pendulum. This is done automatically by \"parsing\" a `.urdf` (Unified Robot Description Format) file.\n",
        "2. Implement the controller you derived in the written part of this exercise. _This is the only piece of code you will need to write._\n",
        "3. Wire up the closed-loop block diagram: connect the controller output with the system input, the system output with the visualizer etc.\n",
        "4. Set up and run a simulation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GIRp_6nZF-9C"
      },
      "source": [
        "# parameters of the harmonic motion of the base\n",
        "# defined globally in the notebook\n",
        "h = 1.\n",
        "omega = np.pi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xeI3GOC-F-9E"
      },
      "source": [
        "## Parse the `.urdf`\n",
        "The first block of our diagram is the vibrating pendulum.\n",
        "No need to write its equations of motion by hand: all the parameters of the system are described in its `.urdf` file.\n",
        "Drake can directly parse this file, and construct a `MultibodyPlant` (i.e. the vibrating-pendulum block in our diagram).\n",
        "The `.urdf` file contains all the physical data of the system, the visualization parameters (shapes and colors of the bodies), etc.\n",
        "Its html-like syntax is very easy to understand, give it a try!\n",
        "\n",
        "Our robot has two bodies:\n",
        "1. The base. This moves on a 1D rail and oscillates according to the harmonic law $h \\sin (\\omega t)$.\n",
        "2. The pendulum. It is connected to the base through a pin. This is the body you will need to control.\n",
        "\n",
        "**Attention!** Since the robot has two bodies, it also has two configuration variables.\n",
        "When writing the controller, we will take care of the first (position of the base) and ensure that it oscillates as required.\n",
        "Then the problem will be reduced to the control of the pendulum only."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0kTVs7f9F-9F"
      },
      "source": [
        "# think of the builder as the construction site of our block diagram\n",
        "builder = DiagramBuilder()\n",
        "\n",
        "# instantiate the vibrating pendulum and the scene graph\n",
        "# the scene graph is a container for the geometries of all the physical systems in our diagram\n",
        "vibrating_pendulum, scene_graph = AddMultibodyPlantSceneGraph(\n",
        "    builder,\n",
        "    time_step=0. # discrete update period , set to zero since system is continuous\n",
        ")\n",
        "\n",
        "# parse the urdf and populate the vibrating pendulum\n",
        "urdf_path = FindResource('models/vibrating_pendulum.urdf')\n",
        "Parser(vibrating_pendulum).AddModelFromFile(urdf_path)\n",
        "vibrating_pendulum.Finalize()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pzwpcaBAF-9G"
      },
      "source": [
        "## Write the Controller\n",
        "In this section we define two controllers:\n",
        "1. An inner controller that makes the base oscillate with the harmonic motion. We wrote this for you.\n",
        "2. The outer controller to make the pendulum spin at constant velocity. You will write part of this.\n",
        "\n",
        "The final diagram will have the following structure:\n",
        "\n",
        "![figure](https://raw.githubusercontent.com/RussTedrake/underactuated/master/figures/exercises/vibrating_pendulum.jpg)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wEm6oLg5F-9I"
      },
      "source": [
        "# this controller enforces the harmonic motion to the base\n",
        "class InnerController(VectorSystem):\n",
        "\n",
        "    def __init__(self, vibrating_pendulum):\n",
        "        # 5 inputs: state of base + pendulum, pendulum torque\n",
        "        # 2 outputs: force on base + torque on pendulum\n",
        "        VectorSystem.__init__(self,5, 2)\n",
        "        self.vibrating_pendulum = vibrating_pendulum\n",
        "\n",
        "    def DoCalcVectorOutput(\n",
        "            self,\n",
        "            context,\n",
        "            controller_input, # state of base + pendulum, pendulum torque\n",
        "            controller_state, # unused input (static controller)\n",
        "            controller_output # force on base + torque on pendulum\n",
        "        ):\n",
        "        \n",
        "        # unpack state\n",
        "        q = controller_input[:2] # base position + pendulum angle\n",
        "        q_dot = controller_input[2:4] # time derivative of q\n",
        "        \n",
        "        # extract manipulator equations: M*a + Cv = tauG + B*u + tauExt\n",
        "        # (for this system B is the identity and the external forces tauExt are zero\n",
        "        # hence, for simplicity, we will just drop them from the code)\n",
        "        M, Cv, tauG, B, tauExt = ManipulatorDynamics(self.vibrating_pendulum, q, q_dot)\n",
        "        Minv = np.linalg.inv(M)\n",
        "        tau = tauG - Cv\n",
        "        \n",
        "        # desired acceleration of the base\n",
        "        # note that this depends on time\n",
        "        t = context.get_time()\n",
        "        a_base = - omega**2 * h * np.sin(omega * t)\n",
        "        \n",
        "        # cancel out the dynamics of the pendulum\n",
        "        # and enforce harmonic motion to the base\n",
        "        # (to fully explain these lines we would need a small math derivation,\n",
        "        # since this is not the goal of the exercise we skip it,\n",
        "        # if you want, you can try your own, it is not hard)\n",
        "        torque = controller_input[-1]\n",
        "        force = - tau[0] # cancel gravity, centrifugal, and Coriolis\n",
        "        force += - (tau[1] + torque) * Minv[0,1] / Minv[0,0] # cancel pendulum effects on the base\n",
        "        force += a_base / Minv[0,0] # enforce desired acceleration\n",
        "        \n",
        "        # control signal\n",
        "        controller_output[:] = [force, torque]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jGvca9vCF-9I"
      },
      "source": [
        "# this is the controller that makes the pendulum spin at constant velocity\n",
        "# you will write the control law in it in the next cell\n",
        "# by defining the function called \"pendulum_torque\"\n",
        "class OuterController(VectorSystem):\n",
        "    \n",
        "    def __init__(self, vibrating_pendulum, pendulum_torque):\n",
        "        # 2 inputs: pendulum state\n",
        "        # 1 output: torque on pendulum\n",
        "        VectorSystem.__init__(self, 2, 1)\n",
        "        self.vibrating_pendulum = vibrating_pendulum\n",
        "        self.pendulum_torque = pendulum_torque\n",
        "        \n",
        "    def DoCalcVectorOutput(\n",
        "            self,\n",
        "            context,\n",
        "            controller_input, # pendulum state\n",
        "            controller_state, # unused input (static controller)\n",
        "            controller_output # torque on pendulum\n",
        "        ):\n",
        "        \n",
        "        # unpack state\n",
        "        theta, theta_dot = controller_input\n",
        "        \n",
        "        # get pendulum parameters\n",
        "        # absolute values to make these parameters positive\n",
        "        pendulum = self.vibrating_pendulum.GetBodyByName('pendulum')\n",
        "        m = pendulum.default_mass()\n",
        "        g = np.abs(self.vibrating_pendulum.gravity_field().gravity_vector()[2])\n",
        "        l = np.abs(pendulum.default_com()[2])\n",
        "\n",
        "        # controller\n",
        "        t = context.get_time()\n",
        "        controller_output[:] = [self.pendulum_torque(m, g, l, theta, theta_dot, t)]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wrn0M1flF-9L"
      },
      "source": [
        "In the following cell, modify the function `pendulum_torque` to implement the control law you derived in the point (b) of the exercise.\n",
        "The function must return the torque to be applied to the pendulum (`int` or `float`).\n",
        "Currently, it just returns zero torque, and the pendulum oscillates freely.\n",
        "\n",
        "The parameters of this function are:\n",
        "- `m`: mass of the bob of the pendulum,\n",
        "- `g`: gravity acceleration ($>0$),\n",
        "- `l`: length of the pendulum rod,\n",
        "- `theta`: angle of the pendulum,\n",
        "- `theta_dot`: angular velocity of the pendulum,\n",
        "- `t`: time.\n",
        "\n",
        "**Very important:**\n",
        "To complete this assignment, you only need to work in the following cell."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "90aI8bv1F-9L"
      },
      "source": [
        "def pendulum_torque(m, g, l, theta, theta_dot, t):\n",
        "    torque = 0 # modify here\n",
        "    return torque"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ge5adoo4F-9M"
      },
      "source": [
        "## Wire up the Block Diagram\n",
        "Now it is time to construct the block diagram: connect the controllers to the system, the sensors to the controllers, etc.\n",
        "We aim to construct the diagram shown above, with the addition of a visualizer, which will be connected with the system state.\n",
        "\n",
        "**Troubleshooting:**\n",
        "Note that we already initialized the `builder` of the block diagram when parsing the `.urdf` file.\n",
        "Hence, by running the following cell multiple times, you would actually try to wire the blocks in the diagram more than once.\n",
        "This is not acceptable, and Drake will raise the error `RuntimeError: Input port is already wired.`\n",
        "If you wish to modify the next cell and rerun the program to see the effects of your modification, you must rerun the cell where the `builder` is initialized first (i.e. the cell with the line `builder = DiagramBuilder()`)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hrcWxoJSF-9N"
      },
      "source": [
        "# instantiate controllers\n",
        "inner_controller = builder.AddSystem(InnerController(vibrating_pendulum))\n",
        "outer_controller = builder.AddSystem(OuterController(vibrating_pendulum, pendulum_torque))\n",
        "\n",
        "# instantiate visualizer\n",
        "visualizer = builder.AddSystem(\n",
        "    PlanarSceneGraphVisualizer(scene_graph, xlim=[-2.5, 2.5], ylim=[-1.5, 2.5], show=False)\n",
        ")\n",
        "\n",
        "# logger that records the state trajectory during simulation\n",
        "logger = LogOutput(vibrating_pendulum.get_state_output_port(), builder)\n",
        "\n",
        "# mux block to squeeze the (base + pendulum) state and\n",
        "# the outer control signal in a single cable\n",
        "mux = builder.AddSystem(Multiplexer([4, 1]))\n",
        "\n",
        "# selector that extracts the pendulum state from the state of the base and the pendulum\n",
        "selector = builder.AddSystem(MatrixGain(np.array([\n",
        "    [0, 1, 0, 0], # selects the angle of the pendulum\n",
        "    [0, 0, 0, 1] # selects the angular velocity of the pendulum\n",
        "])))\n",
        "\n",
        "# (base + pendulum) state, outer control -> mux\n",
        "builder.Connect(vibrating_pendulum.get_state_output_port(), mux.get_input_port(0))\n",
        "builder.Connect(outer_controller.get_output_port(0), mux.get_input_port(1))\n",
        "\n",
        "# mux -> inner controller\n",
        "builder.Connect(mux.get_output_port(0), inner_controller.get_input_port(0))\n",
        "\n",
        "# (base + pendulum) state -> selector\n",
        "builder.Connect(vibrating_pendulum.get_state_output_port(), selector.get_input_port(0))\n",
        "\n",
        "# selector -> outer controller\n",
        "builder.Connect(selector.get_output_port(0), outer_controller.get_input_port(0))\n",
        "\n",
        "# inner controller -> system input\n",
        "builder.Connect(inner_controller.get_output_port(0), vibrating_pendulum.get_actuation_input_port())\n",
        "\n",
        "# scene graph (i.e. all the bodies in the diagram) -> visualizer\n",
        "builder.Connect(scene_graph.get_pose_bundle_output_port(), visualizer.get_input_port(0))\n",
        "\n",
        "# finalize block diagram\n",
        "diagram = builder.Build()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OFUgHmKdF-9N"
      },
      "source": [
        "When connecting all the blocks by hand, it is possible to do some mistakes.\n",
        "To double check your work, you can use the function `plot_system_graphviz`, which plots the overall block diagram you built.\n",
        "You can compare the automatically-generated block diagram with the one above"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QKzrawXPF-9P"
      },
      "source": [
        "# give names to the blocks (just to make the plot nicer)\n",
        "diagram.set_name('Block Diagram for the Control of the Vibrating Pendulum')\n",
        "vibrating_pendulum.set_name('Vibrating Pendulum')\n",
        "inner_controller.set_name('Inner Controller')\n",
        "outer_controller.set_name('Outer Controller')\n",
        "mux.set_name('Mux')\n",
        "selector.set_name('Pendulum-State Selector')\n",
        "visualizer.set_name('Visualizer')\n",
        "logger.set_name('Logger')\n",
        "scene_graph.set_name('Scene Graph')\n",
        "\n",
        "# plot current diagram\n",
        "plt.figure(figsize=(20, 10))\n",
        "plot_system_graphviz(diagram)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oXIU0xSvF-9P"
      },
      "source": [
        "## Closed-Loop Simulation\n",
        "Now we can finally simulate the system in closed loop with the controllers we wrote.\n",
        "In the meanwhile, we will also set up a \"video recording\" with which we will be able to playback the simulation.\n",
        "System trajectories will be stored in the `logger` and plotted in the last cell."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "etDHEpzsF-9Q"
      },
      "source": [
        "# initial state of the pendulum\n",
        "theta = 0.\n",
        "theta_dot = 0.\n",
        "\n",
        "# simulation time\n",
        "sim_time = 10"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y4nf-deZF-9Q"
      },
      "source": [
        "# start recording the video for the animation of the simulation\n",
        "visualizer.start_recording()\n",
        "\n",
        "# set up a simulator\n",
        "simulator = Simulator(diagram)\n",
        "simulator.set_publish_every_time_step(False)\n",
        "\n",
        "# set the initial conditions\n",
        "# do not change the initial conditions of the base\n",
        "# since they must agree with the harmonic motion h*sin(omega*t)\n",
        "context = simulator.get_mutable_context()\n",
        "context.SetTime(0.) # reset current time\n",
        "context.SetContinuousState((\n",
        "    0.,       # initial position of the base, DO NOT CHANGE!\n",
        "    theta,    # initial angle of the pendulum\n",
        "    h*omega,  # initial velocity of the base, DO NOT CHANGE!\n",
        "    theta_dot # initial angular velocity of the pendulum\n",
        "    ))\n",
        "\n",
        "# simulate from zero to sim_time\n",
        "simulator.Initialize()\n",
        "simulator.AdvanceTo(sim_time)\n",
        "\n",
        "# stop the video and build the animation\n",
        "visualizer.stop_recording()\n",
        "ani = visualizer.get_recording_as_animation()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xL3qRKFyF-9R"
      },
      "source": [
        "# play video\n",
        "HTML(ani.to_jshtml())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "77PdUJkYF-9R"
      },
      "source": [
        "We conclude by plotting the position of the base and the angular velocity of the pendulum as functions of time.\n",
        "If we did our job correctly,\n",
        "- the first should coincide with the desired position $h \\sin (\\omega t)$,\n",
        "- the second should coincide with the response of the first-order system $\\ddot \\theta = f(\\dot \\theta)$ you came up with in point (a) of the exercise."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qpLIpgJeF-9S"
      },
      "source": [
        "# base position as a function of time\n",
        "plt.plot(\n",
        "    logger.sample_times(),\n",
        "    logger.data()[0,:],\n",
        "    label='Base position (m)'\n",
        ")\n",
        "\n",
        "# pendulum angular velocity as a function of time\n",
        "plt.plot(\n",
        "    logger.sample_times(),\n",
        "    logger.data()[-1,:],\n",
        "    label='Pendulum angular velocity (rad/s)'\n",
        ")\n",
        "\n",
        "# misc plot settings\n",
        "plt.xlabel('Time (s)')\n",
        "plt.xlim(0, sim_time)\n",
        "plt.grid(True)\n",
        "plt.legend()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bvzkSdtVF-9T"
      },
      "source": [
        "## How Will this Notebook Be Graded?\n",
        "If you are enrolled in the class, this notebook will be graded using [Gradescope](https://www.gradescope.com).\n",
        "We will send you the details of how to access the course page in Gradescope by email.\n",
        "\n",
        "We will replicate your work by running your notebook and checking the relevant functions (pendulum_torque) for correct outputs.\n",
        "We will also test your work by looking at the response of the angular velocity $\\dot \\theta(t)$ from the plot above.\n",
        "You will get full score if *all* the following tests succeed:\n",
        "- The response $\\dot \\theta(t)$ is a nondecreasing function (a first order system, such as $\\ddot \\theta = f(\\dot \\theta)$, cannot oscillate).\n",
        "- The terminal velocity $\\dot \\theta (t=10\\text{ s})$ is less than $1.001 \\text{ rad/s}$.\n",
        "- The terminal velocity $\\dot \\theta (t=10\\text{ s})$ is greater than $0.99 \\text{ rad/s}$.\n",
        "\n",
        "While the first two conditions should always hold (you did things properly), to fulfill the third you might want to increase your controller gains!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UF_mQc70F-9T"
      },
      "source": [
        "## Autograding\n",
        "You can check your work by running the following cell:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LZi1haMeF-9U"
      },
      "source": [
        "from underactuated.exercises.pend.vibrating_pendulum.test_vibrating_pendulum import TestVibratingPendulum\n",
        "from underactuated.exercises.grader import Grader\n",
        "Grader.grade_output([TestVibratingPendulum], [locals()], 'results.json')\n",
        "Grader.print_test_results('results.json')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}