Lesson 7.3: DoubleML in Practice: PLR, IRM and a Real Health Application

Written lesson

The DoubleML workflow

The DoubleML package (Bach, Chernozhukov, Kurz, Spindler) implements cross-fitted, orthogonal estimators. Every analysis follows four steps:

Wrap the data in a DoubleMLData object (outcome, treatment, covariates).
Choose learners for the nuisance functions: ml_l for E[Y|X] and ml_m for E[D|X] (PLR) or ml_g/ml_m (IRM).
Choose the model: PLR (constant effect, any treatment type) or IRM (binary treatment, heterogeneous effects, ATE/ATT, an AIPW-type score).
fit(), then read summary, confint(), and the nuisance losses.

Python · Colab cell

!pip install -q DoubleML lightgbm scikit-learn causaldata
import numpy as np, pandas as pd
import doubleml as dml
from lightgbm import LGBMRegressor, LGBMClassifier
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.linear_model import LassoCV, LogisticRegressionCV
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import PolynomialFeatures, StandardScaler
rng = np.random.default_rng(3)

PLR on the digital-skills simulation

Python · Colab cell

N, P = 3_000, 20
X = rng.normal(0, 1, size=(N, P))
p = 1 / (1 + np.exp(-(0.5 * X[:, 0] - 0.3 * X[:, 1] ** 2 + 0.2 * X[:, 2] * X[:, 4])))
D = rng.binomial(1, p)
Y = 5_000 * D + 2_000 * np.sin(X[:, 0]) + 3_000 * X[:, 3] + 1_000 * X[:, 1] ** 2 + rng.normal(0, 4_000, N)
df = pd.DataFrame(X, columns=[f"X_{i}" for i in range(P)]).assign(D=D, Y=Y)

data = dml.DoubleMLData(df, y_col="Y", d_cols="D")
ml_l = LGBMRegressor(n_estimators=300, learning_rate=0.05, max_depth=4, verbose=-1, random_state=42)
ml_m = LGBMClassifier(n_estimators=300, learning_rate=0.05, max_depth=4, verbose=-1, random_state=42)

plr = dml.DoubleMLPLR(data, ml_l, ml_m, n_folds=5)
plr.fit()
print(plr.summary.round(2))
print("95% CI:\n", plr.confint().round(0))

Comparing learners

DML's guarantee depends on the nuisance models being good. Compare learners by their cross-fitted prediction error, which DoubleML stores for you, and prefer the one with the lowest loss, not the one that gives the estimate you like.

Python · Colab cell

learners = {
    "Lasso (+ squares)": (make_pipeline(PolynomialFeatures(2, include_bias=False), StandardScaler(), LassoCV(cv=3)),
                          make_pipeline(StandardScaler(), LogisticRegressionCV(cv=3, max_iter=2_000))),
    "Random forest":     (RandomForestRegressor(n_estimators=300, min_samples_leaf=5, n_jobs=-1, random_state=1),
                          RandomForestClassifier(n_estimators=300, min_samples_leaf=5, n_jobs=-1, random_state=1)),
    "LightGBM":          (ml_l, ml_m),
}
rows = []
for name, (l, m) in learners.items():
    model = dml.DoubleMLPLR(data, l, m, n_folds=5).fit()
    losses = model.evaluate_learners()
    rows.append((name, model.coef[0], model.se[0], losses["ml_l"].item(), losses["ml_m"].item()))
print(pd.DataFrame(rows, columns=["learner", "theta", "se", "RMSE E[Y|X]", "RMSE E[D|X]"]).round(2))

IRM on NHEFS: smoking cessation and weight

The NHEFS follow-up study (used throughout Hernán and Robins' Causal Inference: What If) records whether adult smokers quit between 1971 and 1982 (qsmk) and their weight change in kg (wt82_71). Quitting was not randomised: age, sex, race, education, smoking intensity, years smoking, exercise, activity level and baseline weight all differ between quitters and non-quitters.

Python · Colab cell

from causaldata import nhefs_complete
nh = nhefs_complete.load_pandas().data
covs = ["sex", "race", "age", "school", "smokeintensity", "smokeyrs", "wt71", "exercise", "active"]
nh = nh[covs + ["qsmk", "wt82_71"]].dropna()
nh = pd.get_dummies(nh, columns=["exercise", "active"], drop_first=True, dtype=float)

naive = nh.wt82_71[nh.qsmk == 1].mean() - nh.wt82_71[nh.qsmk == 0].mean()
nh_data = dml.DoubleMLData(nh, y_col="wt82_71", d_cols="qsmk")
irm = dml.DoubleMLIRM(
    nh_data,
    ml_g=RandomForestRegressor(n_estimators=500, min_samples_leaf=10, max_features=0.5, random_state=0),
    ml_m=RandomForestClassifier(n_estimators=500, min_samples_leaf=10, max_features=0.5, random_state=0),
    n_folds=5, n_rep=3, trimming_threshold=0.01,
)
irm.fit()
print(f"Observations: {len(nh):,} | quitters: {int(nh.qsmk.sum())}")
print(f"Naive difference: {naive:.2f} kg")
print(irm.summary.round(2))

The naive gap (about 2.5 kg) understates the effect: quitters were older and lighter smokers in ways that predict lower weight gain. DML gives roughly 3.3–3.5 kg, close to the IPW and g-formula estimates in Hernán and Robins. n_rep=3 repeats the random fold split and aggregates, so the answer does not depend on one lucky split.

Reporting template

Key takeaway

A DML result is not just a number. Report:

Estimand and model: ATE via IRM (or constant effect via PLR).
Identification: the DAG and the unconfoundedness argument; list the covariates.
Learners: which ones, their tuning, and their cross-fitted losses.
Overlap: propensity score range and trimming threshold.
Stability: estimates across learners and repeated splits (n_rep).
Sensitivity: how strong an unobserved confounder would need to be (Module 14).

Exercise

Re-run the NHEFS IRM with LightGBM learners and with the ATT score (score="ATTE"). Put the three estimates (RF ATE, LightGBM ATE, RF ATT) in one table. Is the conclusion robust?

Show a solution

Python · Colab cell

irm_lgbm = dml.DoubleMLIRM(nh_data, LGBMRegressor(n_estimators=200, learning_rate=0.03, max_depth=3, verbose=-1),
                           LGBMClassifier(n_estimators=200, learning_rate=0.03, max_depth=3, verbose=-1),
                           n_folds=5, trimming_threshold=0.01).fit()
irm_att = dml.DoubleMLIRM(nh_data, RandomForestRegressor(n_estimators=500, min_samples_leaf=10, random_state=0),
                          RandomForestClassifier(n_estimators=500, min_samples_leaf=10, random_state=0),
                          score="ATTE", n_folds=5, trimming_threshold=0.01).fit()
print(pd.DataFrame({"RF ATE": irm.summary.iloc[0, :2], "LGBM ATE": irm_lgbm.summary.iloc[0, :2],
                    "RF ATT": irm_att.summary.iloc[0, :2]}).round(2))

Checklist and recap

I can fit PLR and IRM models in DoubleML and read their output.
I choose nuisance learners by cross-fitted loss.
I can report a DML estimate with its assumptions, learners, overlap and stability.