Differentiable Planning › Code Implementation
World Models HandbookAll chapters Differentiable Planning › Code ImplementationOne section · Sections in this chapter7 / 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 build the planner. Each experiment isolates one point: autograd matches the analytic gradient at a tested action vector, Adam improves an open-loop plan, a Riccati terminal value improves a two-step MPC controller under a fixed optimization budget, and scalar dynamics expose contracting and expanding gradient chains.
The chapter-wide book setup runs before the first executable example so its plotting helpers are available throughout. Seeding controls the common random streams and improves repeatability within a fixed software and hardware environment; it does not by itself guarantee identical results across platforms or library versions.
Next we define the double integrator, the quadratic cost, and the horizon. Keeping the constants in one place makes it easy to vary the horizon later.
import numpy as np
import torch
torch.manual_seed(0)
## Discrete-time double integrator: state z = [position, velocity]
dt = 0.25
A = torch.tensor([[1.0, dt], [0.0, 1.0]])
B = torch.tensor([[0.5 * dt**2], [dt]])
## Quadratic cost with target position 1.0 and target velocity 0.0
Q = torch.diag(torch.tensor([1.0, 0.1]))
R = torch.tensor([[0.01]])
goal = torch.tensor([1.0, 0.0])
z0 = torch.tensor([0.0, 0.0])
H = 12The rollout and the cost are the two ingredients that autograd will differentiate. Notice that the rollout is an ordinary Python loop; the computation graph is built implicitly as the loop runs, exactly as it would be for an RNN. Autograd requires no planning-specific differentiation machinery or hand-derived adjoint, and the planner's autograd path does not assemble explicit Jacobians. The independent check below assembles only for verification. The framework tracks the operations, and a single call to backward on the scalar cost computes and accumulates the full gradient in the action tensor's .grad field.
def rollout(z_init, actions):
"""Roll the model forward. actions: (H,) tensor. Returns states (H+1, 2)."""
b_col = B.squeeze(-1)
states = [z_init]
for t in range(actions.shape[0]):
states.append(A @ states[-1] + b_col * actions[t])
return torch.stack(states)
def trajectory_cost(states, actions, terminal_Q=Q):
"""Stage costs plus an optional quadratic terminal cost."""
stage = sum((z - goal) @ Q @ (z - goal) for z in states[:-1])
ctrl = sum(R.item() * a**2 for a in actions)
terminal = (
torch.tensor(0.0)
if terminal_Q is None
else (states[-1] - goal) @ terminal_Q @ (states[-1] - goal)
)
return stage + ctrl + terminalNow we compute the gradient two ways. Autograd differentiates the rollout; the second calculation assembles the quadratic gradient independently. Agreement at the all-zero action vector is a focused numerical check of this implementation, consistent with the adjoint relation derived above.
## --- Autograd gradient at the all-zero action sequence ---
a_flat = torch.zeros(H, requires_grad=True)
states0 = rollout(z0, a_flat)
J_val = trajectory_cost(states0, a_flat)
J_val.backward()
grad_autograd = a_flat.grad.detach().clone()
## --- Closed-form gradient from the quadratic structure ---
def control_to_state_matrix(k, horizon):
"""G_k such that z_k = A^k z_0 + G_k @ a for a of length `horizon`."""
G = torch.zeros(2, horizon)
for j in range(min(k, horizon)):
G[:, j] = (torch.matrix_power(A, k - 1 - j) @ B).squeeze(-1)
return G
M = torch.zeros(H, H)
b_lin = torch.zeros(H)
for k in range(H + 1):
Gk = control_to_state_matrix(k, H)
M += Gk.T @ Q @ Gk
b_lin += Gk.T @ Q @ (torch.matrix_power(A, k) @ z0 - goal)
M += torch.eye(H) * R.item()
grad_closed_form = 2 * b_lin # evaluated at a = 0, where grad J = 2 M a + 2 bObjective J(a=0) : 13.0000 Max |autograd - closed form| : 0.000e+00 Autograd gradient, first four: [-9. -7.5625 -6.25 -5.0625] Gradient norm : 15.3996
The two gradients agree to numerical precision at this tested point. We check the gradient norm explicitly because agreement between two negligible vectors would be a weak test. Here the all-zero diagnostic vector has a substantial nonzero gradient.
Gradient-based optimization improves the trajectory
Now we hand the action sequence to Adam. The "parameters" are actions rather than model weights; the loop still zeroes gradients, computes the loss, backpropagates, and takes an update.
actions = torch.linspace(-1.5, 1.5, H, requires_grad=True)
optimizer = torch.optim.Adam([actions], lr=0.02)
loss_history = []
for step in range(400):
optimizer.zero_grad()
states = rollout(z0, actions)
loss = trajectory_cost(states, actions)
loss.backward()
optimizer.step()
loss_history.append(loss.item())
with torch.no_grad():
final_states = rollout(z0, actions).clone()
final_actions = actions.clone()We plot the loss curve and resulting phase trajectory. The loss trace shows optimization progress, not a convergence proof. For this Adam run, the optimizer starts from the declared linear ramp from to ; the all-zero vector used in the nearby gradient checks is separate from that run. We use . The deliberately high-loss ramp makes the reduction easy to see, but the 400-step budget does not reach the analytic optimum.
Before looking at the final plan, inspect the gradient by action index. The open-loop double integrator is not asymptotically stable and, because its unit eigenvalue has a nontrivial Jordan block, it is not Lyapunov stable: with nonzero initial velocity, its zero-input position grows linearly. With cumulative state costs, early actions affect more stage costs than late actions, so the all-zero profile is strongest at the start of the horizon. This is an example of the sum of future-cost contributions dominating the profile, not of a vanishing early-action gradient.
## Recompute the gradient at the initial all-zero action sequence to inspect its profile.
probe_actions = torch.zeros(H, requires_grad=True)
probe_states = rollout(z0, probe_actions)
probe_loss = trajectory_cost(probe_states, probe_actions)
probe_loss.backward()
grad_profile_numpy = probe_actions.grad.detach().abs().numpy()
## Compare the earliest and latest action-gradient magnitudes.
grad_ratio_first_last = (
float(grad_profile_numpy[0] / grad_profile_numpy[-1])
if grad_profile_numpy[-1] != 0
else float("nan")
)
grad_span_orders = (
float(np.log10(grad_profile_numpy.max() / grad_profile_numpy.min()))
if grad_profile_numpy.min() > 0
else float("nan")
)
print(
f"|dJ/da_0| / |dJ/da_{{H-1}}| : {grad_ratio_first_last:.4f}".replace(
"_{{H-1}}", "[H-1]"
)
)
print(f"Orders of magnitude spanned : {grad_span_orders:.3f}")
(<Figure size 900x750 with 1 Axes>, <Figure size 900x750 with 1 Axes>)
The objective decreases throughout the plotted run. Under this fixed budget, the trajectory approaches and overshoots the target position, ending near rather than settling exactly at . In this run, the optimizer produces rising and then falling velocity while trading finite-horizon position and velocity tracking against action effort. No separately designed feedback law is used: given the model, cost, initialization, and optimizer, backpropagation through the known dynamics computes the action gradients, and Adam uses those gradients to produce the open-loop action sequence.
Ill-conditioning grows with the horizon in this example
We now assemble for increasing horizons in this fixed double-integrator problem. The true action Hessian is , which has the same condition number as . Horizon growth need not worsen every system, but it does here.
horizons = list(range(2, 41))
condition_numbers = []
for h in horizons:
Mh = torch.zeros(h, h)
for k in range(h + 1):
Gk = control_to_state_matrix(k, h)
Mh += Gk.T @ Q @ Gk
Mh += torch.eye(h) * R.item()
eigvals = torch.linalg.eigvalsh(Mh)
# Curvature accumulates through the action-to-state sensitivity matrices G_k.
condition_numbers.append(float(eigvals[-1] / eigvals[0]))
print(f"Condition number at H=2 : {condition_numbers[0]:.1f}")
print(f"Condition number at H=40: {condition_numbers[-1]:.1f}")
For this system, the condition number climbs by orders of magnitude over the tested range. This eigenvalue spread slows unpreconditioned fixed-step gradient descent. Curvature-aware or preconditioned methods can compensate for part of it, although approximate Hessians can remain inaccurate. Regularization choices and finite precision also matter.
For , the Hessian has largest eigenvalue . Gradient descent with step contracts Euclidean iterate error by at most per step. For and , reaching is therefore guaranteed once
For an integer iteration count, take the ceiling of the right-hand side. When , the step reaches the quadratic optimum in one update instead. This is a worst-case bound, not a prediction that every initialization will attain equality. We compare it with the realized iterate-error threshold for the quadratic using the same learning rate.
## Worst-case steps to reduce Euclidean iterate error to the declared tolerance
## with alpha = 1 / L, where L = 2 * lambda_max(M).
tolerance = 1e-2
condition_numbers_np = np.array(condition_numbers, dtype=float)
bound_steps = np.ceil(
np.log(tolerance) / np.log(1.0 - 1.0 / condition_numbers_np)
)
## Handle the kappa=1 boundary or a value rounded to it.
bound_steps = np.maximum(bound_steps, 1.0)
## Realized iterate-error threshold for H=12 with the same alpha = 1 / L.
quadratic_eigvals = torch.linalg.eigvalsh(M)
quad_lr = 1.0 / (2.0 * float(quadratic_eigvals[-1]))
quad_optimum = torch.linalg.solve(M, -b_lin)
quad_actions = torch.zeros(H)
initial_iterate_error = torch.linalg.norm(quad_actions - quad_optimum)
h12_index = horizons.index(H)
actual_steps_h12 = None
for step in range(int(bound_steps[h12_index]) + 1):
relative_error = (
torch.linalg.norm(quad_actions - quad_optimum) / initial_iterate_error
)
if relative_error <= tolerance:
actual_steps_h12 = step
break
quad_gradient = 2.0 * (M @ quad_actions + b_lin)
quad_actions = quad_actions - quad_lr * quad_gradient
assert actual_steps_h12 is not None
assert actual_steps_h12 <= bound_steps[h12_index]
print(f"Worst-case bound at H=2 : {bound_steps[0]:.0f} steps")
print(f"Worst-case bound at H=12 : {bound_steps[h12_index]:.0f} steps")
print(f"Realized steps at H=12 : {actual_steps_h12} steps")
print(f"Worst-case bound at H=40 : {bound_steps[-1]:.0f} steps")
A Riccati terminal value improves a two-step horizon
We solve the discrete algebraic Riccati equation numerically, then compare three receding-horizon controllers under the same 100-iteration Adam budget per planning call: a declared two-step horizon without a terminal value, the same two-step horizon with the Riccati value, and a twelve-step reference without a terminal value. Finite optimizer budgets mean this is an empirical comparison of these implementations, not an exact Bellman-equivalence test.
import scipy.linalg as sla
P_np = sla.solve_discrete_are(
A.numpy().astype(float),
B.numpy().astype(float),
Q.numpy().astype(float),
R.numpy().astype(float),
)
P = torch.tensor(P_np, dtype=torch.float32)
print(f"Riccati solution P:\n{P.numpy().round(4)}")def plan(z_init, horizon, terminal_Q, iters=100, lr=0.1):
"""Optimize an action sequence of length `horizon` from state z_init."""
a = torch.zeros(horizon, requires_grad=True)
opt = torch.optim.Adam([a], lr=lr)
for _ in range(iters):
opt.zero_grad()
loss = trajectory_cost(rollout(z_init, a), a, terminal_Q=terminal_Q)
loss.backward()
opt.step()
return a.detach()
def mpc_closed_loop(horizon, terminal_Q, steps=25):
"""Receding-horizon control: plan, apply the first action, replan."""
z = z0.clone()
trajectory = [z.clone()]
total = 0.0
for _ in range(steps):
a_plan = plan(z, horizon, terminal_Q)
with torch.no_grad():
a0 = a_plan[0]
total += float((z - goal) @ Q @ (z - goal) + R.item() * a0**2)
z = A @ z + B.squeeze(-1) * a0
trajectory.append(z.clone())
return total, torch.stack(trajectory)
SHORT_HORIZON = 2
cost_short, traj_short = mpc_closed_loop(horizon=SHORT_HORIZON, terminal_Q=None)
cost_term, traj_term = mpc_closed_loop(horizon=SHORT_HORIZON, terminal_Q=P)
cost_long, traj_long = mpc_closed_loop(horizon=12, terminal_Q=None)
short_overshoot = float((traj_short[:, 0] - goal[0]).max())
term_overshoot = float((traj_term[:, 0] - goal[0]).max())
cost_reduction_percent = 100.0 * (cost_short - cost_term) / cost_shortCumulative cost, H=2, no terminal value : 3.568 Cumulative cost, H=2 + Riccati value : 2.824 Cumulative cost, H=12, no terminal value : 3.113 Cost reduction from the terminal value : 20.9% Peak overshoot, no terminal / Riccati : 0.0433 / 0.0052
The two-step controller without a terminal value overshoots the target by about in this deterministic run. Adding the Riccati value reduces peak overshoot to about and lowers the 25-step cumulative cost by about relative to the same two-step controller. The twelve-step curve is a reference, not a claimed trajectory match; with 100 Adam iterations per call its cumulative cost is about higher than the terminal-value controller's, an ordering influenced by finite optimization error. A terminal value is useful when it preserves decision-relevant values or gradients over candidate terminal states.
The gradient chain under stable and unstable dynamics
The final experiment isolates the spectral behavior of the Jacobian chain. For and , the general derivative is
At the code's evaluation point with all actions zero, and this reduces to . We sweep nonnegative and the horizon, recording the first-action gradient magnitude.
def terminal_gradient_profile(rho, horizon):
"""|dJ/da_t| for J = (z_H - 1)^2 with z_{t+1} = rho * z_t + a_t."""
a = torch.zeros(horizon, requires_grad=True)
z = torch.tensor(0.0)
for t in range(horizon):
z = rho * z + a[t]
((z - 1.0) ** 2).backward()
return a.grad.detach().abs().numpy()
rhos = [0.7, 0.9, 1.0, 1.1, 1.25]
sweep_horizons = list(range(1, 41))
first_action_grad = {
rho: [float(terminal_gradient_profile(rho, h)[0]) for h in sweep_horizons]
for rho in rhos
}
profile_stable = terminal_gradient_profile(0.9, 30)
profile_unstable = terminal_gradient_profile(1.1, 30)(<Figure size 900x1050 with 1 Axes>, <Figure size 900x1050 with 1 Axes>)
The left panel shows the predicted exponential decay and growth for this terminal-only objective: for the first-action gradient tends toward zero, for it grows, and at it stays constant. More generally, contraction occurs when . At fixed , the first-action magnitude is times the last-action magnitude; with the ordering reverses. With plain gradient descent using one global learning rate, larger gradient components receive proportionally larger updates.
The update rule strongly affects relative step magnitudes. The next plot is an idealized first-step direction comparison, not a simulation of stateful Adam. It contrasts a normalized plain-gradient magnitude profile with elementwise normalization . This sign-like normalization resembles Adam's first update only when moment state starts at zero, gradients dominate , and clipping, weight decay, or other interventions do not materially alter the coordinatewise first update.
## Gradient profile at a representative stable contraction factor and horizon.
imbalance_profile = terminal_gradient_profile(0.9, 30).astype(float)
plain_step_magnitudes = imbalance_profile / np.linalg.norm(imbalance_profile)
## Idealized first-step elementwise normalization, not stateful Adam.
normalized_coordinates = imbalance_profile / np.maximum(
imbalance_profile, 1e-12
)
normalized_step_magnitudes = normalized_coordinates / np.linalg.norm(
normalized_coordinates
)
print(
"Plain-gradient magnitude ratio, late/early : "
f"{plain_step_magnitudes[-1] / plain_step_magnitudes[0]:.3e}"
)
print(
"Elementwise-normalized ratio, late/early : "
f"{normalized_step_magnitudes[-1] / normalized_step_magnitudes[0]:.3e}"
)