S-Learner for CATE estimation

Source: R/metalearners.R

The S-Learner ("single learner") fits one outcome model \(\mu(x, t) = E[Y | X = x, T = t]\) that includes the treatment indicator as a regular feature. CATEs are obtained by contrasting the predictions with the treatment switched on and off: $$\hat\tau(x) = \hat\mu(x, 1) - \hat\mu(x, 0)$$

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

# S3 method for class 's_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 outcome model, e.g. mlr3::lrn("regr.ranger").

covariates

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

object

A fitted s_learner model.

newdata

A data.frame containing at least the covariate columns.

...

Ignored.

Value

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

Methods (by generic)

  • predict(s_learner): Predict CATEs for new data. Returns a numeric vector \(\hat\tau(x)\) with one element per row of newdata (only the covariate columns are used).

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 <- s_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.474322
# }