Chapter 10 — Prediction Loss: How the Model Learns the Next Step

The big picture

A tiny story

A driving student can make many individual turns accurately and still choose the wrong route across town. LeWM’s prediction loss is like a ruler for one local turn. It asks whether the next latent guess is near the shared encoder’s representation of the next recorded observation.

That ruler is essential. Without it, the predictor has no direct reason to follow recorded transitions. But its coordinates are learned—not meters, degrees, or velocities.

The technical backpack

The frozen implementation uses mean squared error:

observed_states = encode(recorded_frames)
guesses = predict(observed_states_without_last, recorded_actions_without_last)
targets = observed_states_without_first       # still connected
L_pred = mean((guesses - targets)^2)

Squaring prevents positive and negative gaps from cancelling; averaging combines batch, time positions, and latent coordinates. If an important physical variable is absent from the representation, MSE has no coordinate with which to complain.

The target is learned. It is not an oracle state, an EMA teacher, or a stopped answer. Prediction gradients reach the shared encoder through context and target uses. This flexibility also admits the constant-code collapse shortcut, which is why SIGReg has a different job.

With the frozen four-observation default, all shifted pairs count:

context/actions:  z0+a0   z1+a1   z2+a2
predictions:      p1      p2      p3
targets:          z1      z2      z3

They are three one-step targets at different context positions—not a direct one-, two-, and three-step rollout from one start. num_preds=1 sets the offset. With frameskip=5, “one model step” spans five environment actions. The paper reports TwoRoom history one, while the frozen global default is history three.

During teacher forcing, every context state comes from a recorded observation. During planning, predicted states are appended and become inputs. The first small displacement can change the next prediction. Error may grow, shrink, rotate, cancel, or saturate; it need not rise monotonically.

The trick that fools us

Build a dataset with 95 straight transitions and five rare doorway turns. A predictor can handle the majority, ignore the turn action, and still report a pleasing average.

Do not throw away MSE. Unpack it:

1. report held-out teacher-forced one-step error by region and prediction slot;
2. replay recorded actions autoregressively and plot error by horizon;
3. compare teacher-forced and autoregressive contexts at the same targets;
4. hold history fixed and swap supported actions;
5. use privileged-state probes only with clear limits;
6. report cost ranking and closed-loop success separately.

Training MSE and terminal planning cost can both contain squared latent gaps yet use different tensors and reductions. One is a mean over shifted transitions; the other is a summed terminal score per candidate. Their raw values are not interchangeable.

Evidence receipt

LeWorldModel v3, frozen train.py, and frozen JEPA.rollout support teacher-forced one-step training and autoregressive planning.

In one single-seed TwoRoom reproduction, one-step accuracy ranked short-horizon behavior across three checkpoints but did not rank long-horizon planning success. That supports measuring the later links separately under that protocol; it does not establish a universal law or identify the cause.

Allowed claim: low connected one-step MSE shows local latent agreement on evaluated transitions. Do not promote it to meaningful state, reachability, stable rollout, effective CEM, or successful control.

Quick check

1. Why is the encoded next observation useful without being fixed ground truth?
2. What changes between teacher-forced and autoregressive context?
3. Which metric would reveal a rare action-dependent failure hidden by the average?