When a Good Prediction Is Not Enough: Why World Models Must Learn to Distinguish Actions
Imagine a robot learning to grasp a cube. A world model trained on recorded demonstrations predicts what happens next with high accuracy: the gripper closes, the cube lifts, the arm moves. But when the robot must choose between two candidate actions from the same starting state -- approach the cube versus retreat -- a model optimized purely for prediction accuracy may produce nearly identical futures for both options. The prediction is correct, yet the planning decision is uninformative. This is the central problem addressed by a new paper from researchers at Nanjing University, and it points to a fundamental gap in how learned world models are trained for planning.
The paper introduces AD-WM, an action-discriminative world model designed to make latent dynamics preserve the differences between alternative actions. On the OGBench-Cube benchmark, AD-WM improves hard-start success from 3.7% to 52.0% over a matched LeWM baseline. On a real Franka robot, it raises basic pick-and-place success from 42.2% to 71.1% without any lab-specific adaptation. The results suggest that the community has been optimizing world models for the wrong objective.
The Gap Between Prediction and Planning
Latent world models like PlaNet, Dreamer, TD-MPC, and the more recent joint-embedding architectures such as LeWM and V-JEPA learn to predict future states from current observations and actions. They then use these predictions to plan: rolling out candidate action sequences, scoring them by predicted terminal cost, and selecting the best sequence through cross-entropy method (CEM) search.
The problem is that training supervises the outcome of the action actually taken, while planning must compare what would happen under different actions from the same state. A predictor can achieve low factual error by largely preserving the current representation, especially when successive visual frames share substantial common content. But this same predictor may obscure the smaller action-dependent changes that distinguish reaching toward a target from moving away from it.
Think of it this way. If the current state carries most of the information about what the next state will look like, a lazy predictor that simply copies the current representation forward will incur only a small error. Yet this predictor assigns all candidate action sequences the same terminal cost, making it useless for distinguishing between them. The model fits observed transitions but provides no guidance for goal-directed action selection.
This mismatch between factual prediction accuracy and counterfactual action comparison has largely gone unexamined. Joint-embedding predictive architectures emphasize slowly varying features, which are good for representing persistent scene content but can suppress the smaller, action-specific signal that planning depends on.
How AD-WM Preserves Action Information
AD-WM addresses this by combining residual latent prediction with predictor-level action recovery. The design philosophy is straightforward: let the current representation carry shared information, and train the predictor to estimate only the remaining change -- but then add constraints that ensure this change retains action-dependent structure.
The residual predictor estimates latent increments. Given a current representation zt and an action embedding et, it predicts a displacement Δz-tilde that is added to the current state to produce the predicted next state. This is supervised by the encoded successor, matching the predicted increment to the actual encoded displacement. In isolation, residual prediction changes the learning bias toward modeling local changes, but it does not require those changes to retain action information.
This is where the action-recovery objectives come in. AD-WM uses two complementary mechanisms to encourage actions to be recoverable from the current and predicted next representations.
Inverse Dynamics
An auxiliary inverse head takes the current and predicted next representations and attempts to recover the action embedding that was actually applied. The key insight is that by operating on model-generated transitions rather than only observed data, the inverse head constrains the dynamics that the MPC planner will roll out during search. If the predictor produces transitions that make action recovery easy, it must have preserved action information in those transitions.
Normalized Recovery via Conditional Mutual Information
The second objective is motivated by conditional mutual information. A separate head receives the current and predicted representations and defines a Gaussian distribution over standardized action embeddings. The loss has two terms: one that recovers the normalized action embedding, and a KL penalty that shrinks the predicted mean toward zero. This regularizes the predicted transitions to carry action-dependent information while preventing the model from simply memorizing the action encoding.
Both objectives operate during training. At deployment, the auxiliary heads are discarded entirely. The MPC planner uses only the residual dynamics, with no changes to the architecture, encoder, or search procedure.
Results: From Simulation to the Real World
The experiments address four questions: cross-environment performance, ablation of components, planning diagnostics, and real-robot transfer.
Simulation Benchmarks
On OGBench-Cube, the primary benchmark, AD-WM outperforms matched LeWM across all hard-start protocols. LeWM achieves only 3.7% hard-start success across five perturbation levels. AD-WM reaches 52.0%. On the Original protocol (near-distribution starts), AD-WM scores 90.7% versus LeWM's 73.3%. The gap widens under larger perturbations: at P04, the hardest protocol where 72% of starting states fall outside the training distribution, AD-WM reaches 25.3% compared to 10.0% for LeWM and 10.7% for INTACT's Actor-CEM mode.
Across five simulation environments -- Cube, Reacher, TwoRoom, PushT, and Scene -- AD-WM improves mean success over matched LeWM in four of five. Cube goes from 73.3% to 90.7%. Reacher from 76.7% to 83.3%. TwoRoom from 90% to 98%. Scene, which uses balanced hard starts across four subtasks, improves from 35.5% to 39.5%. PushT is the sole exception, declining slightly from 94% to 92%.
Ablation Studies
Controlled ablations isolate the contribution of each component. LeWM alone achieves 3.7% hard-start success. Adding inverse dynamics and MI recovery to absolute prediction raises this to 14.4%. Residual prediction alone reaches 34.7%. Adding inverse dynamics to residual prediction gives 37.1%, while residual prediction plus MI gives 54.7%. The final AD-WM, combining all components, achieves 52.0%.
Two findings stand out. First, residual prediction is the largest single contributor, more than doubling performance on its own. Second, MI recovery adds meaningful gains over inverse dynamics alone. The weight sensitivity analysis shows that MI consistently improves hard-start success across its tested range, peaking at a weight of 0.03, with the pre-specified 0.01 still yielding strong results.
Planning Diagnostics: What Actually Predicts Success
The paper introduces a diagnostic framework that distinguishes factual prediction quality from action-selection quality. On a shared candidate bank of 300 sequences, the authors measure three metrics: one-step factual mean squared error, counterfactual action discriminability (CAD, measuring whole-bank ranking agreement via Spearman correlation), and normalized elite regret (measuring how well the top candidates selected by CEM actually perform in the environment).
The results overturn conventional assumptions. LeWM achieves the lowest factual MSE -- 2.72 -- yet the lowest hard-start success. AD-WM has the highest MSE -- 4.27 -- but the highest success rate. Across 15 model-seed observations, hard-start success correlates with CAD at a negative rho of -0.399, meaning worse whole-bank ranking is associated with better control. By contrast, success correlates strongly with elite regret: rho of 0.86 for best-in-elite regret and 0.81 for elite-mean regret.
This means that a world model can be excellent at predicting what happened, yet poor at selecting what to do next. The metric that actually matters for planning is not how accurately the model predicts observed transitions, but whether its predictions make the right action sequences stand out from the wrong ones.
Zero-Shot Transfer to a Real Robot
Perhaps the most compelling result concerns real-robot transfer. AD-WM is evaluated on a Franka robotic arm using the same frozen V-JEPA 2 ViT-G encoder and matched DROID post-training data as V-JEPA 2-AC. No laboratory images or demonstrations are used for adaptation. The comparison uses identical deployment stacks, cameras, and planner configurations.
On basic pick-and-place across three objects and five start-goal settings, AD-WM achieves 71.1% success versus 42.2% for V-JEPA 2-AC. All three objects improve: the first from 33.3% to 60.0%, the second from 46.7% to 73.3%, and the third from 46.7% to 80.0%. AD-WM also reduces abnormal motion events from 6 out of 10 to 2 out of 10 complex-object trials. On specified-target tasks, it moves the correct target in 21 of 27 trials versus 14 of 27, and completes lift-and-place in 17 of 27 trials versus 9 of 27.
What This Means for the Field
The paper makes a deceptively simple argument with far-reaching consequences. World models for planning should preserve action-dependent differences needed for counterfactual selection, rather than optimizing factual prediction accuracy alone. This does not mean prediction accuracy is irrelevant -- AD-WM still uses prediction supervision as its primary loss. But prediction accuracy is necessary, not sufficient, for useful planning dynamics.
The architectural choice to keep the MPC planner unchanged is significant. AD-WM does not propose a new planning algorithm or a new encoder architecture. The auxiliary action-recovery heads are training-only modifications that are discarded at deployment. This means the gains transfer cleanly: any existing JEPA-based world model pipeline can incorporate these objectives without changing the deployment infrastructure.
The diagnostic framework opens a new evaluation dimension for the community. If elite regret tracks closed-loop success better than factual MSE or whole-bank ranking, then the field should adopt regret-based metrics alongside traditional prediction benchmarks. A model that ranks candidates well in the full bank but fails to retain near-best candidates in the elite set will perform poorly in CEM planning, regardless of its factual accuracy.
Limitations and Open Questions
The authors are candid about the study's constraints. Diagnostics remain limited to Cube; the Scene improvement is unresolved over three seeds, with a paired bootstrap p-value of 0.13. External comparisons differ in inference settings and checkpoint counts, preventing clean attribution of performance differences to the proposed methods alone. The simulation and robot post-training use different auxiliary weights, which limits comparability.
The robot evaluation uses manual subgoals, non-randomized trials, a single site, camera, and backbone. The 45-trial basic pick-and-place protocol provides a meaningful signal but leaves precision gaps for the complex-object and specified-target protocols with only 10 and 27 trials respectively. The authors note that future work should examine broader settings and reduce reliance on manual subgoals.
A deeper question raised by this work concerns the generality of the approach. AD-WM operates in a reward-free, image-goal MPC setting learned from offline data. Extending the action-recovery framework to model-based reinforcement learning with reward supervision, or to online fine-tuning settings, remains an open direction. The normalized recovery objective, motivated by conditional mutual information, assumes a particular factorization that may not hold in all representation spaces.
Nevertheless, the paper establishes an important principle that is likely to shape how future world models are designed and evaluated. The gap between prediction and planning is not a minor nuisance -- it is a structural limitation of models trained exclusively on factual transitions. Closing that gap requires explicit counterfactual supervision, and AD-WM demonstrates that the cost of adding such supervision is modest while the planning gains can be transformative.