CausalVAE as a Plug-in for World Models: Towards Reliable Counterfactual Dynamics (ECCV 2026)
Predicting Is Not Yet Understanding
ECCV 2026 · Causal World Models
Ziyi Ding, Xianxin Lai, Weiyu Chen, Xiao-Ping Zhang, Jiayu Chen
Tsinghua Shenzhen International Graduate School · The University of Hong Kong · INFIFORCE Intelligent Technology
A world model should not only say what happens next. It should also know what would happen if something changed.
Latent world models compress images into compact states and roll them forward under actions. That recipe works well for factual prediction: given what actually happened, retrieve or generate the most likely future. But the same latent can be unreliable under an intervention, because features that are useful for prediction are not necessarily separated into causes and effects.
A counterfactual query is more demanding. If we move one object, change its velocity, or alter a mechanism, the model must update the variables downstream of that change while leaving unrelated factors alone. In an entangled latent space, the intervention often becomes a diffuse perturbation: too much changes, too little changes, or the effect travels along the wrong path.
Our question is simple: can we add explicit causal structure to an existing world model without redesigning the predictive backbone?
Attach causal structure where the world model already reasons
We introduce a CausalVAE-inspired structural branch directly in latent space. The original encoder still maps an observation to a latent state, and the original action-conditioned transition still predicts the next latent. Between them, our plug-in organizes latent factors through a learned DAG:
Plug-in design: observation oₜ → encoder E → latent zₜ → CausalVAE → transition F → future ẑₜ₊₁.
The key is modularity. The same structural branch can sit on AE, VAE, modular, or graph-network dynamics backbones. A differentiable acyclicity penalty keeps the learned adjacency close to a DAG, while a lightweight alignment head anchors latent coordinates to simulator state during training. Neither simulator state nor the alignment head is needed at inference.
A three-stage training strategy
01 · Predict
Train the encoder and transition backbone on standard world-model objectives.
02 · Structure
Freeze the backbone and learn the causal branch with reconstruction, KL, alignment, and DAG constraints.
03 · Refine
Freeze the encoder and CausalVAE, then tune the transition with an α-gated mixture and multi-step supervision.
This staged design lets the backbone first acquire predictive competence, gives the structural branch a stable representation to organize, and finally adapts the transition to the causally structured latent without retraining the entire system from scratch.
Intervene, re-simulate, then retrieve the right future
We evaluate four domains—Physics 3-body, Chemistry, 2D Shapes, and 3D Cubes—using state-, object-, and mechanism-level interventions. Each factual transition is paired with a re-simulated counterfactual target.
The evaluation asks a direct question: among candidate futures, does the model retrieve the one produced by the intervention? Factual H@1 and MRR track ordinary prediction, while CF-H@1 and CF-MRR isolate counterfactual retrieval.
The largest gains appear where ordinary dynamics are intervention-fragile
Across four benchmarks and eight paired backbones, the plug-in generally keeps factual retrieval competitive and improves CF-H@1 on 21 of 32 backbone–benchmark pairs. Physics shows the clearest effect: GNN–NLL rises from 11.0 to 41.0 CF-H@1, a gain of 30.0 points, while AE–NLL rises from 10.5 to 41.0, a gain of 30.5 points.
| Physics 3-body backbone | Baseline CF-H@1 |
|
|---|---|---|
| AE · Contrastive | 10.0 | 24.0 |
| AE · NLL | 10.5 | 41.0 |
| VAE · Contrastive | 8.5 | 8.5 |
| VAE · NLL | 10.5 | 13.0 |
| Modular · Contrastive | 22.5 | 25.0 |
| Modular · NLL | 14.5 | 22.0 |
| GNN · Contrastive | 11.5 | 15.0 |
| GNN · NLL | 11.0 | 41.0 |
In the matched Physics comparison, the plug-in adds approximately 1M parameters.
The learned structure is also interpretable at the mechanism-trend level. On Physics, its strongest edges align with a first-order template derived from the locally linearized dynamics. This is evidence of meaningful physical interaction trends—not a claim that a local DAG exactly recovers the full nonlinear law.
Where the plug-in breaks: the host latent still matters
The gains are not universal. Chemistry exposes a capacity bottleneck: some backbones compress 5^5 = 3,125 global color configurations into a five-dimensional code. Passing that code through another structured encode–decode path can erase class-discriminative directions. In this setting, causal regularization has too little representational room to help.
For AE–Contrastive on Chemistry, increasing the latent dimension changes the CF-H@1 gain monotonically:
| Latent dimension | 5 | 16 | 32 | 128 |
|---|---|---|---|---|
| CF-H@1 change | −16.7 | −9.2 | −1.5 | +3.5 |
The practical lesson is specific: use the plug-in when the backbone latent can still carry the relevant causal factors, and validate it per domain. Long-horizon drift and strongly nonlinear regimes remain open problems.
The takeaway
Structure can be added without starting over. Keep the predictive backbone, organize its latent variables, and train the causal bridge carefully. Where ordinary dynamics confuse correlation with mechanism, explicit causal structure can make the alternative future much easier to find.