Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Structured Simulators and Neural Operators › Worked Example: Planning Through a Differentiable Simulator

Michael BrenndoerferJune 22, 20265 min read

World Models HandbookAll chapters

Part 4 of 12Space, Physics, and Causality

Chapter 22Structured Simulators and Neural Operators › Worked Example: Planning Through a Differentiable Simulator

Reading section 7 of 11 · Worked Example: Planning Through a Differentiable Simulator · 5 minSection 7 of 11 · 5 min

Section 7 of 11 · 5 min read

Structured Simulators and Neural Operators › Worked Example: Planning Through a Differentiable SimulatorOne section · Sections in this chapter7 / 11 ↓

Sections in this chapter11

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

Everything above was about prediction quality. This worked example turns a structured model toward decision-making, which is where world models earn their keep, and it demonstrates the difference between open-loop prediction and closed-loop usefulness that runs through this book. The example is intentionally simple, because its purpose is not to showcase a hard control problem but to expose, in a setting where the answer is known, exactly what differentiability contributes to planning and exactly where the resulting plans are fragile.

The task: a damped pendulum starts at angle $\theta_0 = 1.2$ and we want the final angle near $\theta = 0$ after 2.4 seconds, using a bounded torque $|u| \le 2$ . We use normalized units with effective moment of inertia $I=1$ , so torque enters the angular acceleration as $+u$ ; otherwise the dynamics contain $+u/I$ . We know the physics exactly, so the differentiable simulator is the world model, and the plan is simply a sequence of torques found by gradient descent through the rollout:

\min_{u_{0:T-1}} \quad \theta_T^2 + \rho\, \frac{1}{T}\sum_t u_t^2 \qquad \text{s.t.} \quad x_{t+1} = f(x_t, u_t), \quad |u_t| \le u_{\max}.

where:

$u_{0:T-1}$ : the sequence of $T$ torque actions being optimized, the decision variables of the problem
$\theta_T$ : the pendulum angle after $T$ steps under the chosen action sequence, which we want near zero
$\rho$ : the effort weight penalizing control energy, set to $0.01$ in the experiment below
$\frac{1}{T}\sum_t u_t^2$ : the mean squared normalized torque over the horizon, discouraging aggressive actuation
$x_{t+1} = f(x_t, u_t)$ : the differentiable simulator dynamics as a constraint on every step
$u_{\max}$ : the torque bound, with $|u_t| \le u_{\max}$ enforced via a tanh squashing in the implementation

This is a shooting method from Part II, Chapter 4, with autodiff replacing hand-derived adjoints. The constraint is handled implicitly: rather than enforcing the dynamics as constraints in a constrained optimizer, we simulate them forward and optimize only over the action sequence, which is what "shooting" means. The dynamics are satisfied by construction inside the simulator; the price is that the parameterization of the decision variables, a flat sequence of torques, gives the optimizer no mechanism to correct early mistakes except through the actions themselves.

In[27]:

Code

import numpy as np
import torch

torch.manual_seed(0)

theta_start, omega_start = 1.2, 0.0
n_ctrl = 120
u_max = 2.0


def simulate_controlled(u_seq):
    theta = torch.tensor(float(theta_start))
    omega = torch.tensor(float(omega_start))
    traj = [theta]
    for ut in u_seq:
        alpha = -(g_true / L) * torch.sin(theta) - c * omega + ut
        omega = omega + dt * alpha
        theta = theta + dt * omega
        traj.append(theta)
    return torch.stack(traj)


traj_before = simulate_controlled(torch.zeros(n_ctrl)).detach().numpy()

u_param = torch.zeros(n_ctrl, requires_grad=True)
opt = torch.optim.Adam([u_param], lr=0.05)
ctrl_losses = []
for step in range(500):
    u = u_max * torch.tanh(u_param)  # bounded torque via squashing
    traj = simulate_controlled(u)
    loss = traj[-1] ** 2 + 0.01 * torch.mean(u**2)
    opt.zero_grad()
    loss.backward()
    opt.step()
    ctrl_losses.append(loss.item())

with torch.no_grad():
    traj_after = simulate_controlled(u_max * torch.tanh(u_param)).numpy()

theta_before_final = traj_before[-1]
theta_after_final = traj_after[-1]
omega_before = np.gradient(traj_before, dt)
omega_after = np.gradient(traj_after, dt)
times_ctrl = np.arange(n_ctrl + 1) * dt

Out[28]:

Console

Final angle, no control: 0.777 rad
Final angle, optimized:  0.023 rad
Approx. final angular velocity: -2.159 rad/s
Final cost, optimized:   0.0155

Out[29]:

Visualization

Two angle trajectories with the optimized one reaching near zero at the terminal time. — Pendulum angle over time under zero torque and under the torque sequence optimized by gradients through the simulator. The optimized sequence reaches an angle near zero at the 2.4-second terminal time but does not settle the pendulum there.

Semilog plot of planning cost decreasing over iterations. — Shooting cost during action-sequence optimization. The terminal-plus-effort cost decreases overall across 500 gradient steps.

Gradient descent through the simulator finds a torque sequence that reaches the target angle at the terminal time. The pendulum is still moving, because the objective penalizes terminal angle but not terminal angular velocity or dwell time. Notice what made the optimization work: differentiability of the implemented physics supplied action sensitivities, differentiability through time composed them across 120 steps, and differentiable squashing enforced the torque bound at every iterate. These gradients are correct for the implemented model; mismatch between that model and reality can still mislead the planner.

This plan is open loop: we optimized one torque sequence against the model and executed it. If the real pendulum's friction differed from the model's, or a nudge arrived mid-swing, the plan would drift and the realized terminal error would differ from the prediction. The plan is most fragile where the optimized trajectory passes through sensitive regions of state space, yet a single deterministic rollout does not identify those regions. The standard remedies are to re-plan from the measured state every few steps (receding-horizon control, developed in Part VII) and to model residual uncertainty explicitly. Re-planning replaces one long plan with a chain of shorter optimizations at the cost of solving a problem at each control step. A model accurate enough to optimize through is not automatically safe to trust for the full horizon. This open-loop-versus-closed-loop gap is a central evaluation criterion for the architectures in Part V.