T-Learner for CATE estimation

Source: R/metalearners.R

The T-Learner ("two learners") fits separate outcome models for the control and treated groups, \(\mu_0(x) = E[Y | X = x, T = 0]\) and \(\mu_1(x) = E[Y | X = x, T = 1]\), and estimates the CATE as their difference: $$\hat\tau(x) = \hat\mu_1(x) - \hat\mu_0(x)$$

t_learner(
  data,
  outcome = "y",
  treatment = "t",
  learner,
  learner1 = NULL,
  covariates = NULL
)

# S3 method for class 't_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 used for the control-group model (and, if learner1 is NULL, also for the treated-group model).

learner1

Optional separate mlr3 regression learner for the treated group. 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 t_learner model.

newdata

A data.frame containing at least the covariate columns.

...

Ignored.

Value

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

Methods (by generic)

  • predict(t_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 <- t_learner(synth_train, outcome = "y", treatment = "t",
               learner = lrn("regr.rpart"))
tau_hat <- predict(m, synth_test)
pehe(synth_test$tau, tau_hat)
#> [1] 1.831545
# }