Chapter 20 — Implement a Minimal LeWM Planner from Scratch

The big picture

A tiny story

Before building a toy railway, label every connector. Which wire carries a real observation? Which carries an imagined latent? Which action reaches the motor?

A minimal LeWM planner should do the same. “From scratch” here means building clear interfaces around a trained, frozen checkpoint. It does not mean inventing a new LeWM implementation or pretending tutorial pseudocode is the authors’ exact runtime.

The real rule

Keep these objects separate:

current images   [B, T_context, C, H_image, W_image]
current latents  [B, T_context, D]
goal latent      [B, D]
candidate actions[B, N, H_plan, A_block]
imagined latents [B, N, T_rollout, D]
candidate costs  [B, N]
execution prefix [B, K, A_block]

The separately encoded goal enters the terminal cost, not the dynamics predictor. Correct vectorization may flatten (B\times N) for throughput, but it must preserve each environment, candidate, and goal identity. Test a tiny loop against the vectorized path; permuting candidates should only permute the matching outputs.

Baseline LeWM scores the declared final predicted latent with squared Euclidean distance to the goal. CEM refits an action proposal while the model stays frozen. Record both the best sampled sequence and the final proposal mean. Appendix B and the inspected stable-worldmodel snapshots return the final mean; changing execution to best_seen is a labelled variant.

For the reported setup, (H=5) model blocks, (K=5), and each block contains five environment actions. Only the selected prefix crosses into reality. Predicted latents never become fake measurements in the real observation buffer.

The trick that can fool us

Choose a TwoRoom goal close across a wall. Run the same frozen task with a fixed candidate and real-action budget under four named cells:

• learned dynamics + learned latent cost;
• oracle dynamics + learned cost;
• learned dynamics + oracle cost, only if imagined states have a valid scoring bridge;
• both oracles, with the same finite search.

Add an equal-budget random policy and, separately, a support-constrained proposal. These are different questions. Oracle dynamics tests rollout replacement. Oracle cost tests ranking replacement. Both oracles still do not make CEM optimal. A support constraint changes where search looks.

If oracle dynamics repairs the episode, rollout is implicated. If oracle cost repairs it, ranking is implicated. If both fail, inspect candidate coverage, action representation, reachability, and execution. One final frame cannot identify a cause.

Experiment receipt

At every MPC cycle, save current and goal identities, preprocessing, checkpoint, (H), (K), action-block semantics, all seeds, sampled and elite costs, proposal spread, support warnings, selected normalized actions, raw environment actions, dashed imagined latents, solid returned observations, termination, and timing. A fixed projection may visualize the two paths, but native latent distances remain the measurement; simulator coordinates belong in a separate privileged panel.

The source boundary is LeWorldModel v3, frozen jepa.py and eval.py, commit 8edfeb3. The core did not pin its historical solver dependency, although both inspected solver revisions return the final mean. Oracles and the failure dossier are tutorial diagnostics, not baseline features or deployable pixels-only results.

Quick check

1. Why must the goal latent stay outside baseline dynamics rollout?
2. Which assertion can reveal a vectorization identity mix-up?
3. What different bottlenecks do oracle dynamics and oracle cost test?