policytree: Policy learning via doubly robust empirical welfare maximization over trees

Abstract

The problem of learning treatment assignment policies from randomized or observational data arises in many fields. For example, in personalized medicine, we seek to map patient observables (like age, gender, heart pressure, etc.) to a treatment choice using a data-driven rule. There has recently been a considerable amount of work on statistical methodology for policy learning show that if we only consider policies π restricted to a class Π with finite VC dimension and have access to data from a randomized trial with n samples, then an empirical welfare maximization algorithm achieves regret that scales as √ VC(Π)/n. Athey & Wager (2017) extend this result to observational studies via doubly robust scoring, and Zhou et al. (2018) further consider the case with multiple treatment choices (in particular, the regret will depend on the tree depth, feature space, and number of actions). The package policytree for R (R Core Team, 2020) implements the multi-action doubly robust approach of Zhou et al. (2018) in the case where we want to learn policies π that belong to the class Π of depth-k decision trees. In order to use policytree, the user starts by specifying a set of doubly robust scores for policy evaluation; the software then carries out globally optimal weighted search over decision trees. It is well known that finding an optimal tree of arbitrary depth is NP-hard. However, if we restrict our attention to trees of depth k, then the problem can be solved in polynomial time. Here, we implement the global optimization via an exhaustive (unconstrained) tree search that runs in O(P k N k (log N + D) + P N log N) time, where N is the number of individuals, P the number of characteristics observed for each individual and D is the number of available treatment choices (see details below). If an individual's characteristics only takes on a few discrete values, the runtime can be reduced by a factor of N k. Additionally, an optional approximation parameter lets the user control how many splits to consider.