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

Summary

The paper argues that JEPA-style latent world models trained purely on factual transition prediction can achieve low prediction error while failing to distinguish candidate actions during model predictive control (MPC) — a degenerate predictor that copies the current latent incurs only the expected latent-increment error yet scores all action sequences identically. AD-WM addresses this with two additions: residual latent prediction (the predictor estimates latent increments relative to the current state) and predictor-level action recovery, implemented as an inverse-dynamics loss plus a normalized Gaussian recovery objective motivated by a conditional mutual information (Barber–Agakov) bound; auxiliary heads are discarded at test time, leaving the CEM planner unchanged. On OGBench-Cube, AD-WM raises hard-start success from 3.7% to 52.0% over a matched LeWM reproduction and improves mean success in four of five simulation environments (PushT drops from 94% to 92%). Planning diagnostics show elite regret tracks closed-loop success (correlations 0.863/0.810) far better than factual MSE or whole-bank ranking. With a frozen V-JEPA 2 encoder and matched DROID post-training, zero-shot Franka pick-and-place improves from 42.2% to 71.1% over V-JEPA 2-AC. Limitations: the full method trails its own Res+MI ablation (54.7%), shipped defaults are not optimal, INTACT Direct is stronger in-distribution, the Scene gain is unresolved (p=0.13), and robot evaluation lacks statistical testing.