DR-Learner for CATE estimation

The DR-Learner ("doubly robust learner"; Kennedy, 2023) is a two-stage meta-learner that turns the augmented inverse propensity weighting (AIPW) score used by ate_dr() into a model that predicts individual CATEs.

Cross-fit the nuisance functions: the propensity score $\hat e(x) = P(T = 1 | X = x)$ and the potential-outcome models $\hat\mu_0(x)$, $\hat\mu_1(x)$ from a single outcome model that includes the treatment as a feature.
Form the doubly robust pseudo-outcome $$\psi_i = \hat\mu_1(X_i) - \hat\mu_0(X_i) + \frac{T_i (Y_i - \hat\mu_1(X_i))}{\hat e(X_i)} - \frac{(1 - T_i)(Y_i - \hat\mu_0(X_i))}{1 - \hat e(X_i)}$$ and regress it on the covariates to obtain the final CATE model $\hat\tau(x) = E[\psi | X = x]$.

dr_learner(
  data,
  outcome = "y",
  treatment = "t",
  outcome_learner,
  ps_learner,
  tau_learner = NULL,
  covariates = NULL,
  folds = 5,
  ps_trim = NULL
)

# S3 method for class 'dr_learner'
predict(object, newdata, ...)

Arguments

data: A data.frame with the outcome, treatment and covariates.
outcome: Name of the outcome column. Default "y".
treatment: Name of the binary (0/1) treatment column. Default "t".
outcome_learner: An mlr3 regression learner for the stage-1 outcome model, e.g. mlr3::lrn("regr.ranger"). The treatment indicator is included as a feature and potential outcomes are predicted by setting it to 0 and 1.
ps_learner: An mlr3 classification learner used as the propensity score model, e.g. mlr3::lrn("classif.log_reg"). Its predict type is set to "prob" automatically.
tau_learner: Optional mlr3 regression learner for the second-stage pseudo-outcome regression. Defaults to a clone of outcome_learner.
covariates: Optional character vector of covariate columns. Defaults to all columns except the outcome and the treatment.
folds: Number of cross-fitting folds for the nuisance models. Default 5. Cross-fitting is a core ingredient of the DR-Learner; folds = 1 (in-sample nuisance estimates) is allowed but not recommended with flexible learners.
ps_trim: Optional trimming for the estimated propensity scores. Either a single value a (clips to [a, 1 - a]) or a length-2 range. NULL (default) applies no trimming.
object: A fitted dr_learner model.
newdata: A data.frame containing at least the covariate columns.
...: Ignored.

Value

A fitted CATE model of class c("dr_learner", "cate_learner"). Use predict() to obtain CATE estimates for new data and ate() for their average.

Details

Because the pseudo-outcome has conditional mean equal to the true CATE whenever either nuisance is correctly specified, the second-stage regression targets the CATE directly. Cross-fitting the nuisances (the default folds = 5) makes the pseudo-outcome orthogonal to nuisance estimation error, so flexible learners can be used without overfitting bias. Unlike ate_dr(), which averages $\psi_i$ to a scalar ATE, the DR-Learner keeps the second-stage model and can therefore predict CATEs on new test data.

Methods (by generic)

predict(dr_learner): Predict CATEs for new data.

References

Kennedy, E. H. (2023). Towards optimal doubly robust estimation of heterogeneous causal effects. Electronic Journal of Statistics, 17(2), 3008-3049.

Examples

# \donttest{
library(mlr3)
data(synth_train)
data(synth_test)
set.seed(1)
m <- dr_learner(synth_train, outcome = "y", treatment = "t",
                outcome_learner = lrn("regr.rpart"),
                ps_learner = lrn("classif.rpart"),
                ps_trim = 0.01)
tau_hat <- predict(m, synth_test)
pehe(synth_test$tau, tau_hat)
#> [1] 0.7577948
# }