Latent Dynamics · 2026-10-09
CausalDreamer: Learning Predictive World Models with Latent Disentanglement
Prince Jha, Nils Lukas, Kun Zhang, Salem Lahlou
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
CausalDreamer re-encodes a frozen Dreamer 4/MMBench2 video tokenizer latent into four groups factorized by controllability and reward relevance, routes actions only to the controllable groups, and predicts reward from only the reward-relevant groups. The pretrained dynamics model is fine-tuned to predict this factored representation. On MMBench2, it improves planning scores on clean tasks and their manipulated variants.
Why it matters
For world-model planning, the latent should reflect which aspects respond to actions and which determine reward. The authors address the issue that tokenizer latents trained only for reconstruction lack explicit controllability and reward-structure. By adding factorized structure on top of a frozen tokenizer and fine-tuning only the dynamics to match it, the method aims to improve model-predictive control under distractors (e.g., background/object changes) while keeping the pretrained world-model backbone largely intact.
Method
- Keeps the pretrained MMBench2 video tokenizer frozen and re-encodes its latent into four groups organized by controllability and reward relevance.
- Routes the action only to the two controllable groups; predicts reward using only the two reward-relevant groups, while all four groups reconstruct the tokenizer latent.
- Fine-tunes the pretrained dynamics model (same architecture/objective) to predict the factored representation for planning with model-predictive control.
Limitation
The dynamics model is fine-tuned with its original architecture, so every group attends to the action when predicted, and the reward head reads all four groups rather than only the reward-relevant ones. Within the factored representation, the action routing is architectural with no objective ensuring controllable groups capture the effect of the action, and the training objective does not penalize reward information entering reward-irrelevant groups.
Abstract (from arXiv)
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.
Related papers
- Disentangling Task-Relevant Latent Dynamics for Reward Modeling
- RAP-WAM: Learning World-Action Models through Predictive Rehearsal and Action Perturbations
- Learning Intrinsic Dynamics in Latent World Models
- Can World and Action Modeling Benefit Each Other under Unified Causal Training?
- The Planning Limits of Latent World Models
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First public release: 2026-10-08 · Indexed by Pub-AI: 2026-10-09 UTC
CausalDreamer: Learning Predictive World Models with Latent Disentanglement
Preprint posted: 2026-10-08 · Original record ↗
- arxiv
- 2610.12016
Explicit source versions (1)
- v1 · 2026-10-08
- Preprint available
Source date: 2026-10-08 · First detected: 2026-10-25 UTC · Historical baseline
Source evidence (1)
- Official evidence ↗ · atom:published (official metadata)
Abstract · arxiv
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.