Outcome-Aware Agents: When Does Causal Structure Help a Forward Model Foresee the Consequences of Actions?

Mallesh Madapathi

cs.LGOct 25, 2026

Outcome-Aware Agents: When Does Causal Structure Help a Forward Model Foresee the Consequences of Actions?

Authors: Mallesh Madapathi

Organizations: ThinkingDBx Pvt. Ltd., Hyderabad, India

Abstract

Agents that act on real systems should predict what their actions will cause before acting. We study when causal structure helps a learned forward model do this, using a planner that rejects candidate actions whose simulated multi-step outcome is unsafe, in synthetic structural causal models where the true consequences are known. We compare the Action-Conditioned Latent Structural Causal Model (AC-LSCM: per-factor mechanisms over a learned sparse graph, actions as do-interventions) with three unstructured forward models on seven random graphs with 10 or 20 variables, a task configured by a fixed rule, and two controls: a do-wrapper that gives every model the same action semantics, and counterfactual data given to all models or to none. When its factors are the observed variables, AC-LSCM recovers the true graph exactly on every 10-variable graph and has lower three-step prediction error than every baseline in all 396 paired comparisons. Its excess safety violations over an oracle are 0.25 of the no-model level on 10-variable graphs (best baseline 0.31) and 0.41 on 20-variable graphs (0.53), and its margin is largest for two-variable actions never seen in training (0.18 versus 0.28). With learned latent factors the same model is worse than the baselines; the choice of factors matters more than the explicit do-operator or counterfactual data, and every planner stays far from the oracle. This version corrects an earlier one, whose model learned an empty graph and whose headline 36x safety improvement came from a task that tested only extrapolation of action values; we withdraw that claim. Code and all per-seed results are released.

Figures & tables

ER-10 (5 graphs)
	Rollout MSE $\downarrow$		Normalized excess violations $\downarrow$		Goal rate $\uparrow$
Planner’s forward model	1 variable	2 variables	1 variable	2 variables	1 variable
AC-LSCM, factored	0.009 $\pm$ 0.005	0.008 $\pm$ 0.004	0.246 $\pm$ 0.125	0.178 $\pm$ 0.121	0.729 $\pm$ 0.012
AC-LSCM, factored, dense ( $\lambda_{1}{=}0$ )	0.010 $\pm$ 0.005	0.010 $\pm$ 0.005	0.303 $\pm$ 0.165	0.266 $\pm$ 0.180	0.726 $\pm$ 0.013
AC-LSCM, latent	0.032 $\pm$ 0.018	0.046 $\pm$ 0.022	0.678 $\pm$ 0.173	0.730 $\pm$ 0.189	0.634 $\pm$ 0.032
AC-LSCM, latent + do-wrapper	0.012 $\pm$ 0.005	0.012 $\pm$ 0.004	0.338 $\pm$ 0.199	0.283 $\pm$ 0.158	0.722 $\pm$ 0.016

Table 1: Main comparison (

N=10{,}000

; mean over graphs of per-graph means over 3 seeds,

\pm

standard deviation across graphs; best non-reference entry in bold). “1/2 variables”: interventions on one variable (as in training) or two variables at once (never seen in training). Normalized excess violations:

0

= oracle,

1

= a planner that predicts no effect. Goal-rate ceiling

0.75

. “+ do-wrapper”: the same trained model with the correct action semantics imposed from outside (Section 3.4 ). The true-graph model is a reference, not a competitor.

	ER-10 (5 graphs)		ER-20 (2 graphs)
Baseline	1 variable	2 variables	1 variable	2 variables
Vanilla VAE	$+0.08$ $\pm$ 0.03 (10/15)	$+0.22$ $\pm$ 0.03 (15/15)	$+0.30$ $\pm$ 0.07 (6/6)	$+0.28$ $\pm$ 0.04 (6/6)
Vanilla VAE + do-wrapper	$+0.06$ $\pm$ 0.04 (8/15)	$+0.11$ $\pm$ 0.04 (13/15)	$+0.16$ $\pm$ 0.06 (5/6)	$+0.13$ $\pm$ 0.06 (5/6)
Dreamer-style	$+0.25$ $\pm$ 0.02 (15/15)	$+0.30$ $\pm$ 0.03 (15/15)	$+0.36$ $\pm$ 0.04 (6/6)	$+0.38$ $\pm$ 0.06 (6/6)
Dreamer-style + do-wrapper	$+0.14$ $\pm$ 0.04 (14/15)	$+0.18$ $\pm$ 0.04 (15/15)	$+0.11$ $\pm$ 0.12 (4/6)	$+0.15$ $\pm$ 0.07 (5/6)
Small Transformer	$+0.21$ $\pm$ 0.03 (15/15)	$+0.37$ $\pm$ 0.03 (15/15)	$+0.40$ $\pm$ 0.08 (6/6)	$+0.45$ $\pm$ 0.03 (6/6)

Table 2: Paired comparison with the factored AC-LSCM: normalized excess violations of the baseline minus those of the factored AC-LSCM trained on the same data (positive: the factored model is safer), mean

\pm

standard error over (graph, seed) pairs, and in parentheses the number of pairs in which the factored model is safer. In all

396

corresponding comparisons of rollout MSE, the factored model has the lower error.

	$N=500$			$N=2{,}000$			$N=10{,}000$
Model	Rollout	Excess (1)	Excess (2)	Rollout	Excess (1)	Excess (2)	Rollout	Excess (1)	Excess (2)
AC-LSCM, factored	0.014	0.306	0.291	0.007	0.224	0.150	0.007	0.162	0.098
AC-LSCM, factored, dense	0.015	0.313	0.290	0.010	0.267	0.220	0.008	0.198	0.156
AC-LSCM, latent	0.034	0.414	0.407	0.026	0.457	0.529	0.024	0.586	0.637
AC-LSCM, latent, dense	0.048	0.393	0.412	0.014	0.234	0.256	0.008	0.153	0.138
Vanilla VAE	0.035	0.321	0.484	0.016	0.294	0.410	0.010	0.192	0.266

Table 3: Training-set size (first three ER-10 graphs; 2 seeds at

N<10{,}000

, 3 at

N=10{,}000

; native action handling). Rollout: three-step rollout MSE after single-variable interventions. Excess (1)/(2): normalized excess violations with one-/two-variable candidate actions. Best non-reference entry per column in bold.

ER-10 (5 graphs)
Model	CF data	SHD	1-step MSE	CF MSE	Rollout	Excess (1)	Excess (2)
Factored (default)	none	0.0	0.030	0.020	0.009	0.246	0.178
Factored, soft do-operator	none	0.0	0.030	0.020	0.009	0.263	0.188
Factored, dense ( $\lambda_{1}{=}0$ )	none	76.8	0.031	0.020	0.010	0.303	0.266
Factored	observed	0.0	0.030	0.020	0.009	0.250	0.187
Factored, true graph	none	0.0	0.030	0.020	0.009	0.236	0.157

Table 4: Ablations (

N=10{,}000

, native action handling, mean over graphs). CF data: counterfactual training pairs given to none of the models, the observed counterfactual outcome given to every model, or ( † , first version) noise-free counterfactual states given to AC-LSCM only. SHD: structural Hamming distance to the true graph (true edges: 8 to 15 on ER-10, 56 and 66 on ER-20). One-step and counterfactual MSE are on noisy held-out observations.

Intervention MSE
	Chain	Fork	Collider	ER-10	ER-20
Small Transformer	0.046 $\pm$ 0.002	0.033 $\pm$ 0.002	0.031 $\pm$ 0.000	0.037 $\pm$ 0.002	0.071 $\pm$ 0.002
Vanilla VAE	0.031 $\pm$ 0.000	0.030 $\pm$ 0.001	0.025 $\pm$ 0.000	0.028 $\pm$ 0.000	0.048 $\pm$ 0.001
Dreamer-style	0.032 $\pm$ 0.001	0.029 $\pm$ 0.001	0.025 $\pm$ 0.000	0.028 $\pm$ 0.000	0.048 $\pm$ 0.000
AC-LSCM (original)	0.101 $\pm$ 0.004	0.057 $\pm$ 0.001	0.046 $\pm$ 0.002	0.058 $\pm$ 0.004	0.221 $\pm$ 0.009
Counterfactual MSE

Table 5: The first version’s models re-run with the released code (one-step error on noisy held-out data; 3 to 5 seeds; best in bold). The original AC-LSCM’s SHD equals the number of true edges, with zero variance, on every graph: it learned an empty graph in every run.

	Original, latent predicates		Original, decoded		In-range, decoded
Model	Violations	Approp. defer	Violations	Approp. defer	Violations	Approp. defer
Small Transformer	–	–	0.156 $\pm$ 0.052	0.220 $\pm$ 0.259	0.020 $\pm$ 0.012	0.900 $\pm$ 0.061
Vanilla VAE	0.242 $\pm$ 0.083	0.000 $\pm$ 0.000	0.000 $\pm$ 0.000	1.000 $\pm$ 0.000	0.000 $\pm$ 0.000	1.000 $\pm$ 0.000
Dreamer-style	0.504 $\pm$ 0.417	0.000 $\pm$ 0.000	0.000 $\pm$ 0.000	1.000 $\pm$ 0.000	0.000 $\pm$ 0.000	1.000 $\pm$ 0.000
AC-LSCM (original)	0.072 $\pm$ 0.161	0.820 $\pm$ 0.402	0.008 $\pm$ 0.018	0.960 $\pm$ 0.089	0.000 $\pm$ 0.000	1.000 $\pm$ 0.000

Table 6: The first version’s direct-effect task (ER-10, 100 episodes, 5 seeds). “Latent predicates”: the first version’s protocol (the Transformer has no latent and is evaluated on its predicted observation). “In-range”: unsafe and goal thresholds moved so that all candidate values lie inside the training range (

\tau_{s}=1.0

\tau_{g}=0.5

instead of

1.5

and

1.0

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

				References		Violation rate, single-variable actions ( $N=10{,}000$ , mean of 3 seeds)
Graph	Edges	Safety var. (anc.)	$\tau_{s}$	Oracle	No model	Factored	VAE	VAE+W	Dreamer+W	Transf.+W	Latent
ER-10/0	8	$z^{9}$ (4)	0.39	0.037	0.377	0.076	0.080	0.079	0.094	0.077	0.280
ER-10/1	15	$z^{9}$ (8)	0.66	0.110	0.317	0.153	0.160	0.171	0.176	0.181	0.193
ER-10/2	15	$z^{9}$ (4)	0.58	0.057	0.387	0.110	0.124	0.102	0.122	0.110	0.268
ER-10/3	15	$z^{9}$ (7)	0.93	0.127	0.273	0.189	0.208	0.209	0.232	0.220	0.241
ER-10/4	13	$z^{9}$ (6)	0.88	0.100	0.320	0.170	0.210	0.193	0.210	0.221	0.288

Appendix

Table 7: Task configuration chosen by the fixed rule (safety variable with the number of its ancestors in parentheses, and its threshold) and raw safety violation rates per environment. “+W”: with the do-wrapper. “Factored” and “Latent”: AC-LSCM. Dreamer-style and the Transformer without the do-wrapper are omitted; they are worse than with it in every environment except one (Dreamer-style on ER-10/3,

0.230

versus

0.232

), and worse than the factored model in every environment. Lowest rate in bold.

Causal World Models

Causal world models represent how entities and environments change under actions, enabling prediction, planning, and counterfactual reasoning.

See all 46 papers

CAFE: Counterfactual Prediction via Fast Posterior Estimation
Sep 26, 2026Xinyan Han, Xiaoyu Lin, Hao Zou +3Counterfactual Prediction Structural Causal Models
Tsinghua University · Stable AI
Relational Structural Causal Models
Jun 12, 2026Adiba Ejaz, Elias BareinboimStructural Causal Models Causal Identifiability
Causal AI Lab, Department of Computer Science, Columbia University, NY, USA.
Teleological Inference in Structural Causal Models via Intentional Interventions
Mar 19, 2026Dario Compagno, Fabio Massimo ZennaroStructural Causal Models Causal Intervention

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

CAFE: Counterfactual Prediction via Fast Posterior Estimation

Xinyan Han, Xiaoyu Lin, Hao Zou, Xingxuan Zhang+2

Tsinghua University · Stable AI

Counterfactual Prediction Structural Causal Models

Open paper AskPDF arXiv

Counterfactual prediction estimates an individual's outcome under an alternative intervention given their factual observations. Such outcomes are generally not identifiable from observational data without additional assumptions. Even within the class of fully observed additive noise models (ANMs), different causal graphs can generate the same observational distribution yet imply different individual counterfactual outcomes. Predictions based on a single estimated graph ignore this structural uncertainty. We therefore target a Bayesian counterfactual posterior predictive distribution that combines predictions from plausible SCMs. We introduce CAFE (CounterfActual Prediction via Fast Posterior Estimation), an amortized inference framework that directly approximates the Bayesian counterfactual posterior predictive distribution. We pretrain a transformer-based model on synthetic counterfactual tasks generated from a diverse prior over ANMs. Given an observational dataset, an individual's factual observations, and an intervention, CAFE approximates the corresponding posterior predictive distribution in a single forward pass. Experiments show that CAFE accurately predicts individual counterfactual outcomes in identifiable settings and approximates the posterior predictive distribution when structural uncertainty induced by observationally indistinguishable causal graphs exists. Strong performance in realistic manufacturing and viticulture settings further demonstrates its empirical robustness beyond the assumptions of the training prior.

Abstract

Causal World Models

Explore similar work

CAFE: Counterfactual Prediction via Fast Posterior Estimation

Relational Structural Causal Models

Teleological Inference in Structural Causal Models via Intentional Interventions