PaperlayerView exact revision on arXiv

Source-linked AI summary

CausalDreamer: Learning Predictive World Models with Latent Disentanglement

Prince Jha, Nils Lukas, Kun Zhang, Salem Lahlou

arXiv:2610.12016v1cs.LG

TL;DR

Generative world-model tokenizers are trained to reconstruct frames, not to distinguish controllable from uncontrollable or reward-relevant from reward-irrelevant information. CausalDreamer factors frozen tokenizer latents and fine-tunes the dynamics model; it improves normalized scores on clean and manipulated tasks, but not meaningfully on new environments.

  • Problem

    A reconstruction-trained tokenizer does not explicitly separate controllability or reward relevance, leaving these distinctions entangled in the latent used for planning.

    Source: §1 · p. 1

  • Method

    CausalDreamer re-encodes frozen tokenizer latents into four groups, routes actions only to controllable groups, predicts reward from reward-relevant groups, and fine-tunes the dynamics model.

    Source: §1 · p. 2

  • Results

    CausalDreamer scores 0.199 vs. 0.175 on clean tasks and 0.307 vs. 0.246 on manipulated variants, while neither model plans meaningfully above random in new environments.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Most representation change caused by a background-only scene change stays outside the reward-relevant groups.

    Source: §1 · p. 2

    The learning setup assumes a partially observed Markov decision process and uses clips of T consecutive steps as its learning data.

    Source: §3.1 · pp. 3–4

Abstract

from arXiv · show

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.

1 Introduction

CausalDreamer adds controllability and reward-relevance structure to a pretrained generative world model whose tokenizer does not distinguish these factors. It improves planning scores on clean and manipulated tasks, while separating background changes from reward-relevant content.

1 Introduction

  • A reconstruction-trained tokenizer lacks explicit distinctions between controllable and uncontrollable content or reward-relevant and reward-irrelevant content.The dynamics model must discover these distinctions within an entangled latent.Source: §1 · p. 1
  • The method places four groups, organized by controllability and reward relevance, between a frozen tokenizer and the pretrained dynamics model.Only the controllable groups receive actions; reward is predicted from the reward-relevant groups.Source: §1 · p. 2
  • CausalDreamer scores 0.199 vs. 0.175 on clean tasks and 0.307 vs. 0.246 on manipulated variants; neither model plans meaningfully above random in new environments.These are normalized scores, with random actions at 0 and an expert at 1.Source: §1 · p. 2
  • When only the background changes, most representation change stays outside the reward-relevant groups.This indicates separation of reward-irrelevant changes from content used to predict reward.Source: §1 · p. 2

2 Background

The MMBench2 world model uses a frozen video tokenizer and a block-causal dynamics transformer. Its shortcut-forcing objective predicts clean representations from corrupted clips and conditions generation on actions and denoising parameters.

2 Background

  • The MMBench2 model follows Dreamer 4’s recipe: train a video tokenizer first, then train a dynamics model on its frozen representations.The model has 350M parameters.Source: §2 · p. 3
  • The tokenizer encodes each 224 × 224 frame into 64 latent tokens of 64 channels, and its decoder reconstructs the frame from those tokens.It is trained with mean squared error and LPIPS in a masked-autoencoding setup.Source: §2 · p. 3
  • The block-causal dynamics transformer processes frame tokens alongside an action token, signal-level and step-size tokens, spatial tokens, and learned registers.For control, task-conditioned agent tokens predict rewards and actions for the next 8 steps.Source: §2 · p. 3
  • Shortcut forcing trains the model to predict clean representations from corrupted clips, actions, signal levels, and denoising step sizes.At inference, each frame is generated from noise in 4 steps of size 1/4, conditioned on earlier generated frames.Source: §2 · p. 3

3 Method

CausalDreamer inserts a four-group representation between a frozen tokenizer and its dynamics model. The groups organize information by controllability and reward relevance, while reconstruction and reward prediction train the representation before the dynamics model is fine-tuned to predict it.

3.1 Problem Setting

  • The frozen tokenizer’s representation does not explicitly separate controllable from uncontrollable information or reward-relevant from reward-irrelevant information.The method addresses this gap while keeping the tokenizer frozen throughout.Source: §3 · p. 3, §3.1 · p. 4
  • Four groups encode the combinations of controllability and reward relevance; only controllable groups use the action, and reward-relevant groups predict reward.Each projection reads the full tokenizer representation, with the action concatenated only for controllable groups.Source: §3.1 · p. 4, §3.2 · p. 4

3.2 CausalDreamer

  • A residual decoder maps the four groups back to the tokenizer representation, while a reward predictor reads only the reward-relevant groups.The reconstruction term preserves information in the tokenizer representation, and the reward term encourages reward-relevant groups to carry information for predicting reward.Source: §3.2 · p. 4
  • With the projections and decoder frozen, the pretrained dynamics model is fine-tuned to predict factored representations from earlier representations and actions, without architectural changes.Predictions are rendered through the frozen decoder and tokenizer.Source: §3.2 · pp. 4–5
  • Figure 2 separates evaluation tasks into clean training tasks, unseen manipulated variants, and new environments without training data.The six manipulated variants change a background, manipulated object, or maze layout; the new environments have their own dynamics and rewards.Source: Figure 2 · p. 5

4 Experiments

The experiments compare CausalDreamer with its pretrained model on clean tasks, unseen manipulated variants, and new environments, then test whether background changes enter reward-irrelevant representation groups. CausalDreamer scores higher on clean tasks and manipulated variants, while neither model meaningfully outperforms random actions in new environments.

4.1 Experimental Setup

  • Evaluation covers 10 training-seen clean tasks, six unseen variants changing a background, object, or maze layout, and four unseen new environments, using closed-loop planning without task fine-tuning.Scores are normalized against random actions and expert returns; the study also includes a matched control fine-tuned on tokenizer representations.Source: §4.1 · p. 6

4.2 Planning on Clean Tasks; 4.3 Planning on Manipulated Tasks and New Environments

  • Across ten clean tasks, CausalDreamer’s mean normalized score was 0.199 versus 0.175 for the pretrained model; across six manipulated tasks, it was 0.307 versus 0.246; neither model planned meaningfully better than random on four new environments.Source: §4.2 · p. 6, §4.3 · p. 7

4.2 Planning on Clean Tasks

  • The matched control scores 0.167 on its first 10 seeds, below the pretrained model, and CausalDreamer scores higher on six of ten clean tasks.CausalDreamer is lower on four, most clearly on og-point-maze (0.398 vs. 0.524); mw-push is within 0.01 of random for both models.Source: §4.2 · p. 7

4.3 Planning on Manipulated Tasks and New Environments

  • Under changed maze layouts, CausalDreamer scores 0.020 to 0.120 lower on all three tasks, despite gains on background-change and changed-object variants.On cup-catch-var1 and finger-turn-easy-var1, its scores are 0.175 vs. 0.088 and 0.316 vs. 0.087, respectively.Source: §4.3 · p. 7

4.4 Analysis of the Factored Representation

  • With identical simulator states and rewards, background changes place 79% and 74% of representation change in reward-irrelevant groups, versus 44% and 46% in corresponding tokenizer parts.The percentages are for cup-catch and finger-turn-easy, respectively; Figure 3 compares the two representations across 10 seeds.Source: §4.4 · p. 8, Figure 3 · p. 8

5 Conclusion

CausalDreamer adds a four-group representation on top of a frozen tokenizer and fine-tunes the pretrained dynamics model to predict it. Planning improves on clean tasks and manipulated variants, but not in new environments; background-only changes mostly enter reward-irrelevant groups.

5 Conclusion

  • Planning with CausalDreamer scores higher than the pretrained model on clean and manipulated tasks, with the largest gains for background or object changes; neither model plans meaningfully above random in new environments.The improvement does not extend to environments absent from training.Source: §5 · p. 8
  • The method re-encodes frozen-tokenizer latents into four groups organized by controllability and reward relevance, then fine-tunes the pretrained dynamics model to predict them.Only controllable groups receive the action, and only reward-relevant groups predict reward.Source: §5 · p. 8
  • When only the background changes, most of the resulting representation change falls into reward-irrelevant groups rather than reward-relevant groups.Source: §5 · p. 8

A Limitations and Future Work

The limitations concern how far the factorization is enforced, what its representation analysis tests, and how broadly the reported planning gains hold.

A Limitations and Future Work

  • The dynamics model keeps its original architecture, so every group attends to the action during prediction and its reward head reads all four groups.The factorization is imposed only while learning the representation on top of the tokenizer.Source: A Limitations and Future Work · p. 10
  • The action routing is architectural, but the study does not evaluate whether controllable groups capture the action’s effect.Source: A Limitations and Future Work · p. 10
  • Representation analysis tests reward-irrelevant changes in two tasks, but not whether reward-relevant content is confined to reward-relevant groups.The objective encourages reward information in those groups without penalizing it in reward-irrelevant groups; an independence or mutual-information penalty is proposed as a next step.Source: A Limitations and Future Work · p. 10
  • Gains vary by task: CausalDreamer scores lower when the maze layout changes, and reported 20-seed task means have no significance tests.Source: A Limitations and Future Work · p. 10

B Related Work

Related work addresses scaling world models, planning over frozen representations, and factorizing latent states. CausalDreamer retains a pretrained tokenizer-based model while imposing structure based on action and reward relations.

World models for control

  • DreamerV3 and TD-MPC2 scale world-model planning across many tasks, while Dreamer 4 and MMBench2 extend it to video, with MMBench2 using a 350M-parameter model trained on 427 hours across 210 tasks.The present work keeps the MMBench2 model, data, and planner unchanged and changes the latent predicted by its dynamics model.Source: B Related Work · p. 10

Planning on frozen representations

  • Frozen-encoder planners include DINO-WM, V-JEPA 2, and Causal-JEPA; Causal-JEPA’s object-level masking raises Push-T success from 60.67% to 88.67% over an object-centric baseline.Causal-JEPA uses about 1% of the latent features of the patch-based model; CausalDreamer instead structures latents by action and reward responses.Source: B Related Work · pp. 10–11

Factorized latent states

  • Prior factorized-state work separates controllable from uncontrollable dynamics or groups state by controllability and reward relevance, with differing choices about which groups matter for policy learning.IFactor shows that reward-relevant variables include uncontrollable ones that Denoised MDPs discard.Source: B Related Work · p. 11
  • CausalDreamer adopts four categories organized by controllability and reward relevance, imposing them on a pretrained tokenizer latent rather than learning a stochastic state from scratch.Source: B Related Work · p. 11
Loading 2610.12016v1…