Lab: Doubly Robust Estimation (AIPW)

Overview

In this lab you will estimate the average treatment effect of a job training program on earnings using three approaches: outcome regression (OLS), inverse probability weighting (IPW), and augmented inverse probability weighting (AIPW). You will see that the AIPW estimator is consistent if either the propensity score model or the outcome model is correctly specified — the "doubly robust" property.

What you will learn:

• How to estimate propensity scores and check overlap
• How to construct the IPW estimator from first principles
• How to build the AIPW (doubly robust) estimator
• Why double robustness matters when one model is misspecified
• How to compute standard errors via the influence function

Prerequisites: Familiarity with OLS regression and logistic regression. Understanding of potential outcomes notation is helpful.

Step 1: Simulate Observational Data

We create a dataset where treatment assignment depends on covariates (selection on observables).

1set.seed(42)
2n <- 2000
3
4age <- pmin(pmax(rnorm(n, 35, 10), 18), 65)
5educ <- pmin(pmax(rnorm(n, 12, 3), 6), 20)
6married <- rbinom(n, 1, 0.4)
7prev_earn <- rlnorm(n, 9.5, 0.8)
8
9logit_ps <- -2 + 0.03 * age - 0.1 * educ + 0.5 * married + 0.00005 * prev_earn
10true_ps <- plogis(logit_ps)
11treat <- rbinom(n, 1, true_ps)
12
13y0 <- 20000 + 500 * educ + 100 * age + 3000 * married + 0.3 * prev_earn + rnorm(n, 0, 5000)
14y1 <- y0 + 2500 + 200 * educ - 50 * age
15
16earnings <- treat * y1 + (1 - treat) * y0
17true_ate <- mean(y1 - y0)
18
19df <- data.frame(earnings, treat, age, educ, married, prev_earn)
20
21cat("True ATE:", round(true_ate), "\n")
22cat("Naive diff:", round(mean(earnings[treat == 1]) - mean(earnings[treat == 0])), "\n")

Expected output:

Sample data preview (first 5 rows):

Step 2: Estimate the Propensity Score

1# Estimate propensity score
2ps_model <- glm(treat ~ age + educ + married + prev_earn,
3              data = df, family = binomial)
4df$ps_hat <- predict(ps_model, type = "response")
5
6# Check overlap
7library(ggplot2)
8ggplot(df, aes(x = ps_hat, fill = factor(treat))) +
9geom_histogram(alpha = 0.5, position = "identity", bins = 40) +
10labs(x = "Propensity Score", fill = "Treated") +
11theme_minimal()
12
13cat("PS range (treated):", range(df$ps_hat[df$treat == 1]), "\n")
14cat("PS range (control):", range(df$ps_hat[df$treat == 0]), "\n")

Expected output:

Step 3: Outcome Regression (OLS)

1# Outcome regression
2ols <- lm(earnings ~ treat + age + educ + married + prev_earn, data = df)
3cat("OLS ATE:", coef(ols)["treat"], "\n")
4
5# Separate models for each treatment arm
6mu1_model <- lm(earnings ~ age + educ + married + prev_earn, data = df[df$treat == 1, ])
7mu0_model <- lm(earnings ~ age + educ + married + prev_earn, data = df[df$treat == 0, ])
8mu1_hat <- predict(mu1_model, newdata = df)
9mu0_hat <- predict(mu0_model, newdata = df)
10cat("Outcome regression (separate models):", mean(mu1_hat - mu0_hat), "\n")

Expected output:

Note: The OLS coefficient on treat captures the ATE under the assumption that the linear outcome model is correctly specified. Since the true DGP is linear in the outcome equation, OLS should perform well here.

Step 4: Inverse Probability Weighting (IPW)

1# IPW estimator
2ps <- pmin(pmax(df$ps_hat, 0.05), 0.95)
3
4w1 <- df$treat / ps
5w0 <- (1 - df$treat) / (1 - ps)
6
7# Horvitz-Thompson
8ate_ipw <- mean(w1 * df$earnings) - mean(w0 * df$earnings)
9
10# Hajek (normalized)
11ate_hajek <- sum(w1 * df$earnings) / sum(w1) - sum(w0 * df$earnings) / sum(w0)
12
13cat("Horvitz-Thompson IPW:", round(ate_ipw), "\n")
14cat("Hajek IPW:", round(ate_hajek), "\n")

Expected output:

The Horvitz-Thompson IPW can be unstable due to extreme weights (observations with very low propensity scores receive very large weights). The Hajek (normalized) estimator divides by the sum of weights, which stabilizes the estimate. Trimming propensity scores at [0.05, 0.95] further reduces variance.

Step 5: The AIPW (Doubly Robust) Estimator

1# AIPW by hand
2aipw_scores <- (mu1_hat - mu0_hat
3              + df$treat * (df$earnings - mu1_hat) / ps
4              - (1 - df$treat) * (df$earnings - mu0_hat) / (1 - ps))
5
6ate_aipw <- mean(aipw_scores)
7se_aipw <- sd(aipw_scores) / sqrt(n)
8
9cat("AIPW ATE:", round(ate_aipw), " (SE:", round(se_aipw), ")\n")
10cat("95% CI: [", round(ate_aipw - 1.96 * se_aipw), ",",
11  round(ate_aipw + 1.96 * se_aipw), "]\n")
12
13# Or use the AIPW package
14# library(AIPW)
15# aipw_obj <- AIPW$new(Y = df$earnings, A = df$treat,
16#                       W = df[, c("age","educ","married","prev_earn")])
17# aipw_obj$fit()$summary()

Expected output:

All three estimators should be close to the true ATE in this case because both the propensity score model and the outcome model are correctly specified. The advantage of AIPW becomes clear when one model is deliberately misspecified (Step 6).

Step 6: Double Robustness in Action

Now deliberately misspecify one of the two models to see the doubly robust property.

1# Misspecify propensity score (only age)
2ps_wrong <- pmin(pmax(predict(glm(treat ~ age, family = binomial), type = "response"), 0.05), 0.95)
3
4# Misspecify outcome model (only age)
5mu1_wrong <- predict(lm(earnings ~ age, data = df[df$treat == 1, ]), newdata = df)
6mu0_wrong <- predict(lm(earnings ~ age, data = df[df$treat == 0, ]), newdata = df)
7
8# AIPW: correct outcome + wrong PS
9aipw_good_out <- mean(mu1_hat - mu0_hat + df$treat * (df$earnings - mu1_hat) / ps_wrong
10                   - (1 - df$treat) * (df$earnings - mu0_hat) / (1 - ps_wrong))
11
12# AIPW: wrong outcome + correct PS
13aipw_good_ps <- mean(mu1_wrong - mu0_wrong + df$treat * (df$earnings - mu1_wrong) / ps
14                  - (1 - df$treat) * (df$earnings - mu0_wrong) / (1 - ps))
15
16cat("Correct outcome + wrong PS:", round(aipw_good_out), "\n")
17cat("Wrong outcome + correct PS:", round(aipw_good_ps), "\n")
18cat("True ATE:", round(true_ate), "\n")

Expected output:

This table demonstrates the doubly robust property in action:

• Row 2: Even with a badly misspecified propensity score (using only age), AIPW recovers the correct ATE because the outcome model is correct. The IPW-only estimator, which relies solely on the propensity score, is biased.
• Row 3: Even with a badly misspecified outcome model (using only age), AIPW recovers the correct ATE because the propensity score is correct. The OLS-only estimator, which relies solely on the outcome model, is biased.
• Row 4: When both models are wrong, AIPW offers no guarantee and produces a biased estimate.

Exercises

1. Use machine learning for nuisance models. Replace logistic regression and OLS with random forests for the propensity score and outcome models. Does the AIPW estimate improve?

2. Estimate the ATT. Modify the AIPW formula to target the average treatment effect on the treated. How does it differ from the ATE?

3. Bootstrap confidence intervals. Implement a nonparametric bootstrap for the AIPW estimator and compare the bootstrap SE with the influence-function SE.

4. Vary the degree of overlap. Increase the coefficients in the propensity score model so that overlap deteriorates. At what point does AIPW break down?

Summary

In this lab you learned:

• Outcome regression relies on a correctly specified outcome model; IPW relies on a correctly specified propensity score model
• The AIPW estimator combines both approaches, providing consistency if either model is correct (the doubly robust property)
• Propensity score trimming is essential for practical stability of IPW and AIPW
• The influence function provides analytic standard errors without bootstrapping
• When both models are wrong, no estimator can save you — invest effort in specifying both models well
• AIPW is the foundation for many modern causal inference methods, including targeted learning (TMLE) and double machine learning