X-Learner for CATE estimation

Source: R/metalearners.R

The X-Learner (Künzel et al., 2019) proceeds in three stages:

  1. Fit group-specific outcome models \(\hat\mu_0(x)\) and \(\hat\mu_1(x)\) as in the T-Learner.

  2. Impute individual treatment effects, \(D_0 = \hat\mu_1(X_0) - Y_0\) for controls and \(D_1 = Y_1 - \hat\mu_0(X_1)\) for the treated, and regress them on the covariates to obtain \(\hat\tau_0(x)\) and \(\hat\tau_1(x)\).

  3. Combine the two estimates with propensity score weights: $$\hat\tau(x) = \hat e(x) \hat\tau_0(x) + (1 - \hat e(x)) \hat\tau_1(x)$$

x_learner(
  data,
  outcome = "y",
  treatment = "t",
  learner,
  ps_learner,
  tau_learner = NULL,
  covariates = NULL
)

# S3 method for class 'x_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".

learner

An mlr3 regression learner for the stage-1 outcome models.

ps_learner

An mlr3 classification learner for the propensity score model, e.g. mlr3::lrn("classif.log_reg").

tau_learner

Optional mlr3 regression learner for the stage-2 imputed-effect models. Defaults to a clone of learner.

covariates

Optional character vector of covariate columns. Defaults to all columns except the outcome and the treatment.

object

A fitted x_learner model.

newdata

A data.frame containing at least the covariate columns.

...

Ignored.

Value

A fitted CATE model of class c("x_learner", "cate_learner"); see s_learner().

Details

The X-Learner is particularly effective when the treated and control groups are of very different sizes.

Methods (by generic)

  • predict(x_learner): Predict CATEs for new data.

References

Künzel, S. R., Sekhon, J. S., Bickel, P. J., & Yu, B. (2019). Metalearners for estimating heterogeneous treatment effects using machine learning. Proceedings of the National Academy of Sciences, 116(10), 4156-4165.

Examples

# \donttest{
library(mlr3)
data(synth_train)
data(synth_test)
m <- x_learner(synth_train, outcome = "y", treatment = "t",
               learner = lrn("regr.rpart"),
               ps_learner = lrn("classif.rpart"))
tau_hat <- predict(m, synth_test)
pehe(synth_test$tau, tau_hat)
#> [1] 0.8750192
# }