Learning Counterfactual World Models for Embodied Reasoning under Partial Observability
Todd Zhou ⋅ Daniel Zhang
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 — 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 \geq 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 (Section 8). Here we establish the metric and its relationship to planning success for world models trained from scratch.
Chat is not available.
Successful Page Load