acrobot.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook provides examples to go along with the [textbook](https://underactuated.csail.mit.edu/acrobot.html).  I recommend having both windows open, side-by-side!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import mpld3\n",
    "import numpy as np\n",
    "from IPython.display import HTML, display\n",
    "from pydrake.all import (AddMultibodyPlantSceneGraph, ControllabilityMatrix,\n",
    "                         DiagramBuilder, Linearize, LinearQuadraticRegulator,\n",
    "                         MeshcatVisualizer, Parser, Saturation, SceneGraph,\n",
    "                         Simulator, StartMeshcat, WrapToSystem)\n",
    "from pydrake.examples import (AcrobotGeometry, AcrobotInput, AcrobotPlant,\n",
    "                              AcrobotState, QuadrotorGeometry, QuadrotorPlant,\n",
    "                              StabilizingLQRController)\n",
    "from pydrake.solvers import MathematicalProgram, Solve\n",
    "\n",
    "from underactuated import FindResource, running_as_notebook\n",
    "from underactuated.meshcat_cpp_utils import MeshcatSliders\n",
    "from underactuated.quadrotor2d import Quadrotor2D, Quadrotor2DVisualizer\n",
    "\n",
    "if running_as_notebook:\n",
    "    mpld3.enable_notebook()"
   ]
  },
  {
   "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": {},
   "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": [
    "def acrobot_demo():\n",
    "    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),\n",
    "                                 scene_graph)\n",
    "    meshcat.Delete()\n",
    "    meshcat.Set2dRenderMode(xmin=-4, xmax=4, ymin=-4, ymax=4)\n",
    "    MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\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(), acrobot.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",
    "    # Set the initial conditions (theta1, theta2, theta1dot, theta2dot)\n",
    "    context.SetContinuousState([1, 0, 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",
    "        print('Use the slider in the MeshCat controls to apply elbow torque.')\n",
    "        print(\"Press 'Stop Simulation' in MeshCat to continue.\")\n",
    "        meshcat.AddButton('Stop Simulation')\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",
    "acrobot_demo()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dynamics of the Cart-Pole"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cartpole_demo():\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",
    "    meshcat.Delete()\n",
    "    meshcat.Set2dRenderMode(xmin=-2.5, xmax=2.5, ymin=-1.0, ymax=2.5)\n",
    "    MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "\n",
    "    # Setup slider input\n",
    "    meshcat.AddSlider('u', min=-30., max=30, step=.1, value=0.0)\n",
    "    force_system = builder.AddSystem(MeshcatSliders(meshcat,['u']))\n",
    "    builder.Connect(force_system.get_output_port(),\n",
    "                    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",
    "    # Set the initial conditions (x, theta, xdot, thetadot)\n",
    "    context.SetContinuousState([0, 1, 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",
    "        print('Use the slider in the MeshCat controls to apply elbow torque.')\n",
    "        print(\"Press 'Stop Simulation' in MeshCat to continue.\")\n",
    "        meshcat.AddButton('Stop Simulation')\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",
    "cartpole_demo()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LQR for the Acrobot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First let's check the controllability.  Recall that the rank of the matrix is the number of of non-zero singular values.  So seeing very small singular values suggests \"weak\" controllability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "def acrobot_controllability():\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",
    "    linearized_acrobot = Linearize(acrobot, context)\n",
    "    print(\n",
    "        f\"The singular values of the controllability matrix are: {np.linalg.svd(ControllabilityMatrix(linearized_acrobot), compute_uv=False)}\"\n",
    "    )\n",
    "\n",
    "\n",
    "acrobot_controllability()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def acrobot_balancing_example():\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",
    "    meshcat.Delete()\n",
    "    meshcat.Set2dRenderMode(xmin=-4, xmax=4, ymin=-4, ymax=4)\n",
    "    MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\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",
    "    simulator.set_target_realtime_rate(1.0 if running_as_notebook else 0.0)\n",
    "    duration = 4.0 if running_as_notebook else 0.1\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": [
    "def cartpole_balancing_example():\n",
    "    def UprightState():\n",
    "        state = (0, np.pi, 0, 0)\n",
    "        return state\n",
    "\n",
    "    def Controllability(plant):\n",
    "        context = plant.CreateDefaultContext()\n",
    "        plant.get_actuation_input_port().FixValue(context, [0])\n",
    "        plant.SetPositionsAndVelocities(context, UprightState())\n",
    "\n",
    "        linearized_plant = Linearize(\n",
    "            plant,\n",
    "            context,\n",
    "            input_port_index=plant.get_actuation_input_port().get_index(), output_port_index=plant.get_state_output_port().get_index())\n",
    "        print(linearized_plant.A())\n",
    "        print(linearized_plant.B())\n",
    "        print(\n",
    "            f\"The singular values of the controllability matrix are: {np.linalg.svd(ControllabilityMatrix(linearized_plant), compute_uv=False)}\"\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",
    "        plant.SetPositionsAndVelocities(context, UprightState())\n",
    "\n",
    "        Q = np.diag((10., 10., 1., 1.))\n",
    "        R = np.array([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.1)\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",
    "    meshcat.Delete()\n",
    "    meshcat.Set2dRenderMode(xmin=-2.5, xmax=2.5, ymin=-1.0, ymax=2.5)\n",
    "    MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\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",
    "    plant_context = plant.GetMyMutableContextFromRoot(context)\n",
    "\n",
    "    # Simulate\n",
    "    simulator.set_target_realtime_rate(1.0 if running_as_notebook else 0.0)\n",
    "    duration = 5.0 if running_as_notebook else 0.1\n",
    "    for i in range(5):\n",
    "        context.SetTime(0.)\n",
    "        plant.SetPositionsAndVelocities(\n",
    "            plant_context,\n",
    "            UprightState() + 0.1 * np.random.randn(4,))\n",
    "        simulator.Initialize()\n",
    "        simulator.AdvanceTo(duration)\n",
    "\n",
    "np.set_printoptions(formatter={'float': lambda x: \"{0:0.4f}\".format(x)})\n",
    "cartpole_balancing_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LQR for the Planar Quadrotor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def planar_quadrotor_example():\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",
    "    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=False))\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\n",
    "    visualizer.start_recording()\n",
    "    print(\"simulating...\")\n",
    "    for i in range(5):\n",
    "        context.SetTime(0.)\n",
    "        context.SetContinuousState(np.random.randn(6,))\n",
    "        simulator.Initialize()\n",
    "        simulator.AdvanceTo(duration)\n",
    "    print(\"done.\\ngenerating animation...\")\n",
    "    ani = visualizer.get_recording_as_animation()\n",
    "    display(HTML(ani.to_jshtml()))\n",
    "\n",
    "planar_quadrotor_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LQR for the 3D Quadrotor\n",
    "\n",
    "Art assets courtesy of our friends at Skydio."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "    meshcat.Delete()\n",
    "    meshcat.ResetRenderMode()\n",
    "    meshcat.SetProperty('/Background','visible',False)\n",
    "    MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\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 if running_as_notebook else 0.0)\n",
    "    context = simulator.get_mutable_context()\n",
    "\n",
    "    # Simulate\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(4.0 if running_as_notebook else 0.1)\n",
    "\n",
    "quadrotor_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The 3D Quadrotor again, now using MultibodyPlant\n",
    "\n",
    "The example above used [`QuadrotorPlant`](https://drake.mit.edu/doxygen_cxx/classdrake_1_1examples_1_1quadrotor_1_1_quadrotor_plant.html), in which I've typed in the equations of motion by hand.  This is a great way to work with the system, but is hard to generalize.\n",
    "\n",
    "If you want to add additional forces to the model -- for instance if you want the quadrotor to have contact dynamics with the world (to land on the ground, or collide with a tree) or if you want to add e.g. wind -- then we need a more general solution.  We can also implement the quadrotor dynamics using `MultibodyPlant`.  We just need to manually wire the [`Propeller`](https://drake.mit.edu/doxygen_cxx/classdrake_1_1multibody_1_1_propeller.html) forces into the `Diagram` (because `Propeller` is not a concept supported in URDF nor SDF just yet)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydrake.all import (DiagramBuilder, FindResourceOrThrow, MultibodyPlant,\n",
    "                         Parser, Propeller, PropellerInfo, RigidTransform)\n",
    "from pydrake.examples import QuadrotorPlant\n",
    "\n",
    "from underactuated.scenarios import AddFloatingRpyJoint\n",
    "\n",
    "\n",
    "def MakeMultibodyQuadrotor():\n",
    "    builder = DiagramBuilder()\n",
    "    # The MultibodyPlant handles f=ma, but doesn't know about propellers.\n",
    "    plant = builder.AddSystem(MultibodyPlant(0.0))\n",
    "    model_instance = Parser(plant).AddModelFromFile(\n",
    "        FindResourceOrThrow(\"drake/examples/quadrotor/quadrotor.urdf\"))\n",
    "    # By default the multibody has a quaternion floating base.  To match\n",
    "    # QuadrotorPlant, we can manually add a FloatingRollPitchYaw joint. We set\n",
    "    # `use_ball_rpy` to false because the BallRpyJoint uses angular velocities\n",
    "    # instead of ṙ, ṗ, ẏ.\n",
    "    AddFloatingRpyJoint(plant, plant.GetFrameByName(\"base_link\"),\n",
    "                        model_instance, use_ball_rpy=False)\n",
    "    plant.Finalize()\n",
    "\n",
    "    # Now we can add in propellers as an external force on the MultibodyPlant.\n",
    "    body_index = plant.GetBodyByName(\"base_link\").index()\n",
    "    # Default parameters from quadrotor_plant.cc:\n",
    "    L = 0.15     # Length of the arms (m).\n",
    "    kF = 1.0     # Force input constant.\n",
    "    kM = 0.0245  # Moment input constant.\n",
    "\n",
    "    # Note: Rotors 0 and 2 rotate one way and rotors 1 and 3 rotate the other.\n",
    "    prop_info = [\n",
    "        PropellerInfo(body_index, RigidTransform([L, 0, 0]), kF, kM),\n",
    "        PropellerInfo(body_index, RigidTransform([0, L, 0]), kF, -kM),\n",
    "        PropellerInfo(body_index, RigidTransform([-L, 0, 0]), kF, kM),\n",
    "        PropellerInfo(body_index, RigidTransform([0, -L, 0]), kF, -kM),\n",
    "    ]\n",
    "    propellers = builder.AddSystem(Propeller(prop_info))\n",
    "    builder.Connect(propellers.get_output_port(),\n",
    "                    plant.get_applied_spatial_force_input_port())\n",
    "    builder.Connect(plant.get_body_poses_output_port(),\n",
    "                    propellers.get_body_poses_input_port())\n",
    "    builder.ExportInput(propellers.get_command_input_port(), \"u\")\n",
    "\n",
    "    return builder.Build(), plant\n",
    "\n",
    "# This test demonstrates that the MultibodyPlant version has identical dynamics\n",
    "# to the QuadrotorPlant version (except that the state variables are permuted).\n",
    "# TODO(russt): Move this to Drake as a unit test.\n",
    "def MultibodyQuadrotorExample():\n",
    "    mbp_plant, mbp = MakeMultibodyQuadrotor()\n",
    "\n",
    "    hand_derived_plant = QuadrotorPlant()\n",
    "\n",
    "    # Compare the dynamics at a handful of states.\n",
    "    mbp_plant_context = mbp_plant.CreateDefaultContext()\n",
    "    mbp_context = mbp.GetMyContextFromRoot(mbp_plant_context)\n",
    "    hand_derived_context = hand_derived_plant.CreateDefaultContext()\n",
    "\n",
    "    # Permute mbp <=> hand_derived states.\n",
    "    P = np.array([[1, 0, 0, 0, 0, 0],\n",
    "                  [0, 1, 0, 0, 0, 0],\n",
    "                  [0, 0, 1, 0, 0, 0],\n",
    "                  [0, 0, 0, 0, 0, 1],\n",
    "                  [0, 0, 0, 0, 1, 0],\n",
    "                  [0, 0, 0, 1, 0, 0]])\n",
    "    PP = np.block([[P, np.zeros((6,6))], [np.zeros((6,6)), P]])\n",
    "\n",
    "    rng = np.random.default_rng(seed=1037)\n",
    "    for i in range(5):\n",
    "        u = rng.random((4,))\n",
    "        mbp_x = rng.random((12,))\n",
    "        mbp_plant.get_input_port().FixValue(mbp_plant_context, u)\n",
    "        mbp_context.SetContinuousState(mbp_x)\n",
    "        mbp_xdot = mbp_plant.EvalTimeDerivatives(\n",
    "            mbp_plant_context).CopyToVector()\n",
    "\n",
    "        hand_derived_x = PP @ mbp_x\n",
    "        hand_derived_plant.get_input_port().FixValue(hand_derived_context, u)\n",
    "        hand_derived_context.SetContinuousState(hand_derived_x)\n",
    "        hand_derived_xdot = PP @ hand_derived_plant.EvalTimeDerivatives(\n",
    "            hand_derived_context).CopyToVector()\n",
    "\n",
    "        assert np.allclose(\n",
    "            mbp_xdot, hand_derived_xdot\n",
    "        ), f\"\\nmbp\\t\\t = {mbp_xdot}\\nhand_derived\\t = {hand_derived_xdot}\"\n",
    "\n",
    "\n",
    "MultibodyQuadrotorExample()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Differential Flatness"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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": {
  "interpreter": {
   "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}