A new research paper introduces AD-WM, an action-discriminative joint-embedding world model designed for counterfactual Model Predictive Control (MPC). The AD-WM model addresses a limitation in traditional latent world models, which often prioritize accurate factual predictions but struggle to differentiate between various potential actions from the same state. This can lead to poor performance in control tasks where comparing alternative actions is crucial.
Model Predictive Control is an advanced control method that optimizes control actions for a dynamic system over a finite time horizon, while adhering to specified constraints. MPC relies on dynamic models to predict future system behavior and calculate optimal inputs. At each step, the controller applies the first calculated action and then repeats the optimization process with an updated system state.
AD-WM integrates residual latent dynamics with a predictor-level action-recovery regularization. This regularization uses two primary objectives: inverse dynamics and a normalized recovery objective, both informed by conditional mutual information. These objectives encourage the planning transitions to retain information about the actions taken, which is critical for counterfactual selection. During deployment, the auxiliary heads used for these objectives are discarded, leaving the MPC planner unchanged.
The researchers evaluated AD-WM on the OGBench-Cube benchmark, a suite of tasks designed to assess dexterous cube handling, spatial planning, and sim-to-real transfer in robotics. OGBench Cube tasks involve pick-and-place manipulation of colored cube blocks, requiring object manipulation, sequential generalization, and combinatorial generalization. On the hard-start success metric within OGBench-Cube, AD-WM improved performance from 3.7% to 52.0% when compared to a matched LeWM baseline.
The model also showed improved mean success rates over the reproduced baseline in four out of five simulation environments.
Further testing involved a real Franka robot setup, where AD-WM demonstrated improved zero-shot transfer capabilities. With a frozen V-JEPA 2 encoder and DROID post-training, the model increased basic pick-and-place success from 42.2% to 71.1% without requiring lab-specific adaptations.
Analysis of planning diagnostics revealed that factual prediction error and the overall ranking of actions did not align with the closed-loop success ordering. However, CEM-aligned elite regret, which measures the difference between the best possible action and the chosen action, more closely correlated with success. These findings suggest that for planning tasks, world models should prioritize preserving action-dependent differences necessary for counterfactual selection over solely optimizing factual prediction accuracy.