Notebook

Welcome! If you are new to Google Colab/Jupyter notebooks, you might take a look at this notebook first.

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.

Notebook Setup¶

The following cell will:

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.
import packages used throughout the notebook.

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.

In [ ]:

import importlib
import sys
from urllib.request import urlretrieve

# Install drake (and underactuated).
if 'google.colab' in sys.modules and importlib.util.find_spec('underactuated') is None:
    urlretrieve(f"http://underactuated.csail.mit.edu/scripts/setup/setup_underactuated_colab.py",
                "setup_underactuated_colab.py")
    from setup_underactuated_colab import setup_underactuated
    setup_underactuated(underactuated_sha='560c2adace05eb20ebd78377582015d5b2d3859a', drake_version='0.27.0', drake_build='release')

server_args = []
if 'google.colab' in sys.modules:
  server_args = ['--ngrok_http_tunnel']

# Start a single meshcat server instance to use for the remainder of this notebook.
from meshcat.servers.zmqserver import start_zmq_server_as_subprocess
proc, zmq_url, web_url = start_zmq_server_as_subprocess(server_args=server_args)

    
import numpy as np
import matplotlib.pyplot as plt
import meshcat
import meshcat.transformations as tf
import time

from pydrake.common.containers import namedview
from pydrake.common.value import AbstractValue
from pydrake.math import RigidTransform, RotationMatrix
from pydrake.geometry import FramePoseVector
from pydrake.multibody.plant import MultibodyPlant
from pydrake.multibody.parsing import Parser
from pydrake.systems.framework import (BasicVector, BasicVector_, LeafSystem_,
                                       LeafSystem)
from pydrake.systems.scalar_conversion import TemplateSystem
from pydrake.all import (
    SceneGraph, DiagramBuilder, Parser, ConnectPlanarSceneGraphVisualizer, DrakeVisualizer, 
    ConnectMeshcatVisualizer, DirectCollocation, Solve, PiecewisePolynomial
)

from underactuated import FindResource

# Note: In order to use the Python system with drake's autodiff features, we
# have to add a little "TemplateSystem" boilerplate (for now).  For details,
# see https://drake.mit.edu/pydrake/pydrake.systems.scalar_conversion.html

GliderState = namedview(
    "GliderState", ["x", "z", "pitch", "elevator", "xdot", "zdot", "pitchdot"])

@TemplateSystem.define("GliderPlant_")
def GliderPlant_(T):

    class Impl(LeafSystem_[T]):

        def _construct(self, converter=None):
            LeafSystem_[T].__init__(self, converter)
            # one inputs (elevator_velocity)
            self.DeclareVectorInputPort("elevatordot", BasicVector_[T](1))
            # four positions, three velocities
            self.DeclareContinuousState(4, 3, 0)
            # seven outputs (full state)
            self.DeclareVectorOutputPort("state", BasicVector_[T](7),
                                         self.CopyStateOut)
            self.DeclareVectorOutputPort(
                "forces", BasicVector_[T](2),
                self.OutputForces)

            # parameters based on Rick Cory's "R1 = no dihedral" model.
            self.Sw = 0.0885  # surface area of wing + fuselage + tail.
            self.Se = 0.0147  # surface area of elevator.
            self.lw = 0  # horizontal offset of wing center.
            self.le = 0.022  # elevator aerodynamic center from hinge.
            #self.lh = 0.317  # elevator hinge.
            self.lh = 0.27   # elevator hinge.
            self.inertia = 0.0015  # body inertia.
            self.m = 0.08  # body mass.
            self.rho = 1.204  # air density (kg/m^3).
            self.gravity = 9.81  # gravity

            # TODO(russt): Declare elevator constraints:
            self.elevator_lower_limit = -0.9473
            self.elevator_upper_limit = 0.4463

        def _construct_copy(self, other, converter=None):
            Impl._construct(self, converter=converter)

        def ComputeForces(self, context):
            s = GliderState(
                context.get_mutable_continuous_state_vector().CopyToVector())
            elevatordot = self.EvalVectorInput(context, 0)[0]

            xwdot = s.xdot + self.lw * s.pitchdot * np.sin(s.pitch)
            zwdot = s.zdot + self.lw * s.pitchdot * np.cos(s.pitch)
            vw = np.sqrt(zwdot**2 + xwdot**2)
            fw = -self.rho * self.Sw * (np.sin(s.pitch)*xwdot + np.cos(s.pitch)*zwdot)*vw

            e = s.pitch + s.elevator
            edot = s.pitchdot + elevatordot
            xedot = s.xdot + self.lh * s.pitchdot * np.sin(s.pitch) \
                + self.le * edot * np.sin(e)
            zedot = s.zdot + self.lh * s.pitchdot * np.cos(s.pitch) \
                + self.le * edot * np.cos(e)
            ve = np.sqrt(zedot**2 + xedot**2)
            fe = -self.rho * self.Se * (np.sin(e)*xedot + np.cos(e)*zedot)*ve
            
            return fw, fe

            
        def DoCalcTimeDerivatives(self, context, derivatives):
            s = GliderState(
                context.get_mutable_continuous_state_vector().CopyToVector())
            elevatordot = self.EvalVectorInput(context, 0)[0]
            e = s.pitch + s.elevator
            fw, fe = self.ComputeForces(context)
            
            sdot = GliderState(s[:])
            sdot[0:3] = s[4:7]
            sdot.elevator = elevatordot
            sdot.xdot = (fw * np.sin(s.pitch) + fe * np.sin(e)) / self.m
            sdot.zdot = (fw * np.cos(s.pitch) + fe * np.cos(e)) / self.m - self.gravity
            sdot.pitchdot = (fw * self.lw + fe * (self.lh * np.cos(s.elevator) + self.le)) / self.inertia
            derivatives.get_mutable_vector().SetFromVector(sdot[:])

        def CopyStateOut(self, context, output):
            x = context.get_continuous_state_vector().CopyToVector()
            output.SetFromVector(x)

        def OutputForces(self, context, output):
            fw, fe = self.ComputeForces(context)
            output.SetFromVector([fw, fe])
            
            
    return Impl

GliderPlant = GliderPlant_[None]  # Default instantiation


# To use glider.urdf for visualization, follow the pattern from e.g.
# drake::examples::quadrotor::QuadrotorGeometry.
class GliderGeometry(LeafSystem):

    def __init__(self, scene_graph):
        LeafSystem.__init__(self)
        assert scene_graph

        mbp = MultibodyPlant(1.0)  # Timestep doesn't matter, and this avoids a warning
        parser = Parser(mbp, scene_graph)
        model_id = parser.AddModelFromFile(
            FindResource("models/glider/glider.urdf"))
        mbp.Finalize()
        self.source_id = mbp.get_source_id()
        self.body_frame_id = mbp.GetBodyFrameIdOrThrow(
            mbp.GetBodyIndices(model_id)[0])
        self.elevator_frame_id = mbp.GetBodyFrameIdOrThrow(
            mbp.GetBodyIndices(model_id)[1])

        self.DeclareVectorInputPort("state", BasicVector(7))
        self.DeclareAbstractOutputPort(
            "geometry_pose", lambda: AbstractValue.Make(FramePoseVector()),
            self.OutputGeometryPose)

    def OutputGeometryPose(self, context, poses):
        assert self.body_frame_id.is_valid()
        assert self.elevator_frame_id.is_valid()
        lh = 0.317  # elevator hinge.
        state = GliderState(self.get_input_port(0).Eval(context))
        body_pose = RigidTransform(RotationMatrix.MakeYRotation(state.pitch),
                                   [state.x, 0, state.z])
        elevator_pose = RigidTransform(
            RotationMatrix.MakeYRotation(state.pitch + state.elevator), [
                state.x - lh * np.cos(state.pitch), 0,
                state.z + lh * np.sin(state.pitch)
            ])
        poses.get_mutable_value().set_value(self.body_frame_id, body_pose)
        poses.get_mutable_value().set_value(self.elevator_frame_id,
                                            elevator_pose)

    @staticmethod
    def AddToBuilder(builder, glider_state_port, scene_graph):
        assert builder
        assert scene_graph

        geom = builder.AddSystem(GliderGeometry(scene_graph))
        builder.Connect(glider_state_port, geom.get_input_port(0))
        builder.Connect(geom.get_output_port(0),
                        scene_graph.get_source_pose_port(geom.source_id))

        return geom

class GliderForceVisualizer(LeafSystem):

    def __init__(self, visualizer):
        LeafSystem.__init__(self)

        inspector = visualizer._scene_graph.model_inspector()
        source_name = inspector.GetOwningSourceName(inspector.all_frame_ids()[0])
        self.visualizer = visualizer
        self.vis = visualizer.vis[visualizer.prefix][source_name]['glider']
        draw_period = 1 / 30.
        self.DeclarePeriodicPublish(draw_period, 0.0)
        self.DeclareVectorInputPort("forces", BasicVector(2))
        
    def load(self):
        self.vis["fuselage/force"].set_transform(tf.rotation_matrix(np.pi/2.0, [1, 0, 0]) @ tf.translation_matrix([-.05, 0, 0]))
        self.vis["fuselage/force/scale"].set_transform(tf.scale_matrix(.1, direction=[0, 1, 0]))
        self.vis["fuselage/force/scale/cylinder"].set_object(meshcat.geometry.Cylinder(
                        height=1., radius=.01),
                        meshcat.geometry.MeshLambertMaterial(color=0x33cc33))
        self.vis["fuselage/force/scale/cylinder"].set_transform(tf.translation_matrix([0, .5, 0]))
        self.vis["elevator/force"].set_transform(tf.rotation_matrix(np.pi/2.0, [1, 0, 0]) @ tf.translation_matrix([-.022, 0, 0]))
        self.vis["elevator/force/scale"].set_transform(tf.scale_matrix(.1, direction=[0, 1, 0]))
        self.vis["elevator/force/scale/cylinder"].set_object(meshcat.geometry.Cylinder(
                        height=1., radius=.01),
                        meshcat.geometry.MeshLambertMaterial(color=0x33cc33))
        self.vis["elevator/force/scale/cylinder"].set_transform(tf.translation_matrix([0, .5, 0]))

    def DoPublish(self, context, event):
        LeafSystem.DoPublish(self, context, event)
        
        force = self.get_input_port(0).Eval(context)
        self.vis["fuselage/force/scale"].set_transform(tf.scale_matrix(force[0], direction=[0, 1, 0]))
        self.vis["elevator/force/scale"].set_transform(tf.scale_matrix(force[1], direction=[0, 1, 0]))

        # TODO: Figure out why meshcat animation discards the scaling (only keeps position/orientation).  Can I generalize it?
        if self.visualizer._is_recording:
            with self.visualizer._animation.at_frame(
                    self.vis["fuselage/force/scale"], self.visualizer._recording_frame_num) as f:
                f.set_transform(tf.scale_matrix(force[0], direction=[0, 1, 0]))
                f.set_property('visible','boolean',False)
            with self.visualizer._animation.at_frame(
                    self.vis["elevator/force/scale"], self.visualizer._recording_frame_num) as f:
                f.set_transform(tf.scale_matrix(force[1], direction=[0, 1, 0]))
                f.set_property('visible','boolean',False)
    


def draw_glider(x):
    builder = DiagramBuilder()
    glider = builder.AddSystem(GliderPlant())
    scene_graph = builder.AddSystem(SceneGraph())
    GliderGeometry.AddToBuilder(builder, glider.GetOutputPort("state"), scene_graph)
    meshcat_vis = ConnectMeshcatVisualizer(builder, 
        scene_graph=scene_graph, 
        zmq_url=zmq_url)
#    meshcat_vis.set_planar_viewpoint(xmin=-4, xmax=1, ymin=-1, ymax=1)

    diagram = builder.Build()
    context = diagram.CreateDefaultContext()
    context.SetContinuousState(x)
    meshcat_vis.load()
    diagram.Publish(context)

    return meshcat_vis

In [ ]:

def dircol_perching():
    builder = DiagramBuilder()
    glider = builder.AddSystem(GliderPlant())
    scene_graph = builder.AddSystem(SceneGraph())
    GliderGeometry.AddToBuilder(builder, glider.GetOutputPort("state"), scene_graph)
    visualizer = ConnectMeshcatVisualizer(builder, 
        scene_graph=scene_graph, 
        zmq_url=zmq_url)
    visualizer.vis.delete()
    visualizer.set_planar_viewpoint(xmin=-4, xmax=1, ymin=-1, ymax=1)
    visualizer.vis['/Background'].set_property('visible', False)
    visualizer.vis['/Grid'].set_property('visible', False)
    visualizer.vis['/Axes'].set_property('visible', False)
    
    diagram = builder.Build()
    context = diagram.CreateDefaultContext()
    visualizer.load()
    diagram.Publish(context)

    N = 20
    dircol = DirectCollocation(glider, glider.CreateDefaultContext(), N, 0.01, 2.0/N)
    dircol.AddEqualTimeIntervalsConstraints()

    # Input limits
    u = dircol.input()
    elevator_velocity_limit = 13  # max servo velocity (rad/sec)
    dircol.AddConstraintToAllKnotPoints(-elevator_velocity_limit <= u[0])
    dircol.AddConstraintToAllKnotPoints(u[0] <= elevator_velocity_limit)
    
    # Initial conditions
    s0 = GliderState(np.zeros(7))
    s0.x = -3.5
    s0.z = 0.1
    s0.xdot = 7.0
    dircol.AddBoundingBoxConstraint(s0[:], s0[:], dircol.initial_state())
    context.SetContinuousState(s0[:])
    diagram.Publish(context)

    # Final conditions
    s = GliderState(dircol.final_state())
    dircol.AddBoundingBoxConstraint(0, 0, s.x)
    dircol.AddBoundingBoxConstraint(0, 0, s.z)
    dircol.AddBoundingBoxConstraint(-1.0, -np.pi/6.0, s.pitch)
    dircol.AddBoundingBoxConstraint(-2.0, 2.0, s.xdot)
    dircol.AddBoundingBoxConstraint(-2.0, 2.0, s.zdot)

    # State constraints
    s = GliderState(dircol.state())
    dircol.AddConstraintToAllKnotPoints(s.x >= s0.x)  # TODO: s0.x <= s.x gives an error.
    dircol.AddConstraintToAllKnotPoints(s.x <= 1)
    dircol.AddConstraintToAllKnotPoints(s.z >= -1)
    dircol.AddConstraintToAllKnotPoints(s.z <= 1)
    dircol.AddConstraintToAllKnotPoints(s.pitch >= -np.pi/2.0)
    dircol.AddConstraintToAllKnotPoints(s.pitch <= np.pi/2.0)
    dircol.AddConstraintToAllKnotPoints(s.elevator >= glider.elevator_lower_limit)
    dircol.AddConstraintToAllKnotPoints(s.elevator <= glider.elevator_upper_limit)

    
    # Cost
    dircol.AddRunningCost(10*u**2)
    
    sf_d = GliderState(np.zeros(7))
    sf_d.pitch = -np.pi/4.0
    dircol.AddQuadraticErrorCost(np.diag([10, 10, 10, 0, 1, 1, 1]), sf_d[:], dircol.final_state())
    
    def plot_trajectory(times, states):
        vertices = np.vstack([states[0,:], 0*times, states[1,:]])
        visualizer.vis["dircol"].set_object(meshcat.geometry.Line(meshcat.geometry.PointsGeometry(vertices),meshcat.geometry.LineBasicMaterial(color=0x0000dd)))
        
    dircol.AddStateTrajectoryCallback(plot_trajectory)

    initial_x_trajectory = PiecewisePolynomial.FirstOrderHold(
        [0., .8], np.column_stack((s0[:], sf_d[:])))
    dircol.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)
    
    result = Solve(dircol)
    if not result.is_success():
        infeasible = result.GetInfeasibleConstraints(dircol)
        print("Infeasible constraints:")
        for i in range(len(infeasible)):
            print(infeasible[i])

    x_traj = dircol.ReconstructStateTrajectory(result)
    u_traj = dircol.ReconstructInputTrajectory(result)

    # Animate trajectory
    visualizer.start_recording()
    for t in np.hstack((np.arange(x_traj.start_time(), x_traj.end_time(), visualizer.draw_period), x_traj.end_time())):
        context.SetTime(t)
        context.SetContinuousState(x_traj.value(t))
        diagram.Publish(context)

    visualizer.stop_recording()
    visualizer.publish_recording()

    return x_traj, u_traj

x_traj, u_traj = dircol_perching()

In [ ]:

ts = np.linspace(u_traj.start_time(), u_traj.end_time(), 301)
s = GliderState(x_traj.vector_values(ts))

fig, ax = plt.subplots()
ax.plot(ts, u_traj.vector_values(ts).T, label='elevatordot');
ax.plot(ts, s.pitch, label='pitch')
#ax.plot(ts, s.elevator, label='elevator')
#ax.plot(ts, s.pitchdot, label='pitchdot')

ax.legend();

In [ ]: