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

Forum topic · 小凯 · 2026-09-28

Summary

AD-WM is an action-discriminative joint-embedding world model designed for counterfactual model predictive control (MPC), proposed by Jiabin Qiu et al. (arXiv:2609.30264). Latent world models are usually trained to predict factual transitions, but MPC must compare alternative actions from the same state, so low factual prediction error does not guarantee good action discrimination. 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. The auxiliary heads are discarded at test time, leaving the MPC pipeline unchanged. On OGBench-Cube, AD-WM raises hard-start success from 3.7% to 52.0% over a matched LeWM baseline and beats reproduced baselines on four of five simulated environments. Diagnostics show factual prediction error and full-corpus action ranking do not track closed-loop success, while CEM-aligned elite regret correlates better. With a frozen V-JEPA 2 encoder and identical DROID post-training, AD-WM improves zero-shot transfer to a real Franka platform, lifting base pick-and-place success from 42.2% to 71.1% without lab-specific adaptation.

Paper Overview

Field: Machine Learning
Authors: Jiabin Qiu, Zixuan Chen, Hongye Cao, Jieqi Shi, Jing Huo, Yang Gao
arXiv: 2609.30264
Project page: https://ad-wm.github.io/

Abstract

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. The authors introduce AD-WM, an action-discriminative joint-embedding world model for counterfactual MPC.

Method

AD-WM combines two components:

Residual latent dynamics for state-transition prediction.
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. The auxiliary heads are discarded at test time, leaving the standard MPC procedure unchanged.

Key Results

OGBench-Cube: AD-WM improves hard-start success rate from 3.7% to 52.0% compared to a matched LeWM baseline, and exceeds reproduced baselines' average success rate on four of five simulated environments.
Planning diagnostics: Factual prediction error and full-corpus action ranking do not follow the ordering of closed-loop success rates; CEM-aligned elite regret correlates more closely with success.
Real-world transfer: With a frozen V-JEPA 2 encoder and the same DROID post-training, AD-WM improves zero-shot transfer to a Franka robot platform, raising base pick-and-place success from 42.2% to 71.1% with no lab-specific adaptation.

Takeaway

Planning-oriented world models should retain action-relevant discriminative information to support counterfactual action selection, rather than optimizing solely for factual prediction accuracy.