Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Predictive and Generative Objectives › Code Implementation

Michael BrenndoerferJuly 1, 202611 min read

World Models HandbookAll chapters

Part 6 of 12Learning World Models at Scale

Chapter 31Predictive and Generative Objectives › Code Implementation

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

Section 8 of 12 · 11 min read

Predictive and Generative Objectives › Code ImplementationOne 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

Let us walk through the example in code. The book-wide plotting setup and shared imports are loaded in the hidden setup cell at the top of the chapter.

We now generate synthetic trajectories from the controlled system described in the worked example. We hold out separate training and test sets.

In[9]:

Code

A = np.array([[0.9, 0.2], [-0.1, 0.85]])
B = np.array([0.3, 0.5])


def step_true(x, a):
    x_next = A @ x + B * a
    x_next = x_next + 0.05 * np.array([np.sin(x[0]), np.cos(x[1])])
    return x_next


def rollout_true(x0, actions):
    xs = [x0.copy()]
    x = x0.copy()
    for a in actions:
        x = step_true(x, a)
        xs.append(x)
    return np.array(xs)


def make_dataset(n, T, seed):
    r = np.random.default_rng(seed)
    X, U = [], []
    for _ in range(n):
        x0 = r.normal(scale=0.5, size=2)
        actions = r.normal(scale=1.0, size=T)
        traj = rollout_true(x0, actions)
        X.append(traj)
        U.append(actions)
    return np.array(X), np.array(U)


T = 20
X_train, U_train = make_dataset(200, T, seed=1)
X_test, U_test = make_dataset(50, T, seed=2)

Out[10]:

Console

Train trajectories: (200, 21, 2), Test trajectories: (50, 21, 2)

Each training trajectory is a sequence of states $x_0, x_1, \dots, x_T$ together with actions $a_0, \dots, a_{T-1}$ . The example uses the resulting paired transitions directly; observation-based models may first infer latent states from raw inputs.

Next we define the model. A small MLP with two hidden layers is enough for this system.

In[11]:

Code

torch.manual_seed(0)


class TransitionModel(nn.Module):
    def __init__(self, state_dim=2, action_dim=1, hidden=32):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim + action_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, state_dim),
        )

    def forward(self, x, a):
        return self.net(torch.cat([x, a], dim=-1))

The forward pass takes a state and an action and returns a predicted next state. There is no ambiguity about what we are training; the interesting question is the loss.

Now we train the one-step model. We collapse the trajectories into individual transitions and minimize MSE between predicted and true next states.

In[12]:

Code

torch.manual_seed(0)

x_flat = torch.tensor(X_train[:, :-1].reshape(-1, 2), dtype=torch.float32)
a_flat = torch.tensor(U_train.reshape(-1, 1), dtype=torch.float32)
y_flat = torch.tensor(X_train[:, 1:].reshape(-1, 2), dtype=torch.float32)

model_1step = TransitionModel()
opt = torch.optim.Adam(model_1step.parameters(), lr=1e-2)
one_step_losses = []
for step in range(2000):
    pred = model_1step(x_flat, a_flat)
    loss = ((pred - y_flat) ** 2).mean()
    opt.zero_grad()
    loss.backward()
    opt.step()
    one_step_losses.append(loss.item())

Out[13]:

Console

Final one-step loss: 0.000018

We now train the multi-step model. On each training step we sample a batch of trajectories, take the true initial state, and rollout the model for $K = 5$ steps, penalizing the sum of per-step squared errors:

In[14]:

Code

K = 5
X_seq = torch.tensor(X_train, dtype=torch.float32)
U_seq = torch.tensor(U_train, dtype=torch.float32).unsqueeze(-1)

## Match the one-step model's initial weights. The update data and compute budgets
## remain different, as described above.
torch.manual_seed(0)
model_kstep = TransitionModel()
opt = torch.optim.Adam(model_kstep.parameters(), lr=1e-2)
kstep_losses = []
batch_size = 32
for step in range(2000):
    idx = torch.randint(0, X_seq.shape[0], (batch_size,))
    x = X_seq[idx, 0]
    loss = torch.tensor(0.0)
    for k in range(K):
        x = model_kstep(x, U_seq[idx, k])
        loss = loss + ((x - X_seq[idx, k + 1]) ** 2).mean()
    loss = loss / K
    opt.zero_grad()
    loss.backward()
    opt.step()
    kstep_losses.append(loss.item())

Out[15]:

Console

Final 5-step loss: 0.002434

The two training losses are not directly comparable. The one-step value averages errors over all training transitions. The five-step value averages errors across five dependent rollout horizons for a sampled batch of trajectory prefixes; it is not the error at horizon five alone. We compare both models under the same held-out rollout metric below.

We now evaluate both models on held-out rollouts. Starting from the true initial state of each test trajectory, we unfold both models for the full horizon and record the mean squared error at each step.

In[16]:

Code

def eval_rollout(model, X, U, horizon=T):
    X_t = torch.tensor(X, dtype=torch.float32)
    U_t = torch.tensor(U, dtype=torch.float32).unsqueeze(-1)
    x = X_t[:, 0]
    errors = []
    with torch.no_grad():
        for k in range(horizon):
            x = model(x, U_t[:, k])
            mse = ((x - X_t[:, k + 1]) ** 2).mean().item()
            errors.append(mse)
    return np.array(errors)


err_1step = eval_rollout(model_1step, X_test, U_test)
err_kstep = eval_rollout(model_kstep, X_test, U_test)
horizons = np.arange(1, T + 1)

Out[17]:

Console

Horizon-1 error, one-step model:  0.000021
Horizon-1 error, k-step model:    0.000152
Horizon-20 error, one-step model: 0.000238
Horizon-20 error, k-step model:   0.007867
k-step / one-step error ratio at horizon 1: 7.259
k-step / one-step error ratio at horizon 20: 33.053

The comparison that matters is not who wins at horizon 1, but how the two curves separate as the horizon grows. Now we plot it.

Out[18]:

Visualization

Log-scale line chart over horizons 1 through 20. The one-step model stays below the five-step model; errors are about 0.000021 versus 0.000152 at horizon 1 and 0.000238 versus 0.007867 at horizon 20. — Held-out rollout mean squared error for the one-step and five-step training procedures. The one-step model reaches about 0.000238 at horizon 20, while the five-step model reaches about 0.007867, a ratio of 33.05. The comparison includes unequal transition coverage and compute per update, so it should not be read as an objective-only effect.

The one-step model is lower at every plotted horizon in this run. The five-step model's error is about $7.26\times$ larger at horizon 1 and $33.05\times$ larger at horizon 20. Because the training batches and transition coverage differ, this result describes these procedures under this budget; it does not isolate an objective-only effect.

The horizon curves summarize the error, but the drift is geometric, so rolling both models open-loop in state space makes the divergence between the two objectives visible directly.

In[19]:

Code

## Open-loop rollout of both trained models from one held-out initial state (index 0 of X_test).
traj_index = 0
with torch.no_grad():
    x_init = torch.tensor(X_test[traj_index, 0], dtype=torch.float32).reshape(
        1, 2
    )
    u_col = torch.tensor(U_test[traj_index], dtype=torch.float32).reshape(T, 1)

    @torch.no_grad()
    def open_loop(model):
        states = [x_init.squeeze(0).detach().cpu().numpy()]
        state = x_init
        for k in range(T):
            state = model(state, u_col[k : k + 1])
            states.append(state.squeeze(0).detach().cpu().numpy())
        return np.array(states)


rollout_one_step_traj = open_loop(model_1step)
rollout_k_step_traj = open_loop(model_kstep)
true_traj = X_test[traj_index]

Out[20]:

Visualization

Two-dimensional state-space trajectories over 20 steps from the same held-out initial state. The one-step rollout follows the true path more closely; the five-step rollout peels farther away and ends with the larger Euclidean drift. — Open-loop state trajectories in the (x1, x2) plane from the same held-out initial state. Under these training procedures, the one-step model stays closer to the true 20-step path; final Euclidean drift is about 0.0085 for the one-step model and 0.1789 for the five-step model.

We can also track the training losses side by side to see how quickly each objective converges.

Out[21]:

Visualization

Log-scale training-loss curves over 2000 updates. The one-step full-batch curve decreases smoothly; the five-step minibatch curve decreases overall but fluctuates substantially. — Training losses over 2000 updates. The full-batch one-step curve is comparatively smooth, while the minibatch five-step curve trends downward but remains noisy and spans about 2.39 orders of magnitude after the first 100 updates. Both runs complete without numerical divergence. Their absolute values are not directly comparable because the objectives and batch construction differ.

Neither finite-training curve needs to reach zero. The MLP can approximate the nonlinear transition, so we do not claim a nonzero irreducible error in this deterministic example; the remaining loss reflects finite optimization and approximation error. The multi-step objective also measures errors at several dependent horizons, so its numerical scale is not directly comparable with the one-step objective. What we are checking here is optimization stability and held-out rollout behavior, not equality of the two training losses.

Now let us illustrate the balancing problem directly. Suppose we attach a reward head and a reconstruction-style head on the same trunk and inspect the magnitude of the gradients each task produces on shared parameters.

In[22]:

Code

class MultiHeadModel(nn.Module):
    def __init__(self, state_dim=2, action_dim=1, hidden=32):
        super().__init__()
        self.trunk = nn.Sequential(
            nn.Linear(state_dim + action_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
        )
        self.next_state_head = nn.Linear(hidden, state_dim)
        self.reward_head = nn.Linear(hidden, 1)

    def forward(self, x, a):
        h = self.trunk(torch.cat([x, a], dim=-1))
        return self.next_state_head(h), self.reward_head(h)


torch.manual_seed(1)
bal_model = MultiHeadModel()

In[23]:

Code

## Deterministic proxy targets that mimic the real training loop: next-state is a 2D vector, reward is a scalar. Derived from x_b and a_b below so the diagnostic is reproducible.
x_b = x_flat[:256]
a_b = a_flat[:256]
y_state = y_flat[:256]
reward_scale = 0.1
y_reward = torch.sin(x_b[:, :1]) + reward_scale * a_b


def grad_norm_in_trunk(loss, model):
    model.zero_grad()
    loss.backward(retain_graph=True)
    total = 0.0
    for p in model.trunk.parameters():
        if p.grad is not None:
            total += p.grad.pow(2).sum().item()
    return total**0.5


pred_state, pred_reward = bal_model(x_b, a_b)
loss_state = ((pred_state - y_state) ** 2).mean()
loss_reward = ((pred_reward - y_reward) ** 2).mean()

gn_state = grad_norm_in_trunk(loss_state, bal_model)
gn_reward = grad_norm_in_trunk(loss_reward, bal_model)

Out[24]:

Console

Gradient norm in trunk from state loss:  0.3618
Gradient norm in trunk from reward loss: 0.4305
Ratio (state / reward): 0.84

The ratio is 0.84, so the reward loss produces the larger shared-trunk gradient in this batch: about 0.4305 versus 0.3618 from the state loss. The values depend on initialization, data, reduction, and target scaling. This diagnostic supports measuring shared gradients, but not a claim that reconstruction dominates here. Detaching the reward input in the next cell demonstrates gradient routing; it is not proposed as a fix for these particular numbers.

The two norms are close enough that neither term dominates by an order of magnitude. A useful diagnosis would repeat this measurement over batches and training steps rather than infer a general weighting rule from one sample.

In[25]:

Code

## Trunk gradient norms measured by the diagnostic above; deterministic given the fixed seed and batch.
head_names = ["State head", "Reward head"]
trunk_grad_norms = np.array([gn_state, gn_reward])
dominance_ratio = float(gn_state / gn_reward)

Out[26]:

Visualization

Two-bar chart of shared-trunk gradient norms: 0.3618 from the state loss and 0.4305 from the reward loss, making the state-to-reward ratio 0.84. — L2 gradient norms on the shared trunk for one batch. The reward loss produces a norm of about 0.4305 and the state loss about 0.3618, giving a state-to-reward ratio of 0.84. This single-batch diagnostic shows why gradient influence must be measured rather than inferred from output dimensionality.

In[27]:

Code

## Detach reward head from trunk to isolate the reward gradient
bal_model.zero_grad()
h = bal_model.trunk(torch.cat([x_b, a_b], dim=-1))
pred_state = bal_model.next_state_head(h)
pred_reward = bal_model.reward_head(h.detach())
## The reward loss now only updates the reward head; its gradient does not flow into the shared trunk.
loss_reward_detached = ((pred_reward - y_reward) ** 2).mean()
loss_reward_detached.backward()

In a full world model, each loss also updates head-specific parameters that the other losses do not see. Those private gradients do not tell you which loss controls the shared state. The practical rule is to inspect and deliberately configure gradient flow on shared modules instead of reasoning from scalar losses or head sizes alone.