Lesson 7.1: Doubly Robust Estimation: AIPW

Written lesson

Two routes to the same counterfactual

Under unconfoundedness there are two ways to build the missing potential outcome:

Outcome modelling (regression adjustment): fit μ₁(X) = E[Y | D=1, X] and μ₀(X) = E[Y | D=0, X], then average μ₁(X) − μ₀(X) over everyone.
Treatment modelling (IPW): model e(X) = P(D=1 | X) and reweight (Module 6).

Each is biased if its model is wrong. The augmented IPW (AIPW) estimator combines them:

ψᵢ = μ₁(Xᵢ) − μ₀(Xᵢ) + Dᵢ (Yᵢ − μ₁(Xᵢ)) / e(Xᵢ) − (1 − Dᵢ)(Yᵢ − μ₀(Xᵢ)) / (1 − e(Xᵢ)), ATE = mean(ψᵢ)

Start from the outcome-model prediction, then correct it with the weighted residuals. If the outcome model is right, the residuals average zero and the correction vanishes. If the propensity model is right, the correction removes the outcome model's error. Either one suffices.

Python · Colab cell

!pip install -q numpy pandas scikit-learn statsmodels matplotlib seaborn
import numpy as np, pandas as pd
from sklearn.linear_model import LinearRegression, LogisticRegression
import matplotlib.pyplot as plt, seaborn as sns
rng = np.random.default_rng(3)
sns.set_theme(style="whitegrid")

A Medicaid-style expansion with non-linear confounding

Enrolment in an insurance expansion depends on income and chronic illness, and so does the outcome (a health index). Income enters both equations non-linearly, which is typical in real data.

Python · Colab cell

def simulate(n=4_000):
    income = rng.normal(0, 1, n)
    chronic = rng.binomial(1, 0.3, n)
    logit = -0.3 - 1.2 * income + 0.8 * chronic + 0.6 * income ** 2
    d = rng.binomial(1, 1 / (1 + np.exp(-logit)))
    y = 2.0 * d + 3 * income - 2 * chronic + 1.5 * income ** 2 + rng.normal(0, 2, n)
    X = pd.DataFrame({"income": income, "chronic": chronic})
    return X, d, y

X, d, y = simulate()
TRUE_ATE = 2.0
print(f"Naive difference: {y[d == 1].mean() - y[d == 0].mean():.2f} (truth {TRUE_ATE})")

Double robustness by simulation

We write the three estimators with pluggable features, so we can feed each model either the right specification (with income²) or a wrong one (linear only).

Python · Colab cell

def features(X, correct):
    F = X.copy()
    if correct:
        F["income2"] = F.income ** 2
    return F.to_numpy()

def estimators(X, d, y, outcome_ok, ps_ok):
    Fo, Fp = features(X, outcome_ok), features(X, ps_ok)
    mu1 = LinearRegression().fit(Fo[d == 1], y[d == 1]).predict(Fo)
    mu0 = LinearRegression().fit(Fo[d == 0], y[d == 0]).predict(Fo)
    e = LogisticRegression(C=1e6, max_iter=1_000).fit(Fp, d).predict_proba(Fp)[:, 1]
    e = np.clip(e, 0.01, 0.99)
    ra = np.mean(mu1 - mu0)
    ipw = np.mean(d * y / e) - np.mean((1 - d) * y / (1 - e))
    psi = mu1 - mu0 + d * (y - mu1) / e - (1 - d) * (y - mu0) / (1 - e)
    return ra, ipw, psi.mean(), psi

results = []
for outcome_ok in (True, False):
    for ps_ok in (True, False):
        reps = [estimators(*simulate(), outcome_ok, ps_ok)[:3] for _ in range(100)]
        ra, ipw, aipw = np.mean(reps, axis=0)
        results.append((outcome_ok, ps_ok, ra, ipw, aipw))
table = pd.DataFrame(results, columns=["outcome model OK", "PS model OK", "Regression adj.", "IPW", "AIPW"]).round(2)
print(table, f"\n(truth = {TRUE_ATE})")

Read the table: regression adjustment fails when the outcome model is wrong, IPW fails when the propensity model is wrong, and AIPW is fine as long as one of them is right. When both are wrong, nothing saves you, which is why Lesson 7.2 replaces hand-written models with flexible machine learning.

Standard errors from the influence function

The ψᵢ values are also the estimator's influence function: their standard deviation divided by √n is a valid standard error (when the models are estimated well, as Lesson 7.2 ensures with cross-fitting).

Python · Colab cell

ate, psi = estimators(X, d, y, True, True)[2:]
se = psi.std(ddof=1) / np.sqrt(len(psi))
print(f"AIPW ATE = {ate:.3f}, SE = {se:.3f}, 95% CI [{ate - 1.96 * se:.3f}, {ate + 1.96 * se:.3f}]")

Key takeaway

AIPW has a second, deeper property: Neyman orthogonality. Small errors in μ or e only affect the estimate at second order (through their product). That is what allows machine-learning models, which converge slowly, to be plugged in without ruining inference. It is the foundation of Double Machine Learning.

Exercise

Repeat the simulation table with n = 500 instead of 4,000 and look at the spread (standard deviation) of the 100 estimates for each estimator, not just the mean. Which is most variable, and why?

Show a solution

Python · Colab cell

for outcome_ok, ps_ok in [(True, True), (False, True)]:
    reps = np.array([estimators(*simulate(500), outcome_ok, ps_ok)[:3] for _ in range(100)])
    print(f"outcome_ok={outcome_ok}, ps_ok={ps_ok}: mean {reps.mean(0).round(2)}, sd {reps.std(0).round(2)}")
# IPW is the noisiest: dividing by small propensity scores amplifies noise. AIPW is close to the
# regression estimator's precision when the outcome model is right.

Checklist and recap

I can write regression adjustment, IPW and AIPW.
I can demonstrate double robustness by simulation.
I can compute an AIPW standard error from the influence function.