# CausalVAE as a Plug-in for World Models: Towards Reliable Counterfactual Dynamics

arXiv: https://arxiv.org/abs/2604.07712 · HTML: https://mlsanity.com/paper/2604.07712/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in World models for language agents, one of 11 directions in World models.
- Builds on 16 of its 21 references that are in the map.
- First posted on arXiv 9 Apr 2026.

## Placed in

- World models / [World models for language agents](https://mlsanity.com/wm/world-models-for-language-agents/)

## Builds on (most cited first, 16 of 16)

- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2301.08243/) (2023)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Learning to Simulate Complex Physics with Graph Networks](https://mlsanity.com/paper/2002.09405/) (2020)
- [Object-Centric Learning with Slot Attention](https://mlsanity.com/paper/2006.15055/) (2020)
- [DAGs with NO TEARS: Continuous Optimization for Structure Learning](https://mlsanity.com/paper/1803.01422/) (2018)
- [Revisiting Feature Prediction for Learning Visual Representations from Video](https://mlsanity.com/paper/2404.08471/) (2024)
- [Variational Autoencoders and Nonlinear ICA: A Unifying Framework](https://mlsanity.com/paper/1907.04809/) (2019)
- [MONet: Unsupervised Scene Decomposition and Representation](https://mlsanity.com/paper/1901.11390/) (2019)
- [Multi-Object Representation Learning with Iterative Variational Inference](https://mlsanity.com/paper/1903.00450/) (2019)
- [Neural Relational Inference for Interacting Systems](https://mlsanity.com/paper/1802.04687/) (2018)
- [Systematic Evaluation of Causal Discovery in Visual Model Based Reinforcement Learning](https://mlsanity.com/paper/2107.00848/) (2021)
- [Causal World Models by Unsupervised Deconfounding of Physical Dynamics](https://mlsanity.com/paper/2012.14228/) (2020)
- [Understanding Physical Dynamics with Counterfactual World Modeling](https://mlsanity.com/paper/2312.06721/) (2023)

## Built on by (most cited first, 2 of 2)

- [Latent State Design for World Models under Sufficiency Constraints](https://mlsanity.com/paper/2605.01694/) (2026)
- [Event-Conditioned Diagnostics of Kinematic, Contact, and Object-Permanence Structure in Passive Object-State World Models](https://mlsanity.com/paper/2606.28455/) (2026)

## Abstract (arXiv)

In this work, CausalVAE is introduced as a plug-in structural module for latent world models and is attached to diverse encoder-transition backbones. Across the reported benchmarks, competitive factual prediction is preserved and intervention-aware counterfactual retrieval is improved after the plug-in is added, suggesting stronger robustness under distribution shift and interventions. The largest gains are observed on the Physics benchmark: when averaged over 8 paired baselines, CF-H@1 is improved by +102.5%. In a representative GNN-NLL setting on Physics, CF-H@1 is increased from 11.0 to 41.0 (+272.7%). Through causal analysis, learned structural dependencies are shown to recover meaningful first-order physical interaction trends, supporting the interpretability of the learned latent causal structure.