AD-WM: Teaching World Models to Tell Actions Apart, Not Just Predict the Next Frame

Latent world models have quietly become the substrate for a large share of learned robot control: encode the scene, roll the model forward under candidate actions, score the imagined outcomes, execute the best one. The training signal that makes these models work, however, was inherited from generative modeling rather than from control. It rewards getting the next frame right. Planning, by contrast, is a comparative act — it needs the model to say which of several actions is better, and those two demands are not the same.

Background

The standard recipe trains a world model to predict factual transitions: given the current latent state and the action that was actually taken, predict what came next. The flaw is that a model can minimize this error by learning to ignore action detail and emit a plausible generic successor. When the same model is then handed to model predictive control (MPC), it is asked a different question — compare alternative actions from one state — and a model that is accurate but action-insensitive has nothing useful to say. The paper states the tension plainly: low factual prediction error can coexist with poor candidate-action discrimination.

Core Innovation

AD-WM is an action-discriminative joint-embedding world model for counterfactual MPC, built on two additions to a standard latent-dynamics backbone:

• Residual latent dynamics. The model predicts a residual on top of the current latent rather than the full successor state, which keeps the learned dynamics focused on what changes.
• Predictor-level action-recovery regularization. An inverse-dynamics objective plus a normalized recovery objective, motivated by conditional mutual information, is attached at the predictor. If a predicted transition is to be worth anything for planning, it must still contain enough information to recover the action that produced it. Both objectives push planning-relevant transitions to preserve action information.
• Auxiliary heads are discarded at test time. The regularization lives only in training; at deployment the heads are removed and the MPC loop is unchanged. That is an important engineering property — the method buys better imagination without changing the controller that consumes it.

Results

• OGBench-Cube, hard-start success: 3.7% → 52.0% against a matched LeWM baseline — a jump large enough that it cannot be attributed to tuning noise.
• Mean success improves in four of five simulated environments over the reproduced baseline.
• Planning diagnostics. Factual prediction error and whole-bank action ranking do not follow the closed-loop success ordering, whereas CEM-aligned elite regret tracks success more closely. This is the paper’s most transferable finding: the usual proxy metrics are not the metrics that matter.
• Zero-shot real transfer. With a frozen V-JEPA 2 encoder and matched DROID post-training, AD-WM raises basic pick-and-place success on the authors’ Franka setup from 42.2% to 71.1%, with no lab-specific adaptation.

Limitations

The action-recovery objective adds auxiliary heads and at least one trade-off weight between the recovery and dynamics terms; the abstract reports no sensitivity analysis, so how brittle the balance is remains open. Real-robot evidence is a single task — basic pick-and-place — on a single arm, which is far narrower than the simulated sweep. Recovering the action is a proxy for action-discriminability, not a guarantee that every planning-relevant difference is preserved, and the paper offers no formal bound linking recovery quality to counterfactual ranking quality. The model inherits its perception from a frozen V-JEPA 2 encoder, so part of the zero-shot transfer result belongs to that encoder rather than to AD-WM. The key diagnostic, CEM-aligned elite regret, is specific to cross-entropy-method planners, and it is untested on longer horizons. Finally, no inference or planning-latency figures are given, even though MPC is a real-time loop.

Industry Implications

The practical lesson for anyone building a learned controller is that the metric used to train and select world models should be the one that predicts task success. Teams typically track prediction error because it is cheap and dense; this work shows that it can be flat while closed-loop performance moves, and that an elite-regret-style signal tracks success better. That directly changes how a model checkpoint should be chosen. The second lesson is architectural: action-awareness can be trained in and then compiled out, so a control stack does not need new runtime machinery to benefit. For robot vendors, the strongest signal is the zero-shot transfer number — 42.2% to 71.1% on hardware with no lab-specific adaptation is the kind of gap that separates a research demonstration from something deployable. The remaining work is breadth: more tasks, more arms, and latency numbers before world-model MPC becomes a default in production cells.