Doubly Robust Direct Learning for Estimating Conditional Average Treatment Effect

Inferring the heterogeneous treatment effect is a fundamental problem in the sciences and commercial applications. In this paper, we focus on estimating Conditional Average Treatment Effect (CATE), that is, the difference in the conditional mean outcome between treatments given covariates. Traditionally, Q-Learning based approaches rely on the estimation of conditional mean outcome given treatment and covariates. However, they are subject to misspecification of the main effect model. Recently, simple and flexible one-step methods to directly learn (D-Learning) the CATE without model specifications have been proposed. However, these methods are not robust against misspecification of the propensity score model. We propose a new framework for CATE estimation, robust direct learning (RD-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. As a by-product, we develop a competitive statistical inference tool for the treatment effect, assuming the propensity score is known. We provide theoretical insights to the proposed method using risk bounds 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.