# AD-WM: Action-Discriminative World Models for Counterfactual Model Predictive Control

arXiv: https://arxiv.org/abs/2609.30264 · HTML: https://mlsanity.com/paper/2609.30264/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Latent predictive world models, one of 11 directions in World models.
- Builds on 29 of its 34 references that are in the map.
- First posted on arXiv 24 Sep 2026.

## Placed in

- World models / [Latent predictive world models](https://mlsanity.com/wm/latent-predictive-world-models/)

## Builds on (most cited first, 29 of 29)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Curiosity-driven Exploration by Self-supervised Prediction](https://mlsanity.com/paper/1705.05363/) (2017)
- [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2301.08243/) (2023)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [DROID: A Large-Scale In-The-Wild Robot Manipulation Dataset](https://mlsanity.com/paper/2403.12945/) (2024)
- [V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning](https://mlsanity.com/paper/2506.09985/) (2025)
- [Revisiting Feature Prediction for Learning Visual Representations from Video](https://mlsanity.com/paper/2404.08471/) (2024)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [DINO-WM: World Models on Pre-trained Visual Features enable Zero-shot Planning](https://mlsanity.com/paper/2411.04983/) (2024)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Diffusion for World Modeling: Visual Details Matter in Atari](https://mlsanity.com/paper/2405.12399/) (2024)
- [Deep Visual Foresight for Planning Robot Motion](https://mlsanity.com/paper/1610.00696/) (2016)
- [Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images](https://mlsanity.com/paper/1506.07365/) (2015)
- [LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels](https://mlsanity.com/paper/2603.19312/) (2026)
- [Learning to Poke by Poking: Experiential Learning of Intuitive Physics](https://mlsanity.com/paper/1606.07419/) (2016)
- [DreamDojo: A Generalist Robot World Model from Large-Scale Human Videos](https://mlsanity.com/paper/2602.06949/) (2026)
- [Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models](https://mlsanity.com/paper/2502.14819/) (2025)
- [Learning Markov State Abstractions for Deep Reinforcement Learning](https://mlsanity.com/paper/2106.04379/) (2021)
- [Action-Sufficient State Representation Learning for Control with Structural Constraints](https://mlsanity.com/paper/2110.05721/) (2021)
- [Joint Embedding Predictive Architectures Focus on Slow Features](https://mlsanity.com/paper/2211.10831/) (2022)
- [Sub-JEPA: Subspace Gaussian Regularization for Stable End-to-End World Models](https://mlsanity.com/paper/2605.09241/) (2026)
- [Fast LeWorldModel](https://mlsanity.com/paper/2606.26217/) (2026)
- [Delta-JEPA: Learning Action-Sensitive World Models via Latent Difference Decoding](https://mlsanity.com/paper/2606.31232/) (2026)
- [INTACT: Isomorphic Intent-to-Action Learning for Search-Free World Models](https://mlsanity.com/paper/2607.26056/) (2026)
- [Qantara: Bridge-Flow Training for Multi-Paradigm JEPA Control](https://mlsanity.com/paper/2607.04978/) (2026)

## Abstract (arXiv)

Latent world models are typically trained to predict factual transitions, whereas model predictive control (MPC) must compare alternative actions from the same state. A model can therefore achieve low factual prediction error yet poorly distinguish candidate actions. We introduce AD-WM, an action-discriminative joint-embedding world model for counterfactual MPC. AD-WM combines residual latent dynamics with predictor-level action-recovery regularization, using inverse dynamics and a normalized recovery objective motivated by conditional mutual information. Both objectives encourage planning transitions to preserve action information; their auxiliary heads are discarded at test time, leaving MPC unchanged. On OGBench-Cube, AD-WM improves hard-start success from 3.7% to 52.0% over a matched LeWM baseline and improves mean success over the reproduced baseline in four of five simulation environments. Planning diagnostics show that 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. With a frozen V-JEPA 2 encoder and matched DROID post-training, AD-WM also improves zero-shot transfer to our Franka setup, increasing basic pick-and-place success from 42.2% to 71.1% without lab-specific adaptation. These results suggest that world models for planning should preserve action-dependent differences needed for counterfactual selection, rather than optimize factual prediction accuracy alone. More videos and code are available at https://ad-wm.github.io/.
