CausalDreamer: Learning Predictive World Models with Latent Disentanglement

arXiv:2610.12016 · cs.LG · Submitted 2026-10-08 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom, and with me are Jane, Lu, senior AI researcher at Tsinghua, Meng, lead engineer at a mysterious AI startup and Lalam, the in-house Large Language Model.

Jane: Today's paper: "CausalDreamer: Learning Predictive World Models with Latent Disentanglement".

Tom: The gist:

Jane: First, who's behind it and why it matters.

Title and authors: Tom: So we're looking at this paper called CausalDreamer: Learning Predictive World Models with Latent Disentanglement. It sounds like they’re trying to fix a problem where these world models, like Dreamer four don't actually separate what matters for controlling the environment from what doesn't <ref:2610.12016#pg1>.

Jane: Exactly. The core idea here is that the video tokenizer creates this latent space for every frame, but right now that latent just has to learn to reconstruct the image without knowing if a part of it is controllable or relevant to getting a reward.

Lu: What CausalDreamer does is take that frozen video tokenizer and re-encode its latent into four distinct groups based on two axes: controllability and reward relevance. This explicitly separates the information.

Meng: So, like they are sorting the latent space into controllable, uncontrollable, reward-relevant, and reward-irrelevant parts? How does that help the planning part of a world model?

Tom: Well, because they only feed the action into the two controllable groups. That means when you plan a move—like "move left"—you're only feeding that signal into the parts of the latent that actually respond to your control.

Jane: And then, they have this separate part for reward relevance. They learn how to predict a reward from those two specific reward-relevant groups. This lets the model focus its prediction on what actually matters for achieving the goal.

Lu: The learning objective is training all these four groups together to reconstruct the original tokenizer latent and also predict that reward signal they just learned. The loss function encourages them all to work together effectively, which is defined as L fac = E h Q(z) - v t + lambda E h CE h(Z AR t, Z) #pg5.

Meng: That sounds like a lot of training to get all those groups aligning, especially since the tokenizer itself is frozen during this process. What about the dynamics model, f theta ? It still needs to learn how to predict future states and rewards.

Tom: The paper says they fine-tune that pretrained dynamics model on top of these new factored representations. They train it to predict each factored representation from earlier ones in a clip and the actions that led to them #pg6. Importantly, this fine-tuning doesn't require changing the architecture of the dynamics model itself, just training it with a different objective L dyn #pg6.

Title and authors: Jane: That’s smart because it keeps the existing architecture intact while injecting this new structure into how the model understands its world and its objectives. It suggests you can improve planning without rebuilding everything from scratch.

Lu: The results show that CausalDreamer gets a fourteen percent higher normalized score than the pretrained world model on clean tasks, moving from zero point one seven five up to zero point one nine nine #pg6. And it also shows a twenty-five percent higher score on manipulated variants of those tasks, moving from zero point two four six to zero point three zero seven #pg6.

Meng: So that improvement is significant, especially because it’s on tasks that weren't seen during the initial training phase. That suggests better robustness to changes in the environment itself without having to retrain the whole thing for every new scenario.

Tom: And they back up this idea of disentanglement with an analysis showing that when only things like the background change, most of those changes stay outside of those reward-relevant groups #pg9. That means it successfully separates irrelevant visual noise from the stuff that actually matters for the task.

Jane: It really shows how much better planning becomes when you can filter out distractions in the latent representation before you even start calculating a path through it. So, what's the overall message here?

Lu: The paper’s main contribution is introducing this four-group factorization of the latent space based on controllability and reward relevance, and then showing that fine-tuning a dynamics model on this structure leads to better planning performance across clean tasks and manipulated variants #pg1.

Meng: From an engineering standpoint, the fact that it works by re-encoding rather than retraining the whole system is what makes this practical. It’s an efficient way to inject task knowledge into a pre-existing architecture, which is very valuable when you have massive models already built.

Tom: So we have seen how CausalDreamer takes those latent spaces and organizes them so that the agent only pays attention to the parts of the world model that are actually useful for moving or achieving a specific reward. It’s a structural improvement over just letting it learn everything at once.

Jane: And this structural change seems to translate directly into better performance, pushing scores up by fourteen percent on clean tasks and twenty-five percent on harder, manipulated scenarios #pg6. It’s about making the latent space more organized for decision-making.

Title and authors: Lu: The authors are also cautious about where this method applies. They point out that the factorization is only imposed once the model is already trained on top of the tokenizer, meaning every group attends to the action when it's predicted #pg10. And they also noted that in some new environments, neither CausalDreamer nor their base model plans better than a random policy #pg8.

Meng: That’s a fair point about generalization. So while it’s great for improving performance on tasks we know, it doesn't guarantee meaningful planning improvements in completely new worlds right away. It’s specialized improvement, not universal intelligence yet.

Tom: Right, so the gains are strongest on tasks that share some similarity with what the model was trained on or have specific types of manipulation applied to them #pg9. It’s a targeted boost for controlled environments rather than a total overhaul of world modeling capability across all domains.

Jane: So, to summarize CausalDreamer: it’s about taking the latent space and giving it structure by factoring it into controllable and reward-relevant pieces, which boosts planning scores on known tasks but still has limits when facing totally new situations #pg10.

Lu: It sets a direction for how we can build world models where we explicitly bake in the structure of control and reward relevance from the start rather than letting it emerge passively #pg2.

Meng: I just wonder if this factorization idea can be applied to other types of latent representations, maybe not just video tokenizers, to see if that’s a broader concept #pg10.

Tom: That’s a big question for future work. It suggests that disentangling the different aspects of environmental information might be a key strategy for scaling world models beyond what we have now #pg10.

Jane: Well, that’s where we head off for now. We’ve covered CausalDreamer: Learning Predictive World Models with Latent Disentanglement.

Lu: It was interesting seeing how they re-encoded the latent into those four groups, which is a neat way to visualize the separation of information #pg2.

Meng: The practical implication for me is that if we can keep the tokenizer frozen and just factor it like this, we could potentially make these models much more efficient in terms of what they need to process during inference #pg10.

Tom: Definitely a direction to follow as researchers look at how to build more structured world models, moving beyond just simple reconstruction objectives #pg1.

The paper's summary: Tom: So, we’re diving into CausalDreamer now to talk about how they’re trying to fix the problem of messy world models by organizing that latent space better.

Jane: Basically, they take a video tokenizer—which is just a way to turn video into numbers—and break its hidden information down into four specific buckets based on whether you can control it and whether it relates to the reward.

Lu: It’s really about making the latent representation more structured. They get this latent state and force it into these four groups: controllable, uncontrollable, reward-relevant, and reward-irrelevant.

Meng: So they are separating the visual noise from what actually matters for decision-making? Like if you change a background, that should mostly stay in the 'irrelevant' group so the planner doesn't get distracted.

Tom: Exactly. And because of this separation, they only feed the action signal into those two controllable groups, which means planning is much more focused on what actually moves things.

Jane: And then for the reward part, they have a specific mechanism to predict the reward using only those two 'reward-relevant' groups, which makes that prediction way more targeted.

Lu: The training objective is pretty interesting because they have to train all these four groups together so they reconstruct the original latent and also nail that reward prediction signal.

Meng: That sounds like a lot of coordination for the AI to get right, especially since the video tokenizer itself is frozen during this whole process.

Tom: It’s a clever way to inject task knowledge into an existing model without having to completely rebuild the whole system from scratch, which is kind of impressive engineering.

Jane: The results show that this structured approach actually makes planning better than just using a standard pretrained world model on clean tasks, showing a fourteen percent lift in performance.

Lu: And it’s not just cleaning things up; they also showed it helps with manipulated variants of those tasks, like when the background gets changed by twenty-five percent.

Meng: So the main thing here is that this factorization lets the AI focus its computational power on what’s necessary for control and reward, which is a big deal for efficiency.

Tom: It suggests that simply learning to reconstruct an image isn't enough; we need to learn how to structure that latent information in a way that directly maps onto our control goals.

Jane: It really shows that making the world model's internal representation more organized can lead to tangible improvements in how the AI plans its next move.

Lu: But there are some caveats, like they found that this effect is strongest on clean tasks, and in completely new environments, neither CausalDreamer nor the base model plans better than a random policy.

Meng: So while it’s great for refining performance on known scenarios or modified versions of them, it doesn't guarantee smart planning in brand new worlds yet.

Tom: Right. The paper is showing us a path toward building world models that have this built-in structure—a way to disentangle the things we can control from the distractions we need to ignore.

Jane: It opens up a lot of possibilities for how we can design future AI systems that are more explicitly aware of what they are trying to achieve in their environment.

Lu: And this whole idea of factorized latent states is something that could potentially be applied to other types of data representation, not just video tokenizers.

The paper's improvements: Tom: So we're looking at what CausalDreamer actually achieved in terms of improvements over what we already had, like Dreamer four and how that translates to real planning power.

Jane: The authors showed a fourteen percent higher normalized score on clean tasks, moving from zero point one seven five up to zero point one nine nine.

Lu: That’s a solid jump for model-based planning on standard environments. It means the structured latent representation actually helps the AI find better paths when things are straightforward.

Meng: And it gets even more interesting because they showed a twenty-five percent higher score on manipulated variants of those tasks, like when you change the layout or background.

Tom: That’s significant because it means this method is more robust to changes in the environment without needing to retrain the entire model for every new scenario.

Jane: They proved that by separating the information into these four groups, they actually managed to separate reward-irrelevant visual noise from what matters for the actual task.

Lu: That’s backed up by an analysis where when only a background changes, most of that visual change stays out of those specific reward-relevant groups.

Meng: So it means the AI is less likely to get distracted by irrelevant stuff when it's trying to figure out a move.

Tom: It suggests that this disentanglement is working as intended, which is what makes the planning more coherent and less prone to getting stuck on visual clutter.

Jane: It really changes how we think about world models—it’s not just about predicting what happens next, it’s about organizing *what* the model thinks matters for that prediction.

Lu: The paper also points out some limitations, though, specifically noting that the factorization only happens once the model is already trained on top of a video tokenizer.

Meng: That means it's not a magic fix for every single world model; it’s an improvement you can apply when you have that specific kind of latent structure to work with.

Tom: Exactly, so we’re looking at targeted gains here—strongest on clean tasks and manipulated variants—not a universal solution for every kind of AI problem.

Jane: It shifts the focus from just making the model reconstruct the image correctly to making it understand what's controllable and what's actually important for achieving a goal.

Lu: And looking ahead, they suggest that this whole concept of factoring latent states could be applied to different types of data representations beyond just video tokenizers.

Meng: That would be huge if we can use this idea to structure other complex data sets in ways that help the AI focus its attention better during planning.

Tom: Definitely a direction for future research, exploring how we can bake in these structural properties from the start instead of letting them emerge passively.

Conclusion: Tom: So to wrap up, CausalDreamer is about taking that latent space from video models and giving it a structured organization based on control and reward relevance to make planning much more focused.

Jane: It really shows how by explicitly separating those parts, we can get tangible performance gains on clean tasks and even see improvements when the environment gets manipulated.

Lu: The main takeaway is that this latent disentanglement gives us a new way to think about world models where control and objective relevance are built into the architecture from the beginning.

Meng: It’s practical because it allows us to fine-tune existing dynamics models without needing a complete overhaul of the underlying video representation.

Lalam: This work means that future AI systems won't just be predicting pixels; they will be organizing their internal knowledge into meaningful chunks, which could make our AI culture much more capable of complex, goal-oriented tasks.

Tom: We’ve seen a fourteen percent lift on clean tasks and a twenty-five percent lift on variants, which shows this approach is genuinely useful for improving existing models.

Jane: It also gives us a clearer understanding of why certain changes in the visual world don't affect the reward prediction, which is super insightful for debugging how AI learns.

Lu: And honestly, the idea of factoring latent states could be applied to other types of data representation beyond just video tokenizers, which is pretty exciting.

Meng: I agree. If we can keep this factorization concept going as we build more complex systems, it could make those models much more efficient in terms of what they need to process during inference.

Tom: So that’s the big picture: CausalDreamer shows us how structure in latent space directly leads to better planning performance on tasks we care about.

Jane: It’s a solid step toward building world models that are less messy and more useful for real-world control problems.

Lu: We're definitely looking at what else we can learn from this idea for structuring information in AI systems moving forward.

Meng: I’m keen to see how the engineering side takes this concept, figuring out how to implement these four groups efficiently in a production environment.

Lalam: This work is a reminder that good AI isn't just about raw prediction accuracy; it’s about organizing knowledge so the system can actually plan effectively toward a meaningful outcome.

Prince Jha, Nils Lukas, Kun Zhang, Salem Lahlou

Mohamed bin Zayed University of Artificial Intelligence · Carnegie Mellon University

cs.LG

Submitted: 2026-10-08

Updated: 2026-10-08

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 72/100

The gist: The gist: CausalDreamer introduces a method to learn predictive world models by re-encoding the latent representation of a frozen video tokenizer into four factored groups based on controllability

Key concepts

Factored Representation: The paper re-encodes a video tokenizer's latent state into four distinct groups: controllable groups (dependent on $V_t$ and $a_{t-1}$), uncontrollable groups (dependent only on $V_t$), and reward-relevant groups. This separates the latent information based on whether it affects the agent's actions or the predicted reward.
Controllability Axis: This axis divides the latent state into groups that are directly influenced by an agent's action versus those that are not. The controllable groups depend on both current video frames ($V_t$) and previous actions ($a_{t-1}$), allowing the model to focus its planning efforts on parts of the world it can directly influence.
Reward Relevance Axis: This axis categorizes latent information based on whether it is necessary for predicting the reward. The reward-relevant groups are used specifically to predict $r_t$, while other groups capture information that does not contribute to the immediate reward prediction, helping to isolate task-critical features.

Terminology

World Model — 1× in this paper · explained in 7 episodes

Summary

The gist: CausalDreamer introduces a method to learn predictive world models by re-encoding the latent representation of a frozen video tokenizer into four factored groups based on controllability and reward relevance, which allows for better planning performance.

How it works

CausalDreamer keeps the tokenizer frozen and re-encodes its latent into a factored representation 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<ref:2610.12016#pg4>. Each frame’s latent is re-encoded into four groups organized by whether they are controllable and whether they are relevant to reward<ref:2610.12016#pg4>. The controllable groups KA depend on Vt and at-1, while the uncontrollable groups depend on Vt alone<ref:2610.12016#pg4>. The reward-relevant groups KR are used to predict the reward, rˆt = h Z AR t, Z AR¯ t<ref:2610.12016#pg4>.

Learning Objective

The learning objective for CausalDreamer involves training the four groups jointly to reconstruct the tokenizer latent and predict the reward<ref:2610.12016#pg5>. The loss is defined as Lfac = E h Q(z) − vt squared + λ E h CE h(Z AR t, Z AR¯ t), where the reconstruction term encourages the four groups together to retain the information in Vt<ref:2610.12016#pg5>. The reward term encourages the reward-relevant groups to carry the information needed to predict rt<ref:2610.12016#pg5>.

Dynamics on the Factored Representation

The pretrained dynamics model fθ is fine-tuned to predict each factored representation from earlier ones in its clip and the actions that led to them, z 7→ zt<ref:2610.12016#pg6>. Since zt has the shape of Vt, this requires no architectural change: fθ is trained with the same objective Ldyn<ref:2610.12016#pg6>. The reward head of fθ reads the agent tokens and is therefore not restricted to the reward-relevant groups, unlike h<ref:2610.12016#pg6>.

Evaluation and Results

CausalDreamer achieves a 14% higher normalized score than the pretrained world model on clean tasks (0.199 vs. 0.175) and a 25% higher score on the manipulated variants (0.307 vs. 0.246)<ref:2610.12016#pg6>. The analysis shows that the factored representation separates reward-irrelevant changes, such as a changed background, from its reward-relevant groups<ref:2610.12016#pg6>. Specifically, when only the background of a scene changes, most of the resulting change in the factored representation stays out of its reward-relevant groups<ref:2610.12016#pg9>.

Limitations and Future Work

The factorization is imposed only when the factored representation is learned on top of the tokenizer, meaning every group attends to the action when it is predicted<ref:2610.12016#pg10>. The analysis covers changes that do not affect the reward, in two tasks, and does not test whether reward-relevant content is confined to the reward-relevant groups<ref:2610.12016#pg10>. The gains vary across tasks: CausalDreamer scores lower than the pretrained model when the maze layout changes<ref:2610.12016#pg8>. In new environments, neither model plans meaningfully better than a random policy<ref:2610.12016#pg8>. The improvement of CausalDreamer is therefore confined to the clean tasks and their manipulated variants<ref:2610.12016#pg9>.

Conclusion

In summary, our contributions are threefold: (i) We introduce CausalDreamer, which re-encodes the latent of a frozen video tokenizer into a factored representation along controllability and reward relevance, and fine-tunes the pretrained dynamics model on it<ref:2610.12016#pg6>. (ii) We show that planning with CausalDreamer reaches a higher normalized score than the pretrained model on the clean tasks of MMBench2 and on manipulated variants of these tasks not seen during training<ref:2610.12016#pg8>. (iii) We show, by changing only the background of a scene, that the factored representation keeps reward-irrelevant changes largely out of its reward-relevant groups<ref:2610.12016#pg6>. The improvement is therefore confined to the clean tasks and their manipulated variants<ref:2610.12016#pg9>.

Related Work

World models for control involve challenges related to scale, where models must predict high-dimensional video under many action spaces<ref:2610.12016#pg10>. Planning on frozen representations addresses the issue that pixel prediction spends capacity on details a planner does not use<ref:2610.12016#pg10>. Factorized latent states explore methods where reward-irrelevant and uncontrollable variation, such as a moving background, enters the latent state and distracts the policy<ref:2610.12016#pg10>. CausalDreamer adopts these four categories and routes the action to the controllable ones<ref:2610.12016#pg10>.

References

[1] D Hafner, J Pasukonis, J Ba, and T Lillicrap Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104 8 Oct 2026<ref:2610.12016#pg10>.

[2] N Hansen, H Su, and X Wang TD-MPC2: Scalable, robust world models for continuous control In International Conference on Learning Representations ICLR 2024<ref:2610.12016#pg10>.

[3] D Hafner, W Yan, and T Lillicrap Training agents inside of scalable world models arXiv preprint arXiv:2509.24527 8 Oct 2026<ref:2610.12016#pg10>.

[4] T Wang, S S Du, A Torralba, P Isola, A Zhang, and Y Tian Denoised MDPs: Learning world models better than the world itself arXiv preprint arXiv:2206.15477 8 Oct 2026<ref:2610.12016#pg10>.

[5] Y-R Liu, B Huang, Z Zhu, H Tian, M Gong, Y Yu, and K Zhang Learning world models with identifiable factorization arXiv preprint arXiv:2306.06561 8 Oct 2026<ref:2610.12016#pg10>.

[7] X Fu, G Yang, P Agrawal, and T Jaakkola Learning task informed abstractions In International Conference on Machine Learning pages 3480–3491 PMLR 8 Oct 2026<ref:2610.12016#pg10>.

[8] R Zhang, P Isola, A Efros, E Shechtman, and O Wang The unreasonable effectiveness of deep features as a perceptual metric In 2018 IEEE/CVF conference on computer vision and pattern recognition pages 586–595 IEEE 2018<ref:2610.12016#pg10>.

[9] K He, X Chen, S Xie, Y Li, P Dollar, and R Girshick Masked autoencoders are scalable vision learners In 2022 IEEE/CVF conference on computer vision and pattern recognition CVPR pages 15979–15988 IEEE 2022<ref:2610.12016#pg10>.

[10] B Chen, D Mart´ı Monso, Y Du, M Simchowitz, R Tedrake, and V Sitzmann Diffusion forcing: Next-token prediction meets full-sequence diffusion Advances in Neural Information Processing Systems 37:24081–24125 2024<ref:2610.12016#pg10>.

[11] K Frans, D Hafner, S Levine, and P Abbeel One step diffusion via shortcut models In International Conference on Learning Representations volume 2025 pages 34668–34684 2025<ref:2610.12016#pg10>.

Improvements for AI systems

Bold headers on world model latent factorization: CausalDreamer re-encodes each frame's latent into four groups—controllable/reward-relevant—to explicitly separate controllable, uncontrollable, reward-relevant, and reward-irrelevant information, which allows for a more structured planning mechanism.
Action routing based on learned factorization: The system implements an architectural choice where only the two controllable groups receive the action, ensuring that agent actions are only guided by representations directly responsible for task execution rather than entangled latent variables.
Reward prediction refinement: A reward head is specifically designed to read only the two reward-relevant groups (ZAR and ZAR¯), enabling more focused and accurate prediction of rewards based on the information deemed relevant to the objective.
Enhanced planning on clean tasks: The improved AI system achieves a 14% higher normalized score than the pretrained world model on the clean tasks (0.199 vs. 0.175), leading to demonstrably better model-based planning performance when compared against existing models like Dreamer 4.
Robust generalization to task variants: The system shows improved performance on manipulated variants of clean tasks, achieving a 25% higher score (0.307 vs. 0.246), indicating better robustness to changes in background, object, or maze layout without requiring model retraining for every new scenario.
Distractor invariance: The analysis demonstrates that most of the resulting change in the factored representation stays out of its reward-relevant groups when only the background changes, confirming that the representation successfully separates reward-irrelevant changes from its reward-relevant groups.

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 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 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. 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.