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