Learning Counterfactual World Models for Embodied Reasoning under Partial Observability

Todd Y. Zhou; Daniel Zhang

cs.AISep 5, 2026

Learning Counterfactual World Models for Embodied Reasoning under Partial Observability

Authors: Todd Y. Zhou, Daniel Zhang

Organizations: Harvard University

Abstract

World models promise a general route to embodied intelligence: learn predictive dynamics once, then reason, plan, and act with them. Increasingly, the representations beneath such models are pretrained on large-scale video, interaction, and multimodal corpora, which raises a question prediction quality alone cannot answer: when is a learned representation actually actionable? We identify a failure mode we call counterfactual collapse: a model predicts visually plausible futures while failing to distinguish interventions with different behavioral consequences. This arises whenever a representation is optimized for perceptual similarity rather than intervention structure, which is precisely the objective under which most large-scale pretrained encoders are learned. We introduce Counterfactual Latent World Models (CLWM), which combine a recurrent belief-state encoder, action-conditioned latent dynamics, and a contrastive counterfactual objective that separates futures induced by distinct interventions even when their observations look alike. Across occluded manipulation, aliased navigation, and long-horizon manipulation, CLWM improves planning success over the strongest baseline (65.1% $\to$ 74.6% on Occluded Push and 67.3% $\to$ 78.9% on Aliased Maze) and reduces exploitative planning failures (18.4% $\to$ 9.7% on Deferred Kitchen), with ablations attributing the gains to hard counterfactual negatives, especially perceptual-alias negatives. Finally, our counterfactual separability metric, which tracks planning success across the five baseline model classes ( $r \ge 0.94$ ), is representation-agnostic: given intervention-outcome labels, it can audit any encoder, pretrained or trained from scratch, before a planner trusts it. We do not yet measure it on large-scale pretrained encoders. Here we establish the metric and its relationship to planning success for world models trained from scratch.

Figures & tables

Method	Occluded Push		Aliased Maze		Deferred Kitchen
Method	Success $\uparrow$	CFS $\uparrow$	Success $\uparrow$	CFS $\uparrow$	Return $\uparrow$	Exploit $\downarrow$
RSSM-Recon	47.8 $\pm$ 2.1	0.61	52.4 $\pm$ 1.8	0.58	43.2 $\pm$ 2.7	31.5
Reward-aware RSSM	55.6 $\pm$ 1.9	0.66	58.1 $\pm$ 2.0	0.63	51.7 $\pm$ 2.4	24.8
CPC-RSSM	58.4 $\pm$ 2.4	0.70	61.5 $\pm$ 1.7	0.68	53.9 $\pm$ 2.1	22.7
Dreamer-style	62.9 $\pm$ 1.6	0.69	65.7 $\pm$ 1.9	0.67	59.4 $\pm$ 1.8	20.9
TD-MPC-style	65.1 $\pm$ 1.6	0.71	67.3 $\pm$ 1.5	0.70	62.8 $\pm$ 2.0	18.4

Table 1: Main results on partially observable embodied benchmarks. Success denotes episode success rate. CFS is counterfactual separability AUROC. Exploit rate is the percentage of planning failures caused by high predicted value but low realized value. All success and return values are means over 5 random seeds (seeds 0–4), and

\pm

denotes the standard error of the mean. CFS and exploitation rate are computed over the same seeds. All methods are evaluated from their final checkpoints on 100 fixed test episodes. Aliased Maze results use aliasing intensity

0.0

, the benchmark’s default layout (details in Appendix A ).

Panel of Figure 1 — Figure 1: Counterfactual separability tracks downstream planning success. Across the five baseline model classes—none of which optimizes CFS—the correlation between CFS and success is $r=0.94$ (Occluded Push) and $r=0.95$ (Aliased Maze). The dashed line is a least-squares fit to the ten baseline points. CLWM (rightmost point in each series) directly optimizes counterfactual structure, so its points illustrate consistency with the trend rather than independent validation of the metric.

Variant	Push Succ. $\uparrow$	Maze Succ. $\uparrow$	Avg. CFS $\uparrow$	Kitchen Exploit $\downarrow$
Full CLWM	74.6	78.9	0.85	9.7
No same-belief negatives	70.2	72.4	0.79	13.6
No perceptual-alias negatives	67.5	69.1	0.74	16.8
No value-disagreement negatives	71.3	74.0	0.81	12.9
No outcome head	70.8	73.5	0.80	13.2
Random negatives only	68.1	70.6	0.76	15.4

Table 2: Ablation study. All ablation variants are trained and evaluated identically on all three benchmarks. Each column reports one benchmark’s metric: success on Occluded Push, success on Aliased Maze, CFS averaged over those two tasks (rounded to two decimals, with per-task values for CLWM and the reference in Table 1 ), and exploitation rate on Deferred Kitchen. The bottom row is not an ablation: it repeats the strongest non-counterfactual baseline from Table 1 for reference.

Panel of Figure 2 — Figure 2: Robustness to perceptual aliasing. All results use the same final checkpoints and evaluation protocol as Table 1 (100 test episodes, default maze layout). Intensity $0.0$ corresponds to the Table 1 setting. Points are means over 5 seeds. All methods degrade as local observations become less informative. CLWM remains the strongest at every intensity and retains the largest fraction of its default-setting performance, consistent with latent dynamics trained to preserve action-conditioned distinctions that are not visible in the current observation.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	Value
Seeds & statistics	5 seeds (0–4); $\pm$ is the standard error of the mean over seeds
Belief / latent dimension	200 / 32
Planner	MPC with $N=512$ candidate sequences, horizon $H=12$
Counterfactual rollout horizon	$K=8$
Contrastive temperature	$\tau=0.07$
Loss weights	$\lambda_{\mathrm{cf}}=0.1$ , $\lambda_{\mathrm{inv}}=0.01$ , $\beta_{\mathrm{KL}}=0.5$ (with free bits), $\beta_{r}=1.0$ , $\beta_{v}=0.1$

Appendix

Table 3: Training and evaluation configuration.

Field	Value
Counterfactual branching	$M=64$ intervention sequences per anchor; branch source: 50% policy rollouts, 25% random shooting, 25% replay trajectories
Negative mining	perceptual-alias similarity via LPIPS on image patches over a 5-step window; 16 negatives per anchor; outcome distance threshold $\epsilon_{y}=0.1$
Outcome descriptors $y$	Occluded Push: hidden-object $(x,y)$ pose relative to goal, success threshold 0.05 m; Aliased Maze: discrete topological cell index (0–63); Deferred Kitchen: binary feasibility vector over later subtasks (e.g., stove on, kettle filled)
Simulators	Occluded Push: MuJoCo (MJX) v1.2.1; Aliased Maze: Gymnasium-based grid world v0.1.0; Deferred Kitchen: RoboSuite v2.0 (adapted)
Encoder $e_{\theta}$	4-layer CNN (32/64/128/256 channels, kernel 3, stride 2, ReLU) followed by a fully connected layer to 256 units; proprioceptive/tactile branch: 2-layer MLP (128, 64), concatenated
Decoder	transposed CNN mirroring the encoder

Appendix

Table 4: Implementation details.

Embodied World Models

Embodied world models predict how environments change under an agent’s actions to support planning, learning, simulation, and control.

See all 164 papers

AD-WM: Action-Discriminative World Models for Counterfactual Model Predictive Control
Sep 24, 2026Jiabin Qiu, Zixuan Chen, Hongye Cao +3World Models for Robotics World Model-Based Planning
Nanjing University
The Planning Limits of Latent World Models
Sep 30, 2026Ali Alrasheed, Basim Azam, Naveed AkhtarWorld Model-Based Planning Action-Conditioned World Models
The University of Melbourne, Australia
IMWM: Intuition Models Complement World Models for Latent Planning
Jun 1, 2026Baoqi Gao, Ruize Han, Miao Wang +1World Model-Based Planning Goal-Conditioned Planning
Beihang University · Shenzhen University of Advanced Technology

Sep 24, 2026J/K move · Enter open · S save

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

Jiabin Qiu, Zixuan Chen, Hongye Cao, Jieqi Shi+2

Nanjing University

World Models for Robotics World Model-Based Planning Model Predictive Control Latent World Models

Open paper AskPDF arXiv

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

Abstract

Embodied World Models

Explore similar work

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

The Planning Limits of Latent World Models

IMWM: Intuition Models Complement World Models for Latent Planning