Hyoungseo Son
Matrix Calculus & Autodiff
Contents · Matrix Calculus & Autodiff

Part 1 · Automatic Differentiation

Part 2 · Backpropagation

Part 3 · Differentiable Simulation

  • Differentiating Through a Simulator

Part 3 · Differentiable Simulation

Differentiating Through a Simulator

A simulator is a function iterated. One step maps state to state, xt+1=f(xt,θ)x_{t+1}=f(x_t,\theta), with parameters θ\theta (controls, masses, gains) held fixed across the rollout, and a scalar loss L(xT)L(x_T) reads off the final state. To improve θ\theta you want dL/dθ\mathrm{d}L/\mathrm{d}\theta. Running reverse-mode autodiff through the unrolled trajectory is not a separate trick: it is exactly the adjoint method, and it returns the full gradient in a single backward pass whose cost does not depend on how many parameters θ\theta holds.

The unrolled trajectory

A rollout of TT steps is the composition xT=f(,θ)f(x0,θ)x_T=f(\cdot,\theta)\circ\cdots\circ f(x_0,\theta). The loss depends on θ\theta through every step, so the chain rule threads a product of per-step Jacobians:

dLdθ=t=0T1LxT(k=t+1T1xk+1xk)fθt.\frac{\mathrm{d}L}{\mathrm{d}\theta} =\sum_{t=0}^{T-1} \frac{\partial L}{\partial x_T} \left(\prod_{k=t+1}^{T-1}\frac{\partial x_{k+1}}{\partial x_k}\right) \frac{\partial f}{\partial \theta}\bigg|_{t}.

Forming those products left to right (forward mode) costs one pass per input direction, so θRp\theta\in\mathbb{R}^p needs pp passes. The adjoint reorders the sum to sweep once from the right.

Reverse mode is the adjoint

Define the adjoint λt=L/xt\lambda_t=\partial L/\partial x_t, the sensitivity of the loss to the state at time tt. Seed it at the end with the loss gradient, then propagate backward with the transposed state Jacobian, accumulating a parameter term at each step:

λT=LxT,λt=(fxt) ⁣ ⁣λt+1,dLdθ=t=0T1(fθt) ⁣ ⁣λt+1.\lambda_T=\frac{\partial L}{\partial x_T},\qquad \lambda_t=\left(\frac{\partial f}{\partial x_t}\right)^{\!\top}\!\lambda_{t+1},\qquad \frac{\mathrm{d}L}{\mathrm{d}\theta}=\sum_{t=0}^{T-1}\left(\frac{\partial f}{\partial \theta}\bigg|_t\right)^{\!\top}\!\lambda_{t+1}.

One backward traversal produces every entry of the gradient at once, so the cost is independent of pp: a million-parameter controller and a one-parameter one cost the same backward sweep. That asymmetry, cheap gradient regardless of input dimension, is why reverse mode drives differentiable physics1.

A scalar sim, by hand

Take the smallest nontrivial case: a scalar linear sim xt+1=axt+θx_{t+1}=a\,x_t+\theta with loss L=(xTx)2L=(x_T-x^\star)^2. The two partials are constant,

fxt=a,fθ=1,\frac{\partial f}{\partial x_t}=a,\qquad \frac{\partial f}{\partial \theta}=1,

so the seed is λT=2(xTx)\lambda_T=2(x_T-x^\star), the recurrence is λt=aλt+1\lambda_t=a\,\lambda_{t+1}, and every accumulation term is just λt+1\lambda_{t+1}. Summing the geometric chain gives a closed form:

dLdθ=t=0T1λt+1=2(xTx)k=0T1ak.\frac{\mathrm{d}L}{\mathrm{d}\theta}=\sum_{t=0}^{T-1}\lambda_{t+1} =2\,(x_T-x^\star)\sum_{k=0}^{T-1}a^{k}.

Now plug in a=0.9a=0.9, T=3T=3, x0=0x_0=0, θ=1\theta=1, x=2x^\star=2. The forward rollout is

x1=0.9(0)+1=1,x2=0.9(1)+1=1.9,x3=0.9(1.9)+1=2.71,x_1=0.9(0)+1=1,\quad x_2=0.9(1)+1=1.9,\quad x_3=0.9(1.9)+1=2.71,

with L=(2.712)2=0.5041L=(2.71-2)^2=0.5041. The backward sweep seeds λ3=2(2.712)=1.42\lambda_3=2(2.71-2)=1.42, then

λ2=0.9(1.42)=1.278,λ1=0.9(1.278)=1.1502,\lambda_2=0.9(1.42)=1.278,\qquad \lambda_1=0.9(1.278)=1.1502,

and the parameter gradient is the running sum

dLdθ=λ1+λ2+λ3=1.1502+1.278+1.42=3.8482.\frac{\mathrm{d}L}{\mathrm{d}\theta}=\lambda_1+\lambda_2+\lambda_3 =1.1502+1.278+1.42=3.8482.

The closed form agrees: 2(0.71)(1+0.9+0.81)=1.42×2.71=3.84822(0.71)(1+0.9+0.81)=1.42\times2.71=3.8482. And so does a central finite difference, which never sees the adjoint at all:

L(1.001)L(0.999)0.002=0.5079550.5002590.002=3.848.\frac{L(1.001)-L(0.999)}{0.002}=\frac{0.507955-0.500259}{0.002}=3.848.

The demo runs the same recurrence for T=8T=8: drag aa and θ\theta to set the forward trajectory, then step the adjoint backward and watch each λt\lambda_t get folded into dL/dθ\mathrm{d}L/\mathrm{d}\theta, ending at the finite-difference value.

The adjoint, one backward step at a timeinteractive
forward xₜ →← backward λₜtarget012345678
x_8
3.278
target
3.00
loss (x_T − t)²
0.077
seed λ_8 = 2(x_T − t)
0.556
adjoint backward pass
Seeded λ_8. Step backward to propagate and accumulate.
g (running, 0/8)
0.000
dL/dθ (adjoint)
3.168
dL/dθ (finite diff)
3.168
match

The forward pass (top, left to right) builds the trajectory. The adjoint pass (bottom, right to left) starts from the seed λ_8 and multiplies by a at every step; each λ it visits is added into dL/dθ. One backward sweep produces the exact gradient, which matches the finite difference.

Top: the forward rollout x_t (accent) toward the target (green). Bottom: the adjoint λ_t, seeded at the right and propagated left by one factor of a per step (current step in yellow, backward flow in red). Each λ it visits is added into dL/dθ; the running total lands on the finite-difference gradient once the sweep reaches t = 1.

The continuous limit and solver layers

As the step size h0h\to0 the backward recurrence becomes an ODE for λ(t)\lambda(t) integrated in reverse time, the continuous adjoint behind neural ordinary differential equations2. And when a step is not an explicit update but an equation solve, say a contact quadratic program, you need not unroll the solver at all: the implicit function theorem differentiates its solution directly from the optimality conditions, giving an exact gradient through a fixed number of solver iterations3.

One caveat carries over from the contact-cliff view of differentiable simulation: when contact makes ff nonsmooth, f/xt\partial f/\partial x_t jumps or spikes, and the adjoint faithfully propagates a gradient that is biased or high-variance. The machinery is exact for the model it is given; it cannot repair a discontinuous ff.

Footnotes

  1. Hu et al., DiffTaichi: Differentiable Programming for Physical Simulation (2020).

  2. Chen et al., Neural Ordinary Differential Equations (2018).

  3. Amos & Kolter, OptNet: Differentiable Optimization as a Layer in Neural Networks (2017).