PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Counterfactual World Models for Embodied Reasoning under Partial Observability

Todd Y. Zhou, Daniel Zhang

arXiv:2609.05834v1cs.AI

TL;DR

Embodied world models can predict plausible observations while collapsing interventions with different behavioral consequences, leaving their representations insufficient for action. The paper introduces CLWM, which adds recurrent belief-state modeling and outcome-based contrastive separation of action-conditioned futures. Across partially observable benchmarks, CLWM improves planning, reduces exploitative failures, and supports counterfactual separability as an audit metric, while its evaluation remains limited to world models trained from scratch.

  • Problem

    Observation prediction alone may fail to preserve intervention consequences, especially when different actions produce visually similar futures under partial observability.

  • Method

    CLWM combines recurrent belief inference, action-conditioned latent dynamics, and hard counterfactual negatives defined by intervention outcomes.

  • Results

    CLWM improves planning across all three partially observable benchmarks by up to 11.6 success points, nearly halves exploitative failures, and yields counterfactual separability correlated with planning success across five baselines.

  • Takeaways & Limitations

    Counterfactual separability can audit whether a representation preserves action-sensitive structure before a planner relies on it.

  • Takeaways & Limitations

    The pretrained-representation claim is conceptual rather than empirical, and the metric has so far been calibrated on world models trained from scratch.

Abstract

from arXiv · show

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.

1 Introduction

The paper argues that plausible visual prediction is insufficient for embodied reasoning because representations must preserve how alternative interventions change future outcomes. CLWM addresses this counterfactual collapse with action-conditioned latent modeling and contrastive supervision over intervention-conditioned futures.

  • Motivation: Counterfactual collapse occurs when visually similar futures induced by different actions remain indistinguishable despite different behavioral consequences.This is especially problematic under occlusion, perceptual aliasing, and partial observability.
  • Motivation: Pretrained representations optimized for perceptual or semantic similarity have no explicit pressure to separate futures that look alike but behave differently.The paper proposes counterfactual sufficiency as a property that can also be audited independently of CLWM.
  • Approach: CLWM combines a recurrent belief encoder, stochastic latent dynamics, imagination-based planning, and contrastive supervision over action-conditioned futures.Distinct intervention outcomes should be separated even when their pixels are similar, while equivalent outcomes should align despite nuisance variation.
  • Contributions: CLWM improves planning by up to 11.6 success points and nearly halves exploitative planning failures across three partially observable settings.Ablations identify hard counterfactual negatives, especially perceptual-alias negatives, as the key source of improvement.

2 Related Work

CLWM builds on behavioral rather than perceptual equivalence, adapting that concern to belief-state, intervention-conditioned futures. Its distinctive contribution is hard-negative construction that makes intervention outcomes—not visual similarity or generic contrast—organize the representation.

  • Latent world models: CLWM changes training pressure on action-conditioned latent futures while remaining orthogonal to the planner or actor architecture.The control loop and planner choice are not the contribution.
  • Value equivalence and bisimulation: Counterfactual collapse extends behavioral-equivalence concerns to imagined open-loop futures under alternative interventions from a common belief.The paper requires futures with different outcomes to remain separated, rather than comparing only on-policy states.
  • Contrastive representation learning for control: Unlike generic temporal or augmentation-based contrast, CLWM uses intervention outcomes to define supervised contrastive labels.Its negative construction includes same-belief, perceptual-alias, and value-disagreement negatives.
  • Causal reasoning in reinforcement learning: The method uses interventional language operationally for alternative open-loop action sequences and makes no identifiability claim about recovering a structural causal model.It preserves task-indexed intervention distinctions instead.
  • Model exploitation: The exploitation-rate metric measures deployment failures where a planner selects an imagined future whose predicted value the environment does not realize.This targets model error that can make planning appear successful in imagination while remaining unsafe.

3 Problem Setting

The problem setting is a partially observable decision process in which observations are not Markov, so planning requires a learned history-dependent belief and latent dynamics. The central gap is that observation prediction alone need not preserve distinctions among intervention outcomes.

  • Problem setting: In the POMDP, hidden states generate observations, transitions, and rewards, while the agent maintains history or a learned belief state.The current observation may not contain enough information to determine the relevant state.
  • Problem setting: A latent world model predicts next latent states, observations, and rewards, and a planner rolls out candidate futures in latent space.This supports imagination-based action selection.
  • Problem setting: The usual objective predicts observations, rewards, and continuation values along observed trajectories but does not directly preserve alternative interventional futures.Embodied reasoning requires the family of futures generated by different action sequences.
  • Problem setting: A planning-useful representation preserves distinctions among alternative action-sequence futures whenever those distinctions change control-relevant outcomes.The requirement concerns action consequences rather than observation similarity alone.
  • Counterfactual collapse: Counterfactual collapse is formalized as near-identical imagined latent futures paired with substantially different intervention outcomes.Visual plausibility can coexist with this failure because visual similarity and intervention equivalence are different relations.

4 Counterfactual Latent World Models

CLWM augments recurrent latent world-model learning with open-loop, action-conditioned rollouts and contrastive supervision that preserves intervention-sensitive structure. It uses outcome-based hard negatives to shape the latents consumed by planning, while retaining standard prediction, regularization, reward, and value losses.

  • Architecture: CLWM keeps a recurrent state-space architecture but changes the training signal applied to imagined action-conditioned futures.The recurrent belief is intended to retain hidden information unavailable in the current frame.
  • Architecture: The encoder forms a deterministic recurrent belief from prior belief, previous action, and multimodal observation features.The resulting belief supports inference under partial observability.
  • Counterfactual rollouts: Candidate intervention sequences come from random shooting, policy proposals, model-predictive-control candidates, or replay, then are rolled out in latent space.During imagination, the belief advances open loop using sampled prior latents because future observations are unavailable.
  • Counterfactual rollouts: Imagined rollouts remain fully open loop, with no ground-truth future observation posterior-fed into the rollout.A projection head shapes training representations but is not used as a decoder feature at test time.
  • Counterfactual contrastive objective: Positive pairs share intervention outcomes, while negatives differ in outcome despite originating from the same or similar belief state.Outcome labels may come from privileged simulation state or approximate real-world predicates, trackers, tactile events, or sparse labels.
  • Counterfactual contrastive objective: Hard negatives include same-belief action, perceptual-alias, and value-disagreement cases, preventing visual plausibility from substituting for intervention structure.Perceptual-alias negatives specifically pair visually similar rollouts with different hidden-state outcomes.
  • Full training objective: The full objective combines prediction, KL regularization, reward, value, counterfactual contrastive, and optional intervention-outcome losses.The outcome head helps separated futures correspond to interpretable intervention consequences rather than arbitrary representation spread.
  • Planning with CLWM: At test time, any latent-space planner can use CLWM, and contrastive computation is unnecessary during deployment.The training objective shapes the latent dynamics used by model-predictive control.

5 Experimental Design

The experiments evaluate whether world models preserve intervention structure under partial observability, using three embodied benchmarks, five baselines, and task and counterfactual diagnostics. Counterfactual separability is computed from imagined latent dynamics and intervention outcomes to measure distinctions relevant to planning.

  • Experimental Design: The evaluation tests whether plausible-prediction world models preserve action-contingent futures when short-horizon observations are insufficient for control.The benchmarks target failures caused by visually similar observations with different intervention consequences.
  • Experimental Design: Three partially observable benchmarks measure episode success on Occluded Push and Aliased Maze and normalized return on Deferred Kitchen.Occluded Push includes hidden or visually ambiguous target configurations that can produce different outcomes from the same RGB-D observation.
  • Baselines: Five baselines span reconstruction-based, reward-aware, generic contrastive, imagination-based, and value-guided latent world models.The comparison includes RSSM-Recon, Reward-aware RSSM, CPC-RSSM, Dreamer-style, and TD-MPC-style models.
  • Metrics: Task performance, counterfactual separability, and exploitation rate provide complementary measures of planning quality and intervention sensitivity.CFS measures whether imagined latent rollouts separate intervention outcomes, while exploitation rate captures failures with high predicted value but low realized value.
  • Metrics: CFS is AUROC over distances between imagined latent futures for different intervention sequences, labeled by whether their intervention outcomes differ.It is computed on the imagined dynamics latents with the projection head discarded, matching the geometry used by the planner.

6 Results

CLWM improves planning and reduces exploitative failures, with gains concentrated where perceptual prediction cannot preserve intervention-relevant distinctions. Counterfactual separability tracks planning success across baseline models, while ablations identify hard counterfactual negatives—especially perceptual-alias negatives—as central.

  • Main results: 9.5 points on Occluded Push and 11.6 on Aliased Maze over the strongest non-counterfactual baseline, while Deferred Kitchen return rises from 62.8 to 72.3 and exploitation falls from 18.4% to 9.7%.CLWM achieves the best performance on every benchmark and diagnostic.
  • Main results: Reward-aware RSSM and CPC-RSSM improve over simpler baselines but remain below CLWM because they do not sufficiently separate alternative futures with different downstream outcomes.Reward supervision is sparse and trajectory-local, while generic contrastive prediction does not target same-belief interventions and perceptual aliases.
  • Counterfactual separability: r = 0.94 on Occluded Push and r = 0.95 on Aliased Maze between CFS and planning success across five baselines provide the informative metric-validation trend.CLWM directly optimizes CFS, so its plotted point is consistent with, but does not independently validate, the relationship.
  • Ablations: Removing perceptual-alias negatives drops success to 67.5% and 69.1% and lowers average CFS from 0.85 to 0.74, restoring the collapse of visually similar but behaviorally different futures.This ablation directly removes pressure to distinguish the cases the method targets.
  • Ablations: Removing value-disagreement negatives raises exploitation to 12.9%, while random negatives improve over the reference but remain far below full CLWM.Same-belief and value-disagreement negatives provide complementary benefits, making negative construction quality central.
  • Robustness under perceptual aliasing: 66% success at the highest aliasing level compares with 55% for TD-MPC-style planning and 38% for RSSM-Recon, preserving CLWM’s advantage as observations become less informative.CLWM’s margin over the strongest baseline changes only from 11.6 points at intensity 0.0 to 11.0 at intensity 1.0.

7 Discussion

The results support evaluating embodied world models by whether they preserve action consequences, not only plausible observations. CLWM reduces exploitative failures, and counterfactual separability offers a representation-agnostic audit for pretrained or learned representations.

  • 18.4% → 9.7% model-exploitation rate shows that CLWM reduces failures caused by overoptimistic imagined trajectories.The reduction reflects fewer cases where collapsed latent futures make invalid action sequences appear valuable.
  • The paper’s practical claim is that planners need representations preserving intervention distinctions, without claiming structural causal modeling or identifiability.The scope is deliberately limited to action-relevant representation structure.
  • Counterfactual separability can audit frozen, fine-tuned, pretrained, or policy representations before planners rely on them.It requires action-conditioned futures and intervention-outcome labels, while reward alone is slower and confounded for this property.
  • Across baseline model classes, counterfactual separability tracks planning success with r ≥0.94.The experiments establish this relationship for world models trained from scratch rather than large-scale pretrained encoders.

8 Limitations and Broader Impact

CLWM depends on reliable intervention-outcome signals that are straightforward in simulation but harder to obtain in the real world. The method remains complementary to uncertainty-aware model-based RL and does not address several broader failure modes, while deployment of more autonomous embodied systems requires safeguards.

  • Reliable intervention-outcome signals are available from privileged simulator state but require trackers, tactile events, learned predicates, or human labels in the real world.This creates a practical supervision requirement outside simulation.
  • CLWM does not solve long-horizon compounding error, exploration, distribution shift, or reward misspecification.The authors position it as complementary to uncertainty-aware model-based RL.
  • Pretrained-representation claims remain conceptual because CFS is not measured on large-scale pretrained encoders in these experiments.The paper identifies this as a natural next experiment.
  • More autonomous embodied planning requires conservative deployment practices, including uncertainty estimation, human override, environment-specific validation, and limits on autonomous authority.These safeguards follow from the increased autonomy of physical systems.

9 Conclusion

CLWM targets counterfactual collapse by training latent world models to distinguish behaviorally different futures that look similar observationally. Its training-time machinery improves planning and its counterfactual separability diagnostic can audit representations before deployment.

  • CLWM improves planning on all three partially observable benchmarks by up to 11.6 success points.The reported gains are accompanied by increased counterfactual separability and nearly halved model-exploitation failures.
  • Hard perceptual-alias negatives account for the strongest ablation gains, with the advantage persisting under maximum perceptual aliasing.This supports targeting futures that appear similar in observation space but differ behaviorally.
  • The counterfactual machinery operates only during training, so inference retains the existing planner and control loop without deployment cost.CLWM attaches to latent world models that roll out action-conditioned futures rather than replacing their architecture.
  • CFS can audit frozen, fine-tuned, or pretrained representations before planners trust them, and tracks planning success at r ≥0.94 across scratch-trained world-model classes.The diagnostic requires imagined action-conditioned futures and intervention-outcome labels.

A Experimental Details

The appendix reports configuration and implementation details for the experimental runs. Reported statistics in Tables 1 and 2 are derived from the per-seed values collected in the supplementary material.

  • Tables 3 and 4 provide the training, evaluation, configuration, and implementation details for all reported runs.
  • Per-seed values for every reported number are included in the supplementary material.
  • All statistics in Tables 1 and 2 are derived from the supplementary per-seed table.
Loading 2609.05834v1…