simple_legs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook provides examples to go along with the [textbook](http://underactuated.csail.mit.edu/simple_legs.html).  I recommend having both windows open, side-by-side!\n",
    "\n",
    "[Click here](http://underactuated.csail.mit.edu/drake.html#notebooks) for instructions on how to run the notebook on Deepnote and/or Google Colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import mpld3\n",
    "import numpy as np\n",
    "from IPython.display import display\n",
    "from pydrake.all import (BasicVector, ConstantVectorSource, DiagramBuilder,\n",
    "                         LeafSystem, LogVectorOutput, MeshcatVisualizerCpp,\n",
    "                         PortDataType, PublishEvent, PyPlotVisualizer,\n",
    "                         SceneGraph, Simulator, StartMeshcat,\n",
    "                         UnrestrictedUpdateEvent, VectorLogSink,\n",
    "                         WitnessFunctionDirection, namedview)\n",
    "from pydrake.examples.compass_gait import CompassGait, CompassGaitGeometry\n",
    "from pydrake.examples.rimless_wheel import (RimlessWheel, RimlessWheelGeometry,\n",
    "                                            RimlessWheelParams)\n",
    "\n",
    "from underactuated import running_as_notebook\n",
    "from underactuated.jupyter import AdvanceToAndVisualize\n",
    "\n",
    "if running_as_notebook:\n",
    "    mpld3.enable_notebook()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer (run this cell only once, each instance consumes a port)\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Rimless Wheel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def rimless_wheel(slope=0.08, initial_angle=0, initial_angular_velocity=5.0):\n",
    "  params = RimlessWheelParams()\n",
    "  params.set_slope(slope)\n",
    "\n",
    "  builder = DiagramBuilder()\n",
    "  rimless_wheel = builder.AddSystem(RimlessWheel())\n",
    "  scene_graph = builder.AddSystem(SceneGraph())\n",
    "  RimlessWheelGeometry.AddToBuilder(\n",
    "      builder, rimless_wheel.get_floating_base_state_output_port(), params,\n",
    "      scene_graph)\n",
    "  MeshcatVisualizerCpp.AddToBuilder(builder, scene_graph, meshcat)\n",
    "  meshcat.Set2dRenderMode(xmin=-2, xmax=14, ymin=-2, ymax=3)\n",
    "\n",
    "  diagram = builder.Build()\n",
    "  simulator = Simulator(diagram)\n",
    "\n",
    "  context = simulator.get_mutable_context()\n",
    "\n",
    "  diagram.GetMutableSubsystemContext(\n",
    "      rimless_wheel, context).get_numeric_parameter(0).set_slope(slope)\n",
    "  context.SetAccuracy(1e-4)\n",
    "  context.SetContinuousState([initial_angle, initial_angular_velocity])\n",
    "  simulator.Initialize()\n",
    "  simulator.set_target_realtime_rate(1.0 if running_as_notebook else 0.0)\n",
    "  simulator.AdvanceTo(5.0 if running_as_notebook else 0.1)\n",
    "\n",
    "rimless_wheel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rimless_wheel(initial_angular_velocity=5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rimless_wheel(initial_angular_velocity=10.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rimless_wheel(initial_angular_velocity=0.95)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rimless_wheel(initial_angular_velocity=-5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rimless_wheel(initial_angular_velocity=-4.8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a little interactive plot to allow you to visualize the trajectories of the rimless wheel as you vary the initial velocity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO(russt): Port to meshcat and make it interactive, once again.\n",
    "\n",
    "def rimless_wheel_return_map():\n",
    "    rimless_wheel = RimlessWheel()\n",
    "    simulator = Simulator(rimless_wheel)\n",
    "    context = simulator.get_mutable_context()\n",
    "    params = context.get_numeric_parameter(0)\n",
    "    qmin = params.slope() - rimless_wheel.calc_alpha(params) - 0.1\n",
    "    qmax = params.slope() + rimless_wheel.calc_alpha(params) + 0.1\n",
    "    qdmin = -2\n",
    "    qdmax = 2\n",
    "\n",
    "    context.SetAccuracy(1e-2)\n",
    "    integrator = simulator.get_mutable_integrator()\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    # TODO(russt): make the slope interactive, too.\n",
    "    def simulate(initial_angular_velocity=1.5, duration=1.5):\n",
    "        rimless_wheel.SetDefaultContext(context)\n",
    "        context.SetTime(0.0)\n",
    "        if initial_angular_velocity >= 0:\n",
    "            initial_angle = params.slope() - rimless_wheel.calc_alpha(params)\n",
    "        else:\n",
    "            initial_angle = params.slope() + rimless_wheel.calc_alpha(params)\n",
    "        if initial_angular_velocity == 0:\n",
    "            # Set double_support = True.\n",
    "            context.get_mutable_abstract_state(0).set_value(True)\n",
    "            \n",
    "        context.SetContinuousState([initial_angle, initial_angular_velocity])\n",
    "\n",
    "        integrator.StartDenseIntegration()\n",
    "        simulator.Initialize()\n",
    "        simulator.AdvanceTo(duration)\n",
    "        pp = integrator.StopDenseIntegration()\n",
    "\n",
    "        return pp.vector_values(pp.get_segment_times())\n",
    "\n",
    "    if False:\n",
    "        data = simulate()\n",
    "        line, = ax.plot(data[0,:], data[1,:],'b')\n",
    "        pt, = ax.plot(data[0,0], data[1,0],'b*', markersize=12)\n",
    "        \n",
    "        def update(initial_angular_velocity):\n",
    "            data = simulate(initial_angular_velocity)\n",
    "            line.set_xdata(data[0,:])\n",
    "            line.set_ydata(data[1,:])\n",
    "            pt.set_xdata(data[0,0])\n",
    "            pt.set_ydata(data[1,0])\n",
    "            fig.canvas.draw()\n",
    "        \n",
    "        interact(update, initial_angular_velocity=widgets.FloatSlider(min=qdmin, max=qdmax, step=.1, value=1.1))\n",
    "\n",
    "    else:\n",
    "        data = simulate()\n",
    "        ax.plot(data[0,:], data[1,:],'b')\n",
    "        ax.plot(data[0,0], data[1,0],'b*', markersize=12)\n",
    "\n",
    "    # Plot the energy contours.\n",
    "    nq = 151\n",
    "    nqd = 151\n",
    "    mgl = params.mass() * params.gravity() * params.length()\n",
    "    q = np.linspace(qmin, qmax, nq)\n",
    "    qd = np.linspace(qdmin, qdmax, nqd)\n",
    "    Q, QD = np.meshgrid(q, qd)\n",
    "    Energy = .5 * params.mass() * params.length()**2 * QD**2 + mgl * np.cos(Q)\n",
    "    ax.contour(Q,\n",
    "            QD,\n",
    "            Energy,\n",
    "            alpha=0.5,\n",
    "            linestyles=\"dashed\",\n",
    "            colors=\"black\",\n",
    "            linewidths=0.5)\n",
    "\n",
    "    ax.plot(params.slope() - rimless_wheel.calc_alpha(params)*np.array([1, 1]), np.array([0, qdmax]), 'k--')\n",
    "    ax.plot(params.slope() - rimless_wheel.calc_alpha(params)*np.array([1, 1]), np.array([0, qdmin]), 'k', linewidth=.25)\n",
    "    ax.plot(params.slope() + rimless_wheel.calc_alpha(params)*np.array([1, 1]), np.array([0, qdmin]), 'k--')\n",
    "    ax.plot(params.slope() + rimless_wheel.calc_alpha(params)*np.array([1, 1]), np.array([0, qdmax]), 'k', linewidth=.25)\n",
    "    ax.plot([qmin, qmax], [0, 0], 'k', linewidth=.25)\n",
    "    ax.plot([0, 0], [qdmin, qdmax], 'k', linewidth=.25)\n",
    "    ax.set_xlabel(\"theta\")\n",
    "    ax.set_ylabel(\"thetadot\")\n",
    "    ax.axis([qmin, qmax, qdmin, qdmax])\n",
    "    ax.set_title(\"Trajectories of the Rimless Wheel (w/ contours of \"\n",
    "                \"constant energy)\")\n",
    "    display(mpld3.display())\n",
    "\n",
    "rimless_wheel_return_map()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Compass Gait"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compass_gait():\n",
    "    builder = DiagramBuilder()\n",
    "    compass_gait = builder.AddSystem(CompassGait())\n",
    "\n",
    "    hip_torque = builder.AddSystem(ConstantVectorSource([0.0]))\n",
    "    builder.Connect(hip_torque.get_output_port(0), compass_gait.get_input_port(0))\n",
    "\n",
    "    scene_graph = builder.AddSystem(SceneGraph())\n",
    "    CompassGaitGeometry.AddToBuilder(\n",
    "        builder, compass_gait.get_floating_base_state_output_port(), scene_graph)\n",
    "    MeshcatVisualizerCpp.AddToBuilder(builder, scene_graph, meshcat)\n",
    "    meshcat.Set2dRenderMode(xmin=-1, xmax=5, ymin=-1, ymax=2)\n",
    "\n",
    "    logger = builder.AddSystem(VectorLogSink(14))\n",
    "    builder.Connect(compass_gait.get_output_port(1), logger.get_input_port(0))\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    simulator = Simulator(diagram)\n",
    "    simulator.set_target_realtime_rate(1.0 if running_as_notebook else 0.0)\n",
    "\n",
    "    context = simulator.get_mutable_context()\n",
    "    context.SetAccuracy(1e-4)\n",
    "    context.SetContinuousState([0., 0., 0.4, -2.])\n",
    "\n",
    "    simulator.AdvanceTo(8.0)\n",
    "    log = logger.FindLog(context)\n",
    "    plt.figure()\n",
    "    plt.plot(log.data()[4, :], log.data()[11, :])\n",
    "    plt.xlabel(\"left leg angle\")\n",
    "    plt.ylabel(\"left leg angular velocity\")\n",
    "    display(mpld3.display())\n",
    "\n",
    "compass_gait()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Spring-Loaded Inverted Pendulum (SLIP)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SLIPState = namedview(\n",
    "    \"SLIPState\", [\"x\", \"z\", \"r\", \"theta\", \"xdot\", \"zdot\", \"rdot\", \"thetadot\"])\n",
    "\n",
    "class SpringLoadedInvertedPendulum(LeafSystem):\n",
    "\n",
    "    def __init__(self):\n",
    "        LeafSystem.__init__(self)\n",
    "\n",
    "        self.DeclareVectorInputPort(\"touchdown_angle\", BasicVector(1))\n",
    "        self.DeclareContinuousState(BasicVector(np.zeros(8)), 4, 4, 0)\n",
    "\n",
    "        self.DeclareVectorOutputPort(\"state\", BasicVector(8), self.CopyStateOut)\n",
    "\n",
    "        # Parameters from Geyer05, p.23\n",
    "        self.mass = 80.  # kg\n",
    "        self.r0 = 1.  # m\n",
    "        self.gravity = 9.81  # m/s^2\n",
    "        # Define spring constant in terms of the dimensionless number.\n",
    "        # Definition in section 2.4.3, values in figure 2.4.\n",
    "        # Note: Geyer05 says 10.8 (which doesn't work? -- I get no fixed pts).\n",
    "        dimensionless_spring_constant = 10.7\n",
    "        self.stiffness = (dimensionless_spring_constant * self.mass *\n",
    "                          self.gravity / self.r0)\n",
    "\n",
    "        self.last_apex = None  # placeholder for writing return map result.\n",
    "\n",
    "        self.touchdown_witness = self.MakeWitnessFunction(\n",
    "            \"touchdown\", WitnessFunctionDirection.kPositiveThenNonPositive,\n",
    "            self.foot_height, UnrestrictedUpdateEvent(self.touchdown))\n",
    "        self.takeoff_witness = self.MakeWitnessFunction(\n",
    "            \"takeoff\", WitnessFunctionDirection.kPositiveThenNonPositive,\n",
    "            self.leg_compression, UnrestrictedUpdateEvent(self.takeoff))\n",
    "        self.apex_witness = self.MakeWitnessFunction(\n",
    "            \"apex\", WitnessFunctionDirection.kPositiveThenNonPositive,\n",
    "            self.apex, PublishEvent(self.publish_apex))\n",
    "\n",
    "    def foot_height(self, context):\n",
    "        s = SLIPState(context.get_continuous_state_vector().CopyToVector())\n",
    "        return s.z - self.r0 * np.cos(s.theta)\n",
    "\n",
    "    def touchdown(self, context, event, state):\n",
    "        s = SLIPState(\n",
    "            context.get_mutable_continuous_state_vector().CopyToVector())\n",
    "\n",
    "        # print(\"touchdown\")\n",
    "\n",
    "        # Update rdot and thetadot to match xdot and ydot, using\n",
    "        # x = -r*sin(theta), z = r*cos(theta)\n",
    "        #  => xdot = -rdot*s - r*c*thetadot, zdot = rdot*c - r*s*thetadot\n",
    "        #  => xdot*c + zdot*s = -r*thetadot\n",
    "        # r^2 = x^2 + z^2\n",
    "        #  => 2r*rdot = 2x*xdot + 2z*zdot\n",
    "        #  => rdot = -xdot*sin(theta) + zdot*cos(theta)\n",
    "        # (matches Geyer05 Eq. 2.24 up to the symbol changes)\n",
    "        s.r = self.r0\n",
    "        s.rdot = -s.xdot * np.sin(s.theta) + s.zdot * np.cos(s.theta)\n",
    "        s.thetadot = -(s.xdot * np.cos(s.theta) +\n",
    "                       s.zdot * np.sin(s.theta)) / self.r0\n",
    "        state.get_mutable_continuous_state().get_mutable_vector().SetFromVector(\n",
    "            s[:])\n",
    "\n",
    "    def leg_compression(self, context):\n",
    "        s = SLIPState(context.get_continuous_state_vector().CopyToVector())\n",
    "        return self.r0 - s.r\n",
    "\n",
    "    def takeoff(self, context, event, state):\n",
    "        s = SLIPState(\n",
    "            context.get_mutable_continuous_state_vector().CopyToVector())\n",
    "\n",
    "        # print(\"takeoff\")\n",
    "\n",
    "        # Setup flight state (these lines aren't strictly required, since we\n",
    "        # choose to also integrate x and z in stance below).\n",
    "        s.z = self.r0 * np.cos(s.theta)\n",
    "        s.xdot = -s.rdot * np.sin(s.theta) - self.r0 * s.thetadot * np.cos(\n",
    "            s.theta)\n",
    "        s.zdot = s.rdot * np.cos(s.theta) - self.r0 * s.thetadot * np.sin(\n",
    "            s.theta)\n",
    "\n",
    "        # Update theta to commanded leg angle.\n",
    "        s.theta = self.EvalVectorInput(context, 0).GetAtIndex(0)\n",
    "        s.thetadot = 0\n",
    "        s.r = self.r0\n",
    "        s.rdot = 0\n",
    "\n",
    "        state.get_mutable_continuous_state().get_mutable_vector().SetFromVector(\n",
    "            s[:])\n",
    "\n",
    "    def apex(self, context):\n",
    "        return context.get_continuous_state_vector().GetAtIndex(5)  # zdot\n",
    "\n",
    "    def publish_apex(self, context, event):\n",
    "        # TODO(russt): provide an option to terminate here instead, pending\n",
    "        # resolution of #4447.\n",
    "        # print(\"apex\")\n",
    "        if self.last_apex is None:\n",
    "            s = SLIPState(\n",
    "                context.get_mutable_continuous_state_vector().CopyToVector())\n",
    "            self.last_apex = s.z\n",
    "\n",
    "    def apex_velocity_from_dimensionless_system_energy(self, Etilde, z):\n",
    "        E = Etilde * self.mass * self.gravity * self.r0\n",
    "        # E = 0.5*m*v^2 + m*g*z\n",
    "        xdot = np.sqrt(2. / self.mass * (E - self.mass * self.gravity * z))\n",
    "        return xdot\n",
    "\n",
    "    def energy_flight(self, context):\n",
    "        s = SLIPState(\n",
    "            context.get_mutable_continuous_state_vector().CopyToVector())\n",
    "        return (0.5 * self.mass * (s.xdot**2 + s.zdot**2) +\n",
    "                self.mass * self.gravity * s.z)\n",
    "\n",
    "    def energy_stance(self, context):\n",
    "        s = SLIPState(\n",
    "            context.get_mutable_continuous_state_vector().CopyToVector())\n",
    "        return (0.5 * self.mass * (s.rdot**2 + s.r**2 * s.thetadot**2) +\n",
    "                self.mass * self.gravity * s.r * np.cos(s.theta) +\n",
    "                0.5 * self.stiffness * (self.r0 - s.r)**2)\n",
    "\n",
    "    def CopyStateOut(self, context, output):\n",
    "        x = context.get_continuous_state_vector().CopyToVector()\n",
    "        y = output.SetFromVector(x)\n",
    "\n",
    "    def DoGetWitnessFunctions(self, context):\n",
    "        return [self.touchdown_witness, self.takeoff_witness, self.apex_witness]\n",
    "\n",
    "    def DoCalcTimeDerivatives(self, context, derivatives):\n",
    "        s = SLIPState(context.get_continuous_state_vector().CopyToVector())\n",
    "        sdot = SLIPState(np.zeros(8))\n",
    "        sdot[0:4] = s[4:8]\n",
    "\n",
    "        if (self.foot_height(context) < 0):\n",
    "            # then we're in \"stance\"\n",
    "            sdot.rdot = (self.stiffness / self.mass * (self.r0 - s.r) +\n",
    "                         s.r * s.thetadot**2 - self.gravity * np.cos(s.theta))\n",
    "            sdot.thetadot = (self.gravity / s.r * np.sin(s.theta) -\n",
    "                             2 * s.rdot * s.thetadot / s.r)\n",
    "\n",
    "            # Integrate x and z also, just for the sake of visualization (all\n",
    "            # the integrated values except x will be overwritten in the\n",
    "            # take-off reset).\n",
    "            # x = -r*sin(theta), y = r*cos(theta) =>\n",
    "            sdot.xdot = (-sdot.rdot * np.sin(s.theta) -\n",
    "                         2 * s.rdot * s.thetadot * np.cos(s.theta) +\n",
    "                         s.r * s.thetadot**2 * np.sin(s.theta) -\n",
    "                         s.r * sdot.thetadot * np.cos(s.theta))\n",
    "            sdot.zdot = (sdot.rdot * np.cos(s.theta) -\n",
    "                         2 * s.rdot * s.thetadot * np.sin(s.theta) -\n",
    "                         s.r * sdot.thetadot * np.sin(s.theta) -\n",
    "                         s.r * s.thetadot**2 * np.cos(s.theta))\n",
    "\n",
    "        else:\n",
    "            sdot.xdot = 0\n",
    "            sdot.zdot = -self.gravity\n",
    "            sdot.rdot = 0\n",
    "            sdot.thetadot = 0\n",
    "\n",
    "        derivatives.get_mutable_vector().SetFromVector(sdot[:])\n",
    "\n",
    "\n",
    "class SLIPVisualizer(PyPlotVisualizer):\n",
    "\n",
    "    def __init__(self, ax=None, show=True):\n",
    "        PyPlotVisualizer.__init__(self, ax=ax, show=show)\n",
    "        self.DeclareInputPort(\"slip_state\", PortDataType.kVectorValued, 8)\n",
    "        self.ax.set_aspect(\"equal\")\n",
    "        self.ax.set_xlim(0, 2)\n",
    "        self.ax.set_ylim(-2, 2)\n",
    "\n",
    "        # Draw the ground.\n",
    "        self.ax.plot([-50, 50], [0, 0], \"k\")\n",
    "\n",
    "        a = np.linspace(0, 2 * np.pi, 50)\n",
    "        radius = 0.1\n",
    "        self.hip_fill = self.ax.fill(radius * np.sin(a),\n",
    "                                     radius * np.cos(a),\n",
    "                                     zorder=1,\n",
    "                                     edgecolor=\"k\",\n",
    "                                     facecolor=[.6, .6, .6])\n",
    "        self.hip = copy.copy(self.hip_fill[0].get_path().vertices)\n",
    "\n",
    "        self.leg_line = [self.ax.plot([0, 0], [0, -1], \"k\")[0]]\n",
    "        self.leg_data = [self.leg_line[0].get_xydata().T]\n",
    "        for i in range(1, 13):\n",
    "            self.leg_line.append(\n",
    "                self.ax.plot(\n",
    "                    0.1 * np.array(\n",
    "                        [np.sin((i - 1) * np.pi / 2.),\n",
    "                         np.sin(i * np.pi / 2.)]),\n",
    "                    -.2 - .7 / 13 * np.array([i - 1, i]), \"k\")[0])\n",
    "            self.leg_data.append(self.leg_line[i].get_xydata().T)\n",
    "\n",
    "    def draw(self, context):\n",
    "        state = SLIPState(self.EvalVectorInput(context, 0).CopyToVector())\n",
    "\n",
    "        self.hip_fill[0].get_path().vertices[:, 0] = state.x + self.hip[:, 0]\n",
    "        self.hip_fill[0].get_path().vertices[:, 1] = state.z + self.hip[:, 1]\n",
    "\n",
    "        R = np.array([[np.cos(state.theta), -np.sin(state.theta)],\n",
    "                      [np.sin(state.theta),\n",
    "                       np.cos(state.theta)]])\n",
    "        for i in range(len(self.leg_line)):\n",
    "            self.leg_line[i].set_xdata(state.x +\n",
    "                                       state.r * R[0, :].dot(self.leg_data[i]))\n",
    "            self.leg_line[i].set_ydata(state.z +\n",
    "                                       state.r * R[1, :].dot(self.leg_data[i]))\n",
    "\n",
    "        self.ax.set_title(\"t = {:.1f}\".format(context.get_time()))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def slip_demo():\n",
    "    builder = DiagramBuilder()\n",
    "    plant = builder.AddSystem(SpringLoadedInvertedPendulum())\n",
    "\n",
    "    # Parameters from Geyer05, Figure 2.4\n",
    "    # Note: Geyer uses angle of attack = 90 - touchdown_angle\n",
    "    touchdown_angle = np.deg2rad(30)\n",
    "    Etilde = 1.61\n",
    "\n",
    "    s = SLIPState(np.zeros(8))\n",
    "    s.z = 0.9\n",
    "    s.theta = touchdown_angle\n",
    "    s.r = 1\n",
    "    s.xdot = plant.apex_velocity_from_dimensionless_system_energy(Etilde, s.z)\n",
    "\n",
    "    visualizer = builder.AddSystem(SLIPVisualizer(show=False))\n",
    "    builder.Connect(plant.get_output_port(0), visualizer.get_input_port(0))\n",
    "\n",
    "    log = LogVectorOutput(plant.get_output_port(0), builder)\n",
    "\n",
    "    command = builder.AddSystem(ConstantVectorSource([touchdown_angle]))\n",
    "    builder.Connect(command.get_output_port(0), plant.get_input_port(0))\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    simulator = Simulator(diagram)\n",
    "    context = simulator.get_mutable_context()\n",
    "    context.SetAccuracy(1e-10)\n",
    "\n",
    "    context.get_mutable_continuous_state_vector().SetFromVector(s[:])\n",
    "\n",
    "    AdvanceToAndVisualize(simulator, visualizer, 2.0)\n",
    "\n",
    "slip_demo()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Apex-to-apex map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def slip_map():\n",
    "    plant = SpringLoadedInvertedPendulum()\n",
    "\n",
    "    # Parameters from Geyer05, Figure 2.4\n",
    "    # Note: Geyer uses angle of attack = 90 - touchdown_angle\n",
    "    touchdown_angle = np.deg2rad(30)\n",
    "    Etilde = 1.61\n",
    "\n",
    "    s = SLIPState(np.zeros(8))\n",
    "    s.theta = touchdown_angle\n",
    "    s.r = 1\n",
    "\n",
    "    simulator = Simulator(plant)\n",
    "    context = simulator.get_mutable_context()\n",
    "    plant.get_input_port(0).FixValue(context, [touchdown_angle])\n",
    "    context.SetAccuracy(1e-5)\n",
    "\n",
    "    zs = np.linspace(np.cos(touchdown_angle) + 0.001, 0.95, 25)\n",
    "    zns = 0 * zs\n",
    "    for i in range(len(zs)):\n",
    "        s.z = zs[i]\n",
    "        s.xdot = plant.apex_velocity_from_dimensionless_system_energy(Etilde, s.z)\n",
    "        context.SetTime(0.)\n",
    "        context.get_mutable_continuous_state_vector().SetFromVector(s[:])\n",
    "        simulator.Initialize()\n",
    "        # Note: With this duration, I sometimes get an extra \"touchdown\" after the\n",
    "        # apex, which results in apex-touchdown; touchdown-takeoff-apex on the\n",
    "        # console.  It's not a double reset, the consecutive touchdowns are two\n",
    "        # different sims.\n",
    "        simulator.AdvanceTo(0.6)\n",
    "        zns[i] = plant.last_apex\n",
    "        plant.last_apex = None\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.plot(zs, zns)\n",
    "    ax.plot(zs, zs)\n",
    "    ax.axis(\"equal\")\n",
    "    ax.set_xlim([zs[0], zs[-1]])\n",
    "    ax.set_ylim([zs[0], zs[-1]])\n",
    "    ax.set_xlabel(\"apex height z[n]\")\n",
    "    ax.set_ylabel(\"apex height z[n+1]\")\n",
    "\n",
    "    display(mpld3.display())\n",
    "\n",
    "slip_map()"
   ]
  },
  {
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}