PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jiabin Qiu, Zixuan Chen, Hongye Cao, Jieqi Shi, Jing Huo, Yang Gao

arXiv:2609.30264v1cs.AIcs.RO

TL;DR

Factual prediction accuracy does not ensure that latent world models preserve the action differences MPC needs to compare alternatives from the same state. AD-WM combines residual prediction with predictor-level action recovery, and improves success over matched baselines across simulation and robot evaluations. Its Cube diagnostics find elite regret more closely associated with closed-loop success than factual error or whole-bank ranking.

  • Problem

    Factual prediction can fit observed transitions without preserving the action-dependent differences needed to compare candidate actions for MPC.

    Source: §I · p. 1

  • Method

    AD-WM combines residual latent prediction with inverse and normalized predictor-level action recovery to preserve action information in predicted transitions.

    Source: §I · p. 1, C. Predictor-Level Action Recovery · p. 3

  • Results

    AD-WM improves mean success over matched LeWM in four of five simulation environments, and Cube elite regret is more closely associated with closed-loop success than factual error or whole-bank ranking.

    Source: §I · p. 2, §V · p. 8

  • Takeaways & Limitations

    In the evaluated Cube setting, elite regret tracks closed-loop success more closely than factual prediction error or whole-bank action ranking.

    Source: §I · p. 2, §V · p. 8

    Robot evaluation used manual subgoals and non-randomized trials at one site with one camera and backbone.

    Source: §V · p. 8

Abstract

from arXiv · show

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

I. INTRODUCTION

The introduction identifies a mismatch between factual prediction accuracy and the action comparisons required by MPC. It presents AD-WM as a way to preserve action-dependent differences and evaluates its planning performance across simulation and robot settings.

I. INTRODUCTION

  • Low factual prediction error can mask the action-dependent changes needed to distinguish candidate actions and guide goal-directed selection.Persistent scene content can dominate representations while smaller changes, such as approaching versus moving away from a cube, are obscured.Source: §I · p. 1, Figure 1 · p. 1
  • AD-WM combines residual latent prediction with predictor-level action recovery, using inverse dynamics and normalized recovery motivated by conditional mutual information.The auxiliary heads are discarded at deployment, leaving MPC unchanged.Source: §I · pp. 1–2
  • AD-WM improves mean success over matched LeWM in four of five simulation environments and transfers to a real robot without lab-specific adaptation.The evaluation spans five simulation environments and a real Franka robot.Source: §I · p. 2
  • In controlled Cube experiments, elite regret is more closely associated with closed-loop success than factual prediction error or whole-bank action ranking.The diagnostics evaluate selected elites using environment-realized outcomes on shared candidate sequences.Source: §I · p. 2

II. RELATED WORK

Related work uses learned latent predictions for planning, imagined rollouts, and visual foresight, including feature-predictive JEPA approaches. AD-WM focuses on preserving action-dependent differences in predicted transitions for reward-free, image-goal MPC.

II. RELATED WORK

  • World models use predicted dynamics for latent-space planning, imagined-rollout behavior learning, and visual foresight; this work studies reward-free, image-goal MPC from offline data.In this setting, predicted latent transitions directly determine how candidate action sequences are compared.Source: §II · p. 2
  • JEPA-based planners predict future representations without reconstructing observations, and LeWM is AD-WM’s closest end-to-end simulation baseline.V-JEPA 2-AC provides pretrained dynamics for the real-robot study.Source: §II · p. 2
  • Fast-LeWM, Sub-JEPA, INTACT, Qantara, and Delta-JEPA address rollout efficiency, representation stability, control interfaces, or action sensitivity.AD-WM instead regularizes predicted transitions for counterfactual comparison without modifying the MPC planner.Source: §II · p. 2
  • Prior action-aware representation work motivates retaining action-relevant information, while AD-WM applies inverse and normalized recovery directly to model-generated transitions.Their gradients constrain the transitions MPC rolls out and compares, complementing factual prediction supervision.Source: §II · p. 2

III. METHOD

  • During training, residual prediction and action-recovery objectives shape transitions; at test time, MPC compares candidate sequences after auxiliary heads are discarded.The method section proceeds through task formulation, model and training objectives, planning, and candidate-selection diagnostics.Source: §III · p. 2

A. Problem Formulation

The task is to learn latent dynamics from offline image-action transitions so MPC can select action sequences that reach an encoded visual goal. The formulation emphasizes that low factual error alone may not preserve the action differences needed for that selection.

A. Problem Formulation

  • From offline image-action transitions, the agent encodes the current and goal images and scores candidate action sequences by predicted terminal distance to the goal.MPC executes the first action block and replans without parameter updates.Source: A. Problem Formulation · p. 2
  • An action-independent predictor can achieve low expected error on small latent increments while assigning every candidate sequence the same terminal cost.This motivates dynamics that preserve action-dependent differences for planning.Source: A. Problem Formulation · p. 2

B. Residual Latent Prediction

  • Because successive visual representations share substantial state information, the current representation carries that shared information while the predictor estimates the remaining change.This is the motivation for residual latent prediction.Source: B. Residual Latent Prediction · p. 2
  • The pipeline combines residual prediction with Inv and MI recovery during training, then uses CEM terminal-cost rollouts, executes the first action block, and replans.The auxiliary heads are discarded at deployment.Source: Figure 2 · p. 3
  • Given a learned action embedding, AD-WM predicts the remaining latent change rather than modeling the successor as an unrelated absolute representation.The passage introduces the action-conditioned prediction after specifying that the predictor estimates the remaining change.Source: B. Residual Latent Prediction · pp. 2–3
  • Residual predictions are supervised by the encoded successor representation.The residual formulation equivalently matches predicted latent displacement to encoded successor displacement.Source: B. Residual Latent Prediction · p. 3
  • Residual and absolute prediction share an unconstrained optimum, but the residual parameterization changes learning bias toward local latent changes.Residual prediction alone does not require those changes to retain action information; predictor-level recovery addresses that gap.Source: B. Residual Latent Prediction · p. 3

C. Predictor-Level Action Recovery

Predictor-level action recovery complements successor prediction by encouraging actions to remain recoverable from model-generated transitions. AD-WM uses inverse dynamics and a normalized Gaussian recovery objective motivated by conditional mutual information.

C. Predictor-Level Action Recovery

  • Recovery from the current and predicted next representations directly shapes the dynamics MPC uses, while prediction supervision anchors transitions to observed outcomes.This discourages the predictor from ignoring action information.Source: C. Predictor-Level Action Recovery · p. 3
  • The inverse head recovers the action embedding from predicted transition endpoints, and its gradients train the predictor and action encoder through the predicted endpoint.The recovery target is detached while the predicted endpoint remains differentiable.Source: C. Predictor-Level Action Recovery · p. 3
  • Normalized recovery standardizes action embeddings and uses a Gaussian recovery objective with fixed unit covariance and a KL penalty toward a standard Gaussian.Its recovery log-likelihood is the variational term in a lower bound on conditional mutual information between normalized actions and predicted transitions given the current representation.Source: C. Predictor-Level Action Recovery · p. 3

D. Training and Planning

Simulation training combines factual prediction supervision with auxiliary recovery objectives that encourage predicted transitions to retain action information. At deployment, the auxiliary heads are removed and the encoder and CEM procedure remain unchanged.

D. Training and Planning

  • In simulation, the encoder and predictor are trained jointly.Their objective combines prediction loss with SIGReg, inverse-dynamics, and normalized recovery losses.Source: D. Training and Planning · p. 4
  • Prediction supervision matches observed outcomes, while inverse and normalized recovery objectives encourage predicted transitions to retain action information.The representation regularizer is the same SIGReg used by LeWM; robot post-training freezes the encoder and omits SIGReg.Source: D. Training and Planning · p. 4
  • At deployment, MPC embeds candidate actions, recursively predicts transitions, and uses CEM to retain low-cost elites and refit its sampling distribution.The agent executes the first action block, then replans from the new observation; the encoder architecture and CEM procedure remain unchanged.Source: D. Training and Planning · p. 4

E. Planning-Facing Diagnostics

The diagnostics assess whether predicted transitions support useful candidate selection by comparing predicted costs with environment-realized outcomes on a shared action-sequence bank. They distinguish whole-bank ranking from the quality of CEM-selected elites.

E. Planning-Facing Diagnostics

  • Realized costs come from executing each candidate sequence from the same initial state and encoding the final image; they are diagnostic-only, outside training and MPC.Each model uses its own encoder for predicted and realized costs, so the metrics assess alignment within its planning space.Source: E. Planning-Facing Diagnostics · p. 4
  • Counterfactual action discriminability (CAD) measures Spearman rank agreement across the full candidate bank, including candidates outside the elite set.The shared bank contains candidate action sequences with predicted and realized terminal costs.Source: E. Planning-Facing Diagnostics · p. 4
  • Normalized best-in-elite regret measures the best realized cost among predicted elites against the realized cost range across the bank.Elite-mean regret compares the mean realized cost of predicted elites with that of the realized top-k set; both regrets normalize selection penalties by the realized cost range.Source: E. Planning-Facing Diagnostics · p. 4

IV. EXPERIMENTS

The experiments address cross-task performance, component contributions, planning diagnostics, and transfer to a real robot. They evaluate AD-WM across selected simulation benchmarks using controlled comparisons and task-specific protocols.

IV. EXPERIMENTS

  • The experiments ask how AD-WM performs across tasks and configuration shifts, how residual prediction and action recovery contribute, which planning metrics track success, and whether it transfers to a real robot.Source: §IV · p. 4

A. Experimental Setup

  • Evaluation covers Cube, Reacher, TwoRoom, PushT, and OGBench Scene, with Cube serving as the primary benchmark for controlled ablations and planning diagnostics.The environments span manipulation, reaching, navigation, and planar pushing.Source: A. Experimental Setup · p. 4
  • Controlled ablations hold images, encoder size, training budget, and MPC settings fixed, while varying transition parameterization and auxiliary training choices.The matched baseline is a LeWM reproduction; external Cube methods retain native inference interfaces, so their training and inference are not controlled against the matched pair.Source: A. Experimental Setup · pp. 4–5
  • Cube evaluation includes dataset-state Original starts and hard-start protocols P00–P04; hard-start success averages the five protocols, each evaluated over 50 episodes per checkpoint.P00 differs from Original despite having no perturbation; protocols use three-seed means and population standard deviations unless stated otherwise.Source: A. Experimental Setup · p. 4
  • The Cube success table reports matched-start results, using one checkpoint for Fast-LeWM and Sub-JEPA and three-seed mean±population SD for other methods.Source: Table I · p. 5
  • Reported measures include success, one-step and five-step factual latent MSE, whole-bank action-ranking agreement, and normalized best-in-elite and elite-mean regret for 30 candidates.Spearman rank correlation is used for rank comparisons.Source: A. Experimental Setup · p. 5

B. Simulation Results

AD-WM outperforms matched LeWM across Cube protocols and improves mean success over reproduced LeWM in four of five simulation environments. Hard-start and configuration-shift results qualify this pattern.

Main results

  • On Cube, AD-WM outperforms matched LeWM across all protocols and reaches 25.3% on P04, versus 10.0% for INTACT Direct and 10.7% for Actor-CEM.The authors note that differing training and inference prevent attributing the P04 crossover to search strategy alone.Source: B. Simulation Results · p. 5

Configuration shift

  • For P02–P04, 2.0%/38.7%/72.0% exceed the joint threshold, while 52.7%/80.7%/95.3% exceed the relative threshold.The passages describe P02–P04 as configuration-tail probes within empirical cube xy support.Source: B. Simulation Results · p. 5

Cross-environment performance

  • Mean success rises over reproduced LeWM in four of five environments: Cube 73.3%→90.7%, Reacher 76.7%→83.3%, TwoRoom 90%→98%, and Scene 35.5%→39.5%; PushT falls from 94% to 92%.Scene uses balanced hard starts; the other environments use original protocols.Source: B. Simulation Results · p. 5, Figure 4 · p. 6
  • The overall Cube gain remains unresolved by the paired seed test (p = 0.13).Source: B. Simulation Results · p. 6

C. Ablation Studies

Ablations identify residual prediction and MI as major contributors to Cube hard-start success, while inverse-dynamics gains are smaller and depend on their inputs and weight.

Component contributions

  • Residual prediction alone reaches 34.7% HS, Res+MI reaches 54.7%, and AD-WM reaches 52.0%, compared with 3.7% for LeWM.Adding Inv+MI to absolute prediction reaches 14.4%; Res+Inv reaches 37.1%, and Inv gives a smaller mean gain over Res without improving on Res+MI at its default weight.Source: C. Ablation Studies · p. 6

MI gains across inverse inputs

  • MI adds 14.9/15.6/20.4 mean HS points across predicted endpoints, encoded endpoints, and predicted increments; the latter two reach 60.7/60.3%, versus 52.0% for the default.These comparisons vary inverse-dynamics inputs while MI always uses predicted endpoints, so they test Inv-input robustness rather than MI-input robustness.Source: C. Ablation Studies · p. 6

Weight sensitivity

  • With Inv fixed at 0.1, tested nonzero MI weights yield 50.9–65.2% mean HS versus 37.1%; with MI fixed at 0.01, Inv weights 0.05/0.20 exceed the no-Inv mean, but 0.10 does not.MI peaks at 0.03, while the pre-specified 0.01 remains primary; Inv’s effect is smaller and weight-dependent.Source: C. Ablation Studies · p. 6

D. Planning Diagnostics

Planning diagnostics find elite-regret measures more closely aligned with Cube success than factual error or whole-bank ranking. On Franka, AD-WM has higher basic pick-and-place success than V-JEPA 2-AC under the evaluated transfer setup.

D. Planning Diagnostics

  • Diagnostics compare factual error on 61,999 held-out clips with candidate selection over 64 shared cases, five variants, three seeds, and 300 candidates per case.The default elite size is k = 30, with k = 15 and 60 used to test robustness.Source: D. Planning Diagnostics · p. 6
  • Successive R30 reductions along LeWM–Res–Res+Inv–AD-WM are 0.028 [0.014, 0.043], 0.002 [−0.003, 0.009], and 0.014 [0.009, 0.023]; the Inv-only increment remains unresolved.Res+MI and AD-WM have similar mean regret, and realized-best candidate retention falls from 0.198 to 0.135 from LeWM to AD-WM.Source: D. Planning Diagnostics · p. 6
  • From LeWM to AD-WM, CEM-center distance to the lowest-cost grid point decreases from 0.827 to 0.130, and mean elite-to-center distance falls from 0.448 to 0.082.From LeWM to Res, latent-increment MSE decreases from 0.0067 to 0.0055, while negative-cosine consecutive increments fall from 0.242 to 0.146.Source: D. Planning Diagnostics · p. 6
  • LeWM has the lowest factual MSE but the lowest HS, while AD-WM reaches 52.0% HS despite the highest MSE; HS correlates with CAD at ρ = −0.399 and negative R30 and ¯R30 at 0.863 and 0.810.Across 15 model–seed observations, regret associations persist at k = 15 and k = 60; these fixed-bank associations do not measure adaptive CEM trajectories.Source: D. Planning Diagnostics · p. 6, Table III · p. 7

E. Real-Robot Transfer

  • Franka MPC follows manually supplied grasp, move, and place image goals with 10, 10, and 4 execution steps, evaluating short-horizon control inside a structured subgoal pipeline.The setup uses one side/rear camera; each model has 45 basic pick-and-place, 10 complex-object, and 27 specified-target trials, in separate non-randomized blocks.Source: E. Real-Robot Transfer · pp. 6–7
  • Basic pick-and-place success rises from 42.2% to 71.1% for AD-WM; complex-object and specified-target results also favor AD-WM, though their smaller samples limit precision.The robot comparison uses a matched frozen V-JEPA 2 encoder and post-training setup, without laboratory images or demonstrations for adaptation.Source: E. Real-Robot Transfer · pp. 6–7, Table IV · p. 8

V. CONCLUSION AND LIMITATIONS

AD-WM combines residual prediction and predictor-level action recovery, with auxiliary heads discarded at test time. It improves over matched LeWM in four of five simulation environments and transfers to Franka without lab-specific adaptation, while several evaluation boundaries remain.

V. CONCLUSION AND LIMITATIONS

  • AD-WM improves over matched LeWM in four of five simulation environments and transfers to Franka without lab-specific adaptation; Cube diagnostics find elite regret tracks success more closely than factual error or whole-bank ranking.Cube ablations identify residual prediction and MI as the largest contributors, while inverse dynamics has a smaller, weight-dependent effect on mean success.Source: §V · p. 8
  • Diagnostics remain limited to Cube, and Scene’s gain is unresolved over three seeds; robot evaluation uses manual subgoals and non-randomized trials at one site, with one camera and backbone.The authors identify broader settings and reduced reliance on manual subgoals as future work.Source: §V · p. 8

Related papers

Loading 2609.30264v1…