CausalDreamer: Learning Predictive World Models with Latent Disentanglement

Prince Jha; Nils Lukas; Kun Zhang; Salem Lahlou

cs.LGOct 8, 2026

CausalDreamer: Learning Predictive World Models with Latent Disentanglement

Authors: Prince Jha, Nils Lukas, Kun Zhang, Salem Lahlou

Organizations: Mohamed bin Zayed University of Artificial Intelligence (MBZUAI) · Carnegie Mellon University

Abstract

World models for control must capture which aspects of the environment respond to the agent's actions and which are relevant to reward. Generative world models such as Dreamer 4 consist of a video tokenizer, which encodes each frame into a latent, and a dynamics model, which is pretrained to predict future latents from past latents and actions. Yet the tokenizer is trained with a reconstruction objective, without action or reward supervision, so its latent provides no explicit mechanism to separate controllable, uncontrollable, reward-relevant, and reward-irrelevant information. We propose \textit{CausalDreamer}, which keeps the tokenizer frozen and re-encodes its latent into a factored representation of four groups along two axes: controllability, where only the two controllable groups receive the action, and reward relevance, learned by predicting the reward from the two reward-relevant groups. The pretrained dynamics model is then fine-tuned to predict the factored representation. We evaluate \textit{CausalDreamer} and the pretrained world model it starts from with model-predictive planning on 20 MMBench2 tasks: 10 clean tasks seen during training and 10 unseen tasks, of which 6 are manipulated variants of clean tasks with a changed background, object, or maze layout, and 4 are new environments. We normalize returns so that a policy taking uniformly random actions scores 0 and an expert scores 1. \textit{CausalDreamer} achieves a 14% higher normalized score than the pretrained world model on the clean tasks (0.199 vs.\ 0.175) and a 25% higher score on the manipulated variants (0.307 vs.\ 0.246), while neither model scores meaningfully above the random policy in the new environments. Additionally, our analysis shows that the factored representation separates reward-irrelevant changes, such as a changed background, from its reward-relevant groups.

Figures & tables

Figure 1: Learning the factored representation. A frozen encoder maps each frame to the tokenizer representation. Four projection heads re-encode it into four groups, organized by controllability and reward relevance; only the two controllable heads receive the action. A reward predictor reads only the two reward-relevant groups, and a residual decoder reconstructs the tokenizer representation from all four groups. Grey components are frozen, and the red boxes are the two training losses.

Figure 2: Evaluation tasks. Top: six manipulated tasks, each a variant of a clean task (shown above it) that differs in the background (colors of the background and floor, with the same embodiment and goal), in the manipulated object, or in the maze layout. Bottom: the four clean tasks without a manipulated variant, and four new environments with their own dynamics and rewards. Clean tasks are part of the training corpus; manipulated tasks and new environments have no training data. Each image shows the first frame of an episode.

	Return					Normalized score
Task	Expert	Random	Pretrained	Matched control	CausalDreamer	Pretrained	Matched control	CausalDreamer
cup-catch	1000.0	9.3	128.3	97.7	179.4	0.120	0.089	0.172
finger-turn-easy	1000.0	217.0	269.1	301.4	402.2	0.067	0.108	0.237
mw-push	1820.1	72.7	79.2	183.5	55.1	0.004	0.063	$-0.010$
ms-push-cube	45.8	1.2	2.1	2.7	4.4	0.019	0.034	0.072
lunarlander-hover	469.2	58.6	84.3	53.9	73.7	0.063	$-0.011$	0.037

Table 1: Clean tasks. Mean episode return of the expert reference, a random policy, the pretrained model, the matched control and CausalDreamer, and the normalized score of each model (Section 4.1 ); higher is better. The best normalized score in each row is in bold.

	Return				Normalized score
Task	Expert	Random	Pretrained	CausalDreamer	Pretrained	CausalDreamer
cup-catch-var1	1000.0	9.3	97.0	182.8	0.088	0.175
finger-turn-easy-var1	1000.0	217.0	284.8	464.6	0.087	0.316
ms-push-banana	45.8	0.9	1.1	10.9	0.005	0.222
og-point-var1	88.4	$-97.7$	$-49.4$	$-53.0$	0.260	0.240
og-point-var2	90.7	$-79.7$	$-13.6$	$-17.9$	0.388	0.363

Table 2: Manipulated tasks and new environments. The first six tasks are manipulated variants of a clean task, changing its background, manipulated object or maze layout; the last four are set in new environments. Neither set of tasks has training data. Columns as in Table 1 .

Panel of Figure 3 — Figure 3: Effect of a background change on the representation. (a) Frames of cup-catch and finger-turn-easy rendered from the same simulator state with the clean background and with the background of the manipulated variant. (b) Share of the change between the two renders that falls into the reward-irrelevant groups of the factored representation, and into the corresponding parts of the tokenizer representation; bars show the mean over 10 seeds, dots individual seeds, and the dashed line an equal split.

Explore similar work

Jun 1, 2026cs.LG

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 AI Control Model-Based Planning Latent World Models

Ask PDF

Jul 21, 2026cs.AI

DWM: Separating World Effects from Actions in Latent World Models

Latent world models underpin much of modern model-based control, yet current action-conditioned formulations supervise the next-latent transition with a single, undifferentiated target, forcing a monolithic learning signal to absorb every source of state change. In real world, however, transitions arise from two heterogeneous sources: an action-driven component induced by the agent, and an action-invariant world effect -- the change that would still occur under a null action, dictated by the environment's intrinsic dynamics (e.g., gravity-driven sliding, inertia, contact rebound, and persistent drift). Fusing them into a single target entangles the two inside the latent transition, prevents the model from attributing observed changes to their underlying causes, and undermines the transferability of the learned dynamics. We introduce DWM (Decomposed World Model), a supervision-level framework that operationalizes this decomposition. DWM augments the predictor of a latent world model with an auxiliary world head, regularized by a normalized world-contrastive objective to be action-invariant, while the original pred head is coupled to it via an orthogonality constraint; together, the two signals induce an explicit additive decomposition of the predicted transition into an action-invariant and a complementary action-driven component, without altering the underlying architecture or inference pipeline. To evaluate DWM under persistent world effects, we construct W-variants of three standard control benchmarks -- PushT-W, Reacher-W, and TwoRoom-W -- each instantiating a distinct action-invariant dynamic. DWM matches strong baselines on the flat counterparts and delivers a mean absolute improvement of 13.1% in CEM planning success across the W-variants.

Yi-Ge Zhang, Tianqi Du, Qi Zhang +1

Peking University

World Model Learning Latent Dynamics Modeling World Model-Based Planning Action-Conditioned World Models Model-Based Planning Latent World Models

Ask PDF

Sep 29, 2026cs.LG

Abductive World Modeling via Causal Representation Learning

The central challenge of world modeling is to learn representations that capture how the world evolves. However, existing world models predominantly represent future states without explicitly capturing the latent causes underlying their evolution, limiting their ability to reason about why and how the world changes. To address this limitation, we propose Abductive World Modeling (AWM), a framework that learns structured causal representations by abductively inferring latent causes from predicted futures. Specifically, we realize AWM through the Hierarchical Abductive State Pyramid (HASP), which organizes the inferred world state into three complementary components - Entity, Dynamic, and Relation - capturing what exists, how it changes, and how entities interact, respectively. By jointly reasoning over the current observation and its predicted future, HASP abductively infers these latent factors and integrates them into a structured state representation for downstream reasoning. To the best of our knowledge, AWM is the first framework to introduce abductive state inference into latent-space world modeling for learning structured representations of world dynamics. Experiments across physical prediction, causal reasoning, and action understanding demonstrate the effectiveness of our approach. Compared with V-JEPA, a state-of-the-art latent-space world model, AWM improves physical prediction AUROC by 10.7%, causal reasoning accuracy by 16.8%, and action Top-1 accuracy by 68.0%.

Ziqi Liu, Songhan Yang, Linfan Zhou +4

BAIZ Team

Causal Representation Learning Abductive Reasoning Causal Reasoning Latent World Models State Representation Learning

Ask PDF