Causal-Policy Forest for End-to-End Policy Learning
Masahiro Kato
TL;DR
The paper addresses policy learning under binary treatments by reframing welfare maximization as a CATE estimation problem restricted to binary predictors. It proves an equivalence between empirical welfare maximization and least-squares regression of the CATE when restricting predictions to $\{-1,1\}$, enabling an end-to-end learning approach. The causal-policy forest extends causal forests by producing leafwise policies through a sign rule on CATE estimates, with splits explicitly optimized to improve policy decisions. The method remains scalable and leverages honest forest design, achieving performance close to an oracle in simulations. This work bridges CATE estimation and policy learning, offering a practical, modular tool for end-to-end policy optimization in observational data settings.
Abstract
This study proposes an end-to-end algorithm for policy learning in causal inference. We observe data consisting of covariates, treatment assignments, and outcomes, where only the outcome corresponding to the assigned treatment is observed. The goal of policy learning is to train a policy from the observed data, where a policy is a function that recommends an optimal treatment for each individual, to maximize the policy value. In this study, we first show that maximizing the policy value is equivalent to minimizing the mean squared error for the conditional average treatment effect (CATE) under $\{-1, 1\}$ restricted regression models. Based on this finding, we modify the causal forest, an end-to-end CATE estimation algorithm, for policy learning. We refer to our algorithm as the causal-policy forest. Our algorithm has three advantages. First, it is a simple modification of an existing, widely used CATE estimation method, therefore, it helps bridge the gap between policy learning and CATE estimation in practice. Second, while existing studies typically estimate nuisance parameters for policy learning as a separate task, our algorithm trains the policy in a more end-to-end manner. Third, as in standard decision trees and random forests, we train the models efficiently, avoiding computational intractability.