{ "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/trajopt.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='560c2adace05eb20ebd78377582015d5b2d3859a', 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", "\n", "# Start a single meshcat server instance to use for the remainder of this notebook.\n", "from meshcat.servers.zmqserver import start_zmq_server_as_subprocess\n", "proc, zmq_url, web_url = start_zmq_server_as_subprocess(server_args=server_args)\n", "\n", " \n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import meshcat\n", "import meshcat.transformations as tf\n", "import time\n", "\n", "from pydrake.common.containers import namedview\n", "from pydrake.common.value import AbstractValue\n", "from pydrake.math import RigidTransform, RotationMatrix\n", "from pydrake.geometry import FramePoseVector\n", "from pydrake.multibody.plant import MultibodyPlant\n", "from pydrake.multibody.parsing import Parser\n", "from pydrake.systems.framework import (BasicVector, BasicVector_, LeafSystem_,\n", " LeafSystem)\n", "from pydrake.systems.scalar_conversion import TemplateSystem\n", "from pydrake.all import (\n", " SceneGraph, DiagramBuilder, Parser, ConnectPlanarSceneGraphVisualizer, DrakeVisualizer, \n", " ConnectMeshcatVisualizer, DirectCollocation, Solve, PiecewisePolynomial\n", ")\n", "\n", "from underactuated import FindResource\n", "\n", "# Note: In order to use the Python system with drake's autodiff features, we\n", "# have to add a little \"TemplateSystem\" boilerplate (for now). For details,\n", "# see https://drake.mit.edu/pydrake/pydrake.systems.scalar_conversion.html\n", "\n", "GliderState = namedview(\n", " \"GliderState\", [\"x\", \"z\", \"pitch\", \"elevator\", \"xdot\", \"zdot\", \"pitchdot\"])\n", "\n", "@TemplateSystem.define(\"GliderPlant_\")\n", "def GliderPlant_(T):\n", "\n", " class Impl(LeafSystem_[T]):\n", "\n", " def _construct(self, converter=None):\n", " LeafSystem_[T].__init__(self, converter)\n", " # one inputs (elevator_velocity)\n", " self.DeclareVectorInputPort(\"elevatordot\", BasicVector_[T](1))\n", " # four positions, three velocities\n", " self.DeclareContinuousState(4, 3, 0)\n", " # seven outputs (full state)\n", " self.DeclareVectorOutputPort(\"state\", BasicVector_[T](7),\n", " self.CopyStateOut)\n", " self.DeclareVectorOutputPort(\n", " \"forces\", BasicVector_[T](2),\n", " self.OutputForces)\n", "\n", " # parameters based on Rick Cory's \"R1 = no dihedral\" model.\n", " self.Sw = 0.0885 # surface area of wing + fuselage + tail.\n", " self.Se = 0.0147 # surface area of elevator.\n", " self.lw = 0 # horizontal offset of wing center.\n", " self.le = 0.022 # elevator aerodynamic center from hinge.\n", " #self.lh = 0.317 # elevator hinge.\n", " self.lh = 0.27 # elevator hinge.\n", " self.inertia = 0.0015 # body inertia.\n", " self.m = 0.08 # body mass.\n", " self.rho = 1.204 # air density (kg/m^3).\n", " self.gravity = 9.81 # gravity\n", "\n", " # TODO(russt): Declare elevator constraints:\n", " self.elevator_lower_limit = -0.9473\n", " self.elevator_upper_limit = 0.4463\n", "\n", " def _construct_copy(self, other, converter=None):\n", " Impl._construct(self, converter=converter)\n", "\n", " def ComputeForces(self, context):\n", " s = GliderState(\n", " context.get_mutable_continuous_state_vector().CopyToVector())\n", " elevatordot = self.EvalVectorInput(context, 0)[0]\n", "\n", " xwdot = s.xdot + self.lw * s.pitchdot * np.sin(s.pitch)\n", " zwdot = s.zdot + self.lw * s.pitchdot * np.cos(s.pitch)\n", " vw = np.sqrt(zwdot**2 + xwdot**2)\n", " fw = -self.rho * self.Sw * (np.sin(s.pitch)*xwdot + np.cos(s.pitch)*zwdot)*vw\n", "\n", " e = s.pitch + s.elevator\n", " edot = s.pitchdot + elevatordot\n", " xedot = s.xdot + self.lh * s.pitchdot * np.sin(s.pitch) \\\n", " + self.le * edot * np.sin(e)\n", " zedot = s.zdot + self.lh * s.pitchdot * np.cos(s.pitch) \\\n", " + self.le * edot * np.cos(e)\n", " ve = np.sqrt(zedot**2 + xedot**2)\n", " fe = -self.rho * self.Se * (np.sin(e)*xedot + np.cos(e)*zedot)*ve\n", " \n", " return fw, fe\n", "\n", " \n", " def DoCalcTimeDerivatives(self, context, derivatives):\n", " s = GliderState(\n", " context.get_mutable_continuous_state_vector().CopyToVector())\n", " elevatordot = self.EvalVectorInput(context, 0)[0]\n", " e = s.pitch + s.elevator\n", " fw, fe = self.ComputeForces(context)\n", " \n", " sdot = GliderState(s[:])\n", " sdot[0:3] = s[4:7]\n", " sdot.elevator = elevatordot\n", " sdot.xdot = (fw * np.sin(s.pitch) + fe * np.sin(e)) / self.m\n", " sdot.zdot = (fw * np.cos(s.pitch) + fe * np.cos(e)) / self.m - self.gravity\n", " sdot.pitchdot = (fw * self.lw + fe * (self.lh * np.cos(s.elevator) + self.le)) / self.inertia\n", " derivatives.get_mutable_vector().SetFromVector(sdot[:])\n", "\n", " def CopyStateOut(self, context, output):\n", " x = context.get_continuous_state_vector().CopyToVector()\n", " output.SetFromVector(x)\n", "\n", " def OutputForces(self, context, output):\n", " fw, fe = self.ComputeForces(context)\n", " output.SetFromVector([fw, fe])\n", " \n", " \n", " return Impl\n", "\n", "GliderPlant = GliderPlant_[None] # Default instantiation\n", "\n", "\n", "# To use glider.urdf for visualization, follow the pattern from e.g.\n", "# drake::examples::quadrotor::QuadrotorGeometry.\n", "class GliderGeometry(LeafSystem):\n", "\n", " def __init__(self, scene_graph):\n", " LeafSystem.__init__(self)\n", " assert scene_graph\n", "\n", " mbp = MultibodyPlant(1.0) # Timestep doesn't matter, and this avoids a warning\n", " parser = Parser(mbp, scene_graph)\n", " model_id = parser.AddModelFromFile(\n", " FindResource(\"models/glider/glider.urdf\"))\n", " mbp.Finalize()\n", " self.source_id = mbp.get_source_id()\n", " self.body_frame_id = mbp.GetBodyFrameIdOrThrow(\n", " mbp.GetBodyIndices(model_id)[0])\n", " self.elevator_frame_id = mbp.GetBodyFrameIdOrThrow(\n", " mbp.GetBodyIndices(model_id)[1])\n", "\n", " self.DeclareVectorInputPort(\"state\", BasicVector(7))\n", " self.DeclareAbstractOutputPort(\n", " \"geometry_pose\", lambda: AbstractValue.Make(FramePoseVector()),\n", " self.OutputGeometryPose)\n", "\n", " def OutputGeometryPose(self, context, poses):\n", " assert self.body_frame_id.is_valid()\n", " assert self.elevator_frame_id.is_valid()\n", " lh = 0.317 # elevator hinge.\n", " state = GliderState(self.get_input_port(0).Eval(context))\n", " body_pose = RigidTransform(RotationMatrix.MakeYRotation(state.pitch),\n", " [state.x, 0, state.z])\n", " elevator_pose = RigidTransform(\n", " RotationMatrix.MakeYRotation(state.pitch + state.elevator), [\n", " state.x - lh * np.cos(state.pitch), 0,\n", " state.z + lh * np.sin(state.pitch)\n", " ])\n", " poses.get_mutable_value().set_value(self.body_frame_id, body_pose)\n", " poses.get_mutable_value().set_value(self.elevator_frame_id,\n", " elevator_pose)\n", "\n", " @staticmethod\n", " def AddToBuilder(builder, glider_state_port, scene_graph):\n", " assert builder\n", " assert scene_graph\n", "\n", " geom = builder.AddSystem(GliderGeometry(scene_graph))\n", " builder.Connect(glider_state_port, geom.get_input_port(0))\n", " builder.Connect(geom.get_output_port(0),\n", " scene_graph.get_source_pose_port(geom.source_id))\n", "\n", " return geom\n", "\n", "class GliderForceVisualizer(LeafSystem):\n", "\n", " def __init__(self, visualizer):\n", " LeafSystem.__init__(self)\n", "\n", " inspector = visualizer._scene_graph.model_inspector()\n", " source_name = inspector.GetOwningSourceName(inspector.all_frame_ids()[0])\n", " self.visualizer = visualizer\n", " self.vis = visualizer.vis[visualizer.prefix][source_name]['glider']\n", " draw_period = 1 / 30.\n", " self.DeclarePeriodicPublish(draw_period, 0.0)\n", " self.DeclareVectorInputPort(\"forces\", BasicVector(2))\n", " \n", " def load(self):\n", " self.vis[\"fuselage/force\"].set_transform(tf.rotation_matrix(np.pi/2.0, [1, 0, 0]) @ tf.translation_matrix([-.05, 0, 0]))\n", " self.vis[\"fuselage/force/scale\"].set_transform(tf.scale_matrix(.1, direction=[0, 1, 0]))\n", " self.vis[\"fuselage/force/scale/cylinder\"].set_object(meshcat.geometry.Cylinder(\n", " height=1., radius=.01),\n", " meshcat.geometry.MeshLambertMaterial(color=0x33cc33))\n", " self.vis[\"fuselage/force/scale/cylinder\"].set_transform(tf.translation_matrix([0, .5, 0]))\n", " self.vis[\"elevator/force\"].set_transform(tf.rotation_matrix(np.pi/2.0, [1, 0, 0]) @ tf.translation_matrix([-.022, 0, 0]))\n", " self.vis[\"elevator/force/scale\"].set_transform(tf.scale_matrix(.1, direction=[0, 1, 0]))\n", " self.vis[\"elevator/force/scale/cylinder\"].set_object(meshcat.geometry.Cylinder(\n", " height=1., radius=.01),\n", " meshcat.geometry.MeshLambertMaterial(color=0x33cc33))\n", " self.vis[\"elevator/force/scale/cylinder\"].set_transform(tf.translation_matrix([0, .5, 0]))\n", "\n", " def DoPublish(self, context, event):\n", " LeafSystem.DoPublish(self, context, event)\n", " \n", " force = self.get_input_port(0).Eval(context)\n", " self.vis[\"fuselage/force/scale\"].set_transform(tf.scale_matrix(force[0], direction=[0, 1, 0]))\n", " self.vis[\"elevator/force/scale\"].set_transform(tf.scale_matrix(force[1], direction=[0, 1, 0]))\n", "\n", " # TODO: Figure out why meshcat animation discards the scaling (only keeps position/orientation). Can I generalize it?\n", " if self.visualizer._is_recording:\n", " with self.visualizer._animation.at_frame(\n", " self.vis[\"fuselage/force/scale\"], self.visualizer._recording_frame_num) as f:\n", " f.set_transform(tf.scale_matrix(force[0], direction=[0, 1, 0]))\n", " f.set_property('visible','boolean',False)\n", " with self.visualizer._animation.at_frame(\n", " self.vis[\"elevator/force/scale\"], self.visualizer._recording_frame_num) as f:\n", " f.set_transform(tf.scale_matrix(force[1], direction=[0, 1, 0]))\n", " f.set_property('visible','boolean',False)\n", " \n", "\n", "\n", "def draw_glider(x):\n", " builder = DiagramBuilder()\n", " glider = builder.AddSystem(GliderPlant())\n", " scene_graph = builder.AddSystem(SceneGraph())\n", " GliderGeometry.AddToBuilder(builder, glider.GetOutputPort(\"state\"), scene_graph)\n", " meshcat_vis = ConnectMeshcatVisualizer(builder, \n", " scene_graph=scene_graph, \n", " zmq_url=zmq_url)\n", "# meshcat_vis.set_planar_viewpoint(xmin=-4, xmax=1, ymin=-1, ymax=1)\n", "\n", " diagram = builder.Build()\n", " context = diagram.CreateDefaultContext()\n", " context.SetContinuousState(x)\n", " meshcat_vis.load()\n", " diagram.Publish(context)\n", "\n", " return meshcat_vis" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [] }, "outputs": [], "source": [ "def dircol_perching():\n", " builder = DiagramBuilder()\n", " glider = builder.AddSystem(GliderPlant())\n", " scene_graph = builder.AddSystem(SceneGraph())\n", " GliderGeometry.AddToBuilder(builder, glider.GetOutputPort(\"state\"), scene_graph)\n", " visualizer = ConnectMeshcatVisualizer(builder, \n", " scene_graph=scene_graph, \n", " zmq_url=zmq_url)\n", " visualizer.vis.delete()\n", " visualizer.set_planar_viewpoint(xmin=-4, xmax=1, ymin=-1, ymax=1)\n", " visualizer.vis['/Background'].set_property('visible', False)\n", " visualizer.vis['/Grid'].set_property('visible', False)\n", " visualizer.vis['/Axes'].set_property('visible', False)\n", " \n", " diagram = builder.Build()\n", " context = diagram.CreateDefaultContext()\n", " visualizer.load()\n", " diagram.Publish(context)\n", "\n", " N = 20\n", " dircol = DirectCollocation(glider, glider.CreateDefaultContext(), N, 0.01, 2.0/N)\n", " dircol.AddEqualTimeIntervalsConstraints()\n", "\n", " # Input limits\n", " u = dircol.input()\n", " elevator_velocity_limit = 13 # max servo velocity (rad/sec)\n", " dircol.AddConstraintToAllKnotPoints(-elevator_velocity_limit <= u[0])\n", " dircol.AddConstraintToAllKnotPoints(u[0] <= elevator_velocity_limit)\n", " \n", " # Initial conditions\n", " s0 = GliderState(np.zeros(7))\n", " s0.x = -3.5\n", " s0.z = 0.1\n", " s0.xdot = 7.0\n", " dircol.AddBoundingBoxConstraint(s0[:], s0[:], dircol.initial_state())\n", " context.SetContinuousState(s0[:])\n", " diagram.Publish(context)\n", "\n", " # Final conditions\n", " s = GliderState(dircol.final_state())\n", " dircol.AddBoundingBoxConstraint(0, 0, s.x)\n", " dircol.AddBoundingBoxConstraint(0, 0, s.z)\n", " dircol.AddBoundingBoxConstraint(-1.0, -np.pi/6.0, s.pitch)\n", " dircol.AddBoundingBoxConstraint(-2.0, 2.0, s.xdot)\n", " dircol.AddBoundingBoxConstraint(-2.0, 2.0, s.zdot)\n", "\n", " # State constraints\n", " s = GliderState(dircol.state())\n", " dircol.AddConstraintToAllKnotPoints(s.x >= s0.x) # TODO: s0.x <= s.x gives an error.\n", " dircol.AddConstraintToAllKnotPoints(s.x <= 1)\n", " dircol.AddConstraintToAllKnotPoints(s.z >= -1)\n", " dircol.AddConstraintToAllKnotPoints(s.z <= 1)\n", " dircol.AddConstraintToAllKnotPoints(s.pitch >= -np.pi/2.0)\n", " dircol.AddConstraintToAllKnotPoints(s.pitch <= np.pi/2.0)\n", " dircol.AddConstraintToAllKnotPoints(s.elevator >= glider.elevator_lower_limit)\n", " dircol.AddConstraintToAllKnotPoints(s.elevator <= glider.elevator_upper_limit)\n", "\n", " \n", " # Cost\n", " dircol.AddRunningCost(10*u**2)\n", " \n", " sf_d = GliderState(np.zeros(7))\n", " sf_d.pitch = -np.pi/4.0\n", " dircol.AddQuadraticErrorCost(np.diag([10, 10, 10, 0, 1, 1, 1]), sf_d[:], dircol.final_state())\n", " \n", " def plot_trajectory(times, states):\n", " vertices = np.vstack([states[0,:], 0*times, states[1,:]])\n", " visualizer.vis[\"dircol\"].set_object(meshcat.geometry.Line(meshcat.geometry.PointsGeometry(vertices),meshcat.geometry.LineBasicMaterial(color=0x0000dd)))\n", " \n", " dircol.AddStateTrajectoryCallback(plot_trajectory)\n", "\n", " initial_x_trajectory = PiecewisePolynomial.FirstOrderHold(\n", " [0., .8], np.column_stack((s0[:], sf_d[:])))\n", " dircol.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)\n", " \n", " result = Solve(dircol)\n", " if not result.is_success():\n", " infeasible = result.GetInfeasibleConstraints(dircol)\n", " print(\"Infeasible constraints:\")\n", " for i in range(len(infeasible)):\n", " print(infeasible[i])\n", "\n", " x_traj = dircol.ReconstructStateTrajectory(result)\n", " u_traj = dircol.ReconstructInputTrajectory(result)\n", "\n", " # Animate trajectory\n", " visualizer.start_recording()\n", " for t in np.hstack((np.arange(x_traj.start_time(), x_traj.end_time(), visualizer.draw_period), x_traj.end_time())):\n", " context.SetTime(t)\n", " context.SetContinuousState(x_traj.value(t))\n", " diagram.Publish(context)\n", "\n", " visualizer.stop_recording()\n", " visualizer.publish_recording()\n", "\n", " return x_traj, u_traj\n", "\n", "x_traj, u_traj = dircol_perching()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ts = np.linspace(u_traj.start_time(), u_traj.end_time(), 301)\n", "s = GliderState(x_traj.vector_values(ts))\n", "\n", "fig, ax = plt.subplots()\n", "ax.plot(ts, u_traj.vector_values(ts).T, label='elevatordot');\n", "ax.plot(ts, s.pitch, label='pitch')\n", "#ax.plot(ts, s.elevator, label='elevator')\n", "#ax.plot(ts, s.pitchdot, label='pitchdot')\n", "\n", "ax.legend();" ] }, { "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 }