AD-WM: Action-Discriminative World Models for Counterfactual Model Predictive Control
AD-WM improves counterfactual planning by preserving action information, raising hard-start success from 3.7% to 52%.
AD-WM is an action-discriminative joint-embedding world model for counterfactual model predictive control. It pairs residual latent dynamics with action-recovery regularization from inverse dynamics and a normalized recovery objective; auxiliary heads are dropped at test time. On OGBench-Cube, hard-start success rises from 3.7% to 52.0% versus a matched LeWM baseline, with mean gains in four of five simulation environments. With a frozen V-JEPA 2 encoder and matched DROID post-training, zero-shot Franka pick-and-place success increases from 42.2% to 71.1%.
- Hard-start OGBench-Cube success rises from 3.7% to 52.0% versus LeWM.
- Mean success improves over the baseline in four of five simulated environments.
- CEM-aligned elite regret tracks closed-loop success better than factual prediction error.
- Frozen V-JEPA 2 and DROID post-training lift Franka pick-and-place from 42.2% to 71.1%.
Full article206 words · extracted from arxiv.org · click to collapse
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/.
Text extracted automatically; images, tables and formatting may be missing. Original: https://arxiv.org/abs/2609.30264