World models that better distinguish actions improve robot control success
AD-WM: Action-Discriminative World Models for Counterfactual Model Predictive Control
Artificial IntelligenceRobotics
Summary
Planning in robots means predicting what happens after different actions, but many models focus on just guessing what will happen, not how actions differ. The authors designed a new model called AD-WM that learns to keep track of the differences between possible actions, helping robots pick better choices. This approach improved success rates a lot when trying hard tasks and transferring skills to new robots. Their work suggests future robot planning models should focus on telling actions apart, not just on prediction accuracy.
What this means in practice
- •For robotics engineers: Improve robot manipulation planning by using models that better differentiate possible actions, increasing task success without extra adaptation.
- •For industrial automation teams: Use action-discriminative world models to enhance model predictive control for precise pick-and-place tasks in robotic arms.$Commercial implications: Enables selling advanced robot control software that improves productivity and reliability in manufacturing automation.
Authors
Jiabin Qiu, Zixuan Chen, Hongye Cao, Jieqi Shi, Jing Huo, Yang Gao
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. 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/.