Augmented direct learning for conditional average treatment effect estimation with double robustness

Abstract

Inferring the heterogeneous treatment effect is a fundamental problem in many applications. In this paper, we focus on estimating the Conditional Average Treatment Effect (CATE), that is, the difference in the conditional mean outcome between treatments given covariates. Traditionally, Q-Learning based approaches estimate each conditional mean outcome. However, they are subject to model misspecification. Recently, flexible one-step methods to directly learn (D-Learning) the CATE without outcome model specifications have been proposed. However, they require a specification of the propensity score. We propose robust direct learning (RD-Learning), to augment D-learning, leading to doubly robust estimators of the treatment effect. The consistency for our CATE estimator is guaranteed if either the main effect model or the propensity score model is correctly specified. The framework can be used in both the binary and the multi-arm settings and is general enough to allow different function spaces and incorporate different generic learning algorithms. We conduct a thorough theoretical analysis of the prediction error of our CATE estimator using statistical learning theory under both linear and non-linear settings. The effectiveness of our proposed method is demonstrated by simulation studies and a real data example about an AIDS Clinical Trials study.