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

Sep 24, 2026

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

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

Jiabin Qiu, Zixuan Chen, Hongye Cao +3

Nanjing University

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

Ask PDF

Sep 30, 2026

The Planning Limits of Latent World Models

World models offer a promising way to help robots understand how the physical world evolves and plan complex behaviours through imagination. Yet existing studies mainly demonstrate what these models can accomplish, leaving unclear when their predictions remain useful for planning and where they fail. We study this question using action-conditioned predictors built on five frozen self-supervised visual backbones: V-JEPA 2, V-JEPA 2.1, VideoMAEv2, VideoPrism, and DINOv2. We use frozen backbones to test representations intended to transfer across environments. We evaluate these models on diverse Meta-World manipulation tasks and real-robot interactions from BridgeData V2. We find that a world model guides action selection reliably only when the goal lies within, or slightly beyond, the trajectory it imagines during planning. With five-step rollouts, the length the predictor was trained on, the world model ranks actions reliably only for targets five to ten control steps ahead, whereas task goals lie 16 to 53 steps away. Neither an 81-fold larger predictor nor longer-rollout training extends this range; the encoder affects both range and closed-loop success, with V-JEPA 2.1 performing most consistently. More fundamentally, the limit persists under perfect prediction: using the real simulator, success falls from 92% to 41% as the target moves from five to twenty steps ahead of a five-step rollout. Planning therefore requires either longer imagined trajectories or closer subgoals. For distant goals, pure imagination succeeds in 23% of episodes, planning with feedback (MPC) raises success to 30%, imagining as far as the goal to 47%, and nearby expert subgoals to 76%. Used within its plannable range, a world model can also improve a vision-language-action (VLA) policy: choosing among eight actions the VLA proposes raises its success from 65% to 77% across 16 different tasks.

Ali Alrasheed, Basim Azam, Naveed Akhtar

The University of Melbourne, Australia

World Model-Based Planning Action-Conditioned World Models Robotic Manipulation Model-Based Planning Latent World Models World Model Evaluation

Ask PDF

Jun 1, 2026

IMWM: Intuition Models Complement World Models for Latent Planning

Planning with a learned latent world model is a promising route to control from raw pixels, but a strong world model alone is not enough. We show this experimentally: even with a perfect world model (operationalized by replacing the learned forward predictor with an idealized rollout of the true environment dynamics), a finite-budget sample-based planner still fails on some tasks, indicating that the bottleneck can lie in search rather than in world-model accuracy. Motivated by this gap, we propose IMWM (Intuition Model + World Model), which pairs the world model with an intuition model trained from demonstrations to recognize promising actions. The two models collaborate through three lightweight components: (i) Retrieval Initialization, which initializes the planner's action proposal from a retrieved demonstration; (ii) Hybrid Cost, which combines the intuition score with the world-model rollout cost; and (iii) a Reliability Gate, which adjusts how much the planner trusts intuition in each setting. Across four pixel-based goal-reaching tasks (Two-Room, Reacher, Push-T, and OGBench-Cube), IMWM has higher mean success than the world-model-only planner on all four, with the largest gains on Two-Room (99.2%, +11.5 percentage points) and OGBench-Cube (94.7%, +28.5 percentage points).

Baoqi Gao, Ruize Han, Miao Wang +1

Beihang University · Shenzhen University of Advanced Technology

World Model-Based Planning Goal-Conditioned Planning AI Control Model-Based Planning Latent World Models Latent-Space Planning

Ask PDF

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