Chapter 7 — The Dynamics Predictor: Moving Time Forward in the Mind

The big picture

A tiny story

Place the agent below a doorway. The current image alone does not choose tomorrow. “Left,” “right,” and “stay” can lead to different successors. The encoder supplies a compact noun—the situation. The dynamics predictor supplies a verb—the action-conditioned change.

History helps when one frame hides motion, but it cannot reveal every hidden force, contact mode, or occluded object. The output is a deterministic latent estimate, not a promise that the real future is single-valued.

The technical backpack

The conceptual interface is:

state_history  = encode(observed_frames)     # [B,T,D]
action_history = encode(action_blocks)       # [B,T,D]
next_guesses   = causal_predict(state_history, action_history)

In frozen data configs, frameskip=5: one model step contains five consecutive environment actions. The dense block is flattened and mapped to width D. The action embedder uses a width-one temporal convolution plus a small multilayer map, so neighboring time positions are not mixed there; temporal mixing happens in the causal Transformer.

The paper describes a six-layer predictor with 10% dropout. Actions enter each conditional Transformer block through AdaLN. The action pathway produces shift, scale, and residual-gate controls for both attention and feed-forward branches—six controls in total. Its final modulation layer starts at zero, so action modulation is neutral at initialization and can grow during learning.

AdaLN installs a control cable; it does not prove the trained model listens. A causal mask similarly blocks peeking at later tokens but does not prove causal physics. A future-mutation check should show that changing a later input leaves earlier outputs unchanged.

With the frozen four-observation default, training aligns all three pairs:

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

This is teacher forcing: context comes from recorded encoded observations. It is not a teacher encoder. During planning, the rollout appends p1, then predicts from a history that increasingly contains its own outputs. The frozen code retains only the most recent configured history window.

The trick that fools us

Mute the action channel by replacing every action block with zero while leaving data, model size, and budget matched. A behavior policy that usually takes the same route may still let the muted model predict an average continuation well.

Now choose a supported doorway context where left and right have different recorded successors. Hold the visual history fixed and swap actions. Compare predicted separation with the separation between encoded real successors. Add a free-space control where two tiny actions genuinely behave similarly.

Three metrics must stay separate: teacher-forced one-step error, action-swap sensitivity, and open-loop error by horizon. An action can affect step one yet vanish later; low one-step error can coexist with drift; a good rollout can still be ranked by a bad planner cost.

Evidence receipt

LeWorldModel v3 and frozen module.py support the causal attention and AdaLN path. Frozen jepa.py supports autoregressive append-and-truncate rollout.

Allowed claim: v3 provides an action-conditioning route, causal training context, and autoregressive planning. Do not say AdaLN proves controllability, teacher forcing implies an EMA teacher, every LeWM step means five raw actions, or low one-step error guarantees planning.

Quick check

1. What does AdaLN make possible, and what does it not prove?
2. Why is teacher forcing not a teacher–student architecture here?
3. Which three diagnostics separate local prediction, action use, and rollout stability?