Skip to content

Commit 00b175e

Browse files
RussTedrakeRussTedrake
committed
add quadrotor2d example
1 parent 403ca21 commit 00b175e

File tree

4 files changed

+192
-3
lines changed

4 files changed

+192
-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

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

0 commit comments

Comments
 (0)