Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Physical Dynamics and Constraints › Code Implementation

Michael BrenndoerferJune 20, 202611 min read

World Models HandbookAll chapters

Part 4 of 12Space, Physics, and Causality

Chapter 20Physical Dynamics and Constraints › Code Implementation

Reading section 9 of 12 · Code Implementation · 11 minSection 9 of 12 · 11 min

Section 9 of 12 · 11 min read

Physical Dynamics and Constraints › Code ImplementationOne section · Sections in this chapter9 / 12 ↓

Sections in this chapter12

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

First the hidden setup cell, which locates the shared plotting style and seeds all randomness for the book's figures.

Simulating the Ground-Truth PendulumLink Copied

We sample states from an RK4 trajectory that starts at $\theta_0=2.0$ radians, where the small-angle approximation is inadequate. At each sampled state, the training target is the analytic instantaneous derivative from the pendulum equation. Using the exact vector field avoids asking a continuous-time Hamiltonian model to reproduce a finite-step forward-difference target that need not itself be Hamiltonian.

In[3]:

Code

import numpy as np

G = 9.81
L = 1.0


def pendulum_deriv(x):
    theta, omega = x
    return np.array([omega, -(G / L) * np.sin(theta)])


def rk4_step(f, x, dt):
    k1 = f(x)
    k2 = f(x + 0.5 * dt * k1)
    k3 = f(x + 0.5 * dt * k2)
    k4 = f(x + dt * k3)
    return x + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)


dt = 0.05
n_steps = 600  # 30 seconds of data
x = np.array([2.0, 0.0])
traj = [x.copy()]
for _ in range(n_steps):
    x = rk4_step(pendulum_deriv, x, dt)
    traj.append(x.copy())
traj = np.array(traj)

From this trajectory we build supervised pairs $(x_t, \dot{x}_t)$ using the analytic derivative at each sampled $x_t$ . Rollout behavior is evaluated separately on an initial condition outside the training trajectory's angle range.

In[4]:

Code

X = traj[:-1]
dX = np.array([pendulum_deriv(x) for x in X])
print(
    f"Training pairs: {X.shape[0]}, dt = {dt} s, horizon = {n_steps * dt:.0f} s"
)
print(f"Angle range: [{traj[:, 0].min():.2f}, {traj[:, 0].max():.2f}] rad")

Training the Naive ModelLink Copied

The naive model is a plain multilayer perceptron from state to derivative. Nothing in its architecture requires $\dot{\theta}=\omega$ or energy conservation. The finite network has ample capacity to approximate the field over the bounded training region, but the loss constrains it only at the sampled states.

In[5]:

Code

import torch
import torch.nn as nn

torch.manual_seed(0)
X_t = torch.tensor(X, dtype=torch.float32)
dX_t = torch.tensor(dX, dtype=torch.float32)

naive = nn.Sequential(
    nn.Linear(2, 64),
    nn.Tanh(),
    nn.Linear(64, 64),
    nn.Tanh(),
    nn.Linear(64, 2),
)
opt = torch.optim.Adam(naive.parameters(), lr=1e-3)
for step in range(3000):
    opt.zero_grad()
    loss = ((naive(X_t) - dX_t) ** 2).mean()
    loss.backward()
    opt.step()
print(f"Naive one-step MSE: {loss.item():.2e}")

Training the Hamiltonian ModelLink Copied

The Hamiltonian model learns a scalar Hamiltonian candidate $\mathcal{H}_\theta(\theta,p)$ from the same data. Because mass and length are both one, the conjugate momentum $p=ml^2\omega$ equals $\omega$ numerically. Its derivatives come from autodiff: $\dot{\theta}=\partial\mathcal{H}_\theta/\partial p$ and $\dot p=-\partial\mathcal{H}_\theta/\partial\theta$ . The implementation detaches the inputs and re-wraps them as leaf tensors so that autograd.grad differentiates the scalar with respect to the state. For every setting of the weights, the antisymmetric pairing conserves the learned autonomous $\mathcal{H}_\theta$ under exact continuous flow. It does not guarantee agreement with the analytic pendulum field or exact conservation by a discrete integrator.

In[6]:

Code

class HamiltonianModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, 64),
            nn.Tanh(),
            nn.Linear(64, 64),
            nn.Tanh(),
            nn.Linear(64, 1),
        )

    def h(self, x):
        return self.net(x).squeeze(-1)

    def deriv(self, x):
        # x: (batch, 2) state [theta, p]. Returns time derivative.
        xin = x.detach().clone().requires_grad_(True)
        H = self.h(xin).sum()
        (dH,) = torch.autograd.grad(H, xin, create_graph=True)
        return torch.stack([dH[:, 1], -dH[:, 0]], dim=-1)


hnn = HamiltonianModel()
opt = torch.optim.Adam(hnn.parameters(), lr=1e-3)
for step in range(3000):
    opt.zero_grad()
    loss = ((hnn.deriv(X_t) - dX_t) ** 2).mean()
    loss.backward()
    opt.step()
print(f"Hamiltonian one-step MSE: {loss.item():.2e}")

For seed 0, the final derivative MSE is approximately $3.56\times10^{-5}$ for the MLP and $1.33\times10^{-5}$ for the Hamiltonian model. Both are small on the sampled training states, with the Hamiltonian loss about 2.7 times lower in this run. A phase portrait shows what each fitted field asserts over a wider grid; Hamiltonian circulation is structural, while fidelity to the true arrows remains empirical.

In[7]:

Code

import numpy as np
import torch

## Evaluation grid: a regular 21x21 grid over the phase plane, fixed by design
## (no randomness, no seed dependence). Range spans the training amplitude and
## the held-out rollout amplitude.
theta_grid = np.linspace(-np.pi, np.pi, 21)
omega_grid = np.linspace(-6.0, 6.0, 21)
Theta_grid, Omega_grid = np.meshgrid(theta_grid, omega_grid)
grid_pts = np.stack([Theta_grid.ravel(), Omega_grid.ravel()], axis=1)

## True derivative field on the grid, computed analytically
true_field = np.array([pendulum_deriv(p) for p in grid_pts])

## Learned derivative fields on the grid. The direct network needs no gradient
## tracking, while the Hamiltonian field is itself defined by an input gradient.
grid_t = torch.tensor(grid_pts, dtype=torch.float32)
with torch.no_grad():
    naive_field = naive(grid_t).numpy()
with torch.enable_grad():
    hnn_field = hnn.deriv(grid_t).detach().numpy()

grid_rmse_naive = np.sqrt(np.mean((naive_field - true_field) ** 2))
grid_rmse_hnn = np.sqrt(np.mean((hnn_field - true_field) ** 2))
print(
    f"Grid RMSE  naive: {grid_rmse_naive:.2f}   Hamiltonian: {grid_rmse_hnn:.2f}"
)

Out[8]:

Visualization

Three side-by-side phase-plane panels for the true, naive MLP, and Hamiltonian vector fields. Both learned fields differ visibly from the true arrows away from the sampled trajectory. — True and learned pendulum vector fields on the same evaluation grid. In this seed, grid RMSE is 2.88 for the MLP and 2.74 for the Hamiltonian model. Hamiltonian structure organizes the learned arrows around contours of its learned scalar, but neither field matches the analytic field globally.

Rolling Out and Measuring Energy DriftLink Copied

We roll both models out from $\theta_0=2.5$ radians for 20 seconds, feeding each prediction back as the next input. This angle lies outside the training trajectory's range of about $[-2,2]$ radians. We record state error, the analytic physical energy $\mathcal H_{\text{phys}}$ of each predicted state, and the Hamiltonian model's learned scalar $\mathcal H_\theta$ . The exact pendulum conserves $\mathcal H_{\text{phys}}$ ; the architecture's continuous-time guarantee applies to $\mathcal H_\theta$ instead.

In[9]:

Code

def pendulum_energy(x):
    theta, omega = x[..., 0], x[..., 1]
    return 0.5 * (L**2) * omega**2 + G * L * (1.0 - np.cos(theta))


def rollout_torch(deriv_fn, x0, dt, steps):
    x = torch.tensor(x0, dtype=torch.float32).unsqueeze(0)
    out = [x.detach().numpy().copy()]
    for _ in range(steps):
        x = rk4_step(
            deriv_fn, x, dt
        )  # same RK4 rule for both models; only the learned field differs
        out.append(x.detach().numpy().copy())
    return np.concatenate(out, axis=0)


def rollout_naive(x0, dt, steps):
    x = torch.tensor(x0, dtype=torch.float32).unsqueeze(0)
    out = [x.detach().numpy().copy()]
    with torch.no_grad():
        for _ in range(steps):
            x = rk4_step(naive, x, dt)
            out.append(x.detach().numpy().copy())
    return np.concatenate(out, axis=0)


x0_test = np.array([2.5, 0.0])
roll_steps = 400  # 20 seconds

## Ground truth via RK4 at the same dt
gt = [x0_test.copy()]
x = x0_test.copy()
for _ in range(roll_steps):
    x = rk4_step(pendulum_deriv, x, dt)
    gt.append(x.copy())
gt = np.array(gt)

roll_naive = rollout_naive(x0_test, dt, roll_steps)
roll_hnn = rollout_torch(hnn.deriv, x0_test, dt, roll_steps)

times = np.arange(roll_steps + 1) * dt
err_naive = np.linalg.norm(roll_naive - gt, axis=1)
err_hnn = np.linalg.norm(roll_hnn - gt, axis=1)

E0 = pendulum_energy(x0_test)
E_naive = pendulum_energy(roll_naive) / E0
E_hnn = pendulum_energy(roll_hnn) / E0
with torch.no_grad():
    learned_H = hnn.h(torch.tensor(roll_hnn, dtype=torch.float32)).numpy()
learned_H_drift = learned_H - learned_H[0]

print(
    f"Final error @ {times[-1]:.0f}s  naive: {err_naive[-1]:.2f}   Hamiltonian: {err_hnn[-1]:.2f}"
)
print(
    f"Final energy ratio     naive: {E_naive[-1]:.3f}   Hamiltonian: {E_hnn[-1]:.3f}"
)
print(f"Max |learned H drift|: {np.abs(learned_H_drift).max():.2e}")

Both learned models use the same RK4 rule, which controls one source of variation but does not remove integration error or the different approximation errors in their vector fields. RK4 is nonsymplectic. In this seeded run, the HNN's learned scalar changes by at most about $1.34\times10^{-4}$ , while the analytic physical-energy ratio on the same trajectory ranges from about 1.000 to 1.035. The distinction is the point: nearly conserving $\mathcal H_\theta$ is structural, while agreement with $\mathcal H_{\text{phys}}$ is empirical.

Now consider a minimal hard unilateral constraint. The unconstrained baseline is free fall with no collision handling. The constrained update projects a penetrated position back to $y=0$ and then reflects the normal velocity using a restitution coefficient. It illustrates feasibility plus an impact law; it is not a full complementarity solver because it computes no contact multiplier, impact time, sustained-contact solution, or friction.

In[10]:

Code

def bounce_projected(dt=0.02, steps=200, y0=1.0, e=0.8):
    y, v = y0, 0.0
    ys = [y]
    for _ in range(steps):
        v = v - 9.81 * dt
        y = y + v * dt
        if y < 0.0:  # constraint phi(y) = y >= 0 violated
            y = 0.0  # projection: eliminate penetration
            v = -e * v  # impact impulse: reflect with restitution
        ys.append(y)
    return np.array(ys)


def bounce_unconstrained(dt=0.02, steps=200, y0=1.0):
    y, v = y0, 0.0
    ys = [y]
    for _ in range(steps):
        v = v - 9.81 * dt
        y = y + v * dt
        ys.append(y)
    return np.array(ys)


ball_free = bounce_unconstrained()
ball_proj = bounce_projected()
t_ball = np.arange(len(ball_free)) * 0.02
print(
    f"Min height, unconstrained: {ball_free.min():.2f} m (penetrates the floor)"
)
print(f"Min height, projected:     {ball_proj.min():.3f} m (never below zero)")

The position test implements $\phi(q)\geq0$ by projecting violations back onto the feasible set, and the velocity reflection adds a restitution update. This code enforces nonpenetration exactly at the sampled states. A generic smooth predictor can approximate the bounce map, but it provides no exact nonpenetration guarantee unless the constraint is added explicitly.

Visualizing the ResultsLink Copied

Both figures are built entirely from the arrays computed above. The first compares rollout error and energy drift for the two pendulum models; the second shows the constraint projection keeping the ball on the correct side of the floor.

Out[11]:

Visualization

Log-scale line chart with oscillatory MLP and Hamiltonian state-error curves; the Hamiltonian curve ends lower at 20 seconds, though the curves overlap and cross earlier. — Held-out rollout state error for seed 0. Both errors oscillate and grow; at 20 seconds the MLP error is 7.98 and the Hamiltonian-model error is 4.92. This single run does not isolate an architecture-wide advantage.

Line chart showing the MLP's physical-energy ratio rising as high as about 1.41 while the Hamiltonian trajectory stays close to one. — Analytic physical energy evaluated on each predicted trajectory. The MLP reaches 1.40 times the initial energy, while the Hamiltonian-model trajectory remains between 1.000 and 1.035 in this run. This physical-energy agreement is observed, not guaranteed by the architecture.

Out[12]:

Visualization

Two side-by-side line charts. The left shows a projected bouncing ball staying between zero and one meter with diminishing rebounds. The right shows unconstrained free fall passing through zero and reaching about minus 78 meters. — The projected trajectory remains at or above the floor and shows the diminishing rebounds produced by restitution. The separate free-fall panel keeps its much larger negative range from flattening the constrained motion.

The panels measure different properties. On the sampled training states, both derivative losses are small and the Hamiltonian loss is lower for this seed. On the held-out rollout, the HNN also has the lower 20-second endpoint error, 4.92 rather than 7.98, while neither trajectory stays close to the ground truth throughout. The energy panel shows a clearer structural contrast: the MLP trajectory gains about 40% physical energy, whereas the HNN trajectory remains within roughly 3.5% of its initial physical energy and its learned $\mathcal H_\theta$ changes by only about $1.34\times10^{-4}$ . This one seed is an illustration, not evidence that Hamiltonian models generally win. It shows instead how to audit separate questions: Was the local field fitted? Is the learned invariant conserved? Does that invariant match the physical quantity of interest? Is the state trajectory accurate enough for the intended horizon?