Differentiable Planning › Stability Under Long Rollout Gradients

Michael BrenndoerferJuly 8, 202611 min read
World Models HandbookAll chapters
Part 7 of 12Planning and Agency
Chapter 38Differentiable Planning › Stability Under Long Rollout Gradients
Reading section 5 of 10 · Stability Under Long Rollout Gradients · 11 minSection 5 of 10 · 11 min

Book contents

Section 5 of 10 · 11 min read

Differentiable Planning › Stability Under Long Rollout GradientsOne section · Sections in this chapter5 / 10 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

We now turn to the central technical difficulty. Everything so far assumed that ∇aJ\nabla_aJ is usable. Over long horizons it can fail when accumulated sensitivities are strongly contractive, expansive, ill-conditioned, discontinuous, or noisy. A long horizon alone does not force failure.

The product of Jacobians

Return to the sensitivity formula

∂zk∂at=Fk−1Fk−2⋯Ft+1 Gt.\frac{\partial z_k}{\partial a_t} = F_{k-1} F_{k-2} \cdots F_{t+1} \, G_t.

If the state Jacobians were constant, Fj=FF_j = F for all jj, this product would be Fk−t−1GtF^{k-t-1} G_t. For a scalar F=ρF=\rho, or for a normal matrix whose relevant modes are excited, its asymptotic scale follows the eigenvalue magnitudes:

  • If ∣ρ∣<1|\rho|<1, the product decays geometrically. The influence of early actions on late states vanishes.
  • If ∣ρ∣>1|\rho|>1 and the action excites the scalar mode with a nonzero gain, the sensitivity magnitude grows geometrically. The influence of early actions on late states then becomes unbounded as the separation grows.
  • If ∣ρ∣≈1|\rho|\approx1, the modal magnitude changes little per transition. Over pp transitions it remains near one only when p∣log⁡∣ρ∣∣p|\log|\rho|| is small, and negative ρ\rho alternates sign.

For a general matrix, spectral radius alone is insufficient over a finite horizon: non-normal matrices can exhibit large transient growth even when every eigenvalue lies inside the unit disk. Time-varying nonlinear rollouts require the singular values and directions of the actual product Fk−1⋯Ft+1F_{k-1}\cdots F_{t+1}. The scalar spectral-radius example is therefore intuition, not a complete stability test. A real system's rates come from its physical dynamics. The implemented model's state and action Jacobians supply the planner's rollout-sensitivity factors and match the real rates only to the extent that the model is locally faithful. The full objective gradient also depends on cost or value derivatives and on the coordinates used to parameterize actions. Constraints separately determine feasible or projected directions unless they are encoded through a barrier, penalty, or reparameterization. For nonlinear dynamics, Jacobians depend on the evaluation point and can vary along a rollout.

What goes wrong

Numerical overflow and underflow. Float32's largest finite magnitude is about 3.4×10383.4\times10^{38}, its smallest positive normal value is about 1.2×10−381.2\times10^{-38}, and subnormal values extend to about 1.4×10−451.4\times10^{-45}. With ρ=1.2\rho = 1.2 and H=100H = 100, ρH≈8.3×107\rho^{H} \approx 8.3 \times 10^7; with ρ=1.5\rho = 1.5, the product is about 4.1×10174.1\times10^{17} at H=100H=100 and 1.7×10351.7\times10^{35} at H=200H=200, still finite, but it overflows near H=219H=219. In the other direction, 0.9400≈5.0×10−190.9^{400}\approx5.0\times10^{-19} is still normal; it reaches the normal/subnormal boundary only near H=829H=829. These scalar thresholds do not predict the conditioning of a matrix product, but they show when finite-precision range becomes an additional failure mode. Float64 postpones overflow and underflow without fixing vanishing or exploding sensitivity.

Ill-conditioning. Ill-conditioning can remain severe even when overflow and underflow are absent because it concerns relative directional scales rather than representable range. For a symmetric positive-definite quadratic Hessian whose smallest and largest eigenvalues are μ\mu and LL, respectively, κ2=L/μ\kappa_2=L/\mu; if [μ,L][\mu,L] is only an enclosing interval, then κ2≤L/μ\kappa_2\leq L/\mu. Gradient descent with step 1/L1/L contracts Euclidean iterate error by at most 1−1/κ21-1/\kappa_2 per step when LL is the largest eigenvalue; a constant-factor reduction can therefore take on the order of κ2\kappa_2 steps. In linear-quadratic problems the action Hessian is built from rollout Jacobians and cost curvature. Exact nonlinear Hessians also contain second derivatives of the dynamics and costs. This spread motivates regularized curvature-aware methods such as DDP, but they do not eliminate conditioning problems in general.

Unequal update scales. Consider the terminal-only objective for zt+1=0.85zt+atz_{t+1}=0.85z_t+a_t with H=40H=40. The final-state sensitivities are ∂z40/∂a0=0.8539≈0.00177\partial z_{40}/\partial a_0=0.85^{39}\approx0.00177 and ∂z40/∂a39=1\partial z_{40}/\partial a_{39}=1. Away from a stationary terminal cost, the common terminal derivative cancels in their ratio, so the late-action gradient is about 0.85−39≈5660.85^{-39}\approx566 times larger. With plain gradient descent using one global learning rate, early-action updates are proportionally smaller. Global-norm clipping rescales the entire vector and preserves this ratio; it bounds magnitude but does not repair the imbalance. Per-coordinate moment normalization can make persistent nonzero components more similar in update magnitude, except where optimizer history or ϵ\epsilon dominates, and may over-amplify directions with little physical influence.

In a strongly contracting system, weak early-action influence on the terminal state can be physically correct. Early actions may still matter for stage costs. At a given candidate rollout, relative action-gradient weights depend on the initial state and visited state-action points through the dynamics, costs, horizon, discounting, and parameterization. Those weights may or may not suit the practical objective. If contraction is a model artifact, the sensitivity is misleading. If it is physical, weak terminal sensitivity is faithful, although its optimization consequences still depend on the rest of the formulation.

Gradient variance under stochastic dynamics. Under regularity conditions that permit differentiation under the expectation, reparameterization gives ∇aEϵ[L(a,ϵ)]=Eϵ[∇aL(a,ϵ)]\nabla_a\mathbb{E}_\epsilon[L(a,\epsilon)]=\mathbb{E}_\epsilon[\nabla_aL(a,\epsilon)]. In practice, we estimate it by averaging pathwise gradients from sampled noise realizations. Jacobian chains can amplify or attenuate each sample's influence, so variance may grow, shrink, or plateau with horizon. When measured Monte Carlo variance is large relative to the mean gradient, use more trajectories or validated control variates. For paired comparisons, common random numbers can help when the induced covariance is favorable.

Mitigations

The mitigations fall into five families.

(A) Shorten the effective horizon.

  • Receding-horizon MPC with a finite, task-appropriate planning horizon, the central deployment pattern in Sampling-Based Planning and Model Predictive Control.
  • A sufficiently decision-faithful terminal value, which estimates the omitted tail without another long transition chain.
  • Truncated backpropagation at selected recurrent-state boundaries. One scheme detaches the state KK transitions before the end and differentiates only the final window; this saves work but generally biases the full planning gradient unless the omitted chain-rule contributions are zero.
  • Discounting with 0≤γ<10\leq\gamma<1, which geometrically downweights late terms and shortens the effective horizon. It changes the objective and does not guarantee a better condition number.

(B) Control the gradient scale.

  • Global-norm clipping: rescale the whole vector when its norm exceeds a calibrated threshold. This preserves direction, unlike coordinate-wise clipping, and is preferable when direction preservation is the goal. It does not correct component imbalance.
  • Cost normalization: choose a practical scale from representative loss and gradient measurements, then tune the learning rate, optimizer ϵ\epsilon, and clipping threshold together. An order-one norm is a convention, not a requirement.
  • Higher precision when diagnostics show overflow, underflow, or harmful roundoff. Float64 extends numerical range without fixing true sensitivity decay or conditioning. FP16 mixed-precision optimization may need dynamic loss scaling when backward gradients would otherwise underflow, together with finite-value checks on overflow-prone paths. BF16 and inference-only mixed precision do not inherently require loss scaling; safeguards should match the format and measured numerical range. Sensitive reductions and optimizer state still need adequate precision.

(C) Improve conditioning.

  • Second-order and quasi-Newton methods (DDP, iLQR, Gauss-Newton, L-BFGS) can reduce sensitivity to eigenvalue spread when the local curvature model is accurate enough and suitably regularized or globalized. Approximate or indefinite curvature can produce an unhelpful step, while regularization choices and finite precision can leave the numerical problem ill-conditioned.
  • A compatible, well-chosen preconditioner, for example one based on a positive diagonal approximation to the Gauss-Newton Hessian, can reduce the transformed spectral spread. A poor, stale, or mismatched preconditioner can leave conditioning unchanged or worsen it.
  • Architecture choices inside the transition can help, but must be checked on planner-relevant rollout Jacobians. Layer normalization offers no conditioning guarantee. An upper spectral bound can limit expansion without preventing vanishing directions. A residual form f(z,a)=z+g(z,a)f(z,a)=z+g(z,a) stays near identity only when ∂g/∂z\partial g/\partial z is controlled, not merely when gg has small values.
  • For a linear component known to be asymptotically stable, Schur or diagonal parameterizations can constrain eigenvalues inside the unit disk. Orthogonal recurrences instead preserve norms with eigenvalues on the unit circle. Neither eigenvalue constraint controls finite-horizon transient growth in a non-normal system; monitor singular values or finite-horizon products separately.

(D) Break the chain.

  • Multiple shooting or direct collocation, which lifts state variables and exposes sparse local dynamics constraints instead of fully condensing every dependency into the action objective. The constrained system can still be ill-conditioned.
  • A square-root checkpoint schedule, which reduces activation memory from O(H)O(H) to O(H)O(\sqrt{H}) for a uniform chain at the cost of recomputation. Other schedules make different tradeoffs. Checkpointing addresses activation-memory pressure rather than Jacobian conditioning; it is worthwhile when the value of lower memory use outweighs its recomputation cost.
  • Explicit Jacobian or finite-horizon sensitivity regularizers over planner-relevant trajectories. Multi-step prediction losses expose rollout error but do not by themselves guarantee well-scaled derivative products.

(E) Avoid differentiating through hard components.

  • Straight-through estimators for quantization or discretization: use a hard value in the forward pass and substitute a surrogate backward derivative. In straight-through Gumbel-softmax, that surrogate is the soft relaxation's gradient.
  • Score-function estimators when pathwise reparameterization is unavailable. For Ez∼pa[L(a,z)]\mathbb E_{z\sim p_a}[L(a,z)], use the ordinary ∂aL\partial_a L term together with the likelihood-ratio term when differentiation and integration can be interchanged, the required terms are integrable, and the support is action-independent or its boundary contributions vanish. Nonzero boundary terms need a separate valid estimator. An estimated score must be conditionally unbiased given zz, or must otherwise make the weighted term jointly unbiased. Score-function estimators often have higher variance than a well-behaved pathwise estimator, so control variates are commonly useful.
  • Carefully chosen smooth surrogates: differentiable compliant contact for selected contact transitions, an interior barrier for a differentiable inequality whose initialization and every evaluated iterate remain strictly feasible, or a finite-temperature relaxation of a hard categorical choice. Each changes the original problem and must be validated.

Training models to be differentiable-friendly

Conventional supervised dynamics models are often fit with one-step prediction error, sometimes augmented with multi-step rollout losses. A standard value-only prediction objective does not explicitly control or guarantee the state-Jacobian spectrum. A model can therefore predict well on its training distribution and still be a poor substrate for gradient planning.

Two remedies target different aspects of the problem. Training with a multi-step prediction objective exposes the model to rollout errors caused by its repeated transitions, but it does not by itself guarantee well-conditioned Jacobians. An explicit regularizer can instead penalize selected Jacobian norms, singular values, or finite-horizon products; a guaranteed bound requires an enforced constraint or suitable parameterization. Suppressing physically real expansion may bias the dynamics. These choices connect to Predictive and Generative Objectives.

The deeper point is that one-step prediction quality and closed-loop decision usefulness are different properties. A model can score well on a one-step benchmark and still produce unusable plans. Part XI will revisit this distinction in evaluation. The architecture and learned parameters determine the realized derivatives; data, objective, regularization, initialization, and optimization influence which parameters are learned. Those derivatives are one major factor in how well a gradient planner works, alongside model accuracy, cost design, constraints, horizon, and optimizer.

A practical recipe

The numerical values in the worked example are not task-independent defaults. A measurement-driven starting procedure is:

  1. Nondimensionalize or rescale each state and action coordinate using documented characteristic scales so their numerical magnitudes are comparable, and transform costs and bounds consistently. Then measure representative objective and gradient scales; coordinate normalization alone does not make gradients order one.
  2. Choose a physical lookahead from the system's controllability and task timescales, sampling interval, terminal-value quality, and compute budget. Expressing the horizon only as a number of steps hides these dependencies.
  3. Tune the optimizer and learning rate from measured objective decrease. If you use clipping, set its threshold from the observed unclipped gradient distribution and log how often it activates.
  4. When consecutive receding-horizon problems evolve smoothly, shift the unused suffix, append a terminal action, and check or restore feasibility for the newly measured state and constraints before using the full-length sequence as a warm start. Reset or use additional starts after abrupt disturbances, reference changes, constraint changes, or mode switches.
  5. Log gradients by time step. A large span is a diagnostic trigger, not a verdict: inspect contraction or expansion, action gains, cost placement, cancellations, constraints, units, and horizon. For curvature, use matrix-free Hessian-vector products to estimate the most negative and largest algebraic eigenvalues of a symmetric action Hessian, or Jacobian-vector and vector-Jacobian products to estimate extreme singular values of a nonsymmetric rollout Jacobian or sensitivity product. Form the matrix directly only when the action problem is small enough.
  6. If progress stalls, first check gradient correctness, scale, saturation, active constraints, smoothness, and initialization. Consider regularized iLQR or DDP when curvature is the limiting factor and a useful local quadratic model exists.