Causal Inference IX — Heterogeneous Effects: the Meta-Learner Zoo

Turning any machine-learning regressor into a CATE estimator — S, T, X, and R learners

The average treatment effect answers "does it work on average?"; the conditional average treatment effect (CATE) $\tau(x)=\mathbb{E}[Y(1)-Y(0)\mid X=x]$ answers "for whom, and how much?" — the question behind targeting, personalization, and policy. Meta-learners are recipes that turn any off-the-shelf regressor (the trees, boosting, and nets of the ML arc) into a CATE estimator by combining a few base models. We build the four canonical ones — S, T, X, R — from scratch.

Because CATE accuracy can only be measured against known individual effects (the counterfactual is never observed in real data), we evaluate two ways: first on a known-truth simulation where $\tau(x)$ is set by design, then on the field's standard IHDP semi-synthetic benchmark — the real covariates of the Infant Health and Development Program with outcomes simulated so the true CATE is known. The two together make the honest point: no single meta-learner is best everywhere — the winner depends on arm balance, effect size, and overlap. Python-lead (from-scratch S/T/X/R + econml cross-check); R companion uses grf. This is example 2 of the Heterogeneous-Effects subsection, following the Causal Forest and preceding Double ML.

1. The problem and the data — a known CATE

Judging a CATE estimator requires knowing the truth, so we simulate a study where $\tau(x)$ is set by design. Think of a job-training program evaluated on $n$ workers with five covariates $x$ (age, prior earnings, education, etc., scaled to $[0,1]$). Three ingredients:

• a prognostic function $\mu_0(x)$ — a complex, nonlinear baseline outcome (what a worker earns untreated);
• a heterogeneous effect $\tau(x)=x_1+0.5\,x_2$ — the program helps some workers much more than others;
• an unbalanced, covariate-dependent treatment assignment — only about 30% are treated, and treatment probability depends on $x$ (confounding, as in any observational study).

The outcome is $Y=\mu_0(x)+T\,\tau(x)+\varepsilon$. Because assignment is confounded and the effect is heterogeneous, a naive treated-minus-control comparison is meaningless; we need a method that adjusts for $x$ and lets the effect vary. We hold out a test set and measure CATE recovery by PEHE (Precision in Estimating Heterogeneous Effects — the RMSE of $\hat\tau(x)$ vs the true $\tau(x)$).

import numpy as np, pandas as pd, matplotlib.pyplot as plt, warnings
warnings.filterwarnings("ignore")
from sklearn.ensemble import GradientBoostingRegressor as GBR, GradientBoostingClassifier as GBC
from sklearn.model_selection import cross_val_predict
BLUE="#2b6cb0"; RED="#c53030"; GREEN="#2f855a"; ORANGE="#dd6b20"; GREY="#a0aec0"; PURP="#6b46c1"
def gen(seed, n=4000, prop=None):
    r=np.random.default_rng(seed); X=r.uniform(0,1,(n,5))
    mu0=2*np.sin(np.pi*X[:,0]*X[:,1])+2*(X[:,2]-0.5)**2
    tau=X[:,0]+0.5*X[:,1]
    e=(0.15+0.25*(X[:,0]>0.4)) if prop is None else prop(X)
    T=(r.uniform(size=n)<e).astype(int); Y=mu0+T*tau+r.normal(0,0.5,n)
    return X,T,Y,tau
Xtr,Ttr,Ytr,tautr=gen(0); Xte,Tte,Yte,taute=gen(100)
def pehe(est): return np.sqrt(np.mean((est-taute)**2))
print(f"n_train={len(Xtr)}, treated fraction={Ttr.mean():.2f} (unbalanced, confounded)")
print(f"true CATE tau(x)=x1+0.5*x2 -> ranges [{taute.min():.2f}, {taute.max():.2f}], ATE={taute.mean():.2f}")
fig,ax=plt.subplots(1,2,figsize=(13,4.2))
ax[0].scatter(Xtr[:,0],Ytr,s=6,c=[RED if t else BLUE for t in Ttr],alpha=.3)
ax[0].scatter([],[],c=RED,label="treated"); ax[0].scatter([],[],c=BLUE,label="control")
ax[0].set_xlabel("covariate x1"); ax[0].set_ylabel("outcome Y"); ax[0].set_title("Confounded, heterogeneous — naive comparison is meaningless"); ax[0].legend()
ax[1].scatter(Xte[:,0],taute,s=6,c=PURP,alpha=.4); ax[1].set_xlabel("covariate x1"); ax[1].set_ylabel("true CATE tau(x)")
ax[1].set_title("The target: how the effect varies with x (mostly x1)")
plt.tight_layout(); plt.show()
print("The effect grows with x1 (and x2): a good CATE estimator must recover this shape from confounded, unbalanced data.")

n_train=4000, treated fraction=0.30 (unbalanced, confounded)
true CATE tau(x)=x1+0.5*x2 -> ranges [0.01, 1.49], ATE=0.76

The effect grows with x1 (and x2): a good CATE estimator must recover this shape from confounded, unbalanced data.

2. The S-learner and the T-learner

The two simplest recipes, both wrapping the same base learner (gradient boosting here):

S-learner ("single") fits one model $\hat\mu(x,t)$ with the treatment as just another feature, then $$\hat\tau_S(x)=\hat\mu(x,1)-\hat\mu(x,0).$$ Simple and data-efficient, but it has a known flaw: a regularized base learner may treat the single treatment feature as unimportant and shrink the estimated effect toward zero, especially when the effect is subtle relative to the prognostic signal.

T-learner ("two") fits separate models on the treated and control arms, $\hat\mu_1$ and $\hat\mu_0$, and takes $$\hat\tau_T(x)=\hat\mu_1(x)-\hat\mu_0(x).$$ It never dilutes the treatment, but it splits the data — and when one arm is small (here only ~30% treated), $\hat\mu_1$ is estimated from few observations and becomes noisy, inflating the CATE error. We fit both from scratch and compare their recovered CATEs to the truth.

def base(): return GBR(n_estimators=200,max_depth=3,learning_rate=0.05)
def clf():  return GBC(n_estimators=200,max_depth=3,learning_rate=0.05)
# S-learner
sm=base().fit(np.column_stack([Xtr,Ttr]),Ytr)
tau_S=sm.predict(np.column_stack([Xte,np.ones(len(Xte))]))-sm.predict(np.column_stack([Xte,np.zeros(len(Xte))]))
# T-learner
m1=base().fit(Xtr[Ttr==1],Ytr[Ttr==1]); m0=base().fit(Xtr[Ttr==0],Ytr[Ttr==0])
tau_T=m1.predict(Xte)-m0.predict(Xte)
print(f"S-learner: PEHE {pehe(tau_S):.3f}, corr with truth {np.corrcoef(tau_S,taute)[0,1]:.2f}, ATE {tau_S.mean():.2f}")
print(f"T-learner: PEHE {pehe(tau_T):.3f}, corr with truth {np.corrcoef(tau_T,taute)[0,1]:.2f}, ATE {tau_T.mean():.2f}")
fig,ax=plt.subplots(1,2,figsize=(13,4.4))
for a,est,nm,c in zip(ax,[tau_S,tau_T],["S-learner","T-learner"],[GREEN,ORANGE]):
    a.scatter(taute,est,s=6,c=c,alpha=.3); lim=[taute.min()-.2,taute.max()+.2]; a.plot(lim,lim,"k--",lw=1)
    a.set_xlabel("true CATE tau(x)"); a.set_ylabel("estimated CATE"); a.set_title(f"{nm} (PEHE {pehe(est):.3f})")
plt.tight_layout(); plt.show()
print("Both track the true CATE, but with scatter. The T-learner suffers from the small treated arm (~30%); the S-learner")
print("can compress the effect. The next two learners are designed to fix exactly these weaknesses.")

S-learner: PEHE 0.137, corr with truth 0.93, ATE 0.71
T-learner: PEHE 0.187, corr with truth 0.86, ATE 0.75

Both track the true CATE, but with scatter. The T-learner suffers from the small treated arm (~30%); the S-learner
can compress the effect. The next two learners are designed to fix exactly these weaknesses.

3. The X-learner and the R-learner

X-learner (Künzel et al. 2019) is built for unbalanced arms. It proceeds in stages: (1) fit $\hat\mu_1,\hat\mu_0$ as in the T-learner; (2) impute each unit's individual effect using the other arm's model — for treated units $\tilde D_i=Y_i-\hat\mu_0(x_i)$, for controls $\tilde D_i=\hat\mu_1(x_i)-Y_i$; (3) fit CATE models $\hat\tau_1,\hat\tau_0$ to these imputed effects; (4) combine them with propensity weights, $\hat\tau_X(x)=e(x)\hat\tau_0(x)+(1-e(x))\hat\tau_1(x)$. The trick: when treated units are scarce, it leans on the many control units (via $\hat\tau_0$) to estimate their effects, sharply reducing variance.

R-learner (Nie & Wager 2021) takes the orthogonalization route. Using cross-fitted nuisance estimates $\hat m(x)=\mathbb{E}[Y\mid x]$ and $\hat e(x)=\mathbb{E}[T\mid x]$, it forms residuals $\tilde Y=Y-\hat m(x)$ and $\tilde T=T-\hat e(x)$ and minimizes the Robinson / R-loss $\sum_i(\tilde Y_i-\tau(x_i)\tilde T_i)^2$ — a Neyman-orthogonal objective that is first-order insensitive to nuisance errors. It is the direct conceptual ancestor of the Double-ML estimator in the next notebook — though its pseudo-outcome divides by $\tilde T$, so it is sensitive to extreme propensities.

# X-learner
D1=Ytr[Ttr==1]-m0.predict(Xtr[Ttr==1]); D0=m1.predict(Xtr[Ttr==0])-Ytr[Ttr==0]
tx1=base().fit(Xtr[Ttr==1],D1); tx0=base().fit(Xtr[Ttr==0],D0)
g=np.clip(clf().fit(Xtr,Ttr).predict_proba(Xte)[:,1],.05,.95)
tau_X=g*tx0.predict(Xte)+(1-g)*tx1.predict(Xte)
# R-learner (Robinson): cross-fitted nuisances, orthogonal loss
mhat=cross_val_predict(base(),Xtr,Ytr,cv=5)
ehat=np.clip(cross_val_predict(clf(),Xtr,Ttr,cv=5,method="predict_proba")[:,1],.1,.9)
Yr=Ytr-mhat; Tr=Ttr-ehat
rm=base().fit(Xtr,Yr/Tr,sample_weight=Tr**2); tau_R=rm.predict(Xte)
print(f"X-learner: PEHE {pehe(tau_X):.3f}, corr {np.corrcoef(tau_X,taute)[0,1]:.2f}")
print(f"R-learner: PEHE {pehe(tau_R):.3f}, corr {np.corrcoef(tau_R,taute)[0,1]:.2f}")
from econml.metalearners import XLearner
xl=XLearner(models=base(),propensity_model=clf(),cate_models=base()); xl.fit(Ytr,Ttr,X=Xtr)
tau_econml=xl.effect(Xte)
print(f"econml XLearner (package cross-check): PEHE {pehe(tau_econml):.3f}, corr with from-scratch X {np.corrcoef(tau_X,tau_econml)[0,1]:.2f}")
fig,ax=plt.subplots(1,2,figsize=(13,4.4))
for a,est,nm,c in zip(ax,[tau_X,tau_R],["X-learner","R-learner"],[BLUE,PURP]):
    a.scatter(taute,est,s=6,c=c,alpha=.3); lim=[taute.min()-.2,taute.max()+.2]; a.plot(lim,lim,"k--",lw=1)
    a.set_xlabel("true CATE tau(x)"); a.set_ylabel("estimated CATE"); a.set_title(f"{nm} (PEHE {pehe(est):.3f})")
plt.tight_layout(); plt.show()
print("The X-learner's cross-arm imputation gives the tightest fit here; the from-scratch and econml X-learners agree closely.")
print("The R-learner is orthogonal and principled but its 1/T-residual pseudo-outcome makes it the most propensity-sensitive.")

X-learner: PEHE 0.118, corr 0.94
R-learner: PEHE 0.260, corr 0.77

econml XLearner (package cross-check): PEHE 0.118, corr with from-scratch X 1.00

The X-learner's cross-arm imputation gives the tightest fit here; the from-scratch and econml X-learners agree closely.
The R-learner is orthogonal and principled but its 1/T-residual pseudo-outcome makes it the most propensity-sensitive.

4. Head-to-head, and the imbalance stress test

We line up all four learners by PEHE on the same test data, then run the stress test that motivates the X-learner: vary the treated fraction from balanced (50%) to severe imbalance (~12%) and watch the T- and X-learners' errors. The T-learner degrades as the treated arm shrinks (its $\hat\mu_1$ starves for data); the X-learner, leaning on the large control arm to impute treated effects, stays more accurate — the robustness Künzel et al. designed it for.

ests={"S":tau_S,"T":tau_T,"X":tau_X,"R":tau_R}
fig,ax=plt.subplots(1,2,figsize=(13,4.4))
nm=list(ests); pe=[pehe(ests[k]) for k in nm]
ax[0].bar(nm,pe,color=[GREEN,ORANGE,BLUE,PURP]); ax[0].set_ylabel("PEHE (lower = better)"); ax[0].set_title("CATE error by meta-learner (simulation)")
for i,v in enumerate(pe): ax[0].text(i,v+0.004,f"{v:.3f}",ha="center")
def learn_TX(frac):
    p=lambda X: np.full(len(X),frac)
    Xa,Ta,Ya,ta=gen(1,prop=p); Xb,Tb,Yb,tb=gen(101,prop=p)
    a1=base().fit(Xa[Ta==1],Ya[Ta==1]); a0=base().fit(Xa[Ta==0],Ya[Ta==0]); tT=a1.predict(Xb)-a0.predict(Xb)
    d1=Ya[Ta==1]-a0.predict(Xa[Ta==1]); d0=a1.predict(Xa[Ta==0])-Ya[Ta==0]
    gg=np.clip(clf().fit(Xa,Ta).predict_proba(Xb)[:,1],.05,.95)
    tX=gg*base().fit(Xa[Ta==0],d0).predict(Xb)+(1-gg)*base().fit(Xa[Ta==1],d1).predict(Xb)
    rt=lambda e: np.sqrt(np.mean((e-tb)**2)); return rt(tT),rt(tX)
fracs=[0.5,0.35,0.25,0.15,0.12]; TXs=[learn_TX(f) for f in fracs]
ax[1].plot(fracs,[a for a,_ in TXs],"o-",color=ORANGE,lw=2,label="T-learner")
ax[1].plot(fracs,[b for _,b in TXs],"s-",color=BLUE,lw=2,label="X-learner")
ax[1].invert_xaxis(); ax[1].set_xlabel("treated fraction (-> more imbalanced)"); ax[1].set_ylabel("PEHE"); ax[1].set_title("Under imbalance the X-learner stays ahead of the T-learner"); ax[1].legend()
plt.tight_layout(); plt.show()
print("Ranking (simulation):", ", ".join(f"{k} {pehe(ests[k]):.3f}" for k in sorted(nm,key=lambda k:pehe(ests[k]))))
print("As the treated arm shrinks, the T-learner's error climbs faster than the X-learner's -- the X-learner's reason for existing.")

Ranking (simulation): X 0.118, S 0.137, T 0.187, R 0.260
As the treated arm shrinks, the T-learner's error climbs faster than the X-learner's -- the X-learner's reason for existing.

5. The field benchmark — IHDP (real covariates, known effects)

A CATE method should be tested on more than a self-designed simulation. The standard benchmark is IHDP (Hill 2011): the real covariates of the Infant Health and Development Program — 747 low-birth-weight infants, 25 covariates on the child and mother — with a treatment (specialist home visits) and outcomes simulated from those real covariates, so the true CATE is known and PEHE is computable. (Simulated outcomes are unavoidable here: a real dataset never reveals both potential outcomes, so every honest CATE benchmark is semi-synthetic.) IHDP is deliberately hard — treatment is confounded and badly overlapping (only ~19% treated, with treated propensities near 0.18) — precisely the regime that separates the learners.

Running all four on IHDP flips the simulation's ranking: the S-learner wins, the T- and X-learners are middling, and the R-learner struggles — its inverse-residual weighting is fragile under IHDP's poor overlap. The lesson lands harder than any single number: the best meta-learner is data-dependent, and a method that shone on one design can fail on another.

raw=pd.read_csv("ihdp_npci_1.csv", header=None).values
Ti=raw[:,0].astype(int); Yi=raw[:,1]; mu0i=raw[:,3]; mu1i=raw[:,4]; Xi=raw[:,5:]
tau_ih=mu1i-mu0i
print(f"IHDP benchmark: n={len(Ti)}, {Xi.shape[1]} real covariates, treated {Ti.mean()*100:.0f}% (poor overlap), true ATE {tau_ih.mean():.2f}")
peh=lambda est: np.sqrt(np.mean((est-tau_ih)**2))
# S
s=base().fit(np.column_stack([Xi,Ti]),Yi); tS=s.predict(np.column_stack([Xi,np.ones(len(Ti))]))-s.predict(np.column_stack([Xi,np.zeros(len(Ti))]))
# T
i1=base().fit(Xi[Ti==1],Yi[Ti==1]); i0=base().fit(Xi[Ti==0],Yi[Ti==0]); tT=i1.predict(Xi)-i0.predict(Xi)
# X
d1=Yi[Ti==1]-i0.predict(Xi[Ti==1]); d0=i1.predict(Xi[Ti==0])-Yi[Ti==0]
gg=np.clip(clf().fit(Xi,Ti).predict_proba(Xi)[:,1],.05,.95)
tX=gg*base().fit(Xi[Ti==0],d0).predict(Xi)+(1-gg)*base().fit(Xi[Ti==1],d1).predict(Xi)
# R
mh=cross_val_predict(base(),Xi,Yi,cv=5); eh=np.clip(cross_val_predict(clf(),Xi,Ti,cv=5,method="predict_proba")[:,1],.05,.95)
tt=Ti-eh; tR=base().fit(Xi,(Yi-mh)/tt,sample_weight=tt**2).predict(Xi)
ih={"S":tS,"T":tT,"X":tX,"R":tR}
print("PEHE on IHDP (real covariates, known simulated effect):")
for k in sorted(ih,key=lambda k:peh(ih[k])): print(f"  {k}-learner: PEHE {peh(ih[k]):.3f}")
fig,ax=plt.subplots(1,2,figsize=(13,4.4))
nmi=["S","T","X","R"]; pei=[peh(ih[k]) for k in nmi]
ax[0].bar(nmi,pei,color=[GREEN,ORANGE,BLUE,PURP]); ax[0].set_ylabel("PEHE"); ax[0].set_title("IHDP: S-learner wins, R-learner strains (poor overlap)")
for i,v in enumerate(pei): ax[0].text(i,v+0.03,f"{v:.2f}",ha="center")
ax[1].hist(eh[Ti==0],bins=25,alpha=.55,color=BLUE,label="control"); ax[1].hist(eh[Ti==1],bins=25,alpha=.7,color=RED,label="treated")
ax[1].set_xlabel("estimated propensity e(x)"); ax[1].set_title("IHDP overlap is poor: treated mass sits at low e(x)"); ax[1].legend(fontsize=8)
plt.tight_layout(); plt.show()
print(f"Ranking flips vs the simulation (where X won): on IHDP the simple S-learner is best and the R-learner is worst --")
print(f"poor overlap (treated propensity ~{eh[Ti==1].mean():.2f}) punishes the propensity-sensitive R-learner. No universal winner.")

IHDP benchmark: n=747, 25 real covariates, treated 19% (poor overlap), true ATE 4.02

PEHE on IHDP (real covariates, known simulated effect):
  S-learner: PEHE 0.447
  X-learner: PEHE 0.678
  T-learner: PEHE 0.717
  R-learner: PEHE 2.721

Ranking flips vs the simulation (where X won): on IHDP the simple S-learner is best and the R-learner is worst --
poor overlap (treated propensity ~0.17) punishes the propensity-sensitive R-learner. No universal winner.

6. Summary

Meta-learners turn the flexible regressors of the machine-learning arc into CATE estimators through simple recipes, and the choice of recipe matters:

• S-learner — one model, treatment as a feature; efficient, and hard to beat when overlap is poor, but can shrink the effect toward zero;
• T-learner — separate arm models; unbiased in principle but high-variance when an arm is small;
• X-learner — cross-arm imputation with propensity weighting; robust to unbalanced treatment, the standout under the simulation's imbalance stress test;
• R-learner — a Neyman-orthogonal Robinson loss; principled and the bridge to Double ML, but propensity-sensitive and fragile under poor overlap.

The two evaluations tell the honest story: on the known-truth simulation the X-learner recovered $\tau(x)$ most accurately (and matched econml); on the real-covariate IHDP benchmark the ranking flipped — the S-learner won and the R-learner struggled with IHDP's poor overlap. There is no universally best CATE learner — performance depends on arm balance, effect size, and overlap, so match the method to the design and always evaluate against a benchmark (a known-truth simulation, a semi-synthetic benchmark like IHDP, or held-out policy value).

Cross-links. The base learners are the Trees, Regularized, and Neural-Network models of the ML arc, now aimed at a causal estimand; the R-learner's orthogonalization is the engine of the next notebook, Double/Debiased Machine Learning; the Causal Forest (example 1) is a closely related honest, orthogonalized estimator of $\tau(x)$; and the Policy Learning notebook (example 4) turns these CATE estimates into decisions. All presume unconfoundedness — the back-door assumption of the DAG notebook. The R companion reproduces the learners with grf and scores them on IHDP.