Learning Counterfactual World Models for Embodied Reasoning under Partial Observability

# Learning Counterfactual World Models for Embodied Reasoning under Partial Observability
Source: [https://arxiv.org/html/2609.05834](https://arxiv.org/html/2609.05834)
\\workshoptitle

PTA: From Pretrained Representations to Acting Agents

Daniel ZhangAffiliation: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—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,CLWMimproves planning success over the strongest baseline \(65\.1%→74\.6%65\.1\\%\\\!\\to\\\!74\.6\\%on Occluded Push and67\.3%→78\.9%67\.3\\%\\\!\\to\\\!78\.9\\%on Aliased Maze\) and reduces exploitative planning failures \(18\.4%→9\.7%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≥0\.94r\\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](https://arxiv.org/html/2609.05834#S8)\)\. Here we establish the metric and its relationship to planning success for world models trained from scratch\.

## 1Introduction

World models have become one of the most compelling paradigms for embodied intelligence\. If an agent can learn how the world evolves, it can evaluate possible actions in imagination, plan before acting, and reuse the same learned dynamics across tasks\. This idea underlies a broad family of model\-based reinforcement learning systems, from latent\-space planning from pixels to large\-scale world\-model agents that improve behavior by imagining future trajectories\([Ha and Schmidhuber, 2018](https://arxiv.org/html/2609.05834#bib.bib1);[Hafner et al\., 2019](https://arxiv.org/html/2609.05834#bib.bib2);[Hafner et al\., 2020](https://arxiv.org/html/2609.05834#bib.bib3);[Hafner et al\., 2023](https://arxiv.org/html/2609.05834#bib.bib4);[Hansen et al\., 2024](https://arxiv.org/html/2609.05834#bib.bib5)\)\. In robotics, world models are especially attractive because real\-world interaction is expensive, dangerous, and difficult to reset\([Wu et al\., 2022](https://arxiv.org/html/2609.05834#bib.bib6)\)\.

However, embodied reasoning demands more than plausible prediction: an agent needs to know which aspects of the future would change under a different action\. A model trained to reconstruct or predict observations can be excellent at perceptual forecasting yet insufficient for decision\-making\. In a corridor with aliased landmarks, moving left or right produces nearly identical short\-horizon images while placing the agent in different regions of the map\. In occluded manipulation, a push may leave the visible pixels unchanged while determining whether a hidden object becomes reachable several steps later\. We therefore argue that the next step for embodied world models is*counterfactual sufficiency*: the learned latent state should preserve the consequences of alternative interventions that are relevant to control\. A model is counterfactually insufficient when two histories are close in representation because they predict similar observations even though the agent should act differently in them, or when two candidate actions remain close in imagined latent space because their near\-term images look alike even though their long\-horizon outcomes diverge\. Such failures are particularly common under partial observability, where the current observation is not a Markov state\.

This framing has a consequence beyond our own method\. Counterfactual sufficiency is a property of any representation used for control, including representations pretrained on large\-scale video, image–text, or interaction corpora and then adapted for decision making\. Those corpora are fit with reconstruction, contrastive instance discrimination, or language alignment—objectives that reward perceptual and semantic similarity and contain no pressure to separate two futures that look alike but behave differently\. There is thus no reason to expect a pretrained encoder to arrive counterfactually sufficient, and considerable reason to expect the opposite in exactly the settings—occlusion, aliasing, delayed contact—where embodied agents are deployed\. We return to this in Section[7](https://arxiv.org/html/2609.05834#S7)\. Our counterfactual separability metric is intended to be usable as a standalone audit of an inherited representation, independent of theCLWMtraining objective\.

We introduce*Counterfactual Latent World Models*\(CLWM\), which retain the strengths of modern latent world models—a recurrent belief encoder, stochastic latent dynamics, and imagination\-based planning—but add a contrastive counterfactual objective over action\-conditioned futures\. The key idea is simple: when two interventions from the same belief state induce different downstream consequences, the model should separate their predicted latent futures even if their pixels are similar\. Conversely, when different trajectories lead to the same intervention outcome, the representation should align them despite nuisance\-level visual variation\.

We make four contributions\. First, we formalize counterfactual collapse as a failure mode of latent world models under partial observability\. Second, we proposeCLWM, which augments latent dynamics learning with contrastive counterfactual supervision over intervention\-conditioned futures\. Third, we introduce diagnostics for whether a world model preserves action\-sensitive structure: counterfactual separability, intervention\-outcome prediction, and model\-exploitation rate\. Fourth, we show empirically thatCLWMimproves planning by up to11\.611\.6success points and nearly halves exploitative planning failures across occluded manipulation, aliased navigation, and long\-horizon manipulation, with ablations confirming that hard counterfactual negatives are the key source of the improvement\.

## 2Related Work

#### Latent world models\.

World Models, PlaNet, Dreamer, DreamerV3, and TD\-MPC2 learn compact latent dynamics for imagination and planning\([Ha and Schmidhuber, 2018](https://arxiv.org/html/2609.05834#bib.bib1);[Hafner et al\., 2019](https://arxiv.org/html/2609.05834#bib.bib2);[Hafner et al\., 2020](https://arxiv.org/html/2609.05834#bib.bib3);[Hafner et al\., 2023](https://arxiv.org/html/2609.05834#bib.bib4);[Hansen et al\., 2024](https://arxiv.org/html/2609.05834#bib.bib5)\), and DayDreamer applies the recipe to physical robots\([Wu et al\., 2022](https://arxiv.org/html/2609.05834#bib.bib6)\)\.CLWMis orthogonal to the choice of planner or actor: it changes the training pressure applied to action\-conditioned latent futures, not the architecture or the control loop\.

#### Value equivalence and bisimulation\.

The value equivalence principle\([Grimm et al\., 2020](https://arxiv.org/html/2609.05834#bib.bib7)\), instantiated at scale by MuZero\([Schrittwieser et al\., 2020](https://arxiv.org/html/2609.05834#bib.bib8)\), holds that a model need only be correct about the quantities planning consumes\. Bisimulation metrics\([Ferns et al\., 2004](https://arxiv.org/html/2609.05834#bib.bib9)\)and their deep variants\([Gelada et al\., 2019](https://arxiv.org/html/2609.05834#bib.bib10);[Zhang et al\., 2021](https://arxiv.org/html/2609.05834#bib.bib11)\)organize state abstractions by behavioral rather than perceptual equivalence\. Counterfactual collapse is the belief\-state, intervention\-conditioned form of the same concern: rather than comparing on\-policy states, we require imagined open\-loop futures under alternative interventions from a common belief to remain separated whenever their outcomes differ\. Our contribution is this counterfactual construction and the accompanying hard\-negative mining, not the underlying principle that behavioral distinctions outrank perceptual ones\.

#### Contrastive representation learning for control\.

Contrastive predictive coding\([van den Oord et al\., 2018](https://arxiv.org/html/2609.05834#bib.bib18)\)and CURL\([Laskin et al\., 2020](https://arxiv.org/html/2609.05834#bib.bib19)\)shape representations with temporal or augmentation\-based contrast\. Our objective \(Eq\.[11](https://arxiv.org/html/2609.05834#S4.E11)\) takes the supervised contrastive form of[Khosla et al\. \(2020\)](https://arxiv.org/html/2609.05834#bib.bib12), with intervention outcomes rather than augmentation identity defining the label structure\. What distinguishes it is the negative construction—same\-belief, perceptual\-alias, and value\-disagreement negatives—which our ablations show carries the empirical benefit\.

#### Causal reasoning in reinforcement learning\.

We use interventional language\([Pearl, 2009](https://arxiv.org/html/2609.05834#bib.bib14);[Peters et al\., 2017](https://arxiv.org/html/2609.05834#bib.bib15)\)in the operational sense of executing alternative open\-loop action sequences from a common anchor\. Related work learns causal world models by deconfounding physical dynamics\([Li et al\., 2020](https://arxiv.org/html/2609.05834#bib.bib16)\)or meta\-learns causal reasoning\([Dasgupta et al\., 2019](https://arxiv.org/html/2609.05834#bib.bib17)\)\.CLWMmakes no identifiability claims: it preserves a task\-indexed set of intervention distinctions rather than recovering a structural causal model\.

#### Model exploitation\.

Planners exploit model error, especially outside the training distribution\([Janner et al\., 2019](https://arxiv.org/html/2609.05834#bib.bib13)\)\. Our exploitation\-rate metric quantifies this failure at deployment: episodes in which the planner selects an imagined future whose predicted value the environment does not realize\.

## 3Problem Setting

We consider a partially observable Markov decision process\([Kaelbling et al\., 1998](https://arxiv.org/html/2609.05834#bib.bib20), POMDP;\)ℳ=\(𝒮,𝒜,𝒪,T,O,R,γ\)\\mathcal\{M\}=\(\\mathcal\{S\},\\mathcal\{A\},\\mathcal\{O\},T,O,R,\\gamma\)with hidden statest∈𝒮s\_\{t\}\\in\\mathcal\{S\}, actionat∈𝒜a\_\{t\}\\in\\mathcal\{A\}, observationot∈𝒪o\_\{t\}\\in\\mathcal\{O\}, transition distributionT⁡\(st\+1∣st,at\)T\(s\_\{t\+1\}\\mid s\_\{t\},a\_\{t\}\), observation distributionO⁡\(ot∣st\)O\(o\_\{t\}\\mid s\_\{t\}\), rewardrt=R⁡\(st,at\)r\_\{t\}=R\(s\_\{t\},a\_\{t\}\), and discountγ∈\(0,1\)\\gamma\\in\(0,1\)\. Becauseoto\_\{t\}is not Markov, an agent maintains a historyht=\(o1:t,a1:t−1\)h\_\{t\}=\(o\_\{1:t\},a\_\{1:t\-1\}\)or a learned belief statebt=fθ​\(ht\)b\_\{t\}=f\_\{\\theta\}\(h\_\{t\}\)\. A world model learns latent dynamics

pθ​\(zt\+1∣zt,at\),pθ​\(ot∣zt\),pθ​\(rt∣zt,at\),p\_\{\\theta\}\(z\_\{t\+1\}\\mid z\_\{t\},a\_\{t\}\),\\quad p\_\{\\theta\}\(o\_\{t\}\\mid z\_\{t\}\),\\quad p\_\{\\theta\}\(r\_\{t\}\\mid z\_\{t\},a\_\{t\}\),\(1\)whereztz\_\{t\}denotes a latent state sampled or inferred from the belief, and a planner selects actions by rolling out candidate futures in latent space\.

The usual learning objective encourages the model to predict observations, rewards, and continuation values along observed trajectories\. This is necessary but incomplete\. For embodied reasoning, the relevant object is not onlyp\(ot\+k∣ht,at:t\+k−1\)p\(o\_\{t\+k\}\\mid h\_\{t\},a\_\{t:t\+k\-1\}\)but the family of interventional futures

p\(ot\+k,st\+k,rt:t\+k∣ht,do\(at:t\+k−1=α\)\),p\(o\_\{t\+k\},s\_\{t\+k\},r\_\{t:t\+k\}\\mid h\_\{t\},\\mathrm\{do\}\(a\_\{t:t\+k\-1\}=\\alpha\)\),\(2\)for alternative action sequencesα\\alpha\. A representation useful for planning should preserve distinctions among such futures when those distinctions change control\-relevant outcomes\.

#### Counterfactual collapse\.

LetΦθ​\(ht,α\)\\Phi\_\{\\theta\}\(h\_\{t\},\\alpha\)denote the model’s imagined latent summary after rolling out action sequenceα\\alphafrom historyhth\_\{t\}, and lety⁡\(ht,α\)y\(h\_\{t\},\\alpha\)denote an intervention outcome, such as the hidden object pose, topological cell, contact mode, success predicate, or long\-horizon reward class induced byα\\alpha\. A representation collapses counterfactuals when

dz​\(Φθ​\(ht,α\),Φθ​\(ht,α′\)\)≈0butdy​\(y⁡\(ht,α\),y⁡\(ht,α′\)\)≫0\.d\_\{z\}\(\\Phi\_\{\\theta\}\(h\_\{t\},\\alpha\),\\Phi\_\{\\theta\}\(h\_\{t\},\\alpha^\{\\prime\}\)\)\\approx 0\\quad\\text\{but\}\\quad d\_\{y\}\(y\(h\_\{t\},\\alpha\),y\(h\_\{t\},\\alpha^\{\\prime\}\)\)\\gg 0\.\(3\)This failure can occur even when the model predicts plausible observations, because visual similarity and intervention equivalence are not the same relation\.

## 4Counterfactual Latent World Models

CLWMaugments latent world\-model learning with an objective that preserves intervention\-sensitive structure\. The architecture follows a recurrent state\-space model but changes the training signal applied to action\-conditioned imagined futures\.

### 4\.1Belief\-State Encoder

Given multimodal observationsot=\(xt,qt,ut\)o\_\{t\}=\(x\_\{t\},q\_\{t\},u\_\{t\}\), wherextx\_\{t\}may be image observations,qtq\_\{t\}proprioceptive state, andutu\_\{t\}optional tactile or force signals, the encoder computes a deterministic recurrent belief

bt=fθ​\(bt−1,at−1,eθ​\(ot\)\),b\_\{t\}=f\_\{\\theta\}\(b\_\{t\-1\},a\_\{t\-1\},e\_\{\\theta\}\(o\_\{t\}\)\),\(4\)whereeθe\_\{\\theta\}is a modality\-specific observation encoder\. A stochastic latent state is then inferred by

qθ​\(zt∣bt,ot\)=𝒩⁡\(μθ​\(bt,ot\),Σθ​\(bt,ot\)\)\.q\_\{\\theta\}\(z\_\{t\}\\mid b\_\{t\},o\_\{t\}\)=\\mathcal\{N\}\(\\mu\_\{\\theta\}\(b\_\{t\},o\_\{t\}\),\\Sigma\_\{\\theta\}\(b\_\{t\},o\_\{t\}\)\)\.\(5\)The transition prior is action\-conditioned:

pθ​\(zt\+1∣bt,zt,at\)=𝒩⁡\(μθp​\(bt,zt,at\),Σθp​\(bt,zt,at\)\)\.p\_\{\\theta\}\(z\_\{t\+1\}\\mid b\_\{t\},z\_\{t\},a\_\{t\}\)=\\mathcal\{N\}\(\\mu^\{p\}\_\{\\theta\}\(b\_\{t\},z\_\{t\},a\_\{t\}\),\\Sigma^\{p\}\_\{\\theta\}\(b\_\{t\},z\_\{t\},a\_\{t\}\)\)\.\(6\)This recurrent belief is crucial: under partial observability, the current frame may not reveal whether a target object is behind an occluder, whether a corridor has been visited, or whether a contact event has changed hidden object pose\.

### 4\.2Action\-Conditioned Counterfactual Rollouts

For each anchor historyhth\_\{t\}, we sample a set of candidate intervention sequences

𝒜tcf=\{αi=\(ati,…,at\+K−1i\)\}i=1M,\\mathcal\{A\}^\{\\mathrm\{cf\}\}\_\{t\}=\\\{\\alpha\_\{i\}=\(a\_\{t\}^\{i\},\\ldots,a\_\{t\+K\-1\}^\{i\}\)\\\}\_\{i=1\}^\{M\},\(7\)drawn from random shooting, policy proposals, model\-predictive control candidates, or a replay buffer\. The world model rolls out each intervention in latent space:

zt\+k\+1i∼pθ\(zt\+k\+1∣bt\+ki,zt\+ki,at\+ki\),k=0,…,K−1\.z\_\{t\+k\+1\}^\{i\}\\sim p\_\{\\theta\}\(z\_\{t\+k\+1\}\\mid b\_\{t\+k\}^\{i\},z\_\{t\+k\}^\{i\},a\_\{t\+k\}^\{i\}\),\\quad k=0,\\ldots,K\-1\.\(8\)Equation \([8](https://arxiv.org/html/2609.05834#S4.E8)\) conditions on an imagined beliefbt\+kib^\{i\}\_\{t\+k\}, which Eq\. \([4](https://arxiv.org/html/2609.05834#S4.E4)\) cannot produce directly because no observation exists during imagination\. We therefore advance the belief open loop, feeding the sampled prior latent in place of the encoder features, as in latent imagination for recurrent state\-space models\([Hafner et al\., 2020](https://arxiv.org/html/2609.05834#bib.bib3)\):

bt\+k\+1i=fθ​\(bt\+ki,at\+ki,zt\+k\+1i\),bti=bt\.b^\{i\}\_\{t\+k\+1\}=f\_\{\\theta\}\\big\(b^\{i\}\_\{t\+k\},\\,a^\{i\}\_\{t\+k\},\\,z^\{i\}\_\{t\+k\+1\}\\big\),\\qquad b^\{i\}\_\{t\}=b\_\{t\}\.\(9\)Imagined rollouts are therefore fully open loop: no ground\-truth future observation is posterior\-fed into the rollout\.

A projection headgθg\_\{\\theta\}maps the rollout to a counterfactual embedding

ci=gθ\(zt\+1:t\+Ki,bt,αi\)\.c\_\{i\}=g\_\{\\theta\}\(z\_\{t\+1:t\+K\}^\{i\},b\_\{t\},\\alpha\_\{i\}\)\.\(10\)The embedding is not used as a decoder feature at test time\. It is a training\-time device that pressures the latent dynamics to preserve intervention distinctions\.

### 4\.3Counterfactual Contrastive Objective

We construct positive and negative pairs based on intervention outcomes\. In simulation, the outcome label can be obtained from privileged state during training only\. In real\-world settings, it can be approximated by task predicates, temporal consistency, object trackers, tactile events, or sparse outcome labels\. For each rolloutii, letyiy\_\{i\}be an outcome descriptor\. Positive pairs share an outcome class or are close under an outcome distancedyd\_\{y\}\. Negatives differ in outcome despite originating from the same or similar belief state\. The contrastive loss takes the supervised contrastive form of[Khosla et al\. \(2020\)](https://arxiv.org/html/2609.05834#bib.bib12):

ℒcf=−∑ilog∑j∈𝒫⁡\(i\)exp⁡\(sim⁡\(ci,cj\)/τ\)∑j∈𝒫⁡\(i\)∪𝒩⁡\(i\)exp⁡\(sim⁡\(ci,cj\)/τ\),\\mathcal\{L\}\_\{\\mathrm\{cf\}\}=\-\\sum\_\{i\}\\log\\frac\{\\sum\_\{j\\in\\mathcal\{P\}\(i\)\}\\exp\(\\mathrm\{sim\}\(c\_\{i\},c\_\{j\}\)/\\tau\)\}\{\\sum\_\{j\\in\\mathcal\{P\}\(i\)\\cup\\mathcal\{N\}\(i\)\}\\exp\(\\mathrm\{sim\}\(c\_\{i\},c\_\{j\}\)/\\tau\)\},\(11\)where𝒫⁡\(i\)\\mathcal\{P\}\(i\)and𝒩⁡\(i\)\\mathcal\{N\}\(i\)are positive and negative rollout sets andτ\\tauis a temperature\. The negatives are not arbitrary trajectories but*hard counterfactual negatives*: interventions whose observations may be similar while their downstream outcomes differ\. We use three negative\-mining strategies\.Same\-belief action negativescompare alternative actions from the same inferred belief state\.Perceptual\-alias negativescompare rollouts with high image\-level similarity but different hidden\-state outcomes\.Value\-disagreement negativescompare futures whose predicted values diverge while their reconstruction loss remains low\. Together, these prevent the model from using visual plausibility as a substitute for intervention structure\.

### 4\.4Full Training Objective

The full objective combines standard latent world\-model learning with counterfactual structure preservation:

ℒCLWM=ℒpred\+βKL​ℒKL\+βr​ℒr\+βv​ℒv\+λcf​ℒcf\+λinv​ℒinv\.\\mathcal\{L\}\_\{\\mathrm\{CLWM\}\}=\\mathcal\{L\}\_\{\\mathrm\{pred\}\}\+\\beta\_\{\\mathrm\{KL\}\}\\mathcal\{L\}\_\{\\mathrm\{KL\}\}\+\\beta\_\{r\}\\mathcal\{L\}\_\{r\}\+\\beta\_\{v\}\\mathcal\{L\}\_\{v\}\+\\lambda\_\{\\mathrm\{cf\}\}\\mathcal\{L\}\_\{\\mathrm\{cf\}\}\+\\lambda\_\{\\mathrm\{inv\}\}\\mathcal\{L\}\_\{\\mathrm\{inv\}\}\.\(12\)ℒpred\\mathcal\{L\}\_\{\\mathrm\{pred\}\}is an observation reconstruction or representation\-prediction loss,ℒKL\\mathcal\{L\}\_\{\\mathrm\{KL\}\}regularizes posterior and prior latent states,ℒr\\mathcal\{L\}\_\{r\}andℒv\\mathcal\{L\}\_\{v\}train reward and value heads, andℒinv\\mathcal\{L\}\_\{\\mathrm\{inv\}\}is an optional intervention\-outcome prediction loss:

ℒinv=𝔼\(ht,α,y\)​\[−log⁡pθ​\(y∣Φθ​\(ht,α\)\)\]\.\\mathcal\{L\}\_\{\\mathrm\{inv\}\}=\\mathbb\{E\}\_\{\(h\_\{t\},\\alpha,y\)\}\[\-\\log p\_\{\\theta\}\(y\\mid\\Phi\_\{\\theta\}\(h\_\{t\},\\alpha\)\)\]\.\(13\)The outcome head stabilizes the contrastive objective by ensuring that separated futures correspond to interpretable intervention consequences rather than arbitrary representation spread\.

### 4\.5Planning with CLWM

At test time,CLWMcan be used with any latent\-space planner\. We use model\-predictive control withNNcandidate action sequences and horizonHH, scoring each by predicted reward plus terminal value,J⁡\(α\)=∑k=0H−1γk​r^θ​\(zt\+k,at\+k\)\+γH​V^θ​\(zt\+H\)J\(\\alpha\)=\\sum\_\{k=0\}^\{H\-1\}\\gamma^\{k\}\\hat\{r\}\_\{\\theta\}\(z\_\{t\+k\},a\_\{t\+k\}\)\+\\gamma^\{H\}\\hat\{V\}\_\{\\theta\}\(z\_\{t\+H\}\), executing the first action of the best sequence, and replanning\. BecauseCLWMshapes the latent dynamics during training, no contrastive computation is required at deployment\.

#### Why prediction alone is not enough\.

A sufficiently accurate predictive model would preserve counterfactual structure in the limit of perfect state estimation, full observability, and unlimited capacity—but not under the conditions embodied world models are actually trained in: finite latents, short rollouts, and evaluation by downstream planning rather than observation likelihood\. Under compression, the equivalence relation induced by prediction loss can differ from the one induced by control\. This observation is the partial\-observability, intervention\-conditioned form of the value equivalence principle\([Grimm et al\., 2020](https://arxiv.org/html/2609.05834#bib.bib7);[Schrittwieser et al\., 2020](https://arxiv.org/html/2609.05834#bib.bib8)\)and of the behavioral\-equivalence view underlying bisimulation\([Ferns et al\., 2004](https://arxiv.org/html/2609.05834#bib.bib9);[Gelada et al\., 2019](https://arxiv.org/html/2609.05834#bib.bib10);[Zhang et al\., 2021](https://arxiv.org/html/2609.05834#bib.bib11)\)\. WhatCLWMadds is a training signal that applies this pressure to imagined open\-loop futures under alternative interventions from a belief state, where on\-policy value or bisimulation targets are unavailable or degenerate\. The learned latent need not recover the full hidden state—only enough to answer the counterfactual queries planning uses\.

## 5Experimental Design

Our experiments test the central claim: world models trained only for plausible prediction can fail under partial observability because they collapse action\-contingent futures that are visually similar but behaviorally distinct\. We therefore evaluate not only task performance, but whether the learned latent dynamics preserve intervention structure, and whether gains arise precisely where counterfactual distinctions are necessary for planning\.

### 5\.1Benchmarks

We evaluate on three partially observable embodied benchmarks in which short\-horizon observations are insufficient to identify the control\-relevant state\.Occluded Push: a tabletop manipulator must push a target object into a goal region while the target may be hidden behind an occluder or displaced into visually ambiguous configurations\. The same RGB\-D observation may correspond to a push that contacts the target, misses it, or moves an occluded object into a recoverable pose, so success depends on object permanence and delayed contact reasoning\.Aliased Maze: navigation through repeated corridor textures where topologically distinct locations produce nearly identical egocentric observations but require different actions, testing whether the belief state preserves topological counterfactuals\.Deferred Kitchen: a long\-horizon manipulation routine in which early interactions determine whether later subtasks remain feasible, often by changing hidden preconditions without immediate visual evidence\. We evaluate normalized return and model\-exploitation rate, defined as the percentage of failures where the planner selects a trajectory with high predicted value but low realized value\.

### 5\.2Baselines

We compare against five latent world\-model baselines\.RSSM\-Recon: a recurrent state\-space model trained primarily with observation reconstruction and latent consistency\.Reward\-aware RSSM: adds reward and value prediction, testing whether task supervision alone recovers intervention structure\.CPC\-RSSM: adds a generic contrastive predictive objective, separating the effect of contrastive learning from counterfactual negative construction\.Dreamer\-style: an imagination\-based latent world model trained for actor\-critic learning in latent space\.TD\-MPC\-style: a decoder\-free latent control model optimized for value\-guided planning\.

### 5\.3Metrics

We report three classes of metrics\.Task performance: episode success rate on Occluded Push and Aliased Maze, normalized return on Deferred Kitchen\.Counterfactual separability\(CFS\) measures whether latent rollouts separate intervention outcomes:

CFS=AUROC\(dz\(Φ\(h,α\),Φ\(h,α′\)\),𝟙\[y\(h,α\)≠y\(h,α′\)\]\),\\mathrm\{CFS\}=\\mathrm\{AUROC\}\\left\(d\_\{z\}\(\\Phi\(h,\\alpha\),\\Phi\(h,\\alpha^\{\\prime\}\)\),\\mathbb\{1\}\[y\(h,\\alpha\)\\neq y\(h,\\alpha^\{\\prime\}\)\]\\right\),\(14\)whereΦ⁡\(h,α\)\\Phi\(h,\\alpha\)is the imagined latent future induced by intervention sequenceα\\alphafrom historyhh, andy⁡\(h,α\)y\(h,\\alpha\)is the corresponding intervention outcome\. CFS is computed on the imagined dynamics latentsΦ\\Phiwith the projection headgθg\_\{\\theta\}discarded, so it measures the geometry the planner actually uses rather than the embedding the contrastive loss directly shapes\. The projection head is a training\-time device only\. Outcome labelsyyat evaluation come from privileged simulator state, exactly as during training\.Model\-exploitation ratemeasures planning failures caused by overoptimistic imagined trajectories: a model can produce high predicted values while remaining unsafe for planning if it collapses distinct action\-conditioned futures\.

## 6Results

### 6\.1CLWM Improves Planning in the Regime Where Prediction Is Ambiguous

Table[1](https://arxiv.org/html/2609.05834#S6.T1)shows the main results\.CLWMachieves the best performance on every benchmark and every diagnostic, gaining9\.59\.5points on Occluded Push and11\.611\.6on Aliased Maze over the strongest non\-counterfactual baseline, raising Deferred Kitchen return from62\.862\.8to72\.372\.3, and cutting model exploitation from18\.4%18\.4\\%to9\.7%9\.7\\%\. The pattern is consistent with the paper’s hypothesis: the largest improvements occur in tasks where visually plausible prediction is not enough to support correct action selection\.

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±\\pmdenotes 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 intensity0\.00\.0, the benchmark’s default layout \(details in Appendix[A](https://arxiv.org/html/2609.05834#A1)\)\.Two comparisons are especially informative\. Reward\-aware RSSM improves over RSSM\-Recon but remains far belowCLWM: reward supervision is sparse and trajectory\-local, and does not by itself separate alternative futures before reward is observed\. CPC\-RSSM improves CFS relative to RSSM\-Recon but still underperformsCLWM: generic contrastive prediction helps representation learning, but does not target the hard cases that matter for embodied control—same\-belief interventions and perceptual aliases with different downstream outcomes\. The exploitation reduction is also central: TD\-MPC\-style planning achieves the strongest baseline return on Deferred Kitchen yet still fails exploitatively in18\.4%18\.4\\%of episodes\.CLWMnearly halves this rate, changing not just average performance but the failure modes the planner encounters\.

### 6\.2Counterfactual Separability Tracks Downstream Control

The main results reveal a strong alignment between CFS and planning performance: methods with low CFS also exhibit lower success and higher exploitation\. Two caveats govern how much this alignment can show\. BecauseCLWMdirectly optimizes an objective aligned with CFS, its own point in Figure[1](https://arxiv.org/html/2609.05834#S6.F1)cannot serve as independent validation of the metric\. The informative evidence is the trend across the five baselines, none of which is trained on CFS\. Across those baselines, the correlation between CFS and planning success isr=0\.94r=0\.94on Occluded Push andr=0\.95r=0\.95on Aliased Maze\. This relationship is consistent with CFS capturing a property that reconstruction loss and return alone can obscure: whether imagined latent distance reflects meaningful differences in intervention outcome\. We do not claim that CFS alone characterizes all aspects of world\-model quality\.

0\.550\.550\.60\.60\.650\.650\.70\.70\.750\.750\.80\.80\.850\.850\.90\.9404060608080Counterfactual separability \(AUROC\)Planning success \(%\)Occluded PushAliased Mazetrend \(baselines\)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 isr=0\.94r=0\.94\(Occluded Push\) andr=0\.95r=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\.
### 6\.3Ablations Show That the Gains Come from Counterfactual Structure

Table[2](https://arxiv.org/html/2609.05834#S6.T2)isolates the contribution of each component of the counterfactual objective\. The bottom row repeats the strongest non\-counterfactual baseline from Table[1](https://arxiv.org/html/2609.05834#S6.T1)as a reference\. The gap between fullCLWMand this reference \(9\.59\.5points on Occluded Push,11\.611\.6on Aliased Maze\) is the quantity the ablations decompose\. The most revealing ablation removes perceptual\-alias negatives: success drops to67\.5%67\.5\\%and69\.1%69\.1\\%, with average CFS falling from0\.850\.85to0\.740\.74—precisely the failure mode the method is designed to address\. When the model no longer receives pressure to distinguish futures that look similar but behave differently, it again collapses the cases that matter most for planning\.

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 forCLWMand the reference in Table[1](https://arxiv.org/html/2609.05834#S6.T1)\), and exploitation rate on Deferred Kitchen\. The bottom row is not an ablation: it repeats the strongest non\-counterfactual baseline from Table[1](https://arxiv.org/html/2609.05834#S6.T1)for reference\.Same\-belief and value\-disagreement negatives contribute complementary benefits: removing the former weakens separation of alternative actions from the same inferred state, while removing the latter raises exploitation to12\.9%12\.9\\%\. Random negatives are insufficient—they improve over the reference but remain far below fullCLWM—showing that negative construction quality is central to the method\.

The outcome head contributes least: removingℒinv\\mathcal\{L\}\_\{\\mathrm\{inv\}\}costs3\.83\.8and5\.45\.4points while remaining well above the reference\. Privileged outcome labels enterCLWMboth throughℒinv\\mathcal\{L\}\_\{\\mathrm\{inv\}\}and through the construction of contrastive pairs\. The small cost of removingℒinv\\mathcal\{L\}\_\{\\mathrm\{inv\}\}indicates that their value is realized mainly through the contrastive geometry rather than through direct outcome prediction\. We caution that Table[2](https://arxiv.org/html/2609.05834#S6.T2)does not include a variant that retains the outcome head without the contrastive loss \(λcf=0\\lambda\_\{\\mathrm\{cf\}\}=0\), nor a baseline given the same labels \(e\.g\., TD\-MPC\-style with an outcome head\)\. These supervision\-matched controls are the right way to fully separate label access from representation geometry, and we leave them for the extended version\.

### 6\.4Robustness under Perceptual Aliasing

We next stress\-test the models by increasing perceptual aliasing intensity in Aliased Maze: as aliasing increases, local observations become less informative about the agent’s true topological state, directly testing whether the model relies on surface\-level observation prediction or preserves history\-dependent intervention structure\.

000\.20\.20\.40\.40\.60\.60\.80\.811404060608080Perceptual aliasing intensitySuccess rate \(%\)RSSM\-ReconDreamer\-styleTD\-MPC\-styleCLWMFigure 2:Robustness to perceptual aliasing\. All results use the same final checkpoints and evaluation protocol as Table[1](https://arxiv.org/html/2609.05834#S6.T1)\(100 test episodes, default maze layout\)\. Intensity0\.00\.0corresponds to the Table[1](https://arxiv.org/html/2609.05834#S6.T1)setting\. Points are means over 5 seeds\. All methods degrade as local observations become less informative\.CLWMremains 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\.Figure[2](https://arxiv.org/html/2609.05834#S6.F2)shows that all methods degrade as aliasing increases, butCLWMremains the strongest at every intensity: at the highest aliasing level it maintains66%66\\%success, compared with55%55\\%for TD\-MPC\-style planning and38%38\\%for RSSM\-Recon\. Its margin over the strongest baseline is essentially undiminished at maximum aliasing \(11\.611\.6points at intensity0\.00\.0,11\.011\.0at intensity1\.01\.0\), and it retains the largest fraction of its default\-setting performance \(84%84\\%, versus82%82\\%,76%76\\%, and73%73\\%for TD\-MPC\-style, Dreamer\-style, and RSSM\-Recon\)\. The gain is therefore not an artifact of the easy regime: it persists as local observations become uninformative and prediction alone is least sufficient\.

## 7Discussion

The results support the paper’s central claim: embodied world models should be evaluated by whether they preserve the consequences of possible actions, not only by whether they predict plausible observations\. The gains arise in settings where these criteria diverge—occlusion, perceptual aliasing, hidden contact, and delayed feasibility—and persist essentially undiminished as aliasing intensifies\. Neither task supervision nor generic contrastive learning suffices on its own, and the outcome\-head ablation suggests the privileged labels’ value is realized through the contrastive geometry rather than direct outcome prediction—though fully supervision\-matched controls remain future work\. The exploitation results clarify why this matters for control: if a latent model collapses two futures that differ in hidden feasibility, a planner can assign high value to an invalid action sequence\. By reducing exploitative failures from18\.4%18\.4\\%to9\.7%9\.7\\%,CLWMchanges not only average performance but the way the model fails\. We do not claim thatCLWMlearns a structural causal model or guarantees identifiability\. The claim is practical: a world model for embodied planning should be trained to preserve the intervention distinctions its planner needs\.

Finally, we note an implication for agents built on pretrained representations\. Our experiments train latent dynamics from scratch, but nothing in the definition of counterfactual collapse depends on that choice\. The CFS metric requires only a way to embed action\-conditioned futures and a source of intervention\-outcome labels\. It can be computed for a frozen pretrained encoder, a fine\-tuned one, or a policy’s internal features\. Given that Figure[1](https://arxiv.org/html/2609.05834#S6.F1)shows CFS tracking planning success across baseline model classes, we suggest measuring it before a pretrained representation is handed to a planner, rather than inferring actionability from downstream reward alone: reward is a slow and confounded signal for this property, and a representation can produce adequate average return while remaining exploitable in precisely the ambiguous states where an agent most needs a reliable model\.

## 8Limitations and Broader Impact

CLWMrequires reliable intervention\-outcome signals: available from privileged state in simulation, but requiring trackers, tactile events, learned predicates, or human labels in the real world\. It also increases training cost by sampling multiple interventions per belief state and mining hard negatives, though this can be reduced by reusing planner candidates or applying the objective only in ambiguous states\.CLWMdoes not solve long\-horizon compounding error, exploration, distribution shift, or reward misspecification, and should be viewed as complementary to uncertainty\-aware model\-based RL\. Our claim about pretrained representations is conceptual rather than empirical: we characterize why counterfactual collapse should be expected under standard pretraining objectives and show that CFS is computable for any representation, but we do not measure CFS on large\-scale pretrained encoders here\. We regard that as the natural next experiment\. Configuration details are collected in Appendix[A](https://arxiv.org/html/2609.05834#A1)\.

World models that better preserve action consequences could improve embodied agents in assistive robotics, manipulation, and navigation, and may reduce overconfident planning through visually plausible but physically invalid futures\. At the same time, more capable embodied planning increases the autonomy of physical systems, which requires conservative deployment practices: uncertainty estimation, human override, environment\-specific validation, and clear limits on autonomous authority\.

## 9Conclusion

World models are usually trained to predict what is likely to happen next\. Embodied agents need something stronger: models that preserve what would happen if they acted differently\.CLWMcombines recurrent belief inference, action\-conditioned latent dynamics, and contrastive supervision over intervention outcomes to target counterfactual collapse—the failure to distinguish behaviorally different futures that look similar in observation space\. Empirically,CLWMimproves planning on all three partially observable benchmarks \(by up to11\.611\.6success points\), increases counterfactual separability, and nearly halves model\-exploitation failures, with ablations attributing the gains to hard perceptual\-alias negatives and the advantage persisting essentially undiminished under maximum perceptual aliasing\.

These gains also come at no deployment cost\. The counterfactual machinery—intervention sampling, hard\-negative mining, the projection head—exists only at training time, so the planner pays nothing at inference for the improved latent geometry\. The objective is likewise orthogonal to the control loop: it attaches to any latent world model that rolls out action\-conditioned futures, which makes it a drop\-in training\-time addition to existing world\-model stacks rather than a competing architecture\.

Looking forward, the most portable piece of this work may be the diagnostic rather than the training objective\. CFS requires only imagined action\-conditioned futures and intervention\-outcome labels, so it can audit any inherited representation—frozen, fine\-tuned, or pretrained at scale—before a planner is asked to trust it, and our results provide an initial calibration for that audit: across world\-model classes trained from scratch, counterfactual separability tracked planning success atr≥0\.94r\\geq 0\.94\. A world model useful for planning must preserve not just what is likely to happen, but the intervention consequences of possible actions—and that property can, and should, be measured before it is assumed\.

## References

- Ha and Schmidhuber \[2018\]David Ha and Jürgen Schmidhuber\.Recurrent world models facilitate policy evolution\.In*Advances in Neural Information Processing Systems*, 2018\.
- Hafner et al\. \[2019\]Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson\.Learning latent dynamics for planning from pixels\.In*International Conference on Machine Learning*, 2019\.
- Hafner et al\. \[2020\]Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi\.Dream to control: Learning behaviors by latent imagination\.In*International Conference on Learning Representations*, 2020\.
- Hafner et al\. \[2023\]Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap\.Mastering diverse domains through world models\.arXiv preprint arXiv:2301\.04104, 2023\.
- Hansen et al\. \[2024\]Nicklas Hansen, Hao Su, and Xiaolong Wang\.TD\-MPC2: Scalable, robust world models for continuous control\.In*International Conference on Learning Representations*, 2024\.
- Wu et al\. \[2022\]Philipp Wu, Alejandro Escontrela, Danijar Hafner, Ken Goldberg, and Pieter Abbeel\.DayDreamer: World models for physical robot learning\.arXiv preprint arXiv:2206\.14176, 2022\.
- Grimm et al\. \[2020\]Christopher Grimm, André Barreto, Satinder Singh, and David Silver\.The value equivalence principle for model\-based reinforcement learning\.In*Advances in Neural Information Processing Systems*, 2020\.
- Schrittwieser et al\. \[2020\]Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy Lillicrap, and David Silver\.Mastering Atari, Go, chess and shogi by planning with a learned model\.*Nature*, 588\(7839\):604–609, 2020\.
- Ferns et al\. \[2004\]Norm Ferns, Prakash Panangaden, and Doina Precup\.Metrics for finite Markov decision processes\.In*Uncertainty in Artificial Intelligence*, 2004\.
- Gelada et al\. \[2019\]Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G\. Bellemare\.DeepMDP: Learning continuous latent space models for representation learning\.In*International Conference on Machine Learning*, 2019\.
- Zhang et al\. \[2021\]Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine\.Learning invariant representations for reinforcement learning without reconstruction\.In*International Conference on Learning Representations*, 2021\.
- Khosla et al\. \[2020\]Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan\.Supervised contrastive learning\.In*Advances in Neural Information Processing Systems*, 2020\.
- Janner et al\. \[2019\]Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine\.When to trust your model: Model\-based policy optimization\.In*Advances in Neural Information Processing Systems*, 2019\.
- Pearl \[2009\]Judea Pearl\.*Causality: Models, Reasoning, and Inference*\.Cambridge University Press, 2nd edition, 2009\.
- Peters et al\. \[2017\]Jonas Peters, Dominik Janzing, and Bernhard Schölkopf\.*Elements of Causal Inference: Foundations and Learning Algorithms*\.MIT Press, 2017\.
- Li et al\. \[2020\]Minne Li, Mengyue Yang, Furui Liu, Xu Chen, Zhitang Chen, and Jun Wang\.Causal world models by unsupervised deconfounding of physical dynamics\.arXiv preprint arXiv:2012\.14228, 2020\.
- Dasgupta et al\. \[2019\]Ishita Dasgupta, Jane Wang, Silvia Chiappa, Jovana Mitrovic, Pedro Ortega, David Raposo, Edward Hughes, Peter Battaglia, Matthew Botvinick, and Zeb Kurth\-Nelson\.Causal reasoning from meta\-reinforcement learning\.arXiv preprint arXiv:1901\.08162, 2019\.
- van den Oord et al\. \[2018\]Aaron van den Oord, Yazhe Li, and Oriol Vinyals\.Representation learning with contrastive predictive coding\.arXiv preprint arXiv:1807\.03748, 2018\.
- Laskin et al\. \[2020\]Michael Laskin, Aravind Srinivas, and Pieter Abbeel\.CURL: Contrastive unsupervised representations for reinforcement learning\.In*International Conference on Machine Learning*, 2020\.
- Kaelbling et al\. \[1998\]Leslie Pack Kaelbling, Michael L\. Littman, and Anthony R\. Cassandra\.Planning and acting in partially observable stochastic domains\.*Artificial Intelligence*, 101\(1–2\):99–134, 1998\.

## Appendix AExperimental Details

Tables[3](https://arxiv.org/html/2609.05834#A1.T3)and[4](https://arxiv.org/html/2609.05834#A1.T4)list the 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](https://arxiv.org/html/2609.05834#S6.T1)and[2](https://arxiv.org/html/2609.05834#S6.T2)are derived from that table\.

Table 3:Training and evaluation configuration\.Table 4:Implementation details\.