Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
42,263 characters · 11 sections · 31 citation commands
Policy Learning with Abstention
Policy learning algorithms guide high-stakes decisions in domains like personalized medicine and public policy by constructing individualized treatment rules from observational data. Yet, a critical deficit of existing methods is their failure to abstain when faced with high uncertainty. In safety-critical applications, forcing a decision when evidence is weak can be harmful; the responsible action may be to defer to a human expert or a trusted baseline policy. One way to facilitate this is through allowing learned policies to abstain --- that is, instead of just assigning a unit to treatment or control (denoted “1” or “0” respectively), they can abstain by outputting a special symbol “$\ast$”.
While abstaining algorithms have been developed in the classification literature bousquet2021fast, generalizing these algorithms to the setting of policy learning from observational data is a non-trivial task. The primary difficulty is the counterfactual nature of the problem: the learner only ever observes an outcome under treatment or control for any given unit, never both. Because of this, the learner must transform samples (either via inverse propensity weighting or through a doubly-robust correction) into “pseudo-outcomes” --- quantities representing the unit-specific contrasts in outcomes between treatment and control.
In this work, we present a framework formalizing policy learning with abstention. The goal of the learner is to use observational data $Z = (X, D, Y)$ consisting of covariates $X \in \mathcal{X}$, treatment $D \in \{0, 1\}$, and outcome $Y \in \mathbb{R}$ to learn a high reward abstaining treatment policy $\pi : \mathcal{X} \rightarrow \{0, 1, \ast\}$. Here, $\pi$ is assumed to exists in some pre-defined policy class $\Pi$ that implicitly encodes constraints of the problem (e.g.\ $\Pi$ is the class of all depth 3 decision trees, thus enforcing explainability). Our per-section contributions are as follows:
In addition to the above, we also conduct simulations to compare our safe policy improvement algorithm (Algorithm (ref)) to relevant benchmarks thomas2015high. In sum, our work illustrates the practicality and importance of studying policy learning with abstention, and further provides insight into how abstention can impact other areas of policy learning as well.
\paragraph{Related work on abstention.} Classification with abstention/a reject option traces back to chow1970optimum, who characterize the optimal error–reject trade-off, prescribing abstention whenever the posterior risk of mis-classification exceeds the reject cost. herbei2006reject consider classification with a reject option in a statistical learning setting, and bartlett2008reject introduce a convex “reject-hinge” loss to make abstention trainable in large-scale settings. Other researchers have developed abstaining classification algorithms for SVMs grandvalet2008svmreject, multiclass classification ramaswamy2015consistent, boosting cortes2016boosting, online learning cortes2018online, pasteris2024bandits,neu2020fast, vanderHoeven2022regretvariance, and even deep learning settings geifman2017selective, geifman2019selectivenet. madras2018defer study a setting where the prediction of an abstaining algorithm is replaced by a downstream expert, but don't provide regret/excess risk guarantees. Most closely related to this paper is the work of bousquet2021fast, who develop an empirical risk minimization-based algorithm for abstention that enables fast $O(1/n)$ risk minimization rates without standard margin/realizability assumptions.
\paragraph{Policy Learning.} Our work builds off of the literature on policy learning, particularly works on welfare maximization/regret minimization athey2021policy, hirano2009asymptotics, kitagawa2018should, manski2004statistical. Also related to our work is the literature of safe policy improvement, which aims to produce a policy that, with high-probability, improves over some baseline “safe” policy. Existing approaches include using hypothesis testing to select between the baseline and learned policy thomas2015high, cho2025cspi, bootstrap the baseline policy with a learned policy in regions of low confidence laroche2019safe, simao2019safe, and directly minimizing “negative regret” ghavamzadeh2016safe. Our work provides new perspective onto safe policy improvement, showing how abstaining algorithms can naturally be used to construct improved policies. Other complementary areas that are not directly considered in the work are policy evaluation dudik2011doubly, karampatziakis2021off, offline reinforcement learning moodie2007demystifying, moodie2012q, murphy2003optimal, policy learning under partial identification kallus2018confounding, and inference on values of optimal treatment policies luedtke2016statistical, chen2023inference, whitehouse2025inference.
We assume the learner observes i.i.d.\ draws $Z=(X,D,Y)$ with covariates $X\in{ \mathcal{X}}$, realized action $D\in\{0,1\}$, and outcome $Y\in[0,1]$. Let $(Y(1),Y(0))$ denote potential outcomes, and write \[ g_o(d,x)\coloneqq \mathbb{E}\!\big[Y(d)\mid X=x\big], ~~~ \tau_o(x)\coloneqq g_o(1,x)-g_o(0,x), \] as respectively the expected outcome mapping and conditional average treatment effect (CATE). We denote the propensity score by $p_o(x)\coloneqq \mathbb{P}(D{=}1\mid X{=}x)$. We assume the following of the data generating process.
We let $\hat g, \hat p, \hat \tau$ denote generic estimates of $g_o, p_o , \tau_o$, respectively. For $q\ge 1$ and a distribution $P_X$ on ${ \mathcal{X}}$, define $\|f\|_{P_X,q}\coloneqq\big(\mathbb{E}_{P_X}|f(X)|^q\big)^{1/q}.$
\paragraph{Policies.} A policy $\pi:{ \mathcal{X}}\to\{0,1\}$ prescribes a binary treatment to a unit with covariates $X$. Let $\Pi$ denote a class of feasible treatment policies. We assume the $\Pi$ has bounded complexity, controlled through the VC dimension vapnik2013nature.
The in-class optimal policy is $ \pi^*\in\operatorname*{arg\,max}_{\pi\in\Pi} V(\pi)$, where $V(\pi) := \mathbb{E}_{P_X}[Y(\pi(X))]$ denotes the expected welfare/reward under policy $\pi$. We also consider the Bayes policy, which maximizes value over all measurable binary policies: $\pi^B(x)\;\coloneqq\;\operatorname*{arg\,max}_{d\in\{0,1\}} \mathbb{E}\big[Y(d)\mid X=x\big] \;=\;\mathbbm{1}\!\!\left\{\tau_o(x)\ge 0\right\}$.
\paragraph{Policy value.} For binary policies, we can rewrite the policy value $\mathbb{E}_{P_X}[Y(\pi(X))]$ as
Likewise, we define the conditional value for policy $\pi$ at realized covariates $x$ as
Let $\mathbb{E}_n$ denote the expectation with respect to the empirical sample distribution. We define the inverse propensity weighted (IPW) analogue of Equation (ref) as
Additionally, we define the normalized IPW contribution for a policy $\pi$ at an observation $z=(x,d,y)$:
This quantity is the per-sample term whose average recovers the policy value up to the factor $\kappa$; the extra $\kappa$ factor ensures the function is uniformly bounded. We use $f(z)$ and $f(x,d,y)$ interchangeably. By $Y\in[0,1]$ and $p_o(X)\in[\kappa,1-\kappa]$, we have $f_\pi(Z)\in[0,1]$ a.s., and $ \mathbb{E}\!\big[f_\pi(Z)\big]=\kappa\,V(\pi). $ As a notational shorthand, write $\mathbb{E}[f_\pi]$ and $\mathbb{E}_n[f_\pi]$ for $\mathbb{E}[f_\pi(Z)]$ and $\mathbb{E}_n[f_\pi(Z)]$, respectively. We use this notation throughout. In particular, when we write $\mathbb{E}_n \left|f_{\hat\pi}-f_\pi\right|$, we mean the empirical average of pointwise differences between the normalized IPW scores: $ \tfrac{1}{n}\sum_{i=1}^n \left|\,f_{\hat\pi}(x_i,d_i,y_i)-f_{\pi}(x_i,d_i,y_i)\,\right|. $
We now present our framework for policy learning with abstention. In this framework, the learner can choose to defer on recommending a treatment to any unit based on their observed covariates $X$. For instance, the learner may not want to prescribe treatment when the estimated CATE for a unit is small. When the learner abstains, they receive a small, additive reward over the value of a random guess.
More formally, an abstaining policy is a mapping $\pi : \mathcal{X} \rightarrow \{0,1,*\}$, where $\pi(X) = *$ denotes the deferral of treatment decision. When the learner abstains, they receive reward $\tfrac{Y(1)+Y(0)}{2}+p$, where $p \geq 0$ is some fixed bonus. We define the expected reward mapping under abstention as
Likewise, we define the expected reward/welfare $V^{(p)}(\pi)$ of an abstaining policy with bonus $p$ (and its empirical analogue $V_n^{(p)}(\pi)$) as {
} Note that we have $V^{(p)}(\pi) = V(\pi)$ and $V_n^{(p)}(\pi) = V_n(\pi)$ for any policy $\pi$ that does not abstain.
The goal of the learner is to use either experimental (Section (ref)) or observational (Section (ref)) data to learn an abstaining policy with small abstaining regret, which is defined with respect to a binary policy class $\Pi$ as
where again $\pi^\ast$ is the in-class optimal policy and $\pi$ is some potentially abstaining policy. We define the classical regret just as $\mathrm{Reg}_n(\pi) := \mathrm{Reg}_n^{(0)}(\pi)$. The learner is constrained to returning a policy ${\widetilde{\pi}}$ that aligns with some $\pi \in \Pi$ when it does not abstain, i.e.\ ${\widetilde{\pi}}(x) = \pi(x)$ when ${\widetilde{\pi}}(x) \neq *$.
We first describe an algorithm for learning an abstaining policy when the treatment assignment mechanism (i.e.\ propensity) for the data is known. Our algorithm (Algorithm (ref)) works by first computing an empirical welfare maximizer $\widehat{\pi}$ from $\Pi$ using half of the data (where welfare is determined with respect to the IPW empirical value, $V_n$). Then, it computes a set of “near-optimal” policies whose empirical welfare is close to that of the empirical risk minimizer, modifying these policies to abstain precisely when they disagree $\widehat{\pi}$. Finally, the algorithm returns the empirical abstaining welfare maximizer in this restricted set of policies on the second half of the data. The following theorem shows that Algorithm (ref) obtain fast abstaining regret rates, per the formulation of regret in Equation (ref).
More details and the full proof are provided in Appendix (ref).
.
We now provide some intuition for the proof. Typically, to obtain fast regret rates for policy learning, there must exist a margin on the CATE, i.e.\ the existence of some value $ h > 0$ such that
The abstention bonus $p$ can be viewed as a “synthetic” margin --- if the learner abstains when $|\tau_o(X)| < p$, they automatically accumulate higher reward than any binary treatment assignment policy. Likewise, when $|\tau_o(X)| \geq p$, the learner will be more certain in their decisions, and hence unlikely to abstain. While this intuition is just heuristic, we provide a fully rigorous proof in Appendix (ref). One important thing to note is that, in Algorithm (ref), we never actually need to estimate the CATE. This is of particular importance when the CATE may be a highly complicated function and is thus difficult to capture using statistical learning methods.
When the additive bonus is $p = 0$, our algorithm can no longer be expected to obtain fast $O(1/n)$ rates. The following proposition shows that, in this undiscounted setting, Algorithm (ref) still obtains $O(1/\sqrt{n})$ regret rates outlined, which are known to be generally unimprovable without margin or realizability assumption athey2021policy. In this setting, we can convert an abstaining policy into a binary one by having it assign a treatment uniformly at random in the set $\{0,1\}$. The following proposition provides this slower regret rate, showing there is no added risk of using our algorithm over just empirical welfare maximization even in setting without the additive bonus. We prove this result in Appendix (ref).
The proof is provided in Appendix (ref).
In observational settings, where either the treatment policy is unknown or choice of treatment is endogenous, we can no longer directly run Algorithm (ref). One might think to salvage the algorithm by replacing the propensity $p_o(X)$ in the definition of $V_n$ with an ML estimate $\widehat{p}(X)$. However, unless $\widehat{p}$ converges to $p_o$ in probability at fast, parametric rates\footnote{In particular, one would need $\|p_o - \widehat{p}\|_{P_X, 2} = O_\mathbb{P}(n^{-1/2})$}, Algorithm (ref) will exhibit sub-optimal regret. Instead, one must leverage more sophisticated methods to still enable low-regret learning.
Taking inspiration from the literature on semiparametric estimation chernozhukov2018double, bang2005doubly, foster2023orthogonal, athey2021policy, we introduce a doubly-robust analogue of Algorithm (ref) that furnishes fast regret rates even when the propensity is unknown. Our algorithm uses ML estimates $\widehat{g}, \widehat{p}$ of the regression $g_o(d, x)$ and propensity $p_o(x)$ to “de-bias” observed outcomes $Y$ into more robust pseudo-outcomes. These nuisance estimates are assumed to be independent of the sample. In more detail, we define
This quantity is known as pseudo-outcomes because when $\widehat{g} = g_o$ and $\widehat{p} = p_o$, one can check that $\mathbb{E}[\widehat{\varphi}(X, d, Y) \mid X] = \mathbb{E}[Y(d) \mid X]$. With pseudo-outcomes, we can then define the doubly-robust abstaining welfare/value on the sample, $V_{n, \mathrm{DR}}^{(p)}$, via {
} When $p = 0$, we arrive at the non-abstaining sample welfare $V_{n,\mathrm{DR}} := V^{(p)}_{n,\mathrm{DR}}$. To obtain a doubly-robust analogue of Algorithm (ref), use and expanded $\alpha$ \footnote{We may simply increase the constant in Step (ref) of Algorithm (ref); when the nuisance product error is $o_p(n^{-1/2})$, this yields the same rate as Theorem (ref).} and simply replace every occurrence of $V_n$, $V^{(p)}_n$ by $V_{n, \mathrm{DR}}$ and $V^{(p)}_{n, \mathrm{DR}}$, respectively. See Algorithm (ref) in Appendix (ref).
We now state the main theorem of this subsection.
See Appendix (ref) for the full proof and supporting lemmas.
The above theorem yields a bound on regret that has two terms --- one quickly decaying term followed by another, random term whose decay is non-obvious. We now discuss conditions under which this second term is negligible. Note that one can obtain (via Cauchy-Schwarz) the following upper bound on the product error in terms of individual error rates in nuisance functions, \[ \big\|\hat p-p\big\|_{P_X, 4} \sum_{d \in \{0,1 \}}\big\|\ \hat g(d,\cdot)-g_o(d,\cdot)\big\|_{P_X,4} \] Thus, for the second term in Theorem (ref) to be negligible, we need $\mathrm{Err}_{\mathrm{DR}} \leq c_{\delta}/\sqrt{n}$ with high probability, which will occur if { $\max\{\|\hat p-p\|_{P_X,4},\|\hat g(0,\cdot)-g_o(0,\cdot)\|_{P_X,4},\|\hat g(1,\cdot)-g_o(1,\cdot)\|_{P_X,4}\}$} $\lesssim n^{-1/4}$ with probability at least $1 - \delta$. In particular, this can be accomplished under a variety of learnability assumptions (finite VC-dimension, learnability by trees, etc.) \newline
In the previous section, we established a formal framework for performing policy learning with abstention. We now explore connections between our results in Section (ref) (in particular, Algorithm (ref)) and other aspects of policy learning. This includes developing (non-abstaining) algorithms for policy learning without standard margin assumptions, safely improving policies relative some baseline treatment strategy, and relating abstaining value/welfare $V^{(p)}$ to policy learning in the presence of distribution shift.
In policy learning, fast regret rates are sometimes achievable when the CATE is deterministically bounded away from zero, as outlined in the margin condition of Equation (ref). In particular, prior works have shown that $O(n^{-1})$ rates are possible in the realizable setting, where the Bayes-optimal policy $\pi^B$ lies in the class $\Pi$ kitagawa2018should, Luedtke2020. However, it is generally impossible to know if a policy class $\Pi$ (say, all depth 3 decision trees) contains the optimal treatment policy in advance. Thus, one should aim to obtain fast regret rates under in agnostic settings, where $\Pi$ may not contain the optimal treatment strategy.
In the classification literature, recent works achieve fast learning rates in agnostic settings ben2014sample, bousquet2021fast. These works eschew the realizability assumption, instead imposing a finite combinatorial diameter on the policy class. In words, a class of policies $\Pi$ which bounds the maximal number of points of disagreement between any two policies in the class. More formally, the combinatorial diameter is defined by
Note that $D$ may be infinite even for classes with finite VC dimension.
We prove an analogous result for policy learning (Algorithm (ref)) under the assumption of a finite combinatorial diameter. Further, we consider settings where the policy class does not have a finite combinatorial diameter but instead we have access to a CATE estimation oracle (Algorithm (ref)). The proofs are deferred to Appendix (ref). We start by stating our main theorem for agnostic policy learning under an assumption of finite combinatorial dimension. {\textcolor{black}{
} }
Proof details are provided in Appendix (ref).
Next, suppose we have an oracle for estimator $\widehat{\tau}_o$ that satisfies for any $(X,D,Y) \sim \mathbb{P}$ the following convergence rate with probability greater than $1-\delta$, \[ \left\lVert {\widehat{\tau}} - \tau \right\rVert_{\mathbb{P}, 2} \leq c_\delta n^{-\beta}. \]
Proof details are provided in Appendix (ref).
Theorem (ref) unifies the regression-rate guarantee Luedtke2020 and fast rate in the realizable case kitagawa2018should. In particular, when $\beta = 1/2$, one recovers $n^{-1}$ regret. More generally, if $\beta < 1/2$ but the policy class $\Pi$ is expressive enough that $V(\pi^B(X)) - V(\pi^*(X))$ is small (e.g., $\leq c/n^{-\alpha}$ for some $\alpha>0$), the resulting bound improves upon a direct plug-in classifier ($\pi(X) = \mathbbm{1}\!\!\left\{\hat \tau(X) > 0\right\}$). Thus, either a strong CATE oracle or an “almost” realizable policy class is sufficient to surpass the $1/\sqrt{n}$ regret.
Safe policy improvement (SPI) addresses the deployment problem under uncertainty: one should adopt a new policy only if it can be certified to outperform a baseline policy $\omega$ with high probability. Following recent work on safe policy improvement thomas2015high, cho2025cspi, we use sample splitting (into $\mathcal{D}_{\text{train}}$ and $\mathcal{D}_{\text{test}}$) to separate the tasks of candidate policy selection and safety testing. This splitting enables one-sided lower confidence bounds (LCBs) on the improvement over the baseline: $V(\hat\pi)-V(\omega)$. We treat Algorithm (ref) as a black box to propose a \(\{0,1,*\}\)-valued policy ${\widetilde{\pi}}$ learned on $\mathcal{D}_{train}$. To obtain a deployable binary policy, we impute abstentions with the baseline policy to obtain a candidate policy Equation (ref).
On $D_{\text{test}}$, we estimate $V(\hat\pi)-V(\omega)$ via IPW/DR and compute a one-sided LCB at level $1-\delta$; we generate a sequence of candidate policies by running the abstention learner over a grid of bonuses, $\mathcal{P}=\{p_1<\cdots<p_k\}$, which trades off overlap (more abstention) against improvement (more overrides), and apply a Bonferroni adjustment with $z_{1-\delta/k}=\Phi^{-1}(1-\delta/k)$. Algorithm (ref) then: (i) selects the grid, (ii) learns abstaining policies on $D_{\text{train}}$, (iii) replaces abstentions with $\omega$, and (iv) tests each on $D_{\text{test}}$, returning the first with $\mathrm{LCB}>0$.
\paragraph{Safe Policy Improvement Experiments}
We compare Algorithm (ref) with an Empirical Welfare Maximizer (EWM) and prior work on synthetic data. In particular, we consider the two variants of High-Confidence Policy Improvement (HCPI) described in thomas2015high: one based on finite-sample confidence intervals and one based on a $t$-test statistic. For EWM, we choose the empirical welfare maximizer as the candidate policy and test it against the baseline in exactly the same way as in Algorithm (ref) (with $k=1$).
The design comprises two variants: (1) we use a baseline policy at a fixed a and vary the noise variance $(0.01\text{--}1.0)$ (see (ref)); and (2) we vary the baseline–optimal value gap at a fixed noise level (see (ref)). Each parameter setting is replicated 100 times; policy value is computed by IPW (with known propensities). For Algo (ref), we use a grid of abstention bonuses \(\mathcal{P} = \{0, 0.01, 0.05, 0.10, 0.20\}\). We report value gain, mistake rate (probability of returning a policy worse than the baseline), and improvement rate (probability of returning a policy strictly better than the baseline). Full configuration and implementation details are deferred to Appendix (ref).
We evaluate all methods at significance level \(\delta = 0.05\). Across more than \(2{,}000\) repetitions, our Safe Policy Learner controls the Type-I error at or below \(0.05\). In the very small–sample regime ($n\!\le\!500$), EWM performs better than the alternatives, but once $n\!\ge\!1000$ our method dominates: (i) it achieves the highest improvement rate and the largest mean value gain among accepted policies across baseline–optimal gaps ((ref)), and (ii) it maintains both higher power and competitive mistake rates as noise increases ((ref)). Overall, for moderate-to-large samples, Algorithm (ref) delivers the best safety–power tradeoff, improving more often and by more, while meeting the $\delta$-level error control.
The abstention value $V^{(p)}$ has a natural interpretation as protection against outcome distribution shift. Specifically, consider a setting where the true data distribution of $(Y(0),Y(1))\mid X$ is different from our training data. This would happen when our observational data is outdated and cannot reflect the true effect of the current treatment of interest.
While deterministic policies are optimal without such an outcome distribution shift, they can be problematic otherwise. To mathematically formulate this intuition, we assume the true potential outcome distribution lies in some $W_1$-ball of the distribution that generates our observations, i.e.
The robust objective is then to maximize a policy’s worst–case value over $\mathcal{P}_{\alpha}(\mathbb{P}_{\mathtt{train}})$. Specifically, consider policies $\pi$ taking values in $\{0,1,\tfrac12\}$, where $\tfrac12$ denotes a uniform random choice between the two treatments. In this model, the robust objective coincides with maximizing the abstention value $V^{(p)}$ for a specific bonus level, as formalized below.
Distributional uncertainty acts like an outcome–level penalty of size $\alpha$ when the policy commits to an arm, and a half–penalty under randomization; hence, maximizing the robust objective is equivalent to maximizing $V^{(p)}$ with $p=\alpha/2$.
In this paper, we introduced a framework for policy learning with abstention. Building off of the classification with abstention literature, we formulated a value/welfare that incentivizes abstention by offering a small, additive reward on top of the value of a random guess in regions of uncertainty. We described algorithms that obtain fast abstaining regret rates in both known and unknown propensity settings, and showed how to apply our algorithm to various disparate policy learning problems. There are still many other interesting open directions related to policy learning with abstention. First, our framework only handles binary treatments. There are many settings (such as in medicine) where designing policies to optimize a continuous treatment (or dosage) may be more appropriate. Second, implementing Algorithm (ref) efficiently for popular policy classes such as decision trees is also an interesting open direction. Likewise, we evaluate the performance of algorithms through regret, but one could also consider loss-based modes of policy evaluation as well.
Vasilis Syrgkanis is supported by NSF Award IIS-2337916. Ayush Sawarni is partially supported by NSF Award IIS-2337916.