Differentiable Planning › Backpropagation Through Dynamics

The previous chapter used $s_t$, $a_t$, $f$, $H$, $\ell$, and $\ell_f$. Here we keep $a_t$, $f$, and $H$, but write the state as $z_t$ to emphasize that it may be latent, and write the stage and terminal costs as $c_t$ and $c_H$. Thus $z_t$ corresponds to the previous chapter's $s_t$, while $c_t$ and $c_H$ correspond to $\ell$ and $\ell_f$.

We work with a deterministic, differentiable dynamics model

$$z_{t+1} = f(z_t, a_t), \qquad z_t \in \mathbb{R}^n,\; a_t \in \mathbb{R}^m,$$

where $z_t$ is the model state: it may be a fully observed physical configuration, a latent or belief state, or an observation when that observation is sufficient for the modeled transition. The action $a_t$ is continuous. The initial state $z_0$ is given and fixed. Over a horizon of $H$ transitions we accumulate the cost

$$J(a_0, \dots, a_{H-1}) = \sum_{t=0}^{H-1} c_t(z_t, a_t) + c_H(z_H),$$

where $c_t$ is the stage cost at time $t$ and $c_H$ is an optional terminal cost. The actions are the inputs to a recurrence, and $J$ is a scalar function of those inputs. This is exactly the structure of a recurrent neural network unrolled over time, except that we are differentiating with respect to the inputs rather than the weights.

To make the chain rule concrete, trace the computational paths from one action $a_t$. It has a direct path to the stage cost $c_t$ and a path through the next state $z_{t+1}=f(z_t,a_t)$ to later states and costs. The corresponding derivatives can still be zero: the cost may be locally flat, the action Jacobian may vanish, or later dynamics may erase the perturbation. The structural point is that a current action can contribute directly to the present cost and indirectly to future costs through the intervening transitions.

Define the state Jacobian and action Jacobian at time $t$ as

$$F_t = \frac{\partial f}{\partial z}\bigg|_{(z_t, a_t)} \in \mathbb{R}^{n \times n}, \qquad G_t = \frac{\partial f}{\partial a}\bigg|_{(z_t, a_t)} \in \mathbb{R}^{n \times m}.$$

The sensitivity of a future state to a past action is then a product of state Jacobians:

$$\frac{\partial z_k}{\partial a_t} = F_{k-1} F_{k-2} \cdots F_{t+1} \, G_t, \qquad k > t.$$

Read this product from right to left: the action perturbs $z_{t+1}$ through $G_t$, and each subsequent Jacobian transforms that perturbation. A subsequent Jacobian may mix or rotate directions and may shear a neighborhood of perturbations. Whether it amplifies or attenuates a particular perturbation depends on that direction and the chosen norm. If all relevant $F_j$ have operator norm at most a common $\rho<1$, the sensitivity decays at least geometrically with separation. Growth requires corresponding expansion along the propagated directions; an operator norm above one alone does not guarantee it.

Computing every action gradient by expanding its Jacobian products separately would repeat shared factors. For example, the state sensitivities $\partial z_H/\partial z_{H-1}=F_{H-1}$ and $\partial z_H/\partial z_{H-2}=F_{H-1}F_{H-2}$ share the same suffix. A naive scalar implementation that forms each horizon-length product independently performs $O(H^2)$ transition multiplications; one backward sweep performs $O(H)$. The adjoint (co-state) recursion is reverse-mode autodiff organized to reuse those suffixes.

Define the adjoint vector

$$\lambda_t = \frac{\partial J}{\partial z_t} \in \mathbb{R}^n,$$

interpreted as a column vector. The adjoint is a sensitivity of the objective to the state: it says how much the total cost would change if we could nudge $z_t$ directly. Because $z_t$ feeds into every later state and cost through the dynamics, this sensitivity already accumulates everything that happens downstream of time $t$. At the end of the horizon the adjoint is just the terminal cost gradient:

$$\lambda_H = \frac{\partial c_H}{\partial z_H}.$$

Now walk backwards. The state $z_t$ influences the objective in exactly two ways: through the stage cost $c_t(z_t, a_t)$, and through the next state $z_{t+1} = f(z_t, a_t)$, whose adjoint we already have. The chain rule gives

$$\lambda_t = \frac{\partial c_t}{\partial z_t} + F_t^\top \lambda_{t+1}, \qquad t = H-1, H-2, \dots, 0.$$

The transpose appears because $z_{t+1}=f(z_t,a_t)$ depends on $z_t$ through $F_t$, and $\lambda$ is a column vector. Under the local linearization, $\delta z_{t+1}=F_t\delta z_t+o(\lVert\delta z_t\rVert)$. Pulling a gradient back through the transition therefore multiplies by $F_t^\top$, converting next-state sensitivity into current-state sensitivity. For this discrete recurrence, the recursion is exactly the BPTT backward pass. It is the discrete-time analogue of the continuous ODE adjoint, which instead satisfies a backward differential equation; differentiating a chosen ODE discretization produces a discrete adjoint whose details depend on that discretization.

The action gradient then follows immediately. The action $a_t$ enters the objective only through $c_t(z_t, a_t)$ and through $z_{t+1} = f(z_t, a_t)$:

$$\boxed{\;\frac{\partial J}{\partial a_t} = \frac{\partial c_t}{\partial a_t} + G_t^\top \lambda_{t+1}.\;}$$

where:

• $\partial c_t / \partial a_t$: the direct sensitivity of the stage cost at time $t$ to the action taken at time $t$
• $G_t = \partial f / \partial a \big|_{(z_t, a_t)}$: the action Jacobian, mapping an action perturbation into a next-state perturbation
• $\lambda_{t+1} = \partial J / \partial z_{t+1}$: the adjoint of the next state, which aggregates the sensitivity of every future cost

One backward sweep computes $\lambda_H,\lambda_{H-1},\dots,\lambda_1$ and every action gradient. Its arithmetic cost is a small constant multiple of the forward evaluation. A standard uncheckpointed implementation uses $O(H)$ activation memory when per-step size is fixed; it retains the local information needed for the pullbacks, not necessarily explicit Jacobian matrices. Autodiff computes vector-Jacobian products such as $F_t^\top\lambda_{t+1}$ and $G_t^\top\lambda_{t+1}$ directly. Checkpointing or recomputation can reduce retained activations at additional compute cost.

The next two cells compute and plot analytic scalar multiplication counts for two organizations of the same action-sensitivity calculation. They do not benchmark Python exponentiation. Independently forming $\rho^{H-1-t}$ by repeated multiplication uses $\sum_{r=0}^{H-1}r=H(H-1)/2$ multiplications by $\rho$. With the terminal seed $\lambda_H$ available, the adjoint sweep needs $H-1$ multiplications to produce $\lambda_{H-1},\dots,\lambda_1$; together, $\lambda_1,\dots,\lambda_H$ are the adjoints required by all $H$ action gradients.

This is the cheap gradient principle: reverse mode computes a full $H\times m$ action gradient at a small constant multiple of one scalar objective evaluation, whereas coordinate-wise finite differences require one or more extra evaluations per action coordinate. This evaluation-cost asymmetry is one important reason differentiable planning is attractive in high dimensions. Directional and simultaneous-perturbation estimators have different tradeoffs, so the comparison here is specifically with coordinate-wise finite differences.

Substituting the adjoint recursion into itself gives the explicit form

$$\lambda_t = \sum_{k=t}^{H-1} \Big(\prod_{j=t}^{k-1} F_j^\top\Big) \frac{\partial c_k}{\partial z_k} \;+\; \Big(\prod_{j=t}^{H-1} F_j^\top\Big)\frac{\partial c_H}{\partial z_H},$$

and therefore

$$\frac{\partial J}{\partial a_t} = \frac{\partial c_t}{\partial a_t} + G_t^\top \Bigg[\sum_{k=t+1}^{H-1} \Big(\prod_{j=t+1}^{k-1} F_j^\top\Big)\frac{\partial c_k}{\partial z_k} + \Big(\prod_{j=t+1}^{H-1} F_j^\top\Big)\frac{\partial c_H}{\partial z_H}\Bigg].$$

Every indirect future-cost contribution contains an action Jacobian and a possibly empty product of state-Jacobian transposes applied to a cost gradient. This formula motivates the stability analysis below, and it explains a subtlety that is easy to miss: the gradient with respect to an early action is a sum of many terms, one for each future cost, each weighted by a generally different, time-indexed Jacobian product. Terms can reinforce or cancel. A vanishing product for one future step does not necessarily mean the total gradient vanishes. Conversely, a single large product term can dominate the sum even when the surrounding terms are small, which is one reason the gradient profile across actions can look erratic in practice.

Ordinary reverse mode needs a valid local pullback for every operation on the executed path. At a switch, an autodiff implementation may select a branch or subgradient convention; a planner may instead install a surrogate or custom pullback; and some operations provide no useful derivative. Backpropagation-based planners therefore need a usable sensitivity route for the optimized path: ordinary pullbacks, event-aware or implicit sensitivities, a relaxation or surrogate, or an estimator. Otherwise they must bypass or search the hard component separately.

• Contact and collision. Rigid-contact dynamics are often smooth within a fixed mode but hybrid across impacts. Classical derivatives can jump or fail to exist when contact timing or the active mode changes, while event-aware sensitivities can be valid away from grazing and other degenerate events. Compliant differentiable contact laws or relaxed complementarity conditions can smooth selected transitions, but they change the physical model and do not automatically smooth friction or contact activation. Another local approximation differentiates the active branch while holding the contact set fixed. This is locally exact only while perturbations preserve the full contact and friction mode; across a mode change it omits event sensitivity, although an optimizer may still use it heuristically.
• Discrete actions and argmax operations. With a one-hot encoding, hard $a=\arg\max_i q_i$ has zero Jacobian away from ties and is nondifferentiable at ties, so it supplies no informative ordinary pathwise gradient for changing the selected category. A temperature-softmax relaxes a deterministic argmax; for categorical sampling, Gumbel-softmax relaxes the one-hot Gumbel-max sample. A straight-through estimator keeps the hard forward value but substitutes a surrogate backward derivative. These introduce approximation choices. For background on discrete and latent action representations, see Actions, Inverse Dynamics, and Latent Actions.
• Stochastic dynamics. A sample $z_{t+1} \sim p_\theta(z_{t+1} \mid z_t, a_t)$ does not automatically provide a pathwise derivative with respect to the action. For a reparameterizable continuous distribution, write the sample as a deterministic function of the action and independent noise, for example $z_{t+1} = \mu_\theta(z_t, a_t) + \sigma_\theta(z_t, a_t) \odot \epsilon$ with $\epsilon\sim\mathcal N(0,I)$. Holding the sampled noise fixed then gives a pathwise gradient for a Monte Carlo objective. When reparameterization is unavailable, differentiating $\mathbb E_{z\sim p_a}[L(a,z)]$ yields the ordinary term $\mathbb E[\partial_a L]$ plus the likelihood-ratio term $\mathbb E[L\,\partial_a\log p_a(z)]$, provided differentiation can be interchanged with integration and the required terms are integrable. The likelihood-ratio term also requires action-independent support or vanishing boundary contributions; nonzero boundary terms from action-dependent support need a separate valid estimator. If the score itself is estimated, its conditional expectation given $z$ must equal the true score, or the weighted term must otherwise be jointly unbiased. Score-function estimators often have high variance. See Stochasticity and Uncertainty for the probabilistic background.
• Rendering and perception. At fixed inputs and parameters, a differentiable deterministic image encoder defines a repeatable mathematical gradient, although it may be expensive or badly scaled. Implementation-level repeatability additionally requires deterministic forward and backward kernels in a fixed software and hardware environment. Stochastic rendering, noisy observations, and sampled layers can add gradient variance. Surrogate derivatives can add bias or mismatch; randomized surrogates can also add variance. Hard thresholding, such as a binarized mask or hard attention selection, needs a relaxation, estimator, or custom pullback.

Simulation systems such as DiffTaichi and the published Brax physics engine provide automatic differentiation for substantial parts of a physics computation. So do MJX-JAX and supported NVIDIA Warp kernels. Coverage is backend- and feature-dependent: current MJX-Warp is not differentiable, and contact smoothness depends on the chosen formulation and settings. Planning by backpropagation may instead use analytical dynamics, a learned model, event-aware sensitivities, or an implicitly differentiated solver.