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

Ziyi Ding; Xianxin Lai; Weiyu Chen; Xiao-Ping Zhang; Jiayu Chen

cs.LGApr 9, 2026

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

Authors: Ziyi Ding, Xianxin Lai, Weiyu Chen, Xiao-Ping Zhang, Jiayu Chen

Organizations: Tsinghua Shenzhen International Graduate School, Tsinghua University · The University of Hong Kong · INFIFORCE Intelligent Technology

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.

Figures & tables

Figure 1: Overview of our framework. Top: $o_{t}\!\to\!z_{t}\!\to\!\tilde{z}_{t}\xrightarrow{a_{t}}\hat{z}_{t+1}$ . Bottom: the CausalVAE branch imposes DAG-structured causal constraints in latent space before decoding back to $\tilde{z}_{t}$ .

Figure 2: Three-stage training strategy.

Figure 3: Counterfactual task construction and evaluation pipeline. From factual tuples $(o_{t},a_{t},o_{t+1})$ , we apply interventions $do(s_{t}^{(i,j)}){:=}s_{t}^{(i,j)}+\Delta$ , re-simulate to obtain $o_{t+1}^{cf}$ , and evaluate retrieval on paired factual/counterfactual futures using H@1, MRR, CF-H@1, and CF-MRR (factual H@1/MRR at 1/5/10-step horizons).

		1 Step		5 Steps		10 Steps
Benchmark	Method	H@1 $\uparrow$	MRR $\uparrow$	H@1 $\uparrow$	MRR $\uparrow$	H@1 $\uparrow$	MRR $\uparrow$
Physics 3-body	AE_Contrastive	$89.00$ / $98.00$	$94.25$ / $99.00$	$10.00$ / $9.00$	$27.91$ / $25.72$	$1.00$ / $3.00$	$6.59$ / $9.31$
	AE_NLL	$92.00$ / $94.00$	$95.67$ / $96.83$	$4.00$ / $4.00$	$16.58$ / $16.39$	$0.00$ / $1.00$	$8.93$ / $9.04$
	VAE_Contrastive	$70.00$ / $91.00$	$83.09$ / $95.33$	$1.00$ / $3.00$	$4.96$ / $10.44$	$2.00$ / $2.00$	$5.74$ / $5.73$
	VAE_NLL	$92.00$ / $94.00$	$95.83$ / $97.00$	$1.00$ / $2.00$	$10.19$ / $9.46$	$0.00$ / $0.00$	$3.61$ / $3.64$
	Modular_Contrastive	$90.00$ / $86.00$	$94.42$ / $92.33$	$20.00$ / $8.00$	$36.31$ / $21.23$	$3.00$ / $2.00$	$10.44$ / $6.37$

Table 1: Comparison on factual retrieval metrics. Each cell reports Baseline/+CausalVAE. Bold indicates the best value in each column.

Benchmark	Method	CF-H@1 $\uparrow$	$\Delta$ H@1 $\uparrow$	CF-MRR $\uparrow$	$\Delta$ MRR $\uparrow$
Physics 3-body	AE_Contrastive	10.00/24.00	+14.00	49.33/51.77	+2.44
	AE_NLL	10.50/41.00	+30.50	52.25/66.92	+14.67
	VAE_Contrastive	8.50/8.50	–	46.00/45.40	–
	VAE_NLL	10.50/13.00	+2.50	51.75/51.00	–
	Modular_Contrastive	22.50/25.00	+2.50	59.83/50.62	–
	Modular_NLL	14.50/22.00	+7.50	55.83/57.50	+1.67

Table 2: Comparison on counterfactual retrieval metrics. Each cell reports Baseline/+CausalVAE.

\Delta

indicates the absolute improvement over the baseline.

Figure 4: Comparison between learned structure and first-order physical template. The physical reference matrix is defined as $A_{\mathrm{GT}}=\left|J-I\right|$ from local linearization. The alignment indicates that the model recovers physically meaningful interaction trends in latent space, while remaining a first-order approximation rather than an exact nonlinear law.

Variant	H@1 $\uparrow$	MRR $\uparrow$	CF-H@1 $\uparrow$	CF-MRR $\uparrow$
three-stage training (ours)	91.00	95.16	$31.00$	$52.15$
joint training (w/o stage split)	$90.33$	$94.75$	$28.00$	$51.05$
Baseline (GNN_Contrastive)	$88.00$	$93.37$	$11.50$	52.96
w/o CausalVAE branch	$80.00$	$89.55$	$27.67$	$50.06$
w/o state align loss	$9.00$	$28.02$	$26.67$	$50.57$
single-step rollout	$90.67$	$94.91$	$30.67$	$52.31$

Table 3: Ablations of training strategy, module removal, and hyperparameter sensitivity in the Physics benchmark.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	Default	Scope
Baseline script	train_baselines.py	all baseline / +CausalVAE paired runs
Seed	42	shared within each paired comparison
Optimizer	Adam	shared within each paired comparison
Batch size	1024	train_baselines.py default
Learning rates	slr=lr=transit-lr=finetune-lr =5e-4	train_baselines.py default
Epoch settings	epochs=100 , pretrain-epochs=100 , finetune-epochs=100	train_baselines.py default

Appendix

Table 4: Concrete training defaults used for paired comparisons.

Component	Ke et al . [ 9 ] reference setup	Our implementation defaults
Baseline families	AE, VAE, GNN, Modular	AE, VAE, GNN, Modular (same family split)
Encoder/decoder backbone	Kipf-style CNN + MLP encoder/decoder; medium setting; object-wise latents for structured models [ 12 ] .	encoder=small ; shared world-model backbone in train_baselines.py .
Latent representation	Fixed per-object embedding (reported as 32 per object in reference setup).	embedding-dim-per-object =5, num-objects =5 (domain-adjusted when needed).
Transition parameterization	AE/VAE: MLP transition; GNN: pairwise message passing; Modular: object-wise modular transition.	Model-type-specific transitions selected by --vae , --gnn , --modular .
Training objectives	NLL and contrastive settings; ranking metrics at 1/5/10-step horizons.	NLL (default) and contrastive route ( --contrastive ); factual metrics reported at $h\in\{1,5,10\}$ .
NLL-type optimization budget	Adam, lr $=5\times 10^{-4}$ , batch size 512, 100-epoch settings (reported).	pretrain-epochs =100, epochs =100, finetune-epochs =100.

Appendix

Table 5: Baseline configuration details: reference protocol vs. our implementation defaults.

Model	Latent structure	Transition module	Key implementation defaults
AE	Monolithic latent vector.	MLP transition over latent + action.	hidden-dim =512, embedding-dim-per-object =5, num-objects =5.
VAE	Monolithic latent with posterior $(\mu,\log\sigma^{2})$ .	Same as AE for transition; stochastic latent sampling.	--vae flag; same transition width and action interface as AE.
GNN	Object-wise factored latents.	Graph message passing transition (pairwise interaction bias).	--gnn flag; action-conditioned transition with optional --ignore-action / --copy-action .
Modular	Object-wise factored latents.	Object-wise modular transition (higher-order interaction capacity).	--modular flag; per-object embedding and shared action interface.

Appendix

Table 6: Baseline backbone details used in this work (corresponding to the AE/VAE/GNN/Modular split in Ke et al . [ 9 ] ).

Item	Default value	Where used
Optimizer	Adam	all baseline runs
Base learning rates	slr=lr=transit-lr=finetune-lr =5e-4	baseline NLL and contrastive routes
Batch size	1024	train_baselines.py
Epoch setting A	100	baseline NLL-type training budget
Epoch setting C	100	implementation default retained for reproducibility
Contrastive hinge	1.0 ( --hinge )	decoder-free contrastive route

Appendix

Table 7: Baseline loss/optimization defaults in our implementation.

Benchmark	Query groups (default)	Intervention variables	Notes
Physics (3-body)	Up to 2000 per run.	Object id $i$ , axis (x/y), magnitude $\Delta$ .	Evaluator default includes obj-idx =0, axis =x, delta =3.0; axis/magnitude sweeps follow protocol settings.
2D Shapes / 3D Cubes	Up to 2000 per run.	Object-level pose/push-related factors.	Single- and multi-target variants; intervention index maps to object-slot metadata.
Chemistry	Up to 2000 per run.	Mechanism/factor id and magnitude $\Delta$ .	Mechanism-level intervention metadata when available; otherwise use global structural protocol.

Appendix

Table 8: Detailed construction of counterfactual benchmark queries.

Component	Choice	Variants	Notes
Intervention time	$t_{0}$	sampled over prefix window	consistent across methods
Targets	object / variable	single, multi	physics: pos/vel; pushing: pose/push
Magnitude	$\Delta$	small/med/large	robustness check under intervention strength
Rollout horizon	$H$	short/long	drift stress-test at long $H$
Retrieval set	$N$ , $h$	fixed per benchmark	negatives from other episodes; $h\in\{1,5,10\}$ for factual metrics
Scoring	NLL / similarity	model-dependent	paired comparison protocol

Appendix

Table 9: Evaluation protocol summary.

Explore similar work

Abductive World Modeling via Causal Representation Learning
Sep 29, 2026Ziqi Liu, Songhan Yang, Linfan Zhou +4Causal Representation Learning Abductive Reasoning
BAIZ Team
Do-JEPA: From Masking to Intervention in Latent World Models
Sep 29, 2026Hossein Resani, Javen Qinfeng ShiCausal Effect Estimation Causal World Models
Australian Institute for Machine Learning, Adelaide University · Responsible AI Research Centre
CAFE: Counterfactual Prediction via Fast Posterior Estimation
Sep 26, 2026Xinyan Han, Xiaoyu Lin, Hao Zou +3Counterfactual Prediction Structural Causal Models
Tsinghua University · Stable AI

Sep 29, 2026J/K move · Enter open · S save

Abductive World Modeling via Causal Representation Learning

Ziqi Liu, Songhan Yang, Linfan Zhou, Jiatong Liu+3

BAIZ Team

Causal Representation Learning Abductive Reasoning Causal Reasoning Latent World Models State Representation Learning

Open paper AskPDF arXiv

The central challenge of world modeling is to learn representations that capture how the world evolves. However, existing world models predominantly represent future states without explicitly capturing the latent causes underlying their evolution, limiting their ability to reason about why and how the world changes. To address this limitation, we propose Abductive World Modeling (AWM), a framework that learns structured causal representations by abductively inferring latent causes from predicted futures. Specifically, we realize AWM through the Hierarchical Abductive State Pyramid (HASP), which organizes the inferred world state into three complementary components - Entity, Dynamic, and Relation - capturing what exists, how it changes, and how entities interact, respectively. By jointly reasoning over the current observation and its predicted future, HASP abductively infers these latent factors and integrates them into a structured state representation for downstream reasoning. To the best of our knowledge, AWM is the first framework to introduce abductive state inference into latent-space world modeling for learning structured representations of world dynamics. Experiments across physical prediction, causal reasoning, and action understanding demonstrate the effectiveness of our approach. Compared with V-JEPA, a state-of-the-art latent-space world model, AWM improves physical prediction AUROC by 10.7%, causal reasoning accuracy by 16.8%, and action Top-1 accuracy by 68.0%.

Abstract

Explore similar work

Abductive World Modeling via Causal Representation Learning

Do-JEPA: From Masking to Intervention in Latent World Models

CAFE: Counterfactual Prediction via Fast Posterior Estimation