This notebook provides a short tutorial for automatic differentiation in Drake. It covers the automatic differentiation of
Drake uses Eigen's AutoDiffScalar for automatic differentiation. Any explicit Eigen (and hence Numpy in Python) operations can be automatically differentiated. Let's consider the simple example $\mathbf{y} = \mathbf{a}^{\top} \mathbf{x} = [1,~ 2] \mathbf{x}$. We want to obtain the derivative $\partial \mathbf{y} / \partial \mathbf{x} = \mathbf{a}^{\top}$ using automatic differentiation. Here is how to do it:
import numpy as np
from pydrake.autodiffutils import InitializeAutoDiff, ExtractGradient
a = np.array([1, 2]).reshape([2, -1])
x = np.random.rand(2, 1)
print("double type array:\n", x)
x = InitializeAutoDiff(x)
print("converted to AutoDiffXd scalar type array:\n", x)
y = a.T @ x
dydx = ExtractGradient(y)
print("Gradient:", ExtractGradient(y))
np.testing.assert_allclose(a.T, dydx) # assert dy/dx = a^T
double type array: [[0.43597967] [0.21388789]] converted to AutoDiffXd scalar type array: [[<AutoDiffXd 0.43597967258466797 nderiv=2>] [<AutoDiffXd 0.21388788500968936 nderiv=2>]] Gradient: [[1. 2.]]
Note that for automatic differentiation, we just need two extra steps in addition to the usual explicit Numpy operations: 1) Declare the variable with respect to which we want to differentiate using InitializeAutoDiff(x), and 2) extract the gradient in the end using ExtractGradient(y).
We can also differentiate w.r.t. multiple variables. Consider the example $\mathbf{y} = \mathbf{a}_1^{\top} \mathbf{x}_1 + \mathbf{a}_2^{\top} \mathbf{x}_2$, where we want to obtain $\partial \mathbf{y} / \partial \mathbf{x}_1$ and $\partial \mathbf{y} / \partial \mathbf{x}_2$. Here is how to do it:
import numpy as np
from pydrake.autodiffutils import InitializeAutoDiffTuple, ExtractGradient
a1 = np.array([1, 2]).reshape([2, -1])
x1 = np.random.rand(2, 1)
a2 = np.array([1, 2, 3]).reshape([3, -1])
x2 = np.random.rand(3, 1)
x1, x2 = InitializeAutoDiffTuple(x1, x2)
y = a1.T @ x1+ a2.T @ x2
dydx = ExtractGradient(y)
print("All gradients:", dydx)
dydx1 = dydx[0][0:2]
dydx2 = dydx[0][2:]
print("Gradients calculated by automatic differentiation:")
print("a1^T =", dydx1)
np.testing.assert_allclose(a1.T, [dydx1])
print("a2^T =", dydx2)
np.testing.assert_allclose(a2.T, [dydx2])
All gradients: [[1. 2. 1. 2. 3.]] Gradients calculated by automatic differentiation: a1^T = [1. 2.] a2^T = [1. 2. 3.]
Note that ExtractGradient(y) extracts derivatives of y with respect to all variables declared by InitializeAutodiffTuple.
Drake's most built-in systems' dynamics only involve explicit Eigen operations. Hence, they are all automatically differentiable. Let's consider the simple discrete-time LinearSystem, whose dynamics is given as $$x_{t+1} = x_t + 2u_t, ~~~y_{t} = 3x_t + 4u_t.$$ For general dynamical systems, the derivatives of next state w.r.t. state $\partial x_{t+1}/ \partial x_t$ and input $\partial x_{t+1}/ \partial u_t$, and output w.r.t. state $\partial y_{t}/ \partial x_t$ and input $\partial y_{t}/ \partial u_t$ are frequently wanted. Here we will show how to obtain them via automatic differentiation. Let's first construct the system:
import numpy as np
from pydrake.systems.primitives import LinearSystem
from pydrake.autodiffutils import InitializeAutoDiffTuple, ExtractGradient
A = np.array([[1]])
B = np.array([[2]])
C = np.array([[3]])
D = np.array([[4]])
timestep = 1 # so that the system is discrete-time
system = LinearSystem(A, B, C, D, timestep)
print("A system using double:", system)
A system using double: <pydrake.systems.primitives.LinearSystem object at 0x73b5fd58ffb0>
By default, the system uses double as the scalar type. We need to convert it to use drake::AutoDiffXd:
system_ad = system.ToAutoDiffXd()
print("The system converted to AutoDiffXd:", system_ad)
The system converted to AutoDiffXd: <pydrake.systems.primitives.LinearSystem_[AutoDiffXd] object at 0x73b620f41850>
Let's set $x_t = 1$ and $u_t=1$ (or any real numbers)
context_ad = system_ad.CreateDefaultContext()
x = np.array([1])
u = np.array([1])
x, u = InitializeAutoDiffTuple(x, u)
context_ad.SetDiscreteState(0, x)
system_ad.get_input_port(0).FixValue(context_ad, u)
<pydrake.systems.framework.FixedInputPortValue at 0x73b620f4aaf0>
Then, we calculate the derivatives of next states $\partial x_{t+1}/ \partial x_t=1$ and $\partial x_{t+1}/ \partial u_t=2$
# allocate the state object
x_next_object = system_ad.AllocateContext().get_discrete_state()
# store value to x_next_object without modifying context
system_ad.CalcForcedDiscreteVariableUpdate(context_ad, x_next_object)
# to extract numpy array from the state object
x_next = x_next_object.get_vector(0).CopyToVector()
grad = ExtractGradient(x_next)
dx_next_dx = grad.flatten()[0]
dx_next_du = grad.flatten()[1]
print("Gradients calculated by automatic differentiation:")
print("dx'/dx =", dx_next_dx)
assert dx_next_dx == 1
print("dx'/du =", dx_next_du)
assert dx_next_du == 2
Gradients calculated by automatic differentiation: dx'/dx = 1.0 dx'/du = 2.0
and the derivatives of output $\partial y_{t}/ \partial x_t=3$ and $\partial x_{t}/ \partial u_t=4$
output_object = system_ad.AllocateOutput()
system_ad.CalcOutput(context_ad, output_object)
output_port_index = system_ad.get_output_port(0).get_index()
output = output_object.get_vector_data(output_port_index).CopyToVector()
grad = ExtractGradient(output)
dy_dx = grad.flatten()[0]
dy_du = grad.flatten()[1]
print("Gradients calculated by automatic differentiation::")
print("dy/dx =", dy_dx)
assert dy_dx == 3
print("dy/du =", dy_du)
assert dy_du == 4
Gradients calculated by automatic differentiation:: dy/dx = 3.0 dy/du = 4.0
You can write your own LeafSystem that supports automatic differentiation by using the drake::AutoDiffXd scalar type as the template value. In Python, you can do so using the TemplateSystem utility. Let's consider the simple discrete-time system $$x_{t+1} = x_t + 2u_t, ~~~y_{t} = x^2_t.$$ We will build it with a discrete-time LinearSystem, and a custom system that squares the input as the output. Let's first define the linear system $x_{t+1} = x_t + 2u_t, ~~~y_{t} = x_t$:
import numpy as np
from pydrake.systems.scalar_conversion import TemplateSystem
from pydrake.systems.primitives import LinearSystem
from pydrake.systems.framework import (
BasicVector_,
DiagramBuilder,
LeafSystem_,
)
from pydrake.autodiffutils import InitializeAutoDiffTuple, ExtractGradient
A = np.array([[1]])
B = np.array([[2]])
C = np.array([[1]])
D = np.array([[0]])
timestep = 1
linear_system = LinearSystem(A, B, C, D, timestep)
Now, let's define the templated custom system:
@TemplateSystem.define("SquareSystem_")
def SquareSystem_(T):
class Impl(LeafSystem_[T]):
def _construct(self, dimension: int, converter=None):
LeafSystem_[T].__init__(self, converter=converter)
self.dimension = dimension
self.input_port = self.DeclareVectorInputPort(
"input", BasicVector_[T](dimension)
)
self.output_port = self.DeclareVectorOutputPort(
"output",
BasicVector_[T](dimension),
self.calc_output,
)
def _construct_copy(self, other, converter=None):
Impl._construct(self, other.dimension, converter=converter)
def calc_output(self, context, output):
input_array = self.input_port.Eval(context)
# Element-wise squared input as the output y = x * x
output.set_value(input_array * input_array)
return Impl
SquareSystem = SquareSystem_[None] # The default system that uses double as the scalar
The main difference from what we saw in Modeling Dynamical Systems is the use of template classes LeafSystem_[T] and BasicVector_[T]. The default double-scalar class, defined as SquareSystem = SquareSystem_[None], is nearly identical to a version defined without using template classes. However, by templating the system, Drake can automatically convert it to use drake::AutoDiffXd for automatic differentiation. Now let’s construct the custom system and compose the full diagram:
squared_output = SquareSystem(dimension=1)
builder = DiagramBuilder()
builder.AddSystem(linear_system)
builder.AddSystem(squared_output)
builder.Connect(linear_system.get_output_port(0), squared_output.get_input_port(0))
builder.ExportInput(linear_system.get_input_port(), "input")
builder.ExportOutput(squared_output.get_output_port(), "output")
# The full dynamical system we are considering
system = builder.Build()
print("Default double systems:\n", system.GetSystems())
Default double systems: [<pydrake.systems.primitives.LinearSystem object at 0x73b5fd58ff50>, <pydrake.systems.scalar_conversion.SquareSystem_[float] object at 0x73b5fd365d30>]
Note that although we only construct the default SquareSystem that uses double as the scalar, it can be converted into a system using AutoDiffXd as the scalar when we do ToAutoDiffXd():
system_ad = system.ToAutoDiffXd()
print("AutoDiffXd systems:\n", system_ad.GetSystems())
AutoDiffXd systems: [<pydrake.systems.primitives.LinearSystem_[AutoDiffXd] object at 0x73b5fd366810>, <pydrake.systems.scalar_conversion.SquareSystem_[AutoDiffXd] object at 0x73b5fd366570>]
Now let's calculate the derivatives at $x_t=1$ and $u_t=1$:
context_ad = system_ad.CreateDefaultContext()
x = np.array([1])
u = np.array([1])
x, u = InitializeAutoDiffTuple(x, u)
context_ad.SetDiscreteState(0, x)
system_ad.get_input_port(0).FixValue(context_ad, u)
<pydrake.systems.framework.FixedInputPortValue at 0x73b5fd586a30>
We obtain the correct derivatives $\left. \partial y_{t}/ \partial x_t \right|_{x_t=1}= \left. 2 x_t \right|_{x_t=1} = 2$ and $\partial y_{t}/ \partial u_t=0$:
output_object = system_ad.AllocateOutput()
system_ad.CalcOutput(context_ad, output_object)
output_port_index = system_ad.GetOutputPort("output").get_index()
output = output_object.get_vector_data(output_port_index).CopyToVector()
grad = ExtractGradient(output)
dy_dx = grad.flatten()[0]
dy_du = grad.flatten()[1]
print("Gradients calculated by automatic differentiation:")
print("dy/dx =", dy_dx)
assert dy_dx == 2
print("dy/du =", dy_du)
assert dy_du == 0
Gradients calculated by automatic differentiation: dy/dx = 2.0 dy/du = 0.0
Lastly, let's simulate a continuous-time system, and differentiate the sampled simulation results. We consider the linear system $$ \dot{x} = -x,~~ y = x.$$ Given the initial state $x_0$, its output solution is given as $$ y(t) = e^{-t} x_0,$$ and the output's derivative w.r.t. the initial state is given as $$ \frac{\partial y}{\partial x_0} = e^{-t},$$ which only depends on time. Now we will calculate this derivative through automatic differentiation, and compare the results to the analytical gradients:
import numpy as np
def calculate_analytical_gradient(t):
return np.exp(-t)
Let's construct the linear system
from pydrake.systems.primitives import LinearSystem, LogVectorOutput
from pydrake.systems.framework import DiagramBuilder, SubsystemIndex
from pydrake.autodiffutils import InitializeAutoDiff, ExtractGradient, AutoDiffXd
from pydrake.systems.analysis import Simulator_
A = np.array([[-1]])
B = np.array([[0]])
C = np.array([[1]])
D = np.array([[0]])
timestep = 0 # so that the system is continuous-time
linear_system = LinearSystem(A, B, C, D, timestep)
builder = DiagramBuilder()
builder.AddSystem(linear_system)
builder.ExportInput(linear_system.get_input_port(), "input")
builder.ExportOutput(linear_system.get_output_port(), "output")
logger = LogVectorOutput(linear_system.get_output_port(), builder, publish_period=0.1)
system = builder.Build() # The dynamical system we are considering
and convert it to use AutoDiffXd scalar
system_ad = system.ToAutoDiffXd()
logger_ad = system_ad.get_system(SubsystemIndex(1))
We get the AutoDiffXd version of the logger to extract the simulation results later. Now we construct a Simulator that uses AutoDiffXd scalar, and set an arbitrary initial state $x_0 = 5$:
simulator_ad = Simulator_[AutoDiffXd](system_ad)
context_ad = simulator_ad.get_mutable_context()
system_ad.get_input_port(0).FixValue(context_ad, 0)
x0 = np.array([5])
x0 = InitializeAutoDiff(x0)
context_ad.SetContinuousState(x0)
simulator_ad.Initialize()
<pydrake.systems.analysis.SimulatorStatus at 0x73b5fde87030>
Finally, we simulate for 1 second, and assert that the derivatives calculated by automatic differentiation match the analytical ones:
simulator_ad.AdvanceTo(1)
log = logger_ad.FindLog(simulator_ad.get_context())
# convert AutoDiffXd back to double
sample_times = np.array([t.value() for t in log.sample_times()])
# calcualte gradients analytically and via autodiff
analytical_gradients = calculate_analytical_gradient(sample_times)
autodiff_gradients = ExtractGradient(log.data()).flatten()
# Let's print the data at an arbitrary sample time
sample_index = 3
print("dy/dx0 at t =", sample_times[sample_index]),
print("Analytical gradient:", analytical_gradients[sample_index])
print("Gradient calculated by autodiff:", autodiff_gradients[sample_index])
# We assert that the autodiff gradients are correct at all sample times
np.testing.assert_allclose(autodiff_gradients, analytical_gradients, atol=1e-6)
dy/dx0 at t = 0.30000000000000004 Analytical gradient: 0.7408182206817179 Gradient calculated by autodiff: 0.7408182212987496
System Scalar Types and Conversions in Drake
double, AutoDiffXd, and symbolic::Expression. Automatic Differentiation with Eigen
Automatic Differentiation with Drake’s Hydroelastic Contact Model