Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Taxonomies of World Models › Implementing the Comparison

Michael BrenndoerferJune 4, 20269 min read

World Models HandbookAll chapters

Part 1 of 12The World-Model Idea

Chapter 4Taxonomies of World Models › Implementing the Comparison

Reading section 8 of 12 · Implementing the Comparison · 9 minSection 8 of 12 · 9 min

Section 8 of 12 · 9 min read

Taxonomies of World Models › Implementing the ComparisonOne section · Sections in this chapter8 / 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

We begin with the shared substrate: the simulator, the dataset, and the nonlinear sensor. Everything downstream consumes these arrays.

In[6]:

Code

import numpy as np

rng = np.random.default_rng(11)

DT = 0.05  # integration step (seconds)
NOISE_STD = 0.08  # process noise on angular velocity


def pendulum_step(theta, omega, u):
    """One Euler step of a damped pendulum with torque input u."""
    omega_next = omega + DT * (-np.sin(theta) - 0.15 * omega + u)
    omega_next += NOISE_STD * np.sqrt(DT) * rng.normal()
    return theta + DT * omega_next, omega_next


N_TRAIN_EPISODES, N_TEST_EPISODES, EP_LEN = 30, 10, 120

episodes = []
for _ in range(N_TRAIN_EPISODES + N_TEST_EPISODES):
    theta, omega = rng.uniform(-np.pi, np.pi), 0.0
    th, om, ac = [theta], [omega], []
    for _ in range(EP_LEN):
        u = float(np.clip(rng.normal(0.0, 1.2), -2.0, 2.0))
        theta, omega = pendulum_step(theta, omega, u)
        th.append(theta)
        om.append(omega)
        ac.append(u)
    episodes.append((np.array(th), np.array(om), np.array(ac)))

train_episodes = episodes[:N_TRAIN_EPISODES]
test_episodes = episodes[N_TRAIN_EPISODES:]

Each episode holds 121 states and 120 actions, giving 120 paired transitions per episode and 3,600 training transitions in total. The actions are independent Gaussian torque samples clipped to a realistic range. They provide broad exploratory coverage without the deliberately designed excitation sequence that a classical system-identification experiment would use. This is deliberate: we want to see how far generic regression gets with ordinary exploratory data rather than idealized persistently exciting inputs.

Next we build the sensor. The key design decision is the $\tanh$ : without it, a linear map of linear-in-feature dynamics would still be linear in observation coordinates, and Contestant A would unfairly escape its fate. Real sensors, cameras above all, are emphatically nonlinear, and the $\tanh$ stands in for that fact.

In[7]:

Code

OBS_DIM = 16


def features(theta, omega):
    """The physically sufficient statistics of the pendulum state."""
    return np.stack([np.sin(theta), np.cos(theta), omega], axis=-1)


proj = rng.normal(size=(OBS_DIM, 3)) / np.sqrt(3.0)


def observe(theta, omega):
    """A fixed nonlinear 'sensor': random projection followed by tanh."""
    return np.tanh(features(theta, omega) @ proj.T)


obs_episodes = [(observe(th, om), ac, th, om) for th, om, ac in train_episodes]
test_obs_episodes = [
    (observe(th, om), ac, th, om) for th, om, ac in test_episodes
]

To see what the sensor does to the state, we plot two views of the same training transitions: raw observation coordinates and the sufficient features, both colored by the true angle. In feature coordinates the state organizes into a clean ring parameterized by $\theta$ ; in observation coordinates the same data is a tangle with no visible angular structure. This is the coordinate problem Contestant A must solve with nothing but a linear map.

Out[9]:

Visualization

Scatter plot of tangled observation coordinates colored by angle. — The first two dimensions of the sixteen-dimensional observation vector, colored by the true angle. The nonlinear sensor tangles the state into a shape with no visible angular structure.

Scatter plot of sufficient features forming a ring, colored by angle. — The sufficient features (sin of angle versus angular velocity), colored by the same angles. In these coordinates the dynamics-relevant structure is a clean ring.

Now Contestant A: ordinary least squares from concatenated observation-and-action to next observation. This is the simplest thing that deserves the name "world model," and it is worth internalizing how little machinery it involves.

In[10]:

Code

def stack(episode_list):
    X = np.concatenate([o[:-1] for o, _, _, _ in episode_list])
    U = np.concatenate([ac for _, ac, _, _ in episode_list])[:, None]
    Y = np.concatenate([o[1:] for o, _, _, _ in episode_list])
    return X, U, Y


X_train, U_train, Y_train = stack(obs_episodes)
print(
    f"Training transitions: {X_train.shape[0]}, observation dimension: {X_train.shape[1]}"
)

phi_train = np.concatenate([X_train, U_train], axis=1)
W_obs, *_ = np.linalg.lstsq(phi_train, Y_train, rcond=None)

resid_one_step = Y_train - phi_train @ W_obs
print(
    f"Observation-space linear model, one-step RMSE: {np.sqrt((resid_one_step**2).mean()):.4f}"
)

Contestant B learns a compact state estimator in PyTorch. Its three outputs are trained to match the sufficient features $(\sin\theta, \cos\theta, \dot{\theta})$ on the simulator's training trajectories. This supervision is not a general recipe for world-model learning because a real agent rarely receives the true state. It is useful here because it creates a controlled middle case: observation-space dynamics must first pass through an imperfect learned encoder, while the transition itself is held fixed across B and C.

In[11]:

Code

import torch
import torch.nn as nn

torch.manual_seed(0)
LATENT_DIM = 3

F_train = np.concatenate(
    [features(th[:-1], om[:-1]) for _, _, th, om in obs_episodes]
)
F_next_train = np.concatenate(
    [features(th[1:], om[1:]) for _, _, th, om in obs_episodes]
)

encoder = nn.Sequential(
    nn.Linear(OBS_DIM, 64),
    nn.SiLU(),
    nn.Linear(64, 64),
    nn.SiLU(),
    nn.Linear(64, LATENT_DIM),
)
optimizer = torch.optim.Adam(encoder.parameters(), lr=1e-3)

x_t = torch.tensor(X_train, dtype=torch.float32)
feature_t = torch.tensor(F_train, dtype=torch.float32)

loss_history = []
for epoch in range(600):
    z = encoder(x_t)
    loss = ((z - feature_t) ** 2).mean()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    loss_history.append(loss.item())

with torch.no_grad():
    z_train = encoder(x_t).numpy()
print(f"Final latent state-estimation loss: {loss_history[-1]:.6f}")

The training curve below shows state-estimation loss descending smoothly. Low supervised loss confirms that the sensor can be untangled on the training distribution, but rollout evaluation is still needed because small initial-state errors can persist or amplify under dynamics.

Out[12]:

Visualization

Log-scale line chart of latent state-estimation loss decreasing over epochs. — Supervised latent state-estimation loss over 600 full-batch epochs. The encoder learns to recover dynamics-relevant features from nonlinear observations, but this training metric does not measure rollout error.

Contestant C receives the true features directly. Both B and C use the deterministic pendulum update with the stochastic term removed, so their difference comes only from learned versus oracle initialization. We also fit a one-step linear feature model solely to estimate the scale of process noise from residuals for the stochastic extension.

In[13]:

Code

g_train = np.concatenate([F_train, U_train], axis=1)
W_feat, *_ = np.linalg.lstsq(g_train, F_next_train, rcond=None)

feat_residuals = F_next_train - g_train @ W_feat
noise_cov = np.cov(feat_residuals.T)
print("Feature-space linear model, residual std per dimension:")
print(np.sqrt(np.diag(noise_cov)).round(4))

Evaluation happens in terms of the physical angle $\theta$ . The observation model therefore needs a small regularized linear readout from sensor coordinates to $(\sin\theta, \cos\theta)$ . Contestants B and C already represent those quantities explicitly, so they require no reporting model.

In[14]:

Code

theta_train = np.concatenate([th[:-1] for _, _, th, om in obs_episodes])
sc_targets = np.stack([np.sin(theta_train), np.cos(theta_train)], axis=1)
readout_design = np.column_stack([X_train, np.ones(len(X_train))])
readout_penalty = np.eye(readout_design.shape[1])
readout_penalty[-1, -1] = 0.0
W_readout = np.linalg.solve(
    readout_design.T @ readout_design + 1e-6 * readout_penalty,
    readout_design.T @ sc_targets,
)

Now the evaluation harness. Each contestant receives a rollout function that starts from an initial observation (or feature vector) and applies a held-out action sequence entirely inside its own imagination. Errors are computed on the wrapped angle difference so that crossings of $\pm\pi$ do not masquerade as huge mistakes.

In[15]:

Code

HORIZON = 60


def wrap(angle):
    return np.arctan2(np.sin(angle), np.cos(angle))


def rollout_obs_linear(o0, actions_h):
    o = o0.copy()
    thetas = []
    for u in actions_h:
        o = np.concatenate([o, [u]]) @ W_obs
        sc = np.concatenate([o, [1.0]]) @ W_readout
        thetas.append(np.arctan2(sc[0], sc[1]))
    return np.array(thetas)


def structured_feature_step(f, u, omega_noise=0.0):
    theta = np.arctan2(f[0], f[1])
    omega = f[2]
    omega_next = omega + DT * (-np.sin(theta) - 0.15 * omega + u) + omega_noise
    theta_next = theta + DT * omega_next
    return features(theta_next, omega_next)


def rollout_latent(o0, actions_h):
    with torch.no_grad():
        z = encoder(torch.tensor(o0, dtype=torch.float32)).numpy()
        thetas = []
        for u in actions_h:
            z = structured_feature_step(z, u)
            thetas.append(np.arctan2(z[0], z[1]))
    return np.array(thetas)


def rollout_oracle_feature(f0, actions_h):
    f = f0.copy()
    thetas = []
    for u in actions_h:
        f = structured_feature_step(f, u)
        thetas.append(np.arctan2(f[0], f[1]))
    return np.array(thetas)


rmse_curves = {}
example_predictions = {}
contestants = [
    ("observation linear", rollout_obs_linear, "observation"),
    ("learned latent", rollout_latent, "observation"),
    ("oracle feature", rollout_oracle_feature, "feature"),
]
for name, rollout, initializer in contestants:
    sq_errs = []
    for idx, (o_seq, ac, th, om) in enumerate(test_obs_episodes):
        initial_state = (
            features(th[0], om[0]) if initializer == "feature" else o_seq[0]
        )
        pred = rollout(initial_state, ac[:HORIZON])
        truth = th[1 : HORIZON + 1]
        sq_errs.append(wrap(pred - truth) ** 2)
        if idx == 0:
            example_predictions[name] = pred
    rmse_curves[name] = np.sqrt(np.mean(sq_errs, axis=0))

true_example = test_episodes[0][0][1 : HORIZON + 1]
check_step = HORIZON // 2
for name in rmse_curves:
    print(
        f"{name:>18s}: angle RMSE at step {check_step} = {rmse_curves[name][check_step - 1]:.3f} rad"
    )

Finally, the stochastic extension. We estimate the angular-velocity noise scale from Contestant C's one-step residuals, inject Gaussian innovations into the structured transition, and sample two hundred rollouts from one held-out initial state with zero torque. For contrast we also simulate true noisy pendulum trajectories from the same initial state.

In[16]:

Code

N_SAMPLES, STOCH_HORIZON = 200, 80
estimated_omega_noise = max(float(np.std(feat_residuals[:, 2])), 1e-6)

_, _, th_ref, om_ref = test_obs_episodes[1]
f_init = features(th_ref[0], om_ref[0])

samples = np.zeros((N_SAMPLES, STOCH_HORIZON))
for s in range(N_SAMPLES):
    f = f_init.copy()
    for h in range(STOCH_HORIZON):
        f = structured_feature_step(
            f,
            0.0,
            omega_noise=estimated_omega_noise * rng.standard_normal(),
        )
        samples[s, h] = np.arctan2(f[0], f[1])

deterministic_rollout = rollout_oracle_feature(f_init, np.zeros(STOCH_HORIZON))
bands = np.percentile(samples, [5, 25, 50, 75, 95], axis=0)

true_noisy = np.zeros((5, STOCH_HORIZON))
for k in range(5):
    th_s, om_s = th_ref[0], om_ref[0]
    for h in range(STOCH_HORIZON):
        th_s, om_s = pendulum_step(th_s, om_s, 0.0)
        true_noisy[k, h] = th_s

All computation is done, and what follows are three views of the same underlying arrays.