{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "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/acrobot.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": {},
   "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='ffe2b28ed89637889c04405e5d7d2d98be3df5b6', drake_version='0.27.0', drake_build='release')\n",
    "\n",
    "server_args = []\n",
    "if 'google.colab' in sys.modules:\n",
    "  server_args = ['--ngrok_http_tunnel']\n",
    "# Start two meshcat server instances to use for the remainder of this notebook.\n",
    "from meshcat.servers.zmqserver import start_zmq_server_as_subprocess\n",
    "proc_planar, zmq_url_planar, web_url_planar = start_zmq_server_as_subprocess(server_args=server_args)\n",
    "proc, zmq_url, web_url = start_zmq_server_as_subprocess(server_args=server_args)\n",
    "\n",
    "import numpy as np\n",
    "from ipywidgets import FloatSlider, ToggleButton\n",
    "from IPython.display import display, SVG\n",
    "import pydot\n",
    "\n",
    "import pydrake.all\n",
    "from pydrake.all import (\n",
    "    DiagramBuilder, LinearQuadraticRegulator, Saturation, SceneGraph, Simulator, \n",
    "    WrapToSystem, AddMultibodyPlantSceneGraph, Parser\n",
    ")\n",
    "from pydrake.examples.acrobot import AcrobotPlant, AcrobotGeometry\n",
    "from pydrake.systems.jupyter_widgets import WidgetSystem\n",
    "from pydrake.common.containers import namedview\n",
    "from underactuated import FindResource\n",
    "from underactuated.jupyter import running_as_notebook\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dynamics of the Acrobot\n",
    "\n",
    "The acrobot is a core example in Drake.  We could certainly load it from a .urdf file, but Drake offers an Acrobot implementation that makes it convenient to manipulate the parameters (and visualize the system with different parameters)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "builder = DiagramBuilder()\n",
    "acrobot = builder.AddSystem(AcrobotPlant())\n",
    "\n",
    "# Setup visualization\n",
    "scene_graph = builder.AddSystem(SceneGraph())\n",
    "AcrobotGeometry.AddToBuilder(builder, acrobot.get_output_port(0), scene_graph)\n",
    "visualizer = pydrake.systems.meshcat_visualizer.ConnectMeshcatVisualizer(\n",
    "    builder, \n",
    "    scene_graph=scene_graph, \n",
    "    zmq_url=zmq_url_planar)\n",
    "visualizer.vis.delete()\n",
    "visualizer.set_planar_viewpoint(xmin=-4, xmax=4, ymin=-4, ymax=4)\n",
    "\n",
    "# Setup slider input\n",
    "slider = FloatSlider(value=0.0, min=-5., max=5., step=0.1, description='u', continuous_update=True)\n",
    "torque_system = builder.AddSystem(WidgetSystem([slider]))\n",
    "builder.Connect(torque_system.get_output_port(0), acrobot.get_input_port(0))\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",
    "stop_button = ToggleButton(value=False, description='Stop Simulation')\n",
    "display(stop_button)\n",
    "\n",
    "# Set the initial conditions\n",
    "context.SetContinuousState([1., 0, 0, 0]) # theta1, theta2, theta1dot, theta2dot\n",
    "context.SetTime(0.0)\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",
    "    while not stop_button.value:\n",
    "        simulator.AdvanceTo(simulator.get_context().get_time() + 1.0)\n",
    "    stop_button.value = False\n",
    "else:\n",
    "    simulator.AdvanceTo(0.1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dynamics of the Cart-Pole"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "builder = DiagramBuilder()\n",
    "plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
    "file_name = FindResource(\"models/cartpole.urdf\")\n",
    "Parser(plant).AddModelFromFile(file_name)\n",
    "plant.Finalize()\n",
    "\n",
    "# Setup visualization\n",
    "visualizer = pydrake.systems.meshcat_visualizer.ConnectMeshcatVisualizer(\n",
    "    builder, \n",
    "    scene_graph=scene_graph, \n",
    "    zmq_url=zmq_url_planar)\n",
    "visualizer.vis.delete()\n",
    "visualizer.set_planar_viewpoint(xmin=-2.5, xmax=2.5, ymin=-1.0, ymax=2.5)\n",
    "\n",
    "# Setup slider input\n",
    "slider = FloatSlider(value=0.0, min=-30., max=30., step=0.1, description='u', continuous_update=True)\n",
    "force_system = builder.AddSystem(WidgetSystem([slider]))\n",
    "builder.Connect(force_system.get_output_port(0), plant.get_actuation_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",
    "stop_button = ToggleButton(value=False, description='Stop Simulation')\n",
    "display(stop_button)\n",
    "\n",
    "# Set the initial conditions\n",
    "context.SetContinuousState([0, 1., 0, 0]) # x, theta, xdot, thetadot\n",
    "context.SetTime(0.0)\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",
    "    while not stop_button.value:\n",
    "        simulator.AdvanceTo(simulator.get_context().get_time() + 1.0)\n",
    "    stop_button.value = False\n",
    "else:\n",
    "    simulator.AdvanceTo(0.1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LQR for the Acrobot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.examples.acrobot import (AcrobotInput, AcrobotGeometry, \n",
    "                                      AcrobotPlant, AcrobotState)\n",
    "\n",
    "def acrobot_balancing_example():\n",
    "    def UprightState():\n",
    "        state = AcrobotState()\n",
    "        state.set_theta1(np.pi)\n",
    "        state.set_theta2(0.)\n",
    "        state.set_theta1dot(0.)\n",
    "        state.set_theta2dot(0.)\n",
    "        return state\n",
    "\n",
    "\n",
    "    def BalancingLQR():\n",
    "        # Design an LQR controller for stabilizing the Acrobot around the upright.\n",
    "        # Returns a (static) AffineSystem that implements the controller (in\n",
    "        # the original AcrobotState coordinates).\n",
    "\n",
    "        acrobot = AcrobotPlant()\n",
    "        context = acrobot.CreateDefaultContext()\n",
    "\n",
    "        input = AcrobotInput()\n",
    "        input.set_tau(0.)\n",
    "        acrobot.get_input_port(0).FixValue(context, input)\n",
    "\n",
    "        context.get_mutable_continuous_state_vector()\\\n",
    "            .SetFromVector(UprightState().CopyToVector())\n",
    "\n",
    "        Q = np.diag((10., 10., 1., 1.))\n",
    "        R = [1]\n",
    "\n",
    "        return LinearQuadraticRegulator(acrobot, context, Q, R)\n",
    "\n",
    "\n",
    "    builder = DiagramBuilder()\n",
    "    acrobot = builder.AddSystem(AcrobotPlant())\n",
    "\n",
    "    saturation = builder.AddSystem(Saturation(min_value=[-10], max_value=[10]))\n",
    "    builder.Connect(saturation.get_output_port(0), acrobot.get_input_port(0))\n",
    "    wrapangles = WrapToSystem(4)\n",
    "    wrapangles.set_interval(0, 0, 2. * np.pi)\n",
    "    wrapangles.set_interval(1, -np.pi, np.pi)\n",
    "    wrapto = builder.AddSystem(wrapangles)\n",
    "    builder.Connect(acrobot.get_output_port(0), wrapto.get_input_port(0))\n",
    "    controller = builder.AddSystem(BalancingLQR())\n",
    "    builder.Connect(wrapto.get_output_port(0), controller.get_input_port(0))\n",
    "    builder.Connect(controller.get_output_port(0), saturation.get_input_port(0))\n",
    "\n",
    "    # Setup visualization\n",
    "    scene_graph = builder.AddSystem(SceneGraph())\n",
    "    AcrobotGeometry.AddToBuilder(builder, acrobot.get_output_port(0), scene_graph)\n",
    "    visualizer = pydrake.systems.meshcat_visualizer.ConnectMeshcatVisualizer(\n",
    "        builder, \n",
    "        scene_graph=scene_graph, \n",
    "        zmq_url=zmq_url_planar)\n",
    "    visualizer.vis.delete()\n",
    "    visualizer.set_planar_viewpoint(xmin=-3.0, xmax=3.0, ymin=-3.0, ymax=4.0)\n",
    "\n",
    "    diagram = builder.Build()\n",
    "\n",
    "    # Set up a simulator to run this diagram\n",
    "    simulator = Simulator(diagram)\n",
    "    simulator.set_target_realtime_rate(1.0)\n",
    "    context = simulator.get_mutable_context()\n",
    "\n",
    "    # Simulate\n",
    "    duration = 4.0 if running_as_notebook else 0.1 # sets a shorter duration during testing\n",
    "    for i in range(5):\n",
    "        context.SetTime(0.)\n",
    "        context.SetContinuousState(UprightState().CopyToVector() +\n",
    "                                0.05 * np.random.randn(4,))\n",
    "        simulator.Initialize()\n",
    "        simulator.AdvanceTo(duration)\n",
    "\n",
    "acrobot_balancing_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LQR for the Cart-Pole"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def cartpole_balancing_example():\n",
    "    def UprightState():\n",
    "        state = (0, np.pi, 0, 0)\n",
    "        return state\n",
    "\n",
    "\n",
    "    def BalancingLQR(plant):\n",
    "        # Design an LQR controller for stabilizing the CartPole around the upright.\n",
    "        # Returns a (static) AffineSystem that implements the controller (in\n",
    "        # the original CartPole coordinates).\n",
    "\n",
    "        context = plant.CreateDefaultContext()\n",
    "        plant.get_actuation_input_port().FixValue(context, [0])\n",
    "\n",
    "        context.get_mutable_continuous_state_vector().SetFromVector(UprightState())\n",
    "\n",
    "        Q = np.diag((10., 10., 1., 1.))\n",
    "        R = [1]\n",
    "\n",
    "        # MultibodyPlant has many (optional) input ports, so we must pass the\n",
    "        # input_port_index to LQR.\n",
    "        return LinearQuadraticRegulator(\n",
    "            plant,\n",
    "            context,\n",
    "            Q,\n",
    "            R,\n",
    "            input_port_index=plant.get_actuation_input_port().get_index())\n",
    "\n",
    "\n",
    "    builder = DiagramBuilder()\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
    "    file_name = FindResource(\"models/cartpole.urdf\")\n",
    "    Parser(plant).AddModelFromFile(file_name)\n",
    "    plant.Finalize()\n",
    "\n",
    "    controller = builder.AddSystem(BalancingLQR(plant))\n",
    "    builder.Connect(plant.get_state_output_port(), controller.get_input_port(0))\n",
    "    builder.Connect(controller.get_output_port(0),\n",
    "                    plant.get_actuation_input_port())\n",
    "\n",
    "    # Setup visualization\n",
    "    visualizer = pydrake.systems.meshcat_visualizer.ConnectMeshcatVisualizer(\n",
    "        builder, \n",
    "        scene_graph=scene_graph, \n",
    "        zmq_url=zmq_url_planar)\n",
    "    visualizer.vis.delete()\n",
    "    visualizer.set_planar_viewpoint(xmin=-2.5, xmax=2.5, ymin=-1.0, ymax=2.5)\n",
    "\n",
    "    diagram = builder.Build()\n",
    "\n",
    "    # Set up a simulator to run this diagram\n",
    "    simulator = Simulator(diagram)\n",
    "    simulator.set_target_realtime_rate(1.0)\n",
    "    context = simulator.get_mutable_context()\n",
    "\n",
    "    # Simulate\n",
    "    duration = 8.0 if running_as_notebook else 0.1 # sets a shorter duration during testing\n",
    "    for i in range(5):\n",
    "        context.SetTime(0.)\n",
    "        context.SetContinuousState(UprightState() + 0.1 * np.random.randn(4,))\n",
    "        simulator.Initialize()\n",
    "        simulator.AdvanceTo(duration)\n",
    "\n",
    "cartpole_balancing_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LQR for the Planar Quadrotor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Consider switching this to meshcat, too.\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from underactuated.quadrotor2d import Quadrotor2D, Quadrotor2DVisualizer\n",
    "from underactuated.jupyter import AdvanceToAndVisualize, SetupMatplotlibBackend\n",
    "\n",
    "\n",
    "def planar_quadrotor_example():\n",
    "    plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "    def QuadrotorLQR(plant):\n",
    "        context = plant.CreateDefaultContext()\n",
    "        context.SetContinuousState(np.zeros([6, 1]))\n",
    "        plant.get_input_port(0).FixValue(context, plant.mass * plant.gravity / 2. * np.array([1, 1]))\n",
    "\n",
    "        Q = np.diag([10, 10, 10, 1, 1, (plant.length / 2. / np.pi)])\n",
    "        R = np.array([[0.1, 0.05], [0.05, 0.1]])\n",
    "\n",
    "        return LinearQuadraticRegulator(plant, context, Q, R)\n",
    "\n",
    "\n",
    "    builder = DiagramBuilder()\n",
    "    plant = builder.AddSystem(Quadrotor2D())\n",
    "\n",
    "    controller = builder.AddSystem(QuadrotorLQR(plant))\n",
    "    builder.Connect(controller.get_output_port(0), plant.get_input_port(0))\n",
    "    builder.Connect(plant.get_output_port(0), controller.get_input_port(0))\n",
    "\n",
    "    # Setup visualization\n",
    "    visualizer = builder.AddSystem(Quadrotor2DVisualizer(show=plt_is_interactive))\n",
    "    builder.Connect(plant.get_output_port(0), visualizer.get_input_port(0))\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",
    "    # Simulate\n",
    "    duration = 4.0 if running_as_notebook else 0.1 # sets a shorter duration during testing\n",
    "    if not visualizer._show:\n",
    "        visualizer.start_recording()\n",
    "    for i in range(5):\n",
    "        context.SetTime(0.)\n",
    "        context.SetContinuousState(np.random.randn(6,))\n",
    "        simulator.Initialize()\n",
    "        if i<4:\n",
    "            simulator.AdvanceTo(duration)\n",
    "        else:\n",
    "            AdvanceToAndVisualize(simulator, visualizer, duration)\n",
    "\n",
    "planar_quadrotor_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LQR for the 3D Quadrotor\n",
    "\n",
    "This example requires a 3D visualizer, and does not work [yet](https://github.com/rdeits/meshcat/issues/67#issuecomment-632453885) on Colab (but should work on Binder).  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.examples.quadrotor import (QuadrotorPlant,\n",
    "                                        StabilizingLQRController,\n",
    "                                        QuadrotorGeometry)\n",
    "\n",
    "def quadrotor_example():\n",
    "\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant = builder.AddSystem(QuadrotorPlant())\n",
    "\n",
    "    controller = builder.AddSystem(StabilizingLQRController(plant, [0, 0, 1]))\n",
    "    builder.Connect(controller.get_output_port(0), plant.get_input_port(0))\n",
    "    builder.Connect(plant.get_output_port(0), controller.get_input_port(0))\n",
    "\n",
    "    # Set up visualization in MeshCat\n",
    "    scene_graph = builder.AddSystem(SceneGraph())\n",
    "    QuadrotorGeometry.AddToBuilder(builder, plant.get_output_port(0), scene_graph)\n",
    "    visualizer = pydrake.systems.meshcat_visualizer.ConnectMeshcatVisualizer(\n",
    "        builder, \n",
    "        scene_graph=scene_graph, \n",
    "        zmq_url=zmq_url)\n",
    "    visualizer.vis.delete()\n",
    "    # end setup for visualization\n",
    "\n",
    "    diagram = builder.Build()\n",
    "\n",
    "    # Set up a simulator to run this diagram\n",
    "    simulator = Simulator(diagram)\n",
    "    simulator.set_target_realtime_rate(1.0)\n",
    "    context = simulator.get_mutable_context()\n",
    "\n",
    "    # Simulate\n",
    "    duration = 4.0 if running_as_notebook else 0.1 # sets a shorter duration during testing\n",
    "    for i in range(5):\n",
    "        context.SetTime(0.)\n",
    "        context.SetContinuousState(0.5*np.random.randn(12,))\n",
    "        simulator.Initialize()\n",
    "        simulator.AdvanceTo(duration)\n",
    "\n",
    "quadrotor_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Balancing a 2D Segway (aka \"ballbot\")\n",
    "\n",
    "I'll also use this example to illustrate a little more about how one might build a new model in drake (using multibody plant).  I've included the URDF in this notebook directly as a string so that you can see it and easily play with it.\n",
    "\n",
    "https://en.wikipedia.org/wiki/Ballbot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "ballbot_urdf = \"\"\"\n",
    "<?xml version=\"1.0\"?>\n",
    "\n",
    "<robot xmlns=\"http://drake.mit.edu\"\n",
    " xmlns:xsi=\"http://www.w3.org/2001/XMLSchema-instance\"\n",
    " name=\"BallBot\">\n",
    "\n",
    "  <link name=\"ground\">\n",
    "    <visual>\n",
    "      <origin xyz=\"0 0 -5\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <box size=\"1000 1000 10\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\"0.93 .74 .4 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "  </link>\n",
    "\n",
    "  <joint name=\"ground_weld\" type=\"fixed\">\n",
    "    <parent link=\"world\" />\n",
    "    <child link=\"ground\" />\n",
    "  </joint>\n",
    "\n",
    "  <link name=\"ball\">\n",
    "    <inertial>\n",
    "      <origin xyz=\"0 0 0\" rpy=\"0 0 0\" />\n",
    "      <mass value=\"5\" />\n",
    "      <inertia ixx=\".02\" ixy=\"0\" ixz=\"0\" iyy=\"0.02\" iyz=\"0\" izz=\"0.02\" />\n",
    "    </inertial>\n",
    "    <visual>\n",
    "      <origin xyz=\"0 0 0\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <sphere radius=\".1\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\"0.25 0.52 0.96 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "  </link>\n",
    "\n",
    "  <link name=\"bot\">\n",
    "    <inertial>\n",
    "      <origin xyz=\"0 0 .05\" rpy=\"0 0 0\" />\n",
    "      <mass value=\"4\" />\n",
    "      <inertia ixx=\"0.018\" ixy=\"0\" ixz=\"0\" iyy=\"0.018\" iyz=\"0\" izz=\"0.0288\" />\n",
    "    </inertial>\n",
    "    <visual>\n",
    "      <origin xyz=\"0 0 .05\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "         <cylinder length=\".1\" radius=\".12\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\".61 .63 .67 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "  </link>\n",
    "  \n",
    "  <joint name=\"x\" type=\"prismatic\">\n",
    "    <parent link=\"world\" />\n",
    "    <child link=\"ball\" />\n",
    "    <origin xyz=\"0 0 .1\" />\n",
    "    <axis xyz=\"1 0 0\" />\n",
    "    <dynamics damping=\"0.1\" />\n",
    "  </joint>\n",
    "\n",
    "  <joint name=\"theta\" type=\"continuous\">\n",
    "    <parent link=\"ball\" />\n",
    "    <child link=\"bot\" />\n",
    "    <origin xyz=\"0 0 0\" />\n",
    "    <axis xyz=\"0 1 0\" />\n",
    "    <dynamics damping=\"0.1\" />\n",
    "  </joint>\n",
    "\n",
    "  <transmission type=\"SimpleTransmission\" name=\"ball_torque\">\n",
    "    <actuator name=\"torque\" />\n",
    "    <joint name=\"theta\" />\n",
    "    <mechanicalReduction>1</mechanicalReduction>\n",
    "  </transmission>\n",
    "\n",
    "  <transmission type=\"SimpleTransmission\" name=\"ball_force\">\n",
    "    <actuator name=\"force\" />\n",
    "    <joint name=\"x\" />\n",
    "    <mechanicalReduction>.1</mechanicalReduction>\n",
    "  </transmission>\n",
    "</robot>\n",
    "\"\"\"\n",
    "\n",
    "def ballbot_example():\n",
    "    def UprightState():\n",
    "        state = (0, 0, 0, 0)\n",
    "        return state\n",
    "\n",
    "\n",
    "    def BalancingLQR(plant):\n",
    "        # Design an LQR controller for stabilizing the CartPole around the upright.\n",
    "        # Returns a (static) AffineSystem that implements the controller (in\n",
    "        # the original CartPole coordinates).\n",
    "\n",
    "        context = plant.CreateDefaultContext()\n",
    "        plant.get_input_port().FixValue(context, [0])\n",
    "\n",
    "        context.get_mutable_continuous_state_vector().SetFromVector(UprightState())\n",
    "\n",
    "        Q = np.diag((10., 10., 1., 1.))\n",
    "        R = [1]\n",
    "\n",
    "        # MultibodyPlant has many (optional) input ports, so we must pass the\n",
    "        # input_port_index to LQR.\n",
    "        return LinearQuadraticRegulator(\n",
    "            plant,\n",
    "            context,\n",
    "            Q,\n",
    "            R)\n",
    "\n",
    "\n",
    "    def MakeBallBot():\n",
    "        builder = DiagramBuilder()\n",
    "        plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
    "        Parser(plant).AddModelFromString(ballbot_urdf, \"urdf\")\n",
    "        plant.Finalize()\n",
    "        # Applying a torque between the ball and the bot torques the bot, but also causes a \n",
    "        # force at the ground (the radius of the ball is entered as a gear reduction in the transmission).\n",
    "        B = np.array([[1],[1]])\n",
    "        gain = builder.AddSystem(pydrake.systems.primitives.MatrixGain(B))\n",
    "        gain.set_name(\"Actuator Mapping\")\n",
    "        builder.Connect(gain.get_output_port(), plant.get_actuation_input_port())\n",
    "        builder.ExportInput(gain.get_input_port(), \"torque\")\n",
    "        builder.ExportOutput(plant.get_state_output_port(), \"continuous_state\")\n",
    "        builder.ExportOutput(scene_graph.get_query_output_port(), \"query\")\n",
    "        return builder.Build()\n",
    "\n",
    "    builder = DiagramBuilder()\n",
    "    ballbot = builder.AddSystem(MakeBallBot())\n",
    "\n",
    "    controller = builder.AddSystem(BalancingLQR(ballbot))\n",
    "    controller.set_name(\"LQR Controller\")\n",
    "    builder.Connect(ballbot.GetOutputPort(\"continuous_state\"), controller.get_input_port())\n",
    "    builder.Connect(controller.get_output_port(),\n",
    "                    ballbot.get_input_port())\n",
    "\n",
    "    # Setup visualization\n",
    "    visualizer = pydrake.systems.meshcat_visualizer.ConnectMeshcatVisualizer(\n",
    "        builder, \n",
    "        output_port=ballbot.GetOutputPort(\"query\"),\n",
    "        zmq_url=zmq_url)\n",
    "    visualizer.vis.delete()\n",
    "    visualizer.vis['/Background'].set_property('visible', False)\n",
    "    #visualizer.set_planar_viewpoint(xmin=-2.5, xmax=2.5, ymin=-1.0, ymax=2.5)\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    \n",
    "    # For reference, let's draw the diagram we've assembled:\n",
    "    display(SVG(pydot.graph_from_dot_data(diagram.GetGraphvizString())[0].create_svg()))\n",
    "\n",
    "    # Set up a simulator to run this diagram\n",
    "    simulator = Simulator(diagram)\n",
    "    simulator.set_target_realtime_rate(1.0)\n",
    "    context = simulator.get_mutable_context()\n",
    "    context.SetContinuousState(UprightState() + 0.0 * np.random.randn(4,) + np.array([1, 0, 0 ,0]))\n",
    "\n",
    "    # Simulate\n",
    "    duration = 10.0 if running_as_notebook else 0.1 # sets a shorter duration during testing\n",
    "    simulator.AdvanceTo(duration)\n",
    "\n",
    "ballbot_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But there is subtle point to be made here.  I've modeled the wheels as a prismatic joint at the base.  Let's instead model the robot with a floating joint, and using collisions elements to simulate the interaction with the robot and the ground..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ballbot_floating_base_urdf = \"\"\"\n",
    "<?xml version=\"1.0\"?>\n",
    "\n",
    "<robot xmlns=\"http://drake.mit.edu\"\n",
    " xmlns:xsi=\"http://www.w3.org/2001/XMLSchema-instance\"\n",
    " name=\"BallBot\">\n",
    "\n",
    "  <link name=\"ground\">\n",
    "    <visual>\n",
    "      <origin xyz=\"0 0 -5\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <box size=\"1000 1000 10\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\"0.93 .74 .4 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "    <collision>\n",
    "      <origin xyz=\"0 0 -5\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <box size=\"1000 1000 10\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\"0.93 .74 .4 1\" />\n",
    "      </material>\n",
    "    </collision>\n",
    "  </link>\n",
    "\n",
    "  <joint name=\"ground_weld\" type=\"fixed\">\n",
    "    <parent link=\"world\" />\n",
    "    <child link=\"ground\" />\n",
    "  </joint>\n",
    "\n",
    "  <link name=\"ball\">\n",
    "    <inertial>\n",
    "      <origin xyz=\"0 0 0\" rpy=\"0 0 0\" />\n",
    "      <mass value=\"5\" />\n",
    "      <inertia ixx=\".02\" ixy=\"0\" ixz=\"0\" iyy=\"0.02\" iyz=\"0\" izz=\"0.02\" />\n",
    "    </inertial>\n",
    "    <visual>\n",
    "      <origin xyz=\"0 0 0\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <sphere radius=\".1\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\"0.25 0.52 0.96 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "    <collision>\n",
    "      <origin xyz=\"0 0 0\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <sphere radius=\".1\" />\n",
    "      </geometry>\n",
    "    </collision>\n",
    "  </link>\n",
    "\n",
    "  <link name=\"bot\">\n",
    "    <inertial>\n",
    "      <origin xyz=\"0 0 .05\" rpy=\"0 0 0\" />\n",
    "      <mass value=\"4\" />\n",
    "      <inertia ixx=\"0.018\" ixy=\"0\" ixz=\"0\" iyy=\"0.018\" iyz=\"0\" izz=\"0.0288\" />\n",
    "    </inertial>\n",
    "    <collision>\n",
    "      <origin xyz=\"0 0 .05\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "         <cylinder length=\".1\" radius=\".12\" />\n",
    "      </geometry>\n",
    "    </collision>\n",
    "    <visual>\n",
    "      <origin xyz=\"0 0 .05\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "         <cylinder length=\".1\" radius=\".12\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\".61 .63 .67 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "  </link>\n",
    "  \n",
    "  <drake:joint name=\"floating_base\" type=\"planar\">\n",
    "    <parent link=\"world\" />\n",
    "    <child link=\"ball\" />\n",
    "    <origin rpy=\"1.57 0 0\" xyz=\"0 0 .1\" />\n",
    "  </drake:joint>\n",
    "\n",
    "  <joint name=\"theta2\" type=\"continuous\">\n",
    "    <parent link=\"ball\" />\n",
    "    <child link=\"bot\" />\n",
    "    <origin rpy=\"-1.57 0 0\" xyz=\"0 0 0\" />\n",
    "    <axis xyz=\"0 1 0\" />\n",
    "    <dynamics damping=\"0.1\" />\n",
    "  </joint>\n",
    "\n",
    "  <transmission type=\"SimpleTransmission\" name=\"ball_torque\">\n",
    "    <actuator name=\"torque\" />\n",
    "    <joint name=\"theta2\" />\n",
    "    <mechanicalReduction>1</mechanicalReduction>\n",
    "  </transmission>\n",
    "\n",
    "</robot>\n",
    "\"\"\"\n",
    "\n",
    "def ballbot_floating_base_example():\n",
    "    builder = DiagramBuilder()\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    Parser(plant).AddModelFromString(ballbot_floating_base_urdf, \"urdf\")\n",
    "    plant.Finalize()\n",
    "\n",
    "    # Just use zero instead of a controller to start\n",
    "    command = builder.AddSystem(pydrake.systems.primitives.ConstantVectorSource([0.0]))\n",
    "    builder.Connect(command.get_output_port(), plant.get_actuation_input_port())\n",
    "\n",
    "    # Setup visualization\n",
    "    visualizer = pydrake.systems.meshcat_visualizer.ConnectMeshcatVisualizer(\n",
    "        builder, \n",
    "        scene_graph=scene_graph, \n",
    "        zmq_url=zmq_url)\n",
    "    visualizer.vis['/Background'].set_property('visible',False)\n",
    "    visualizer.vis.delete()\n",
    "    #visualizer.set_planar_viewpoint(xmin=-2.5, xmax=2.5, ymin=-1.0, ymax=2.5)\n",
    "\n",
    "    diagram = builder.Build()\n",
    "\n",
    "    # Set up a simulator to run this diagram\n",
    "    simulator = Simulator(diagram)\n",
    "    simulator.set_target_realtime_rate(1.0)\n",
    "    context = simulator.get_mutable_context()\n",
    "\n",
    "    State = namedview(\n",
    "    \"State\", [\"x\", \"z\", \"theta1\", \"theta2\", \"xdot\", \"zdot\", \"theta1dot\", \"theta2dot\"])\n",
    "    x0 = State(np.zeros(8))\n",
    "    x0.z = .15\n",
    "    x0.theta1 = 0.05\n",
    "    plant_context = plant.GetMyContextFromRoot(context)\n",
    "    plant.SetPositionsAndVelocities(plant_context, x0[:])\n",
    "\n",
    "    # Simulate\n",
    "    duration = 3.0 if running_as_notebook else 0.1 # sets a shorter duration during testing\n",
    "    simulator.AdvanceTo(duration)\n",
    "\n",
    "ballbot_floating_base_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's the real test... you should be able to run the controller designed for the simple model on the floating base version of the ballbot.  Give it a shot!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ballbot_floating_base_linearization_example():\n",
    "    builder = DiagramBuilder()\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    Parser(plant).AddModelFromString(ballbot_floating_base_urdf, \"urdf\")\n",
    "    plant.Finalize()\n",
    "    builder.ExportInput(plant.get_actuation_input_port(), \"command\")\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    context = diagram.CreateDefaultContext()\n",
    "    diagram.get_input_port().FixValue(context, [0])\n",
    "\n",
    "    # TODO: Would be better to solve for the fixed point and then call Linearize directly.\n",
    "    sys = pydrake.systems.primitives.FirstOrderTaylorApproximation(diagram, context)\n",
    "    print(sys.A())\n",
    "    print(sys.B())\n",
    "    print(sys.f0())\n",
    "\n",
    "    # TODO: implement discrete-time controllability matrix.\n",
    "\n",
    "ballbot_floating_base_linearization_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Differential Flatness"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import math\n",
    "import numpy as np\n",
    "\n",
    "from pydrake.solvers.mathematicalprogram import MathematicalProgram, Solve\n",
    "from pydrake.trajectories import PiecewisePolynomial\n",
    "\n",
    "from underactuated.quadrotor2d import Quadrotor2DVisualizer\n",
    "\n",
    "# TODO(russt): Use drake.trajectories.PiecewisePolynomialTrajectory\n",
    "#  instead (currently missing python bindings for the required constructor),\n",
    "#  or port this class to C++.\n",
    "class PPTrajectory():\n",
    "\n",
    "    def __init__(self, sample_times, num_vars, degree, continuity_degree):\n",
    "        self.sample_times = sample_times\n",
    "        self.n = num_vars\n",
    "        self.degree = degree\n",
    "\n",
    "        self.prog = MathematicalProgram()\n",
    "        self.coeffs = []\n",
    "        for i in range(len(sample_times)):\n",
    "            self.coeffs.append(\n",
    "                self.prog.NewContinuousVariables(num_vars, degree + 1, \"C\"))\n",
    "        self.result = None\n",
    "\n",
    "        # Add continuity constraints\n",
    "        for s in range(len(sample_times) - 1):\n",
    "            trel = sample_times[s + 1] - sample_times[s]\n",
    "            coeffs = self.coeffs[s]\n",
    "            for var in range(self.n):\n",
    "                for deg in range(continuity_degree + 1):\n",
    "                    # Don't use eval here, because I want left and right\n",
    "                    # values of the same time\n",
    "                    left_val = 0\n",
    "                    for d in range(deg, self.degree + 1):\n",
    "                        left_val += (coeffs[var, d] * np.power(trel, d - deg) *\n",
    "                                     math.factorial(d) /\n",
    "                                     math.factorial(d - deg))\n",
    "                    right_val = (self.coeffs[s + 1][var, deg] *\n",
    "                                 math.factorial(deg))\n",
    "                    self.prog.AddLinearConstraint(left_val == right_val)\n",
    "\n",
    "        # Add cost to minimize highest order terms\n",
    "        for s in range(len(sample_times) - 1):\n",
    "            self.prog.AddQuadraticCost(np.eye(num_vars), np.zeros(\n",
    "                (num_vars, 1)), self.coeffs[s][:, -1])\n",
    "\n",
    "    def eval(self, t, derivative_order=0):\n",
    "        if derivative_order > self.degree:\n",
    "            return 0\n",
    "\n",
    "        s = 0\n",
    "        while s < len(self.sample_times) - 1 and t >= self.sample_times[s + 1]:\n",
    "            s += 1\n",
    "        trel = t - self.sample_times[s]\n",
    "\n",
    "        if self.result is None:\n",
    "            coeffs = self.coeffs[s]\n",
    "        else:\n",
    "            coeffs = self.result.GetSolution(self.coeffs[s])\n",
    "\n",
    "        deg = derivative_order\n",
    "        val = 0 * coeffs[:, 0]\n",
    "        for var in range(self.n):\n",
    "            for d in range(deg, self.degree + 1):\n",
    "                val[var] += (coeffs[var, d] * np.power(trel, d - deg) *\n",
    "                             math.factorial(d) / math.factorial(d - deg))\n",
    "\n",
    "        return val\n",
    "\n",
    "    def add_constraint(self, t, derivative_order, lb, ub=None):\n",
    "        \"\"\"Adds a constraint of the form d^deg lb <= x(t) / dt^deg <= ub.\"\"\"\n",
    "        if ub is None:\n",
    "            ub = lb\n",
    "\n",
    "        assert derivative_order <= self.degree\n",
    "        val = self.eval(t, derivative_order)\n",
    "        self.prog.AddLinearConstraint(val, lb, ub)\n",
    "\n",
    "    def generate(self):\n",
    "        self.result = Solve(self.prog)\n",
    "        assert self.result.is_success()\n",
    "\n",
    "\n",
    "tf = 3\n",
    "zpp = PPTrajectory(sample_times=np.linspace(0, tf, 6),\n",
    "                   num_vars=2,\n",
    "                   degree=5,\n",
    "                   continuity_degree=4)\n",
    "zpp.add_constraint(t=0, derivative_order=0, lb=[0, 0])\n",
    "zpp.add_constraint(t=0, derivative_order=1, lb=[0, 0])\n",
    "zpp.add_constraint(t=0, derivative_order=2, lb=[0, 0])\n",
    "zpp.add_constraint(t=1, derivative_order=0, lb=[2, 1.5])\n",
    "zpp.add_constraint(t=2, derivative_order=0, lb=[4, 1])\n",
    "zpp.add_constraint(t=tf, derivative_order=0, lb=[6, 0])\n",
    "zpp.add_constraint(t=tf, derivative_order=1, lb=[0, 0])\n",
    "zpp.add_constraint(t=tf, derivative_order=2, lb=[0, 0])\n",
    "zpp.generate()\n",
    "\n",
    "if False:  # Useful for debugging\n",
    "    t = np.linspace(0, tf, 100)\n",
    "    z = np.zeros((2, len(t)))\n",
    "    knots = np.zeros((2, len(zpp.sample_times)))\n",
    "    fig, ax = plt.subplots(zpp.degree + 1, 1)\n",
    "    for deg in range(zpp.degree + 1):\n",
    "        for i in range(len(t)):\n",
    "            z[:, i] = zpp.eval(t[i], deg)\n",
    "        for i in range(len(zpp.sample_times)):\n",
    "            knots[:, i] = zpp.eval(zpp.sample_times[i], deg)\n",
    "        ax[deg].plot(t, z.transpose())\n",
    "        ax[deg].plot(zpp.sample_times, knots.transpose(), \".\")\n",
    "        ax[deg].set_xlabel(\"t (sec)\")\n",
    "        ax[deg].set_ylabel(\"z deriv \" + str(deg))\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "t = np.linspace(0, tf, 100)\n",
    "z = np.zeros((2, len(t)))\n",
    "\n",
    "for i in range(len(t)):\n",
    "    z[:, i] = zpp.eval(t[i])\n",
    "ax.plot(z[0, :], z[1, :])\n",
    "\n",
    "for t in np.linspace(0, tf, 7):\n",
    "    x = zpp.eval(t)\n",
    "    xddot = zpp.eval(t, 2)\n",
    "    theta = np.arctan2(-xddot[0], (xddot[1] + 9.81))\n",
    "    v = Quadrotor2DVisualizer(ax=ax)\n",
    "    context = v.CreateDefaultContext()\n",
    "    v.get_input_port(0).FixValue(context, [x[0], x[1], theta, 0, 0, 0])\n",
    "    v.draw(context)\n",
    "\n",
    "# Draw the (imaginary) obstacles\n",
    "ax.fill(2 + np.array([-.1, -.1, .1, .1, -.1]),\n",
    "        1.25 * np.array([0, 1, 1, 0, 0]),\n",
    "        facecolor=\"darkred\",\n",
    "        edgecolor=\"k\")\n",
    "ax.fill(2 + np.array([-.1, -.1, .1, .1, -.1]),\n",
    "        1.75 + 1.25 * np.array([0, 1, 1, 0, 0]),\n",
    "        facecolor=\"darkred\",\n",
    "        edgecolor=\"k\")\n",
    "ax.fill(4 + np.array([-.1, -.1, .1, .1, -.1]),\n",
    "        .75 * np.array([0, 1, 1, 0, 0]),\n",
    "        facecolor=\"darkred\",\n",
    "        edgecolor=\"k\")\n",
    "ax.fill(4 + np.array([-.1, -.1, .1, .1, -.1]),\n",
    "        1.25 + 1.75 * np.array([0, 1, 1, 0, 0]),\n",
    "        facecolor=\"darkred\",\n",
    "        edgecolor=\"k\")\n",
    "ax.set_xlim([-1, 7])\n",
    "ax.set_ylim([-.25, 3])\n",
    "ax.set_title(\"\");\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}