Skip to content

Back to the catalog

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

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

AD-WM is an action-discriminative joint-embedding world model designed for counterfactual model predictive control, addressing the problem that world models trained only on factual transition prediction often fail to distinguish between alternative candidate actions during planning. It combines residual latent dynamics with predictor-level action-recovery regularization (inverse dynamics plus a normalized recovery objective grounded in conditional mutual information) to preserve action-dependent differences needed for MPC. On OGBench-Cube it raises hard-start success from 3.7% to 52.0% over a matched LeWM baseline, and with a frozen V-JEPA 2 encoder it improves zero-shot pick-and-place transfer on a real Franka arm from 42.2% to 71.1%.

TL;DR A world model trained to preserve action-discriminative information (not just factual accuracy) dramatically improves counterfactual MPC success on OGBench-Cube and real Franka pick-and-place tasks.
Open original ↗
Document

What it claims

Beyond the headline transfer results the summary already gives, the page establishes a diagnostic finding about which planning metrics predict closed-loop control success — and it turns out the obvious ones do not. AD-WM's argument is that a world model can have low factual prediction error yet still fail at MPC, because MPC compares alternative actions from the same state rather than predicting one factual rollout.

How it works

Two auxiliary objectives sit on top of residual latent dynamics: an inverse-dynamics head and a normalized recovery objective motivated by conditional mutual information, both pushing planning transitions to preserve action information. The page adds a detail the summary omits: both auxiliary heads are discarded at test time, so MPC itself is left unchanged — the regularization shapes the learned representation but adds no inference-time cost or planner modification. Planning uses CEM (cross-entropy method).

Evidence

  • OGBench-Cube hard-start: 3.7% → 52.0% over a matched LeWM baseline (in summary).
  • Breadth: improves mean success over the reproduced baseline in four of five simulation environments — the summary's single number understates the spread of results.
  • Planning diagnostics (the paper's distinctive claim): 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. This is the concrete basis for the "don't optimize factual accuracy alone" thesis, and it is not in the summary.
  • Real robot: zero-shot Franka pick-and-place 42.2% → 71.1%, with a frozen V-JEPA 2 encoder and matched DROID post-training, no lab-specific adaptation (the DROID post-training detail is new against the summary).

Limits

  • This is an abstract-only page; no ablation tables, hyperparameters, per-environment numbers, or the identity of the one simulation environment where it did not improve are visible. The four-of-five claim means one environment regressed or was flat — unspecified here.
  • Real-robot evidence is a single task family (basic pick-and-place) on one Franka setup at 71.1% — far from reliable, and "zero-shot transfer" rests on the V-JEPA 2 + DROID pipeline, not on the world model alone.
  • Code and videos are promised via a link not resolvable in the fetched text; verify availability before planning to build on it.

Organisations named here

  • Franka
  • OGBench-Cube

Technologies named here

  • AD-WM
  • LeWM
  • V-JEPA 2
  • Droid

Techniques named here

  • Counterfactual model predictive control
  • Action-discriminative joint-embedding world model
  • Residual latent dynamics
  • Predictor-level action-recovery regularization
  • Inverse dynamics objective
  • Normalized recovery objective
  • CEM-aligned elite regret

Related in research