Skip to content

Commit 3f69b32

Browse files
RussTedrakeRussTedrake
committed
add quadrotor2d example
1 parent 403ca21 commit 3f69b32

File tree

4 files changed

+189
-3
lines changed

4 files changed

+189
-3
lines changed

‎src/CMakeLists.txt

+12-3
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,16 @@
22
find_package(PythonInterp REQUIRED)
33
find_program(PYCODESTYLE_EXECUTABLE pycodestyle REQUIRED) # pip install pycodestyle
44

5+
function(add_pystyle_test pyfile)
6+
7+
get_filename_component(pydir ${CMAKE_CURRENT_SOURCE_DIR} NAME)
8+
9+
add_test(NAME ${pydir}_${pyfile}_pycodestyle
10+
COMMAND "${PYCODESTYLE_EXECUTABLE}" --count --show-source ${pyfile}
11+
WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
12+
13+
endfunction(add_pystyle_test)
14+
515
function(add_pytest pyfile)
616

717
get_filename_component(pydir ${CMAKE_CURRENT_SOURCE_DIR} NAME)
@@ -16,9 +26,7 @@ function(add_pytest pyfile)
1626
"PYTHONPATH=${PYTHONPATH}"
1727
)
1828

19-
add_test(NAME ${pydir}_${pyfile}_pycodestyle
20-
COMMAND "${PYCODESTYLE_EXECUTABLE}" --count --show-source ${pyfile}
21-
WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
29+
add_pystyle_test(${pyfile})
2230

2331
endfunction(add_pytest)
2432

@@ -42,6 +50,7 @@ add_subdirectory(double_integrator)
4250
add_subdirectory(double_pendulum)
4351
add_subdirectory(lyapunov)
4452
add_subdirectory(pendulum)
53+
add_subdirectory(quadrotor2d)
4554
add_subdirectory(rimless_wheel)
4655
add_subdirectory(simple)
4756
add_subdirectory(van_der_pol)

‎src/quadrotor2d/CMakeLists.txt

+4
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
2+
add_pystyle_test(quadrotor2d.py)
3+
4+
add_pytest(lqr.py --duration 0.1 --trials 1)

‎src/quadrotor2d/lqr.py

+55
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
2+
import argparse
3+
import numpy as np
4+
5+
from pydrake.systems.analysis import Simulator
6+
from pydrake.systems.controllers import LinearQuadraticRegulator
7+
from pydrake.systems.framework import DiagramBuilder
8+
9+
from quadrotor2d import Quadrotor2D, Quadrotor2DVisualizer
10+
11+
12+
parser = argparse.ArgumentParser()
13+
parser.add_argument("-N", "--trials",
14+
type=int,
15+
help="Number of trials to run.",
16+
default=5)
17+
parser.add_argument("-T", "--duration",
18+
type=float,
19+
help="Duration to run each sim.",
20+
default=4.0)
21+
args = parser.parse_args()
22+
23+
24+
def QuadrotorLQR(plant):
25+
context = plant.CreateDefaultContext()
26+
context.SetContinuousState(np.zeros([6, 1]))
27+
context.FixInputPort(0, plant.mass*plant.gravity/2.*np.array([1, 1]))
28+
29+
Q = np.diag([10, 10, 10, 1, 1, (plant.length/2./np.pi)])
30+
R = np.array([[0.1, 0.05], [0.05, 0.1]])
31+
32+
return LinearQuadraticRegulator(plant, context, Q, R)
33+
34+
35+
builder = DiagramBuilder()
36+
37+
plant = builder.AddSystem(Quadrotor2D())
38+
visualizer = builder.AddSystem(Quadrotor2DVisualizer())
39+
builder.Connect(plant.get_output_port(0), visualizer.get_input_port(0))
40+
41+
controller = builder.AddSystem(QuadrotorLQR(plant))
42+
builder.Connect(controller.get_output_port(0), plant.get_input_port(0))
43+
builder.Connect(plant.get_output_port(0), controller.get_input_port(0))
44+
45+
diagram = builder.Build()
46+
47+
simulator = Simulator(diagram)
48+
simulator.set_target_realtime_rate(1.0)
49+
context = simulator.get_mutable_context()
50+
51+
52+
for i in range(args.trials):
53+
context.set_time(0.)
54+
context.SetContinuousState(np.random.randn(6,))
55+
simulator.StepTo(args.duration)

‎src/quadrotor2d/quadrotor2d.py

+118
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
2+
import numpy as np
3+
4+
from underactuated import PyPlotVisualizer
5+
from pydrake.systems.framework import BasicVector_, LeafSystem_, PortDataType
6+
from pydrake.systems.scalar_conversion import TemplateSystem
7+
8+
9+
# Note: In order to use the Python system with drake's autodiff features, we
10+
# have to add a little "TemplateSystem" boilerplate (for now). For details,
11+
# see https://drake.mit.edu/pydrake/pydrake.systems.scalar_conversion.html
12+
@TemplateSystem.define("Quadrotor2D_")
13+
def Quadrotor2D_(T):
14+
15+
class Impl(LeafSystem_[T]):
16+
def _construct(self, converter=None):
17+
LeafSystem_[T].__init__(self, converter)
18+
# two inputs (thrust)
19+
self._DeclareVectorInputPort("u", BasicVector_[T](2))
20+
# six outputs (full state)
21+
self._DeclareVectorOutputPort("x", BasicVector_[T](6),
22+
self._CopyStateOut)
23+
# three positions, three velocities
24+
self._DeclareContinuousState(3, 3, 0)
25+
26+
# parameters based on [Bouadi, Bouchoucha, Tadjine, 2007]
27+
self.length = 0.25 # length of rotor arm
28+
self.mass = 0.486 # mass of quadrotor
29+
self.inertia = 0.00383 # moment of inertia
30+
self.gravity = 9.81 # gravity
31+
32+
def _construct_copy(self, other, converter=None):
33+
Impl._construct(self, converter=converter)
34+
35+
def _CopyStateOut(self, context, output):
36+
x = context.get_continuous_state_vector().CopyToVector()
37+
y = output.SetFromVector(x)
38+
39+
def _DoCalcTimeDerivatives(self, context, derivatives):
40+
x = context.get_continuous_state_vector().CopyToVector()
41+
u = self.EvalVectorInput(context, 0).CopyToVector()
42+
q = x[:3]
43+
qdot = x[3:]
44+
qddot = np.array([-np.sin(q[2])/self.mass*(u[0] + u[1]),
45+
np.cos(x[2])/self.mass*(u[0] + u[1])
46+
- self.gravity,
47+
self.length/self.inertia*(-u[0] + u[1])])
48+
derivatives.get_mutable_vector().SetFromVector(
49+
np.concatenate((qdot, qddot)))
50+
51+
return Impl
52+
53+
54+
Quadrotor2D = Quadrotor2D_[None] # Default instantiation
55+
56+
57+
class Quadrotor2DVisualizer(PyPlotVisualizer):
58+
def __init__(self):
59+
PyPlotVisualizer.__init__(self)
60+
self._DeclareInputPort(PortDataType.kVectorValued, 6)
61+
self.ax.set_aspect('equal')
62+
self.ax.set_xlim(-2, 2)
63+
self.ax.set_ylim(-1, 1)
64+
65+
self.length = .25 # moment arm (meters)
66+
67+
self.base = np.vstack((1.2*self.length*np.array([1, -1, -1, 1, 1]),
68+
0.025*np.array([1, 1, -1, -1, 1])))
69+
self.pin = np.vstack((0.005*np.array([1, 1, -1, -1, 1]),
70+
.1*np.array([1, 0, 0, 1, 1])))
71+
a = np.linspace(0, 2*np.pi, 50)
72+
self.prop = np.vstack((self.length/1.5*np.cos(a), .1+.02*np.sin(2*a)))
73+
74+
self.base_fill = self.ax.fill(self.base[0, :], self.base[1, :],
75+
zorder=1, edgecolor='k',
76+
facecolor=[.6, .6, .6])
77+
self.left_pin_fill = self.ax.fill(self.pin[0, :], self.pin[1, :],
78+
zorder=0, edgecolor='k',
79+
facecolor=[0, 0, 0])
80+
self.right_pin_fill = self.ax.fill(self.pin[0, :], self.pin[1, :],
81+
zorder=0, edgecolor='k',
82+
facecolor=[0, 0, 0])
83+
self.left_prop_fill = self.ax.fill(self.prop[0, :], self.prop[0, :],
84+
zorder=0, edgecolor='k',
85+
facecolor=[0, 0, 1])
86+
self.right_prop_fill = self.ax.fill(self.prop[0, :], self.prop[0, :],
87+
zorder=0, edgecolor='k',
88+
facecolor=[0, 0, 1])
89+
90+
def draw(self, context):
91+
x = self.EvalVectorInput(context, 0).CopyToVector()
92+
R = np.array([[np.cos(x[2]), -np.sin(x[2])],
93+
[np.sin(x[2]), np.cos(x[2])]])
94+
95+
p = np.dot(R, self.base)
96+
self.base_fill[0].get_path().vertices[:, 0] = x[0] + p[0, :]
97+
self.base_fill[0].get_path().vertices[:, 1] = x[1] + p[1, :]
98+
99+
p = np.dot(R, np.vstack((-self.length + self.pin[0, :],
100+
self.pin[1, :])))
101+
self.left_pin_fill[0].get_path().vertices[:, 0] = x[0] + p[0, :]
102+
self.left_pin_fill[0].get_path().vertices[:, 1] = x[1] + p[1, :]
103+
p = np.dot(R, np.vstack((self.length + self.pin[0, :],
104+
self.pin[1, :])))
105+
self.right_pin_fill[0].get_path().vertices[:, 0] = x[0] + p[0, :]
106+
self.right_pin_fill[0].get_path().vertices[:, 1] = x[1] + p[1, :]
107+
108+
p = np.dot(R, np.vstack((-self.length + self.prop[0, :],
109+
self.prop[1, :])))
110+
self.left_prop_fill[0].get_path().vertices[:, 0] = x[0] + p[0, :]
111+
self.left_prop_fill[0].get_path().vertices[:, 1] = x[1] + p[1, :]
112+
113+
p = np.dot(R, np.vstack((self.length + self.prop[0, :],
114+
self.prop[1, :])))
115+
self.right_prop_fill[0].get_path().vertices[:, 0] = x[0] + p[0, :]
116+
self.right_prop_fill[0].get_path().vertices[:, 1] = x[1] + p[1, :]
117+
118+
self.ax.set_title('t = ' + str(context.get_time()))

0 commit comments

Comments
 (0)