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

TL;DR

CausalVAE enhances latent world models with causal structure, improving counterfactual prediction robustness and interpretability across benchmarks, especially in physics simulations.

Contribution

Introduces CausalVAE as a plug-in module that boosts counterfactual retrieval and interpretability in latent world models.

Findings

Significant improvement in CF-H@1 on Physics benchmark (+102.5%)

Enhanced counterfactual prediction robustness under interventions

Learned causal dependencies recover physical interaction trends

Abstract

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.