# Learning Counterfactual World Models for Embodied Reasoning under Partial Observability

arXiv: https://arxiv.org/abs/2609.05834 · HTML: https://mlsanity.com/paper/2609.05834/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Latent predictive world models, one of 11 directions in World models.
- Builds on 16 of its 19 references that are in the map.
- First posted on arXiv 5 Sep 2026.

## Placed in

- World models / [Latent predictive world models](https://mlsanity.com/wm/latent-predictive-world-models/)

## Builds on (most cited first, 16 of 16)

- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Supervised Contrastive Learning](https://mlsanity.com/paper/2004.11362/) (2020)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [DayDreamer: World Models for Physical Robot Learning](https://mlsanity.com/paper/2206.14176/) (2022)
- [DeepMDP: Learning Continuous Latent Space Models for Representation Learning](https://mlsanity.com/paper/1906.02736/) (2019)
- [The Value Equivalence Principle for Model-Based Reinforcement Learning](https://mlsanity.com/paper/2011.03506/) (2020)
- [Causal Reasoning from Meta-reinforcement Learning](https://mlsanity.com/paper/1901.08162/) (2019)
- [Causal World Models by Unsupervised Deconfounding of Physical Dynamics](https://mlsanity.com/paper/2012.14228/) (2020)

## Abstract (arXiv)

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% → 74.6% on Occluded Push and 67.3% → 78.9% on Aliased Maze) and reduces exploitative planning failures (18.4% → 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 ≥ 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.