Chapter 4 — Understand LeWM in One Diagram

The big picture

A tiny story

In the daytime control room, LeWM studies recorded journeys. It may look at the real next frame, compare its guess with that answer, and update its encoder and predictor. At night, the trained model helps choose an action. The unexecuted future is blank, so every candidate future must come from the model itself.

Mixing the rooms creates plausible nonsense: a training graph with a goal and CEM, a planner that sees the true future, or a planner that bends model weights until an impossible route looks cheap.

The real rule

The two lanes are short enough to memorize:

TRAIN
recorded frames + actions -> encode -> predict
encoded real next frames -> connected targets
prediction loss + weighted SIGReg -> update the learned components

PLAN
current + goal -> frozen encoder
candidate actions -> autoregressive latent rollouts -> terminal cost
CEM changes the action proposal -> execute K -> observe -> repeat

The first lane has an answer key but no task goal. The second has a task goal but no answer key for an action that has not happened.

Symbolically, images move from [B,T,C,H,W] to latent states [B,T,D]. In the cited experiments images are 224×224, frame skip is five, and sampled subtrajectories contain four observation positions in the frozen default. The v3 appendix reports history one for TwoRoom and three for PushT and OGBench-Cube; the frozen global config uses history_size: 3 and TwoRoom does not override it. Preserve both labels.

During training, real encoded states provide causal context—“teacher forcing” here does not mean a teacher network. The shifted target stays connected to the shared encoder. During planning, predictions are appended and reused. Logging or inference .detach() calls do not turn v3 training into stop-gradient.

SIGReg acts on training embeddings. It makes the exact constant-code shortcut costly, but it is not a planning cost and does not certify physics. The paper method default for its weight is 0.1; the frozen YAML uses 0.09. Neither is a universal optimum.

The trick that fools us

Suppose planning updates both candidate actions and world-model weights. The terminal cost falls quickly. Did the action improve? Maybe not—the optimizer may simply bend the predictor until its endpoint resembles the goal.

The baseline invariant is:

model_parameter_change = 0
candidate_action_change = allowed

CEM refines only the action distribution. The goal latent bypasses dynamics and meets imagined endpoints at terminal squared Euclidean cost. The paper’s Appendix B prose returns the final proposal mean, while Algorithm 2 also allows the best sampled sequence or the first action of that mean. The repository delegates CEM to an unpinned stable-worldmodel dependency, so record the resolved return convention.

MPC plans H model steps, executes K, then observes. The reported setup uses H=5, K=5, and five environment steps per model step: 25 environment steps before feedback. Replanning can limit stale imagination; it cannot repair a missing wall or bad goal geometry.

Evidence receipt

LeWorldModel v3, frozen train.py, and frozen jepa.py support the two-lane method, connected training graph, vectorized latent rollout, and terminal criterion.

Allowed claim: LeWM jointly trains an encoder and action-conditioned predictor with prediction plus SIGReg, then freezes the model while CEM/MPC search and execute actions. Do not claim that CEM finds a global optimum, latent distance equals reachability, or MPC always executes one action.

Quick check

1. Which information exists in training but not in an unexecuted planning branch?
2. What changes during CEM, and what must stay frozen?
3. Why does a goal detach during frozen planning say nothing about target detach during training?