rlasso: R-learner, implemented via glmnet (lasso)

In xnie/rlearner: R-learner for Heterogeneous Treatment Effect Estimation

Description

R-learner, as proposed by Nie and Wager (2017), implemented via glmnet (lasso)

Usage

rlasso(
  x,
  w,
  y,
  alpha = 1,
  k_folds = NULL,
  foldid = NULL,
  lambda_y = NULL,
  lambda_w = NULL,
  lambda_tau = NULL,
  lambda_choice = c("lambda.min", "lambda.1se"),
  rs = FALSE,
  p_hat = NULL,
  m_hat = NULL,
  penalty_factor = NULL
)

Arguments

x	the input features
w	the treatment variable (0 or 1)
y	the observed response (real valued)
alpha	tuning parameter for the elastic net
k_folds	number of folds for cross-fitting
foldid	user-supplied foldid. Must have length equal to length(w). If provided, it overrides the k_folds option.
lambda_y	user-supplied lambda sequence for cross validation in learning E[y|x]
lambda_w	user-supplied lambda sequence for cross validation in learning E[w|x]
lambda_tau	user-supplied lambda sequence for cross validation in learning the treatment effect E[y(1) - y(0) | x]
lambda_choice	how to cross-validate for learning the treatment effect tau; choose from "lambda.min" or "lambda.1se"
rs	whether to use the RS-learner (logical).
p_hat	user-supplied estimate for E[W|X]
m_hat	user-supplied estimte for E[Y|X]
penalty_factor	user-supplied penalty factor, a vector of length the same as the number of covariates in x.

Value

an rlasso object

Examples

## Not run: 
n = 100; p = 10

x = matrix(rnorm(n*p), n, p)
w = rbinom(n, 1, 0.5)
y = pmax(x[,1], 0) * w + x[,2] + pmin(x[,3], 0) + rnorm(n)

rlasso_fit = rlasso(x, w, y)
rlasso_est = predict(rlasso_fit, x)

## End(Not run)

xnie/rlearner documentation built on April 11, 2021, 12:49 a.m.