R-Learner for CATE estimation

The R-Learner (Nie & Wager, 2021) is the heterogeneous-effect generalisation of the partially linear DML estimator (ate_dml()). It builds on the Robinson decomposition of the model $Y = \tau(X) T + g(X) + \varepsilon$:

Cross-fit the two nuisance functions, the conditional outcome mean $\hat m(x) = E[Y | X = x]$ (the treatment is excluded from the features) and the propensity score $\hat e(x) = P(T = 1 | X = x)$, and form the residuals $\tilde Y = Y - \hat m(X)$ and $\tilde T = T - \hat e(X)$.
Estimate the CATE by minimising the R-Loss $$\hat\tau = \arg\min_\tau \frac{1}{n} \sum_i \left[ \tilde Y_i - \tilde T_i \, \tau(X_i) \right]^2$$ which is fitted as a weighted regression of the pseudo-outcome $\tilde Y_i / \tilde T_i$ on the covariates with weights $\tilde T_i^2$.

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

# S3 method for class 'r_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 conditional outcome mean $m(x) = E[Y | X]$ (the treatment is excluded from the features), e.g. mlr3::lrn("regr.ranger").
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 weighted 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 r_learner model.
newdata: A data.frame containing at least the covariate columns.
...: Ignored.

Value

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

Details

The objective minimised here is exactly the r_loss() used elsewhere for model selection, so the R-Learner is the estimator that directly targets that score. When the treatment effect is constant it reduces to ate_dml(); unlike that scalar estimator, the R-Learner keeps the second-stage model and can predict CATEs on new test data.

The weighted formulation requires a tau_learner that supports observation weights ("weights" %in% learner$properties). If it does not, the R-Learner falls back to an unweighted regression of the pseudo-outcome and issues a warning. As with ate_dml(), small $\tilde T_i$ inflate the pseudo-outcome, so ps_trim is recommended with flexible propensity learners.

Methods (by generic)

predict(r_learner): Predict CATEs for new data.

References

Nie, X., & Wager, S. (2021). Quasi-oracle estimation of heterogeneous treatment effects. Biometrika, 108(2), 299-319.

Examples

# \donttest{
library(mlr3)
data(synth_train)
data(synth_test)
set.seed(1)
m <- r_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.4553936
# }