DoubleML: Partially Linear Models

53.1 The DoubleML package (30 min)

DoubleML (Bach, Chernozhukov, Kurz & Spindler) implements Day 52's estimator and its relatives with a common interface:

1. Describe the data: DoubleMLData(df, y_col, d_cols, x_cols).
2. Pick a model class: DoubleMLPLR (partially linear, today), DoubleMLIRM and DoubleMLIIVM (interactive, tomorrow), and more.
3. Pass any scikit-learn-style learners for the nuisance functions: ml_l for E[Y|X], ml_m for E[D|X].
4. .fit(), then read .coef, .se, .confint().

One practical detail: DoubleML draws its cross-fitting folds from numpy's global random generator, so each script calls np.random.seed(...) first. Without it, every run gives slightly different numbers.

53.2 Three kinds of nuisance learners (70 min)

Create day53/plr.py (about 5 minutes; the random forest is slow):

# day53/plr.py   (python -m day53.plr)
"""
The DoubleML package: the partially linear regression (PLR) model  spend = theta * coupon + g(X) + error,
estimated by orthogonal scores and cross-fitting, with three kinds of nuisance learners, and repeated
cross-fitting (n_rep) so the answer doesn't depend on one random split.
"""
import doubleml as dml
import numpy as np
import polars as pl
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LassoCV
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

from causalkit.coupons import X_COLS, coupons, public
from causalkit.learners import reg, variance_weighted

np.random.seed(53)  # DoubleML draws its cross-fitting folds from numpy's global RNG

full = coupons()
df = public(full).to_pandas()
data = dml.DoubleMLData(df[X_COLS + ["coupon", "spend"]], y_col="spend", d_cols="coupon", x_cols=X_COLS)
print(data)

learners = {
    "lasso": make_pipeline(StandardScaler(), LassoCV(cv=3)),
    "random forest": RandomForestRegressor(300, min_samples_leaf=20, max_features=0.5, n_jobs=-1, random_state=0),
    "LightGBM": reg(),
}
rows = []
for name, learner in learners.items():
    plr = dml.DoubleMLPLR(data, ml_l=learner, ml_m=learner, n_folds=5, n_rep=3)
    plr.fit()
    rmse = plr.evaluate_learners()
    rows.append({"learners": name, "estimate": round(plr.coef[0], 2), "SE": round(plr.se[0], 2),
                 "95% CI": f"[{plr.confint().iloc[0, 0]:.1f}, {plr.confint().iloc[0, 1]:.1f}]",
                 "RMSE of l(X)": round(float(np.mean(rmse["ml_l"])), 1),
                 "RMSE of m(X)": round(float(np.mean(rmse["ml_m"])), 3)})
with pl.Config(tbl_width_chars=120, fmt_str_lengths=30):
    print(pl.DataFrame(rows))
print(f"Target: {variance_weighted(full['_tau'].to_numpy(), full['_ps'].to_numpy()):.2f}")

python -m day53.plr

│ learners      ┆ estimate ┆ SE   ┆ 95% CI       ┆ RMSE of l(X) ┆ RMSE of m(X) │
│ lasso         ┆ 13.4     ┆ 0.6  ┆ [12.2, 14.6] ┆ 38.6         ┆ 0.459        │
│ random forest ┆ 16.33    ┆ 0.58 ┆ [15.2, 17.5] ┆ 38.0         ┆ 0.457        │
│ LightGBM      ┆ 16.16    ┆ 0.56 ┆ [15.1, 17.3] ┆ 36.5         ┆ 0.457        │
Target: 16.39

• With lasso nuisance models, DML gives exactly OLS's answer ($13.4): linear models can't capture the confounding, and no amount of orthogonality fixes a nuisance model that's systematically wrong.
• Random forest and LightGBM both land within half a standard error of the target.
• evaluate_learners() reports each nuisance model's out-of-fold RMSE. LightGBM predicts spend best (36.5), and that's the honest way to choose between learners: compare nuisance prediction accuracy, never the effect estimates themselves (choosing the learner that gives the answer you like is fishing, Day 26).
• n_rep=3 repeats cross-fitting with three different splits; more in 53.5.

53.3 What PLR's single number means (50 min)

Concept. PLR assumes one θ for everyone. When effects differ (they do: −$26 to +$69), what does θ̂ estimate? Not the ATE: a weighted average of customers' effects, weighting each by e(X)(1 − e(X)), the variance of their coupon given their features. Customers whose coupon was a coin flip count most; customers the targeting model was nearly certain about count least.

Create day53/what_plr_estimates.py:

# day53/what_plr_estimates.py   (python -m day53.what_plr_estimates)
"""
When effects differ across customers, PLR's single theta is NOT the ATE. It weights each customer's
effect by e(X)(1 - e(X)): customers whose coupon was hardest to predict count most, customers the
targeting model was nearly sure about (e near 0 or 1) count almost nothing.
"""
import numpy as np
import polars as pl

from causalkit.coupons import coupons, truth
from causalkit.learners import variance_weighted

full = coupons()
tau, ps = full["_tau"].to_numpy(), full["_ps"].to_numpy()
t = truth(full)
print(f"ATE (everyone):                     {t['ATE']:.2f}")
print(f"ATT (customers who got a coupon):   {t['ATT']:.2f}")
print(f"PLR target (variance-weighted):     {variance_weighted(tau, ps):.2f}")

w = ps * (1 - ps)
bands = (full.with_columns(weight=pl.Series(w / w.sum()), ps_band=pl.Series(np.digitize(ps, [0.2, 0.4, 0.6, 0.8])))
         .group_by("ps_band").agg(customers=pl.len(), share_of_weight=pl.col("weight").sum().round(3),
                                  average_effect=pl.col("_tau").mean().round(1)).sort("ps_band")
         .with_columns(propensity=pl.col("ps_band").replace_strict(
             {0: "< 0.2", 1: "0.2-0.4", 2: "0.4-0.6", 3: "0.6-0.8", 4: ">= 0.8"})).drop("ps_band"))
print(bands.select("propensity", "customers", "share_of_weight", "average_effect"))

python -m day53.what_plr_estimates

ATE (everyone):                     16.51
ATT (customers who got a coupon):   20.26
PLR target (variance-weighted):     16.39
│ propensity ┆ customers ┆ share_of_weight ┆ average_effect │
│ < 0.2      ┆ 1360      ┆ 0.042           ┆ -2.2           │
│ 0.2-0.4    ┆ 5397      ┆ 0.275           ┆ 8.5            │
│ 0.4-0.6    ┆ 6146      ┆ 0.369           ┆ 16.9           │
│ 0.6-0.8    ┆ 5014      ┆ 0.256           ┆ 23.9           │
│ >= 0.8     ┆ 2083      ┆ 0.058           ┆ 30.4           │

Here the PLR target happens to be close to the ATE ($16.39 vs $16.51), because the down-weighted customers at both ends (effects −2.2 and +30.4) roughly cancel. That's luck, not a property of PLR. The customers the targeting model was surest about (propensity ≥ 0.8) have the biggest effects (+30.4) and only 6% of the weight. If the question is "what did the coupons we sent achieve?" (the ATT, $20.26), PLR is the wrong tool. Tomorrow's IRM targets the ATE or ATT directly. PLR is the right tool when the effect is plausibly constant, or the treatment is continuous (Days 58–60).

53.4 ML can't find a confounder that isn't there (35 min)

Try it. Drop deal_share from the features and rerun PLR with LightGBM. Create day53/missing_confounder.py.

# day53/missing_confounder.py   (python -m day53.missing_confounder)
import doubleml as dml
import numpy as np

from causalkit.coupons import X_COLS, coupons, public
from causalkit.learners import clf, reg, variance_weighted

full = coupons()
df = public(full).to_pandas()
for name, cols in [("all features", X_COLS), ("without deal_share", [c for c in X_COLS if c != "deal_share"])]:
    np.random.seed(53)
    plr = dml.DoubleMLPLR(dml.DoubleMLData(df[cols + ["coupon", "spend"]], "spend", "coupon", cols),
                          ml_l=reg(), ml_m=clf(), n_folds=5).fit()
    print(f"{name:<20} {plr.coef[0]:.2f} (SE {plr.se[0]:.2f})")
print(f"Target {variance_weighted(full['_tau'].to_numpy(), full['_ps'].to_numpy()):.2f}")

all features         16.12 (SE 0.56)
without deal_share   19.87 (SE 0.54)
Target 16.39

53.5 Day 53 mini-project: does the answer depend on the split? (65 min)

Cross-fitting splits customers into folds at random. Run PLR with 20 different splits and look at the spread. Create day53/split_robustness.py (about a minute and a half):

# day53/split_robustness.py   (python -m day53.split_robustness)
"""
Mini-project: cross-fitting splits the data at random, so the estimate depends a little on the split.
Run PLR with 20 different splits (n_rep=20) and look at the spread. DoubleML reports the median
across splits, with a standard error that includes the between-split variation.
"""
import doubleml as dml
import numpy as np

from causalkit.coupons import X_COLS, coupons, public
from causalkit.learners import clf, reg

np.random.seed(53)  # DoubleML draws its cross-fitting folds from numpy's global RNG

df = public(coupons()).to_pandas()
data = dml.DoubleMLData(df[X_COLS + ["coupon", "spend"]], "spend", "coupon", X_COLS)
plr = dml.DoubleMLPLR(data, ml_l=reg(), ml_m=clf(), n_folds=5, n_rep=20)
plr.fit()
each = plr.all_coef.ravel()
print(f"20 splits: estimates from {each.min():.2f} to {each.max():.2f} (SD {each.std():.2f})")
print(f"Median across splits: {plr.coef[0]:.2f}  (SE {plr.se[0]:.2f}, compared with a typical single-split SE "
      f"{np.median(plr.all_se):.2f})")

python -m day53.split_robustness

20 splits: estimates from 15.98 to 16.34 (SD 0.09)
Median across splits: 16.16  (SE 0.56, compared with a typical single-split SE 0.56)

Different splits move the estimate by up to $0.36, with an SD of $0.09: small compared with the standard error ($0.56), so the split isn't driving the answer. DoubleML reports the median over splits and adds the between-split spread to the standard error (here it's too small to change the second decimal). With smaller samples or noisier learners, the spread can be much larger; always use n_rep of 5 or more for a number you'll report.

Report: "DML with LightGBM, 5 folds, median over 20 splits". Anyone rerunning your analysis with a different seed should get the same answer to the precision you report.

53.6 Save your work with Git (20 min)

git add .
git commit -m "Day 53: DoubleML PLR, learners, what PLR estimates, repeated cross-fitting"
git push

Day 53 checklist

• I can fit DoubleMLPLR with different nuisance learners and compare them by nuisance RMSE
• I can explain what PLR's single coefficient means when effects differ across customers
• I know that ML removes regularisation bias, not omitted-variable bias
• I use repeated cross-fitting and report the median over splits