Chapter 8 — A Complete Forward Pass: Follow One Batch from Start to Finish

The big picture

A tiny story

A factory can produce boxes of the right size with the wrong labels inside. A LeWM batch can do the same. Three predictions and three targets may match in shape even if each prediction is compared with the present instead of the next state, only the final pair is trained, or the target was detached.

This chapter’s simple promise is provenance: we can point to the observation, action block, prediction, and target at every slot.

The technical backpack

Frozen defaults make the whole forward pass concrete: batch B=128, history_size=3, num_preds=1, latent width D=192, and frameskip=5.

pixels                 [B,4,C,H,W]
flattened pixels       [B×4,C,H,W]
observation latents    [B,4,192]
raw action blocks      [B,4,5×A]
encoded actions        [B,4,192]
state/action context   [B,3,192]
connected targets      [B,3,192]
predictions            [B,3,192]
SIGReg input           [4,B,192]
losses                 scalar

The paper reports history one for TwoRoom and three for PushT/OGBench-Cube; the frozen global default is three. These shapes describe the frozen default, not every LeWM.

The shortest faithful forward card is:

z = encode_each_frame(images)              # [B,4,D]
u = encode_each_action_block(actions)       # [B,4,D]
z_pred = causal_predict(z[:,:3], u[:,:3])  # [B,3,D]
z_next = z[:,1:]                            # [B,3,D], do not detach
L_pred = mean_square_gap(z_pred, z_next)
L_sig = sigreg(time_first(z))               # [4,B,D]
L_total = L_pred + lambda * L_sig

All three shifted one-step pairs contribute. num_preds=1 means a one-position target offset, not one final output. A model step spans five environment actions in these data configs.

The prediction branch directly trains the visual encoder, encoder projector, action encoder, causal predictor, and prediction projector. Because z_next stays connected, prediction also returns through the target use of the same visual encoder. SIGReg directly trains only the visual encoder and its projector. One optimizer updates the connected learned system; there is no target-encoder optimizer or EMA update.

SIGReg sees each relative time position across the batch. It must not flatten B×T into one population, because temporal change is not the same as across-example diversity. Frozen code uses 1,024 random projections and 17 frequency knots from 0 through 3. The paper appendix presents quadrature differently. The paper method weight is 0.1; frozen YAML is 0.09. Keep every source label.

The trick that fools us

Make a deliberately wrong branch:

wrong targets: [z0,z1,z2]
right targets: [z1,z2,z3]

Both are [B,3,D]. The wrong branch can learn to copy the present and report a pleasing loss. Use visible time checksums and print prediction_slot / last_visible_context / target_source. Correct rows are 0/0/1, 1/1/2, 2/2/3.

Then test the graph. Detach only the correct target slice in a disposable probe and confirm that the target-side gradient route disappears while the context route remains. Merely finding some encoder gradient is not enough.

Five cheap checks catch most plumbing errors: exact shapes, shift provenance, causal future-mutation, finite/non-null gradients in every expected module, and no NaN or infinity. Passing them proves mechanics—not useful dynamics or planning.

Evidence receipt

LeWorldModel v3, frozen train.py, jepa.py, and module.py support the frozen alignment and connected graph. The paper’s displayed training pseudocode has a malformed F.mse_loss call, so executable train.py is the authority for exact indexing.

Do not call the next target a planning goal, detach it by JEPA habit, mix batch and time in SIGReg, or treat matching shapes as correctness.

Quick check

1. Which three prediction–target pairs come from four observations?
2. Through which two visual-encoder uses can prediction loss flow?
3. Why can a wrong temporal target pass every ordinary shape assertion?