Lesson 8.2: Causal Forests with EconML · Applied Causal Inference

Written lesson

What a causal forest does differently

A regression forest splits to predict Y well. A causal forest (Wager and Athey; Athey, Tibshirani and Wager's generalised random forests) splits to find subgroups whose treatment effects differ. Two ingredients make its estimates trustworthy:

Honesty: one half of each tree's sample chooses the splits, the other half estimates the effects in the leaves, so the same noise is not used twice.
Orthogonalisation: in EconML's CausalForestDML, Y and D are first residualised with flexible models (exactly as in Module 7), and the forest works on the residuals. This removes confounding before looking for heterogeneity.

The result is a CATE estimate τ̂(x) with valid confidence intervals.

Python · Colab cell

!pip install -q econml lightgbm scikit-learn matplotlib seaborn
import numpy as np, pandas as pd
from econml.dml import CausalForestDML
from lightgbm import LGBMRegressor, LGBMClassifier
import matplotlib.pyplot as plt, seaborn as sns
rng = np.random.default_rng(11)
sns.set_theme(style="whitegrid")

Fitting CausalForestDML

We reuse the retraining-stipend world from Lesson 8.1. EconML distinguishes:

X: variables along which the effect may vary (effect modifiers), and
W: additional confounders used only in the first-stage models.

Python · Colab cell

def make_data(n):
    X = pd.DataFrame({
        "months_unemployed": rng.gamma(2, 4, n),
        "years_education": rng.integers(8, 21, n).astype(float),
        "age": rng.uniform(20, 60, n),
        "female": rng.binomial(1, 0.5, n).astype(float),
    })
    W = pd.DataFrame({"prior_earnings": rng.lognormal(10.3, 0.4, n) / 1_000, "region": rng.integers(0, 5, n)})
    tau = 500 + 250 * np.minimum(X.months_unemployed, 24) - 150 * (X.years_education - 12)
    p = 1 / (1 + np.exp(-(-0.5 + 0.05 * X.months_unemployed - 0.03 * (X.age - 40) + 0.02 * (W.prior_earnings - 30))))
    d = rng.binomial(1, p)
    y = 15_000 + 400 * W.prior_earnings + 300 * (X.years_education - 12) - 120 * X.months_unemployed \
        + 500 * W.region + d * tau + rng.normal(0, 3_000, n)
    return X, W, d, y.to_numpy(), tau.to_numpy()

X, W, d, y, tau = make_data(8_000)
X_test, W_test, _, _, tau_test = make_data(2_000)

cf = CausalForestDML(
    model_y=LGBMRegressor(n_estimators=300, learning_rate=0.05, verbose=-1),
    model_t=LGBMClassifier(n_estimators=300, learning_rate=0.05, verbose=-1),
    discrete_treatment=True, n_estimators=1_000, min_samples_leaf=20, cv=5, random_state=0,
)
cf.fit(y, d, X=X, W=W)
cate = cf.effect(X_test)
lo, hi = cf.effect_interval(X_test, alpha=0.05)
print(f"ATE estimate: {cf.ate(X):,.0f} (true {tau.mean():,.0f})")
print(f"PEHE: {np.sqrt(np.mean((cate - tau_test) ** 2)):,.0f}")
print(f"Coverage of 95% CATE intervals: {np.mean((tau_test >= lo) & (tau_test <= hi)):.2f}")

Intervals and drivers of heterogeneity

feature_importances_ shows how often each modifier was used to split, weighted by the heterogeneity it explained. It is a guide, not a causal claim about the modifier itself.

Python · Colab cell

imp = pd.Series(cf.feature_importances_, index=X.columns).sort_values()
imp.plot.barh(figsize=(7, 3), title="Causal forest: drivers of effect heterogeneity"); plt.show()

order = np.argsort(X_test.months_unemployed.to_numpy())
plt.figure(figsize=(10, 4))
plt.fill_between(X_test.months_unemployed.to_numpy()[order], lo[order], hi[order], alpha=.15, label="95% interval")
plt.scatter(X_test.months_unemployed, cate, s=4, alpha=.4, label="Estimated CATE")
plt.scatter(X_test.months_unemployed, tau_test, s=4, alpha=.4, label="True CATE")
plt.xlabel("Months unemployed"); plt.ylabel("Effect on earnings ($)"); plt.legend(); plt.show()

Is the heterogeneity real?

Forests always find some differences. A simple, honest check: split the test sample into groups by predicted CATE and estimate the average effect in each group with a doubly robust score. If the forest found real heterogeneity, the groups' effects should rise with the predicted rank. (This is the idea behind the "best linear predictor" and "sorted GATE" tests of Chernozhukov et al.)

Python · Colab cell

from sklearn.model_selection import cross_val_predict
X_t, W_t, d_t, y_t, tau_t = make_data(4_000)                       # fresh sample, never seen by the forest
F = pd.concat([X_t, W_t], axis=1)
e_hat = cross_val_predict(LGBMClassifier(n_estimators=200, verbose=-1), F, d_t, cv=5, method="predict_proba")[:, 1].clip(0.02, 0.98)
# Outcome models by treatment arm, cross-fitted, for doubly robust (AIPW) scores
mu1, mu0 = np.zeros(len(y_t)), np.zeros(len(y_t))
from sklearn.model_selection import KFold
for tr, te in KFold(5, shuffle=True, random_state=0).split(F):
    m1 = LGBMRegressor(n_estimators=300, verbose=-1).fit(F.iloc[tr][d_t[tr] == 1], y_t[tr][d_t[tr] == 1])
    m0 = LGBMRegressor(n_estimators=300, verbose=-1).fit(F.iloc[tr][d_t[tr] == 0], y_t[tr][d_t[tr] == 0])
    mu1[te], mu0[te] = m1.predict(F.iloc[te]), m0.predict(F.iloc[te])
dr = mu1 - mu0 + d_t * (y_t - mu1) / e_hat - (1 - d_t) * (y_t - mu0) / (1 - e_hat)

quartile = pd.qcut(cf.effect(X_t), 4, labels=["Q1 (lowest)", "Q2", "Q3", "Q4 (highest)"])
check = pd.DataFrame({"quartile": quartile, "doubly robust effect": dr, "true": tau_t})
print(check.groupby("quartile", observed=True).mean().round(0))

The why

Effects that rise steadily across predicted quartiles, estimated on data the forest never saw, are evidence that the heterogeneity is real and usable for targeting. Flat group effects mean the forest is fitting noise, and targeting on it would be arbitrary.

Checklist and recap

I can explain honesty and orthogonalisation in causal forests.
I can fit CausalForestDML with modifiers X and confounders W, and report intervals.
I validate heterogeneity on held-out data with doubly robust scores by predicted quartile.