Causal Forests — Heterogeneous Treatment Effects

From-scratch honest causal tree · econml CausalForestDML · a return-to-work policy

This notebook opens the Causal Inference arc, and it builds directly on the random forest from the ML-Trees arc. The predictive forest answered "what is $E[Y\mid X]$?" — predict an outcome. A causal forest answers a different question:

$$\tau(x) \;=\; E\big[\,Y(1) - Y(0)\;\big|\;X = x\,\big],$$

the conditional average treatment effect (CATE) — how the effect of a treatment varies from one individual to the next. That is the quantity behind targeting: which clients an intervention actually helps, which subgroups a policy moves most, which regime a signal works in.

The fundamental problem of causal inference. $\tau(x)$ can never be observed for any single unit: we see the outcome with treatment $Y(1)$ or without it $Y(0)$, never both. So we cannot fit a tree to $\tau$ the way we fit CART to $y$. Two ideas get us there anyway, and we implement both from scratch before calling the reference packages:

1. Heterogeneity splitting — split to make leaves that differ in their effect, not leaves pure in $y$.
2. Honesty — decide where to split and estimate the leaf effects on disjoint halves of the data, which is what makes the estimates unbiased and the confidence intervals valid.

We then hand the same problem to econml's CausalForestDML (the debiased-ML causal forest, cross-checked against R's grf in the companion notebook) and validate it against a known ground-truth $\tau(x)$.

The identifying assumption — stated once, up front. None of this manufactures causation from correlation. Everything below assumes unconfoundedness: conditional on the covariates $X$, treatment is as-good-as-random ($\{Y(0),Y(1)\}\perp W \mid X$). A randomized experiment guarantees it by design; observational data requires no unmeasured confounders. Given that, the forest estimates how the effect varies — it does not solve identification.

1. A policy experiment with a known ground truth

To validate a CATE estimator we need a world where we know the true $\tau(x)$ for every person — so we simulate one, styled on a return-to-work / job-training program (the policy family behind welfare-to-work and SSA work-incentive evaluations). Each of 6,000 working-age adults is described by five covariates:

covariate	meaning
age	age in years, 22–60
educ	years of schooling, 8–18
prior	earnings the year before ($1,000s) — a strong predictor of later earnings
health	self-rated health index, 0–10
female	sex indicator (a gender earnings gap is built in)

Three mechanisms are wired in on purpose:

• Selection into the program (confounding). Enrollment $W$ is not random: people with higher prior earnings, better health, and younger are more likely to enroll. Because prior earnings and health also raise later earnings, a naïve treated-vs-control comparison will be biased.
• A heterogeneous true effect. The program raises earnings most for the young and the low-prior-earnings — those with the most human capital to gain — and barely at all (even slightly negatively) for older, higher-earning workers: $$\tau(x) = 6 - 0.08\,(\text{age}-22) - 0.12\,\text{prior}.$$ So the effect modifiers are age and prior; educ, health, female shift the baseline but not the effect.
• Outcome. $Y = \underbrace{8 + 0.7\,\text{prior} + 0.5\,\text{educ} + 0.4\,\text{health} - 3\,\text{female}}_{\text{baseline }b(x)} + W\cdot\tau(x) + \varepsilon$, with $\varepsilon\sim N(0,3)$.

The goal: recover $\tau(x)$ — who benefits, and by how much — from observational data, and turn it into a targeting rule.

import numpy as np, matplotlib.pyplot as plt
rng = np.random.default_rng(0)
n = 6000
age    = rng.uniform(22, 60, n)
educ   = np.clip(np.round(rng.normal(13, 2, n)), 8, 18)
prior  = np.clip(rng.normal(20, 10, n), 0, None)                    # prior earnings, $1000s
health = rng.uniform(0, 10, n)
female = rng.integers(0, 2, n)
X = np.column_stack([age, educ, prior, health, female]); names = ["age","educ","prior","health","female"]

# selection into the program (confounding): higher prior/health, younger -> more likely to enroll
propensity = 1/(1+np.exp(-(-0.5 + 0.04*(prior-20) + 0.15*(health-5) - 0.02*(age-40))))
W = (rng.uniform(0,1,n) < propensity).astype(int)

# heterogeneous TRUE effect (unobservable in real life; kept only to grade the estimators)
tau  = 6 - 0.08*(age-22) - 0.12*prior
base = 8 + 0.7*prior + 0.5*educ + 0.4*health - 3*female
Y    = base + W*tau + rng.normal(0, 3, n)

print(f"n={n}   enrolled P(W=1)={W.mean():.2f}   true ATE E[tau]={tau.mean():.3f}")
print(f"true tau ranges from {tau.min():.2f} (older, high earners) to {tau.max():.2f} (young, low earners)")

fig, ax = plt.subplots(1, 3, figsize=(13.5, 3.6))
ax[0].hist(tau, bins=40, color="#3182ce", edgecolor="white")
ax[0].axvline(tau.mean(), color="#c53030", lw=2, label=f"ATE={tau.mean():.2f}"); ax[0].axvline(0, color="k", lw=1, ls=":")
ax[0].set_title("True effect $\\tau(x)$ varies across people"); ax[0].set_xlabel("earnings gain from program ($1000s)"); ax[0].legend()
sc = ax[1].scatter(age, prior, c=tau, cmap="RdYlBu_r", s=6); ax[1].set_xlabel("age"); ax[1].set_ylabel("prior earnings")
ax[1].set_title("$\\tau$ is highest for young & low-prior"); plt.colorbar(sc, ax=ax[1], label="$\\tau(x)$")
ax[2].scatter(prior, propensity, s=5, alpha=.3, color="#805ad5"); ax[2].set_xlabel("prior earnings"); ax[2].set_ylabel("P(enroll)")
ax[2].set_title("Selection: who enrolls (confounding)")
plt.tight_layout(); plt.show()

n=6000   enrolled P(W=1)=0.38   true ATE E[tau]=2.068
true tau ranges from -2.49 (older, high earners) to 5.96 (young, low earners)

2. Why the obvious answers fail — and the econometric benchmark

Before any machine learning, three classical estimates — the natural comparisons a trained econometrician reaches for first:

1. Difference in means, $\bar Y_{W=1}-\bar Y_{W=0}$. Unbiased only under randomization. Here enrollees have higher prior earnings and better health, so it is confounded upward.
2. OLS with a treatment dummy, $Y = \alpha + \delta W + X\beta$. Adjusting for the covariates removes the confounding and recovers the average effect $\delta\approx\text{ATE}$ — but delivers a single number, no heterogeneity.
3. OLS with interactions, $Y = \alpha + \delta W + X\beta + W\!\cdot\!(\text{age},\text{prior})\,\gamma$. This is the textbook way to let the effect vary: the implied $\hat\tau(x)=\delta+\gamma_1\text{age}+\gamma_2\text{prior}$. It works beautifully when you guessed the right interactions — and here the truth is linear in age and prior, so it will look great. The causal forest's job is to find that structure without being told which variables interact or in what shape.

import statsmodels.api as sm
naive = Y[W==1].mean() - Y[W==0].mean()

Xc = sm.add_constant(np.column_stack([W, X]))                       # [const, W, age,educ,prior,health,female]
ols = sm.OLS(Y, Xc).fit(); ate_ols = ols.params[1]

# OLS with treatment x (age, prior) interactions -> parametric CATE
Xi = sm.add_constant(np.column_stack([W, X, W*age, W*prior]))
oi = sm.OLS(Y, Xi).fit()
d, g_age, g_prior = oi.params[1], oi.params[-2], oi.params[-1]
tau_ols = d + g_age*age + g_prior*prior                            # implied parametric CATE

print(f"1. difference in means        {naive:6.3f}   (true ATE {tau.mean():.3f}) -> bias {naive-tau.mean():+.3f}  CONFOUNDED")
print(f"2. OLS, treatment dummy        {ate_ols:6.3f}   adjusts for X -> recovers the average, but one number only")
print(f"3. OLS w/ age,prior interaction  ATE {tau_ols.mean():.3f};  implied tau corr w/ true = {np.corrcoef(tau_ols,tau)[0,1]:.3f}")
print("   -> the interaction model nails it HERE because the true tau really is linear in age & prior.")
print("      The causal forest must discover that shape with no such guess.")

1. difference in means         4.757   (true ATE 2.068) -> bias +2.689  CONFOUNDED
2. OLS, treatment dummy         1.837   adjusts for X -> recovers the average, but one number only
3. OLS w/ age,prior interaction  ATE 1.894;  implied tau corr w/ true = 1.000
   -> the interaction model nails it HERE because the true tau really is linear in age & prior.
      The causal forest must discover that shape with no such guess.

3. From scratch — an honest causal tree

causal.py implements a single honest causal tree (Athey & Imbens, PNAS 2016). Two departures from the predictive CART of the ML-Trees arc:

• Split criterion. Instead of variance reduction in $y$, we score a split by the treatment-effect heterogeneity it creates, $$\text{gain} = \frac{n_L\,n_R}{n_L+n_R}\,\big(\hat\tau_L-\hat\tau_R\big)^2,\qquad \hat\tau_{\text{child}}=\bar Y_{W=1}-\bar Y_{W=0}\ \text{in that child}.$$ The tree actively hunts for the cut that most separates high-effect from low-effect people.

• Honesty. The training rows are split in two: a structure half chooses the splits, a disjoint estimate half sets each leaf's $\hat\tau$. Reusing one sample for both — as ordinary CART does — biases leaf effects toward the very subgroups the tree searched for; separating them is what earns valid inference.

We hold out 30% as a test set, grow one honest tree, and print the subgroups it discovers.

from causal import HonestCausalTree, HonestCausalForest
split = rng.uniform(0,1,n) < 0.7
Xtr,Ytr,Wtr,tau_tr = X[split],Y[split],W[split],tau[split]
Xte,Yte,Wte,tau_te = X[~split],Y[~split],W[~split],tau[~split]

tree = HonestCausalTree(max_depth=3, min_leaf=80, min_treat=25).fit(Xtr,Ytr,Wtr, np.random.default_rng(2))

def show(node, names, depth=0, path="root"):
    pad = "   "*depth
    if node.leaf:
        print(f"{pad}[leaf n={node.n:4d}]  tau_hat = {node.tau:+.2f}")
    else:
        print(f"{pad}if {names[node.feat]} <= {node.thr:.1f}:")
        show(node.left, names, depth+1); print(f"{pad}else ({names[node.feat]} > {node.thr:.1f}):")
        show(node.right, names, depth+1)
show(tree.root, names)

pt = tree.predict(Xte)
print(f"\nhonest tree (test): corr(est, true) = {np.corrcoef(pt,tau_te)[0,1]:.3f}   RMSE = {np.sqrt(np.mean((pt-tau_te)**2)):.3f}")
print("The tree splits on age and prior earnings -- exactly the true effect modifiers -- with no hint of which variables matter.")

if age <= 35.8:
   if prior <= 31.6:
      if health <= 4.6:
         [leaf n= 322]  tau_hat = +4.53
      else (health > 4.6):
         [leaf n= 344]  tau_hat = +4.97
   else (prior > 31.6):
      [leaf n=  96]  tau_hat = +1.78
else (age > 35.8):
   if prior <= 24.2:
      if prior <= 15.9:
         [leaf n= 461]  tau_hat = +3.47
      else (prior > 15.9):
         [leaf n= 409]  tau_hat = +1.59
   else (prior > 24.2):
      if educ <= 14.5:
         [leaf n= 367]  tau_hat = +1.63
      else (educ > 14.5):
         [leaf n= 122]  tau_hat = +2.33

honest tree (test): corr(est, true) = 0.744   RMSE = 1.434
The tree splits on age and prior earnings -- exactly the true effect modifiers -- with no hint of which variables matter.

4. From a tree to a forest — and the confounding catch

A single honest tree gives blocky, high-variance effects. Averaging many honest trees over random subsamples — a causal forest — smooths $\hat\tau(x)$, the same variance-reduction bagging buys a predictive forest. We compare two forests:

• HonestCausalForest (from scratch): honest trees only. It captures the shape of $\tau(x)$ well — but its level is biased, because honesty alone does nothing about confounded selection into treatment.
• econml's CausalForestDML (Athey–Wager, via double/debiased ML): first residualizes both the outcome and the treatment against $X$ with nuisance ML models (the "DML" step, and the first appearance in this arc of the double-ML idea), then runs a causal forest on the residuals. Orthogonalizing away the confounding fixes the level and yields valid confidence intervals.

The scatter of estimated vs. true CATE on the test set is the moment of truth — a perfect estimator lands on the 45° line.

from econml.dml import CausalForestDML
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier

hf = HonestCausalForest(n_trees=150, max_depth=4, min_leaf=40, min_treat=12, subsample=0.5).fit(
        Xtr,Ytr,Wtr, np.random.default_rng(3))
pf = hf.predict(Xte)

cf = CausalForestDML(model_y=RandomForestRegressor(n_estimators=300, min_samples_leaf=20, random_state=0),
                     model_t=RandomForestClassifier(n_estimators=300, min_samples_leaf=20, random_state=0),
                     discrete_treatment=True, n_estimators=800, min_samples_leaf=10, random_state=0, cv=4)
cf.fit(Ytr, Wtr, X=Xtr, W=None)
pe = cf.effect(Xte); lb, ub = cf.effect_interval(Xte, alpha=0.1)
cover = np.mean((tau_te>=lb) & (tau_te<=ub))

def line(m):
    lo,hi = tau_te.min(), tau_te.max(); return [lo,hi],[lo,hi]
fig, ax = plt.subplots(1, 2, figsize=(11.5, 4.8))
for a,(p,t,c) in zip(ax, [(pf,"from-scratch honest forest","#dd6b20"), (pe,"econml CausalForestDML","#2b6cb0")]):
    a.scatter(tau_te, p, s=6, alpha=.35, color=c); a.plot(*line(0), "k--", lw=1)
    a.set_xlabel("true $\\tau(x)$"); a.set_ylabel("estimated $\\hat\\tau(x)$"); a.set_title(t)
    a.text(.05,.9,f"corr {np.corrcoef(p,tau_te)[0,1]:.3f}\nRMSE {np.sqrt(np.mean((p-tau_te)**2)):.2f}\nATEhat {p.mean():.2f}",
           transform=a.transAxes, va="top", fontsize=9, bbox=dict(fc="white",ec="0.7"))
plt.tight_layout(); plt.show()

print(f"true ATE {tau_te.mean():.3f}")
print(f"from-scratch forest ATE {pf.mean():.3f}  (shape good, LEVEL biased by confounding -- no orthogonalization)")
print(f"econml CFDML        ATE {pe.mean():.3f}  (DML removes the confounding bias)")
print(f"econml 90% CI empirical coverage of the true tau: {cover:.3f}")
print()
for nm, pp in (("from-scratch honest forest", pf), ("econml CausalForestDML", pe)):
    print(f"  {nm:<28} corr with true tau {np.corrcoef(pp,tau_te)[0,1]:.3f}   "
          f"RMSE {np.sqrt(np.mean((pp-tau_te)**2)):.3f}")
print()
print("Those two numbers say different things and both matter. The correlations are high, so both")
print("forests recover the SHAPE of the effect -- who benefits more than whom -- which is what a")
print("targeting rule needs. The from-scratch forest's RMSE is the worse of the two because its LEVEL")
print("is off: without orthogonalization it inherits the confounding, so it ranks correctly and")
print("mis-states magnitudes. For ranking that is survivable; for reporting an effect size it is not.")
print()
print(f"AND THE INTERVAL IS THE WEAK PART. A nominal 90% interval covering {cover:.1%} is short by")
print(f"{0.90-cover:.1%}, and this is not sampling noise -- it is the known behaviour of forest confidence")
print("intervals for a smooth effect surface, where the honest-splitting variance estimate does not")
print("capture the smoothing bias. The point estimates are trustworthy enough to target with. The")
print("intervals should be read as indicative rather than as calibrated 90% statements, and this")
print("notebook's own ranking results do not depend on them.")

true ATE 2.061
from-scratch forest ATE 3.350  (shape good, LEVEL biased by confounding -- no orthogonalization)
econml CFDML        ATE 1.946  (DML removes the confounding bias)
econml 90% CI empirical coverage of the true tau: 0.854

  from-scratch honest forest   corr with true tau 0.930   RMSE 1.421
  econml CausalForestDML       corr with true tau 0.955   RMSE 0.460

Those two numbers say different things and both matter. The correlations are high, so both
forests recover the SHAPE of the effect -- who benefits more than whom -- which is what a
targeting rule needs. The from-scratch forest's RMSE is the worse of the two because its LEVEL
is off: without orthogonalization it inherits the confounding, so it ranks correctly and
mis-states magnitudes. For ranking that is survivable; for reporting an effect size it is not.

AND THE INTERVAL IS THE WEAK PART. A nominal 90% interval covering 85.4% is short by
4.6%, and this is not sampling noise -- it is the known behaviour of forest confidence
intervals for a smooth effect surface, where the honest-splitting variance estimate does not
capture the smoothing bias. The point estimates are trustworthy enough to target with. The
intervals should be read as indicative rather than as calibrated 90% statements, and this
notebook's own ranking results do not depend on them.

5. What drives the heterogeneity?

Two reads on which covariates move the effect:

• Forest effect-importance — how often, and how usefully, each variable is split on for the effect (not the outcome). It should light up age and prior and ignore the rest.
• Effect partial dependence — sweep one covariate across its range, hold the others at their medians, and plot $\hat\tau$. Overlaid on the true $\tau$, this shows the forest recovering the built-in slopes ($-0.08$ in age, $-0.12$ in prior) without being told the functional form — the payoff over the parametric interaction model, which only worked because we hand-coded those interactions.

imp = cf.feature_importances_
order = np.argsort(imp)
fig, ax = plt.subplots(1, 3, figsize=(14, 4))
ax[0].barh([names[i] for i in order], imp[order], color="#2b6cb0"); ax[0].set_title("Causal-forest effect importance"); ax[0].set_xlabel("importance")

med = np.median(X, axis=0)
for a, j, lab in [(ax[1], 0, "age"), (ax[2], 2, "prior earnings ($1000s)")]:
    grid = np.linspace(np.percentile(X[:,j],2), np.percentile(X[:,j],98), 60)
    G = np.tile(med, (60,1)); G[:,j] = grid
    est = cf.effect(G); l,u = cf.effect_interval(G, alpha=0.1)
    true_line = 6 - 0.08*(G[:,0]-22) - 0.12*G[:,2]
    a.plot(grid, est, color="#2b6cb0", lw=2, label="causal forest $\\hat\\tau$")
    a.fill_between(grid, l, u, color="#2b6cb0", alpha=.15, label="90% CI")
    a.plot(grid, true_line, color="#c53030", lw=2, ls="--", label="true $\\tau$")
    a.set_xlabel(lab); a.set_ylabel("$\\hat\\tau$"); a.set_title(f"Effect vs {lab.split(' ')[0]}"); a.legend(fontsize=8)
plt.tight_layout(); plt.show()
print("Importance concentrates on prior earnings and age; the partial-dependence lines track the true downward slopes.")

Importance concentrates on prior earnings and age; the partial-dependence lines track the true downward slopes.

6. The payoff — who should the program enroll?

A CATE estimate becomes a policy the moment we rank people by $\hat\tau(x)$ and treat from the top down. With the true $\tau$ in hand we can grade that rule honestly. The targeting curve below plots the average realized effect among the treated as we expand the program from the highest-$\hat\tau$ workers outward, against two references:

• Perfect targeting (rank by the true $\tau$) — the unbeatable ceiling.
• Random / treat-everyone — the flat line at the ATE, what a program with no targeting delivers.

The gap between the forest curve and the flat line is the value of personalization: real earnings gains redirected from workers the program barely helps to those it helps most. The vertical marker is the break-even frontier — enrol only where $\hat\tau>0$.

o_hat  = np.argsort(-pe)        # rank test units by ESTIMATED effect (the deployable rule)
o_true = np.argsort(-tau_te)    # rank by TRUE effect (oracle ceiling)
frac = np.arange(1, len(pe)+1)/len(pe)
cum_hat  = np.cumsum(tau_te[o_hat])/np.arange(1, len(pe)+1)      # realized avg effect among treated
cum_true = np.cumsum(tau_te[o_true])/np.arange(1, len(pe)+1)

fig, ax = plt.subplots(figsize=(12.5, 5.6))
plt.plot(frac*100, cum_hat,  color="#2b6cb0", lw=2.6, label="target by causal forest $\\hat\\tau$")
plt.plot(frac*100, cum_true, color="#2f855a", lw=2.0, ls="--", label="perfect targeting (oracle)")
plt.axhline(tau_te.mean(), color="#c53030", lw=1.8, ls=":", label=f"treat everyone / random (ATE={tau_te.mean():.2f})")
share_pos = np.mean(pe>0)*100
plt.axvline(share_pos, color="0.4", lw=1.2, ls="-.")
ax.text(share_pos+1, 0.06, f"$\\hat\\tau>0$ for {share_pos:.0f}% of workers", fontsize=10, color="0.25", transform=ax.get_xaxis_transform())
plt.xlabel("% of workers enrolled (highest $\\hat\\tau$ first)"); plt.ylabel("average realized effect among enrolled ($1000s)")
plt.title("Targeting curve — the value of knowing who benefits", fontsize=12)
plt.tick_params(labelsize=10); plt.margins(x=0.01)
top20v = cum_hat[int(0.2*len(pe))-1]
plt.plot([20],[top20v], "o", color="#2b6cb0", ms=9, zorder=5)
plt.annotate(f"top 20% -> {top20v:.2f}  ({top20v/tau_te.mean():.1f}x the ATE)", xy=(20, top20v),
             xytext=(27, top20v+0.30), fontsize=10, color="#2b6cb0",
             arrowprops=dict(arrowstyle="->", color="#2b6cb0", lw=1.4))
plt.legend(fontsize=10, loc="upper right", framealpha=.95); plt.tight_layout(); plt.show()

top = o_hat[:int(0.2*len(pe))]
print(f"Treat the top 20% by forest tau: realized avg effect {tau_te[top].mean():.2f}  vs  treat-all {tau_te.mean():.2f}"
      f"  ->  {tau_te[top].mean()/tau_te.mean():.1f}x the per-enrollee gain.")
print(f"Enrolling only where tau_hat>0 covers {share_pos:.0f}% of workers and avoids the ~{np.mean(tau_te<0)*100:.0f}% the program would hurt.")

Treat the top 20% by forest tau: realized avg effect 4.07  vs  treat-all 2.06  ->  2.0x the per-enrollee gain.
Enrolling only where tau_hat>0 covers 90% of workers and avoids the ~9% the program would hurt.

7. Summary

Predictive CART splits to reduce outcome error; causal CART splits to reveal where a treatment's effect differs — and, made honest and ensembled, becomes the causal forest that estimates heterogeneous effects with confidence intervals.

What the notebook showed, end to end:

step	result
Naïve difference in means	badly confounded (biased up by selection into the program)
OLS with covariates	recovers the average effect — one number, no heterogeneity
OLS with hand-picked interactions	great CATE when you guess the right interactions (the econometric benchmark)
From-scratch honest causal tree	discovers age & prior as effect modifiers unaided; blocky effects
From-scratch honest forest	recovers the effect shape (corr ≈ 0.9) but a biased level under confounding
econml CausalForestDML	DML orthogonalization fixes the level and gives valid CIs; recovers the shape of $\tau(x)$ (correlation printed below)
Targeting curve	turns $\hat\tau(x)$ into a policy worth multiples of the per-enrollee gain of treating everyone

Where it sits. Causal forests are the ML-powered members of the causal-inference family: they take the random forest you already built and swap the splitting objective (effect heterogeneity) and add honesty + DML. Two threads continue from here in this arc:

• the identification side — potential outcomes, matching/propensity, IV, difference-in-differences — which is what licenses the unconfoundedness assumption we simply asserted;
• the double/debiased ML side — the residualization inside CausalForestDML, generalized (Chernozhukov et al.) to estimate a single well-identified effect with any ML nuisance models.

The companion causal_forests_R.ipynb fits the same problem with grf (causal_forest), the reference implementation from the authors of the method.