Observationally Informed Adaptive Causal Experimental Design

Erdun Gao; Liang Zhang; Jake Fawkes; Aoqi Zuo; Wenqin Liu; Haoxuan Li; Mingming Gong; Dino Sejdinovic

stat.MLMar 4, 2026

Observationally Informed Adaptive Causal Experimental Design

Authors: Erdun Gao, Liang Zhang, Jake Fawkes, Aoqi Zuo, Wenqin Liu, Haoxuan Li, Mingming Gong, Dino Sejdinovic

Organizations: Australian Institute for Machine Learning, Adelaide University · Hong Kong University of Science and Technology (Guangzhou) · University College London · The University of Sydney ·

Abstract

Randomized Controlled Trials (RCTs) represent the gold standard for causal inference yet remain a scarce resource. While large-scale observational data is often available, it is utilized only for retrospective fusion, and remains discarded in prospective trial design due to bias concerns. We argue this "tabula rasa" data acquisition strategy is fundamentally inefficient. In this work, we propose Active Residual Learning, a new paradigm that leverages the observational model as a foundational prior. This approach shifts the experimental focus from learning target causal quantities from scratch to efficiently estimating the residuals required to correct observational bias. To operationalize this, we introduce the R-Design framework. Theoretically, we establish two key advantages: (1) a structural efficiency gap, proving that estimating smooth residual contrasts admits strictly faster convergence rates than reconstructing full outcomes; and (2) information efficiency, where we quantify the redundancy in standard parameter-based acquisition (e.g., BALD), demonstrating that such baselines waste budget on task-irrelevant nuisance uncertainty. We propose R-EPIG (Residual Expected Predictive Information Gain), a unified criterion that directly targets the causal estimand, minimizing residual uncertainty for estimation or clarifying decision boundaries for policy. Experiments on synthetic and semi-synthetic benchmarks demonstrate that R-Design significantly outperforms baselines, confirming that repairing a biased model is far more efficient than learning one from scratch.

Figures & tables

Panel of Figure 1 — Figure 1: R-Design intuition on 1D synthetic data. The observational prior is biased (Left) but captures high-frequency structure (Middle-Left). Treating it as a fixed offset, R-Design learns a simpler residual function (Middle-Right), exploiting the complexity gap to recover the true CATE with few samples (Right), while the naive observational model fails due to uncorrected bias.

Panel of Figure 2 — Figure 2: Causal diagrams of data-generating processes.

Panel of Figure 3 — Figure 3: Comparison of acquisition functions using the average rank metric and relative performance improvements for effect estimation and policy learning tasks over eight combinations of base effect functions, bias functions, and CATE functions.

Panel of Figure 4 — Figure 4: Comparison of $\sqrt{\text{PEHE}}$ on simulation dataset of two trial estimators (CMGP and NSGP). (1,3) CMGP and NSGP, (2,4) CMGP and NSGP with heavy covariate shift.

Dimension	Variations Evaluated	Ref.
1. Active Learning Configuration
Batch Size	$n_{b}\in\{2,3,5,10,20,30\}$	§ 18
Pool Size	$n_{P}\in\{500,1000,1500,2000,2500\}$	§ 16
2. Model Agnosticism (Robustness)
Stage 1 (Prior)	TabPFN, CMGP, NSGP, CausalPFN	§ 15
Stage 2 (Trial)	CMGP, NSGP, BART, BCF, CMDE	§ 19

Table 1: Index of additional ablation studies.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Scenario	Confounding $\eta({\bm{x}})$	Baseline $\mu({\bm{x}},0)$	Treatment Effect $\tau({\bm{x}})$
1	Simple	Simple	Simple
2	Simple	Simple	Complex
3	Simple	Complex	Simple
4	Simple	Complex	Complex
5	Complex	Simple	Simple
6	Complex	Simple	Complex

Appendix

Table 2: Univariate simulation scenarios defined by component complexity.

Aspect	R-EPIG- $\tau$	R-EPIG- $\mu$
Target Estimand	Residual Contrast $\tau_{\delta}$	Residual Vector $\bm{\delta}$
Information Metric	Scalar Mutual Information	Vector Mutual Information
Complexity	${\mathcal{O}}(n_{P}\cdot n_{O})$	${\mathcal{O}}(n_{P}\cdot n_{O})$ (with small constant overhead)
Intuition	Maximizes correlation with CATE	Maximizes correlation with Outcomes

Appendix

Table 3: Comparison of R-EPIG variants.

Aspect	R-EPIG- $\tau$	R-EPIG- $\mu$	R-EPIG- $\pi$
Target Quantity	CATE Magnitude	Outcomes	Policy
Objective	Minimize PEHE	Minimize PEHE	Minimize APE
Information Metric	Gaussian MI	Gaussian Vector MI	Bernoulli MI
Computation	Analytic (Closed-form)	Analytic (Closed-form)	Gauss-Hermite Quad.
Focus Region	High Epistemic Unc. (CATE)	High Epistemic Unc. (Outcome)	Decision Boundary

Appendix

Table 4: R-EPIG variants: CATE estimation vs. policy optimization.

Panel of Figure 9 — Figure 9: Comparison of $\sqrt{\text{PEHE}}$ univaraite simulation setup. The first row shows the results of simulation setups indexed from 1-4 and the second row shows the results of 5-8.

Panel of Figure 10 — Figure 10: Average Policy Error (APE) Convergence. Comparison across eight univariate simulation benchmarks. The top row displays results for scenarios 1–4, and the bottom row for scenarios 5–8. Lower APE indicates higher accuracy in identifying the optimal treatment arm.

Panel of Figure 11 — Figure 11: Policy Regret Convergence. Comparison across eight univariate simulation benchmarks. The top row displays results for scenarios 1–4, and the bottom row for scenarios 5–8. R-EPIG- $\pi$ achieves the fastest convergence to zero regret in the majority of tasks.

Figure 12: Comparison of PEHE over dim 6 datasets with different trail models. From left to right: CMGP on normal dataset, CMGP on heavy covaraite shift, NSGP on normal dataset, NSGP on heavy covariate shift.

Figure 13: Comparison of decision making performance over the dim 6 dataset with different trail models. From left to right: CMGP (APE), CMGP (AR), NSGP (APE), NSGP (AR).

Figure 14: Comparison of CATE estimation performance with different kernels.

Figure 15: Comparison of CATE estimation performance with different Observational models.

Figure 16: Comparison of CATE estimation performance with different pool sizes.

Figure 17: Comparison of CATE estimation performance with different observational sizes.

Figure 18: Comparison of CATE estimation performance with different batch sizes.

Figure 19: Comparison of CATE estimation performance with different trial models.

Figure 20: Comparison of CATE estimation performance on IHDP dataset.

Figure 21: Comparison of CATE estimation performance on ACTG-175 dataset.

Figure 22: Comparison of policy estimation performance on IHDP dataset.

Figure 23: Comparison of policy estimation performance on ACTG-175 dataset.

Figure 24: Comparison of CATE estimation performance with two different strategies.

Method	Observational Data Size ( $n_{O}$ )
	$n_{O}=20$	$n_{O}=50$	$n_{O}=100$	$n_{O}=200$	$n_{O}=500$	$n_{O}=1000$
UMT (1-Stage)	1.08	1.08	0.91	0.79	0.59	0.53
TSR (2-Stage)	1.53	1.50	1.08	0.81	0.63	0.50

Appendix

Table 5: Comparison of normalized PEHE (Test) AUC for 1-Stage (UMT) and 2-Stage (TSR) methods across varying observational data sizes (

n_{O}

). Both methods utilize the R-EPIG-

\mu

acquisition function.

Active Experimental Design

See all 17 papers

Apr 28, 2026

TrialCalibre: A Fully Automated Causal Engine for RCT Benchmarking and Observational Trial Calibration

Real-world evidence (RWE) studies that emulate target trials increasingly inform regulatory and clinical decisions, yet residual, hard-to-quantify biases still limit their credibility. The recently proposed BenchExCal framework addresses this challenge via a two-stage Benchmark, Expand, Calibrate process, which first compares an observational emulation against an existing randomized controlled trial (RCT), then uses observed divergence to calibrate a second emulation for a new indication causal effect estimation. While methodologically powerful, BenchExCal is resource intensive and difficult to scale. We introduce TrialCalibre, a conceptualized multiagent system designed to automate and scale the BenchExCal workflow. Our framework features specialized agents such as the Orchestrator, Protocol Design, Data Synthesis, Clinical Validation, and Quantitative Calibration Agents that coordi-nate the the overall process. TrialCalibre incorpo-rates agent learning (e.g., RLHF) and knowledge blackboards to support adaptive, auditable, and transparent causal effect estimation.

Amir Habibdoust, Xing Song

Institute for Data Science and Informatics, University of Missouri, Columbia, MO, USA. · Department of Biomedical Informatics, Biostatistics, and Medical Epidemiology, University of Missouri, Columbia, MO, USA.

Causal Effect Estimation Benchmark Design Scientific Workflow Automation Autonomous Research Agents

Ask PDF

May 7, 2026

DARTS: Targeting Prognostic Covariates in Budget-Constrained Sequential Experiments

Randomized controlled trials typically assume that prognostic covariates are known and available at no cost. In practice, obtaining high-dimensional pretreatment data is costly, forcing a trade-off between covariate-adaptive precision and a measurement budget. We introduce Dynamic Adaptive Rerandomization via Thompson Sampling (DARTS), which treats covariate acquisition as a sequential optimization problem embedded within a design-based causal inference task. A budgeted combinatorial Thompson sampler learns which covariates are most prognostic across successive batches; selected covariates then drive rerandomization and regression adjustment to reduce batch-level average treatment effect variance. Our primary theoretical contribution is a decoupling result: adaptive covariate selection based on past batches preserves batch-level randomization validity, and the cumulative inverse-variance weighted estimator achieves at least nominal asymptotic coverage. We further derive a Bayes risk bound for the acquisition layer that matches the minimax lower bound up to logarithmic factors. Empirically, DARTS systematically concentrates the budget on informative features, significantly closing the efficiency gap to oracle designs while maintaining strict inferential validity.

Kateryna Husar, Alexander Volfovsky

Department of Statistical Science Duke University Durham, NC 27705

Thompson Sampling Causal Effect Estimation Feature Selection Bayesian Experimental Design

Ask PDF

Nov 10, 2025

Robust Experimental Design via Generalised Bayesian Inference

Bayesian optimal experimental design is a principled framework for conducting experiments that leverages Bayesian inference to quantify how much information one can expect to gain from selecting a certain design. However, accurate Bayesian inference relies on the assumption that one's statistical model of the data-generating process is correctly specified. If this assumption is violated, Bayesian methods can lead to poor inference and estimates of information gain. Generalised Bayesian (or Gibbs) inference is a more robust probabilistic inference framework that replaces the likelihood in the Bayesian update by a suitable loss function. In this work, we present Generalised Bayesian Optimal Experimental Design (GBOED), an extension of generalised Bayesian inference to the experimental design setting which achieves robustness in both design and inference. Using a generalisation of Bayesian decision theory, we derive a new acquisition function, the Gibbs expected information gain. Our empirical results demonstrate that GBOED can enhance robustness to outliers and incorrect assumptions about the outcome noise distribution.

Yasir Zubayr Barlas, Sabina J. Sloman, Samuel Kaski

The University of Manchester · Aalto University

Generalized Bayesian Inference Bayesian Experimental Design Bayesian Inference

Ask PDF

Abstract

Active Experimental Design

Explore similar work

TrialCalibre: A Fully Automated Causal Engine for RCT Benchmarking and Observational Trial Calibration

DARTS: Targeting Prognostic Covariates in Budget-Constrained Sequential Experiments

Robust Experimental Design via Generalised Bayesian Inference