Learning Treatment Allocations with Risk Control Under Partial Identifiability
Sofia Ek, Dave Zachariah
TL;DR
This work addresses learning treatment allocations under partial identifiability by focusing on constraining treatment risk $T(\pi)$ while minimizing population risk $R(\pi)$. It develops a certifiable, finite-sample policy-learning framework that uses miscalibration bounds with $\Gamma\ge1$ to obtain upper bounds on $R(\pi)$ and $T(\pi)$ via weights $\overline{W}^{\Gamma}$, and enforces a high-probability guarantee $\mathbb{P}(T(\pi)\le \tau) \ge 1-\alpha$ across feasible $\Gamma$. The method splits data into $\mathcal{D}_m$ to learn a nominal policy $\pi(X;t)$ under a constraint, and $\mathcal{D}_n$ to construct an upper confidence bound $\overline{T}^{\alpha}_n(t)$ and select $t_n$ to certify the risk bound; an interpretable fast-and-frugal-tree implementation is used for practical deployment. Experiments on synthetic data and the STAR dataset show that the certified policies achieve the desired treatment-risk control with high probability while reducing population risk, demonstrating applicability to precision medicine and other safety-critical domains.
Abstract
Learning beneficial treatment allocations for a patient population is an important problem in precision medicine. Many treatments come with adverse side effects that are not commensurable with their potential benefits. Patients who do not receive benefits after such treatments are thereby subjected to unnecessary harm. This is a `treatment risk' that we aim to control when learning beneficial allocations. The constrained learning problem is challenged by the fact that the treatment risk is not in general identifiable using either randomized trial or observational data. We propose a certifiable learning method that controls the treatment risk with finite samples in the partially identified setting. The method is illustrated using both simulated and real data.