World-action models: planning in a learned latent

The missing question

A vision-language-action policy is a map from observation to action. Nothing in that signature lets you ask what an action does. So a VLA cannot rank two candidate plans, cannot detect that it is about to enter an unrecoverable state, and cannot say "I don't know" — all three require a prediction of the future, and a policy does not have one.

A world-action model learns exactly that prediction. The object of interest is a latent transition model

$$f_\psi : (z_t,\, a_t) \;\longmapsto\; z_{t+1},$$

paired with an encoder $E_\psi : o_t \mapsto z_t$. This is the PlaNet/Dreamer lineage [1, 2, 3] — learn dynamics in a compact latent, then plan or train inside it. Once you have $f_\psi$, the downstream capability is search: an action sequence can be evaluated before it is executed, which is the same principle MuZero applied in discrete domains with a model trained only to support value prediction [4].

Why the latent, and not pixels

The tempting formulation is to predict the next frame. It is the wrong objective, for a reason that is easy to state quantitatively. Pixel reconstruction loss is dominated by the highest-variance content in the image, which for a fixed workspace camera is background texture, lighting, and shadow — not the $7$ dimensions that matter. Optimizing

$$\mathcal{L}_{\mathrm{pix}} = \bigl\lVert \hat{o}_{t+1} - o_{t+1} \bigr\rVert_2^2$$

allocates capacity in proportion to pixel count, and the gripper occupies a small fraction of it. The Dreamer line makes reconstruction work through careful balancing [3], but the decoder-free alternatives — TD-MPC and TD-MPC2, which train a latent purely to support reward and value prediction [5, 6] — reach comparable or better control performance without ever rendering a pixel.

The joint-embedding predictive architecture states the principle directly: predict in a representation trained to be predictable rather than complete [7]. V-JEPA 2 is the strongest current evidence that this scales to embodied planning, with an action-conditioned head used for zero-shot robot manipulation [8]. With an online encoder $E_\psi$ and an EMA target encoder $E_{\bar\psi}$, the latent-dynamics objective over a rollout of length $K$ is

$$\mathcal{L}_{\mathrm{dyn}} = \sum_{k=1}^{K} \Bigl\lVert\, f_\psi^{(k)}(z_t, a_{t:t+k-1}) \;-\; \mathrm{sg}\bigl[E_{\bar\psi}(o_{t+k})\bigr] \,\Bigr\rVert_2^2 ,$$

where $\mathrm{sg}[\cdot]$ is the stop-gradient and $f^{(k)}$ denotes $k$-step composition. The stop-gradient is what makes this well-posed: without it the encoder can minimize the loss by collapsing, $E_\psi(\cdot) \equiv c$, which is perfectly predictable and perfectly useless. BYOL [9] and SimSiam [10] are the canonical analyses of why the EMA-target-plus-stop-gradient pair avoids that fixed point.

Collapse remains the central practical failure. Two properties are worth monitoring directly rather than inferred from loss curves:

• Rank of the latent batch. Compute the singular values of the $B \times d$ matrix of encoded observations; RankMe [11] formalizes effective rank as a label-free quality measure. A healthy encoder keeps effective rank on the order of $d$; collapse shows up as rank falling to single digits while $\mathcal{L}_{\mathrm{dyn}}$ improves.
• Action sensitivity. The Jacobian $\partial z_{t+1} / \partial a_t$ should be non-negligible. A model that has learned to ignore its action input is a video predictor, not a world model, and it will score every plan identically.

The horizon limit

Rollout error compounds. If the one-step latent error is bounded by $\epsilon$ and $f_\psi$ is $L$-Lipschitz in $z$, the $k$-step error obeys

$$\bigl\lVert \hat{z}_{t+k} - z_{t+k} \bigr\rVert \;\le\; \epsilon \sum_{j=0}^{k-1} L^{j} \;=\; \epsilon\,\frac{L^{k} - 1}{L - 1} \quad (L \neq 1).$$

For $L > 1$ this is exponential, which sets a hard usable horizon: given an error tolerance $\delta$, planning is meaningful only for

$$k \;\lesssim\; \frac{\log\bigl(1 + \delta (L-1)/\epsilon\bigr)}{\log L}.$$

The consequence is architectural, not a matter of more training. MBPO reached the same conclusion from the RL side and drew the practical lesson: use the learned model only for short branched rollouts, and let real data carry the rest [12]. On real manipulation data, honest usable horizons are on the order of one to a few seconds of latent rollout. A WAM is a local search tool. Long-horizon behavior has to come from somewhere else — a symbolic or language-level task decomposition that hands the WAM one subgoal at a time, each inside the horizon. Any system claiming minute-long latent planning is either not measuring compounding error or has an $L < 1$ contraction that has smoothed away the dynamics it needed to model.

Composing a WAM with a VLA

The useful arrangement is not replacement. A VLA is a strong proposal distribution — it has absorbed internet-scale semantics and knows what a plausible next motion looks like. A WAM is a scorer. Together they form sampling-based planning where the expensive part is already cheap; TD-MPC2's use of a learned policy prior to seed its MPPI search over a learned model is the same structure at smaller scale [6]:

$$a^{\star}_{t:t+H} = \arg\max_{a^{(i)} \sim \pi_{\mathrm{VLA}}} \; \sum_{k=1}^{H} \gamma^{k}\, R_\psi\bigl(\hat{z}_{t+k}^{(i)}\bigr) \;-\; \lambda\, C\bigl(\hat{z}^{(i)}_{t:t+H}\bigr),$$

over $N$ chunks sampled from the VLA, with $\hat{z}^{(i)}$ rolled out through $f_\psi$, a learned reward or goal-distance $R_\psi$, and a cost $C$ penalizing predicted collision or joint-limit states. Note the cost enters before execution, which is the whole point.

Because the rollouts are in latent space and $f_\psi$ is small relative to the backbone, the scoring pass is cheap: $N = 32$ chunks over $H = 16$ steps is a batched forward pass through a small network, and it fits inside the VLA's own inference budget rather than adding to it. This is the axis on which the latent approach separates from the video-generation approach to world modelling — UniPi [13], UniSim [14], Genie [15], GAIA-1 [16], iVideoGPT [17] and Cosmos [18] all produce far richer and more inspectable rollouts, but generating pixels for $N$ candidate plans per control step is orders of magnitude off the budget.

Three capabilities fall out of the same machinery:

Recovery. When the executed transition disagrees with the prediction, $\lVert z_{t+1} - f_\psi(z_t, a_t) \rVert$ spikes. That residual is a direct, model-native signal that something unmodeled happened — the object slipped, the grasp failed — and it triggers replanning without a hand-written detector.

Calibrated refusal. An ensemble of $M$ dynamics heads gives a disagreement estimate

$$u(z_t, a_t) = \frac{1}{M}\sum_{m=1}^{M} \Bigl\lVert f_{\psi_m}(z_t,a_t) - \bar{f}(z_t,a_t) \Bigr\rVert^2 ,$$

which is high off the training support. This is PETS' epistemic-vs-aleatoric decomposition [19] on top of deep ensembles [20]. Unlike policy entropy, it separates "several actions are fine" from "I have no idea what happens next" — and only the second should stop the robot.

Counterfactual evaluation. Held-out plans can be scored offline, against a model, without occupying hardware. DayDreamer is the existence proof that a learned world model can carry the sample-efficiency load on physical robots rather than in simulation [21], and it is the part that changes how fast an embodied system can be iterated on.

What is still unresolved

Latent rollouts are not verifiable. $\hat{z}_{t+k}$ has no units and no inspectable content, so a plan that scores well cannot be audited by a human the way a predicted trajectory in task space can. Decoding the latent for inspection reintroduces the pixel objective we removed, and a decoder trained only for visualization gives no guarantee that what it shows is what the planner scored.

Reward and cost learning is the second gap, and it has a name: objective mismatch — a model trained for predictive accuracy is not thereby trained to be useful for control [22]. Worse, $R_\psi$ and $C$ are themselves learned from limited labels, and a planner that maximizes over $N$ samples will find their errors. The scaling behaviour of that failure is now well characterized in the reward-model setting [23], and there is no reason to expect embodied reward models to be better behaved.

Both are open. What is not open is the direction: a system that cannot predict consequences cannot be trusted to act among people, and prediction is the capability a policy alone does not have.

References

1. Hafner et al. Learning Latent Dynamics for Planning from Pixels (PlaNet). 2018. arXiv:1811.04551
2. Hafner et al. Mastering Atari with Discrete World Models (DreamerV2). 2020. arXiv:2010.02193
3. Hafner et al. Mastering Diverse Domains through World Models (DreamerV3). 2023. arXiv:2301.04104
4. Schrittwieser et al. Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model (MuZero). 2019. arXiv:1911.08265
5. Hansen et al. Temporal Difference Learning for Model Predictive Control (TD-MPC). 2022. arXiv:2203.04955
6. Hansen et al. TD-MPC2: Scalable, Robust World Models for Continuous Control. 2023. arXiv:2310.16828
7. Assran et al. Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture (I-JEPA). 2023. arXiv:2301.08243
8. Assran et al. V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning. 2025. arXiv:2506.09985
9. Grill et al. Bootstrap Your Own Latent: A New Approach to Self-Supervised Learning. 2020. arXiv:2006.07733
10. Chen & He. Exploring Simple Siamese Representation Learning (SimSiam). 2020. arXiv:2011.10566
11. Garrido et al. RankMe: Assessing the Downstream Performance of Pretrained Self-Supervised Representations by Their Rank. 2022. arXiv:2210.02885
12. Janner et al. When to Trust Your Model: Model-Based Policy Optimization (MBPO). 2019. arXiv:1906.08253
13. Du et al. Learning Universal Policies via Text-Guided Video Generation (UniPi). 2023. arXiv:2302.00111
14. Yang et al. Learning Interactive Real-World Simulators (UniSim). 2023. arXiv:2310.06114
15. Bruce et al. Genie: Generative Interactive Environments. 2024. arXiv:2402.15391
16. Hu et al. GAIA-1: A Generative World Model for Autonomous Driving. 2023. arXiv:2309.17080
17. Wu et al. iVideoGPT: Interactive VideoGPTs are Scalable World Models. 2024. arXiv:2405.15223
18. NVIDIA et al. Cosmos World Foundation Model Platform for Physical AI. 2025. arXiv:2501.03575
19. Chua et al. Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models (PETS). 2018. arXiv:1805.12114
20. Lakshminarayanan et al. Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles. 2016. arXiv:1612.01474
21. Wu et al. DayDreamer: World Models for Physical Robot Learning. 2022. arXiv:2206.14176
22. Lambert et al. Objective Mismatch in Model-Based Reinforcement Learning. 2020. arXiv:2002.04523
23. Gao et al. Scaling Laws for Reward Model Overoptimization. 2022. arXiv:2210.10760