Causal Inference IX — Meta-Learners (R companion)¶
S/T/X-learners with grf regression forests, the orthogonalized causal_forest, and the IHDP benchmark¶
R's flagship for heterogeneous treatment effects is grf (generalized random forests, Athey-Tibshirani-Wager). This companion uses grf::regression_forest as the common base learner to build the S-, T-, and X-learners from scratch, brings in grf::causal_forest (an orthogonalized, honest estimator of $\tau(x)$ that plays the R-learner's role), and scores them both on the known-truth simulation and the real-covariate IHDP benchmark. As in Python, the winner differs between the two — no meta-learner is best everywhere.
1. The known-truth data¶
The same design as the Python notebook: five covariates, a nonlinear prognostic function $\mu_0(x)$, a heterogeneous effect $\tau(x)=x_1+0.5x_2$, and unbalanced, covariate-dependent treatment (~30% treated). Because $\tau(x)$ is known, we score each learner by PEHE on a held-out test set.
suppressMessages(library(grf))
gen<-function(seed,n=3000){set.seed(seed); X<-matrix(runif(n*5),n,5)
mu0<-2*sin(pi*X[,1]*X[,2])+2*(X[,3]-0.5)^2; tau<-X[,1]+0.5*X[,2]
e<-0.15+0.25*(X[,1]>0.4); T<-rbinom(n,1,e); Y<-mu0+T*tau+rnorm(n,0,0.5)
list(X=X,T=T,Y=Y,tau=tau,e=mean(T))}
tr<-gen(0); te<-gen(100)
pehe<-function(est,truth) sqrt(mean((est-truth)^2))
cat(sprintf("n_train=%d, treated fraction=%.2f (unbalanced); true ATE=%.2f, tau in [%.2f, %.2f]\n",
length(tr$T), tr$e, mean(te$tau), min(te$tau), max(te$tau)))
Warning message: "package 'grf' was built under R version 4.6.1"
n_train=3000, treated fraction=0.30 (unbalanced); true ATE=0.75, tau in [0.02, 1.49]
2. S- and T-learners with regression_forest¶
Both wrap the same base learner. The S-learner fits one regression_forest on $[X,T]$ and differences the prediction at $T=1$ vs $T=0$. The T-learner fits a separate forest per arm; it is hurt by the small treated arm.
sf<-regression_forest(cbind(tr$X,tr$T), tr$Y)
tS<-predict(sf, cbind(te$X,1))$predictions - predict(sf, cbind(te$X,0))$predictions
m1<-regression_forest(tr$X[tr$T==1,], tr$Y[tr$T==1]); m0<-regression_forest(tr$X[tr$T==0,], tr$Y[tr$T==0])
tT<-predict(m1,te$X)$predictions - predict(m0,te$X)$predictions
cat(sprintf("S-learner PEHE %.3f (corr %.2f)\n", pehe(tS,te$tau), cor(tS,te$tau)))
cat(sprintf("T-learner PEHE %.3f (corr %.2f)\n", pehe(tT,te$tau), cor(tT,te$tau)))
S-learner PEHE 0.168 (corr 0.93)
T-learner PEHE 0.147 (corr 0.90)
3. X-learner from scratch, and grf::causal_forest¶
The X-learner imputes each unit's effect from the opposite arm's forest, models those pseudo-effects, and combines them with propensity weights. Then causal_forest estimates $\tau(x)$ directly with built-in orthogonalization (the Robinson logic of the R-learner) and honesty (valid confidence intervals). test_calibration's differential coefficient near 1 with a tiny p-value certifies real heterogeneity.
d1<-tr$Y[tr$T==1]-predict(m0,tr$X[tr$T==1,])$predictions
d0<-predict(m1,tr$X[tr$T==0,])$predictions-tr$Y[tr$T==0]
tx1<-regression_forest(tr$X[tr$T==1,],d1); tx0<-regression_forest(tr$X[tr$T==0,],d0)
ps<-pmin(pmax(predict(regression_forest(tr$X,tr$T),te$X)$predictions,.05),.95)
tX<-ps*predict(tx0,te$X)$predictions+(1-ps)*predict(tx1,te$X)$predictions
cf<-causal_forest(tr$X,tr$Y,tr$T); tC<-predict(cf,te$X)$predictions
cat(sprintf("X-learner PEHE %.3f (corr %.2f)\n", pehe(tX,te$tau), cor(tX,te$tau)))
cat(sprintf("causal_forest PEHE %.3f (corr %.2f)\n", pehe(tC,te$tau), cor(tC,te$tau)))
cat("\ncausal_forest test_calibration (differential coef ~1, tiny p = real heterogeneity):\n")
print(round(test_calibration(cf),3))
X-learner PEHE 0.090 (corr 0.96)
causal_forest PEHE 0.100 (corr 0.96)
causal_forest test_calibration (differential coef ~1, tiny p = real heterogeneity):
Best linear fit using forest predictions (on held-out data)
as well as the mean forest prediction as regressors, along
with one-sided heteroskedasticity-robust (HC3) SEs:
Estimate Std. Error t value Pr(>t)
mean.forest.prediction 0.976 0.030 32.814 < 2.2e-16 ***
differential.forest.prediction 1.149 0.083 13.769 < 2.2e-16 ***
---
Signif. codes: 0 '***' 0.001 '**' 0.01 '*' 0.05 '.' 0.1 ' ' 1
4. Head-to-head on the simulation¶
Ranking by PEHE reproduces the Python finding: the X-learner and orthogonalized causal forest give the tightest CATE recovery, the T-learner trails (starved treated arm), the S-learner sits between.
ests<-list(S=tS,T=tT,X=tX,`causal forest`=tC); pe<-sapply(ests,function(e) pehe(e,te$tau))
options(repr.plot.width=13, repr.plot.height=4.6); par(mfrow=c(1,2))
cols<-c("#2f855a","#dd6b20","#2b6cb0","#6b46c1")
barplot(pe, col=cols, ylab="PEHE (lower = better)", main="CATE error by learner (simulation)"); text(seq(0.7,by=1.2,length.out=4), pe+0.005, sprintf("%.3f",pe))
best<-names(which.min(pe))
plot(te$tau, ests[[best]], pch=19, col=rgb(.17,.42,.69,.25), cex=.4, xlab="true CATE tau(x)", ylab="estimated CATE",
main=sprintf("Best (sim): %s (PEHE %.3f)", best, min(pe))); abline(0,1,lty=2)
par(mfrow=c(1,1))
cat("Ranking (sim):", paste(sprintf("%s %.3f", names(sort(pe)), sort(pe)), collapse=" | "), "\n")
Ranking (sim): X 0.090 | causal forest 0.100 | T 0.147 | S 0.168
5. The field benchmark — IHDP (real covariates, known effects)¶
The standard CATE benchmark is IHDP (Hill 2011): the real covariates of the Infant Health and Development Program (747 infants, 25 covariates) with simulated outcomes so the true CATE is known — the only honest way to get a CATE benchmark, since a real dataset never reveals both potential outcomes. IHDP is deliberately hard: treatment is confounded and poorly overlapping (~19% treated). We run the same grf learners and score PEHE against the true effect mu1 - mu0. As in Python, the ranking flips from the simulation — the simple learners (T and S) lead while the more sophisticated X-learner and causal forest trail under the poor overlap — underscoring that the best learner is data-dependent.
ih<-read.csv("ihdp_npci_1.csv", header=FALSE)
Ti<-ih$V1; Yi<-ih$V2; mu0i<-ih$V4; mu1i<-ih$V5; Xi<-as.matrix(ih[,6:30]); tau_ih<-mu1i-mu0i
cat(sprintf("IHDP: n=%d, %d real covariates, treated %.0f%% (poor overlap), true ATE %.2f\n", nrow(ih), ncol(Xi), 100*mean(Ti), mean(tau_ih)))
# S
sfi<-regression_forest(cbind(Xi,Ti),Yi); iS<-predict(sfi,cbind(Xi,1))$predictions-predict(sfi,cbind(Xi,0))$predictions
# T
im1<-regression_forest(Xi[Ti==1,],Yi[Ti==1]); im0<-regression_forest(Xi[Ti==0,],Yi[Ti==0]); iT<-predict(im1,Xi)$predictions-predict(im0,Xi)$predictions
# X
id1<-Yi[Ti==1]-predict(im0,Xi[Ti==1,])$predictions; id0<-predict(im1,Xi[Ti==0,])$predictions-Yi[Ti==0]
itx1<-regression_forest(Xi[Ti==1,],id1); itx0<-regression_forest(Xi[Ti==0,],id0)
ips<-pmin(pmax(predict(regression_forest(Xi,Ti),Xi)$predictions,.05),.95)
iX<-ips*predict(itx0,Xi)$predictions+(1-ips)*predict(itx1,Xi)$predictions
# causal forest
icf<-causal_forest(Xi,Yi,Ti); iC<-predict(icf)$predictions
ihp<-c(S=pehe(iS,tau_ih), T=pehe(iT,tau_ih), X=pehe(iX,tau_ih), `causal forest`=pehe(iC,tau_ih))
cat("PEHE on IHDP (real covariates, known simulated effect):\n")
for(k in names(sort(ihp))) cat(sprintf(" %-14s %.3f\n", k, ihp[k]))
options(repr.plot.width=8, repr.plot.height=4.4)
bp<-barplot(ihp, col=cols, ylab="PEHE", main="IHDP: ranking flips vs the simulation"); text(bp, ihp+0.03, sprintf("%.2f",ihp))
cat(sprintf("\nRanking flips vs the sim (where X/causal-forest led): on IHDP's poor overlap the simple learners (T, S) lead and the\nsophisticated X-learner/causal-forest trail. No universal winner -- the best learner is data-dependent.\n"))
IHDP: n=747, 25 real covariates, treated 19% (poor overlap), true ATE 4.02
PEHE on IHDP (real covariates, known simulated effect):
T 0.357 S 0.365 X 0.516 causal forest 0.573
Ranking flips vs the sim (where X/causal-forest led): on IHDP's poor overlap the simple learners (T, S) lead and the sophisticated X-learner/causal-forest trail. No universal winner -- the best learner is data-dependent.
6. Summary¶
Using grf::regression_forest as a common base learner, the S-, T-, and X-learners and grf::causal_forest reproduced the Python notebook across both evaluations: on the known-truth simulation the X-learner and causal forest gave the tightest CATE recovery (with test_calibration certifying real heterogeneity), while on the real-covariate IHDP benchmark the ranking flipped — the simple T- and S-learners led while the sophisticated X-learner and causal forest trailed under IHDP's poor overlap. The value R adds is that grf ships the orthogonalization and honest confidence intervals as a tested package.
The workflow to remember: pick a flexible base learner, wrap it in a meta-learner suited to the design (X-learner for imbalance), or use causal_forest for a self-contained orthogonalized estimate — and always validate the CATE against known truth or a benchmark like IHDP. Cross-links: causal_forest is the subject of example 1; the base learners are the ML arc's forests; the orthogonalization inside both causal_forest and the R-learner is the engine of Double/Debiased Machine Learning (example 3); and Policy Learning (example 4) turns these estimates into decisions.