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

Jiabin Qiu, Zixuan Chen, Hongye Cao, Jieqi Shi, Jing Huo, Yang Gao · N/A · 2026

Summary

Abstract Summary

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/.

Key Points

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.

Method

Fig. 2 illustrates the training and planning pipeline of AD-WM . During training, residual latent prediction and action-recovery objectives jointly shape action-discriminative transitions. At test time, MPC uses the learned dynamics to compare candidate action sequences, with the auxiliary heads discarded. We first formulate the task, then detail the model, training objectives, and planning procedure.

Finally, we introduce diagnostics for evaluating candidate selection.

Method — AD-WM: Action-Discriminative World Models for Counterfactual Model Predictive Control — Fig. 3 : The simulation environments.

Experiments

How well does AD-WM perform across simulation tasks and configuration shifts? Q2: How do residual prediction and action recovery contribute to control performance? Q3: Which prediction and selection metrics best reflect closed-loop success? Q4: Can AD-WM transfer to a real robot without lab-specific adaptation? Benchmarks and Baselines.

We follow LeWM’s simulation environment selection [ 11 ] , evaluating on Cube, Reacher, TwoRoom, and PushT, and additionally include Scene from OGBench [ 32 ] . As shown in Fig.

Experiments — AD-WM: Action-Discriminative World Models for Counterfactual Model Predictive Control — Fig. 5 : Real-robot evaluation. (a) Franka setup with a single camera. (b) Basic pick-and-place, complex-object, and specified-target protocols use 45, 10, and 27 trials per model. (c) Matched examples illustrate grasping and target selection. Both models share image goals, execution stages, gripper handling, and action clipping.

Sources and Demonstrations

Figures are reproduced from Qiu et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.