Policy Learning with Abstention

AI summary

Overview

Research area: Offline policy learning / causal inference from observational data, sitting at the intersection of statistical learning theory, semiparametric estimation, and individualized treatment rule design.

Technical level: Advanced. The paper is a theory-driven contribution built on VC-dimension arguments, inverse propensity weighting, doubly robust estimation, and high-probability regret bounds, though the motivating intuition is stated in accessible terms.

Scope: The paper introduces a formal framework in which learned treatment policies may decline to decide ("abstain") in exchange for a small additive reward, proves regret guarantees for a two-stage learner under known and unknown propensities, and shows how the same machinery improves several other policy learning problems.

What This Paper Is About

Standard policy learning algorithms always output a treatment decision — assign treatment 1 or 0 — even when the evidence for that unit is weak. In high-stakes settings such as medicine or public policy that is risky; the safer action may be to defer to a human expert or a trusted default. This paper formalizes what it means for a learned policy to abstain (output a special symbol for "defer"), gives the abstaining decision a small bonus p on top of the value of a random guess, and builds an algorithm that provably earns near-optimal reward quickly. The authors then show that this abstention machinery is not a narrow trick but a reusable tool for other policy learning problems.

Key Contributions

1. A formal framework for policy learning with abstention. The paper defines abstaining policies as mappings from covariates to the set {0, 1, }, where a decision to abstain yields reward (Y(1)+Y(0))/2 + p for a fixed bonus p ≥ 0. It defines a "safe" or abstaining regret, Reg_n^(p)(π) := V(π*) − V^(p)(π), measured against the best binary policy π in a class Π, and requires that any non-abstaining decision of the learned policy agrees with some policy in Π.

2. A two-stage learner with fast O(1/n) regret guarantees (Algorithm 1). The algorithm first computes an empirical welfare maximizer on one half of the data, then builds a set of near-optimal policies, and lets each of those abstain exactly where it disagrees with the maximizer. Under known propensities, Theorem 3.1 shows regret bounded roughly by (d log(n/d) + log(1/δ)) / (p n κ²) with probability at least 1 − δ, where d is the VC dimension of Π and κ is the overlap parameter.

3. A doubly robust extension for unknown propensities (Theorem 3.4). Replacing the IPW value with a doubly robust value V_{n,DR}^{(p)} — built from ML estimates ĝ and p̂ of the outcome regression and propensity — yields regret bounded by the same fast term plus an additive term Err_DR² / (p κ²), where Err_DR measures the product error between propensity and outcome-regression estimates. The paper notes that this second term is negligible when the nuisance estimates converge at rate n^(−1/4) in the relevant L4 norms.

4. Three downstream applications of the same black box. Algorithm 1 is used to obtain improved margin-condition guarantees without the usual realizability assumption (Theorems 4.1 and 4.2), to construct a safe policy improvement procedure that certifies improvement over a baseline with high probability (Algorithm 3), and to connect abstention with robustness to distribution shift (Proposition 4.3).

Main Findings

• Fast abstaining rates when propensities are known. Theorem 3.1 establishes that Algorithm 1 achieves regret on the order of (d log(n/d) + log(1/δ)) / (p n κ²) with probability at least 1 − δ, an O(1/n)-type rate. The bonus p acts as a "synthetic margin": abstaining when |τ_o(X)| < p automatically yields more reward than any binary assignment, while abstention is unlikely when |τ_o(X)| ≥ p.

• No added risk when the bonus is zero. Proposition 3.3 shows that running Algorithm 1 with any fixed p ∈ (0,1) and then evaluating its value at p = 0 still gives standard regret on the order of (1/κ) sqrt((d log(n/d) + log(1/δ))/n) — an O(1/√n) rate described as matching rates due to Athey and Wager (2021) and as generally unimprovable without margin or realizability assumptions. So using the abstaining learner costs nothing relative to plain empirical welfare maximization.

• Doubly robust guarantees under unknown propensities. With nuisances ĝ and p̂ independent of the sample, Theorem 3.4 gives regret bounded by the fast O(1/n) term plus Err_DR²/(p κ²). The paper shows the additive term vanishes if the propensity and outcome regression errors each decay at roughly n^(−1/4) in L4 norm with probability at least 1 − δ. Appendix B contains the full algorithm (Algorithm 4) and proof.

• Faster rates without realizability under a margin condition. Algorithm 2 first runs the abstention learner with bonus p = h/2, then refines on the region where it abstained. When the policy class Π has finite combinatorial diameter D and the margin condition P(|τ_o(X)| ≥ h) = 1 holds, Theorem 4.1 gives regret bounded roughly by (D + d log(n/d) + log(1/δ)) / (κ² h n) with probability at least 1 − δ. Since D can be infinite even for finite-VC classes, the paper offers a CATE-oracle alternative.

• Unifying regression-rate and realizable-case guarantees. With a CATE regression oracle of exponent β (error ‖τ̂ − τ‖_{P,2} ≤ c_δ n^(−β)), Theorem 4.2 decomposes regret into three terms with rates n^(−1), n^(−2β), and n^(−1/2−β). At β = 1/2 this recovers n^(−1) regret, unifying the regression-rate guarantee of Luedtke and Chambaz (2020) with the fast rate of the realizable case (Kitagawa and Tetenov, 2018). The authors note the bound improves on a direct plug-in classifier when the CATE oracle is strong or the class is "almost" realizable.

• Safe policy improvement that controls Type-I error. Algorithm 3 learns abstaining policies on a training split across a grid of bonuses, replaces abstentions with the baseline policy ω, and tests each candidate on a held-out split using a one-sided lower confidence bound with a Bonferroni adjustment z_{1−δ/k}. Across more than 2,000 repetitions at significance level δ = 0.05, the safe policy learner controlled Type-I error at or below 0.05, achieved the highest improvement rate and largest mean value gain among accepted policies across baseline–optimal gaps, and maintained higher power with competitive mistake rates as noise increased. In the very small-sample regime (n ≤ 500), EWM performed better than the alternatives, but for n ≥ 1000 the proposed method dominated.

• Experimental design details. The synthetic study compares Algorithm 3 against two variants of High-Confidence Policy Improvement (HCPI) from Thomas et al. (2015) and against Safe EWM (with k = 1). Two designs are used: a fixed baseline with noise variance varied from 0.01 to 1.0 (Figure 1), and a varying baseline–optimal value gap at fixed noise (Figure 2). Each parameter setting is replicated 100 times, policy value is computed by IPW with known propensities, and the abstention grid is P = {0, 0.01, 0.05, 0.10, 0.20}.

• Connection to distribution shift. The abstract and introduction state that abstention provides a natural connection to distributionally robust policy learning, acting as a hedge against small data shifts (Proposition 4.3). The provided text of Section 4.3 is truncated, so the formal statement of Proposition 4.3 and its guarantees are not reported here.

Methodology in Plain English

The core difficulty the paper confronts is that in observational data you only ever see the outcome under the treatment a unit actually received, never both treatments. So the learner must convert raw samples into "pseudo-outcomes" that stand in for the treatment–control contrast, either by inverse propensity weighting (when the treatment assignment probabilities are known) or by a doubly robust correction that combines estimated outcome regressions with estimated propensities (when they are not).

Given those de-biased scores, the algorithm works in two stages. First, it splits the data in half, finds the policy in the class with the best empirical welfare on one half, and collects all policies whose empirical welfare is within a statistical tolerance of that maximizer. Second, for each such near-optimal policy it forms an abstaining variant that copies the maximizer where they agree and abstains where they disagree. The method then picks the best of these abstaining variants on the second half of the data.

The intuition for why this is fast: disagreement between two near-optimal policies tends to happen exactly in the regions where the treatment effect is small and the decision is genuinely uncertain. Abstaining there captures the bonus p, so the learner ends up collecting extra reward precisely where a forced binary decision would have been a coin flip. The paper emphasizes that the algorithm never needs to estimate the conditional average treatment effect directly, which matters when that function is complicated. For the applications, the same learner is reused as a black box: to fill in the abstention regions with a refined policy in the margin-condition setting, to generate a sequence of candidate policies that a safety test can accept or reject, and to characterize behavior under distribution shift.

Why This Matters

Impact on research. The paper imports the classification-with-reject-option literature — tracing back to Chow (1970) and developed by Herbei and Wegkamp (2006), Bartlett and Wegkamp (2008), and most closely Bousquet and Zhivotovskiy (2021) — into the counterfactual setting of policy learning, which is non-trivial because outcomes are only partially observed. It also reframes safe policy improvement (previously approached by hypothesis testing, bootstrap ensembling, or negative-regret minimization) as a natural consequence of abstention, and shows abstention can substitute for standard margin or realizability assumptions.

Real-world applications.

• Personalized medicine: deferring a treatment recommendation to a clinician when the estimated benefit for a patient is small or uncertain, rather than committing to a possibly harmful prescription.
• Online advertising and marketing: falling back to a default campaign in segments where the model cannot confidently distinguish which creative or intervention wins.
• Public policy targeting: declining to allocate a scarce program slot in cases where the evidence does not support a clear decision, and routing those cases to human review.
• Sequential deployment of new models: using the safe policy improvement procedure to certify, with a high-probability guarantee, that a newly learned policy beats the incumbent before it is rolled out.

Industry relevance. Any organization that routinely deploys treatment-assignment models on observational or logged data — healthcare systems, ad platforms, recommendation systems, A/B-testing platforms — faces exactly the failure mode this paper targets: an algorithm that must output something even when its confidence is low. The doubly robust variant matters especially for practitioners, since real systems rarely know their propensities exactly, and the safety certification procedure directly addresses an operational question (when is it safe to ship this model?) that is usually answered with far weaker heuristics.

Future Directions

• Filling out the distribution-shift analysis. The paper claims abstention acts as a hedge against small data shifts via Proposition 4.3, but the details are not part of the text reviewed here; a natural next step is developing the full theory of when abstention meaningfully improves robustness and quantifying the size of the shifts it tolerates.

• Extending beyond the binary treatment, binary outcome setting. The framework assumes D ∈ {0,1} and Y ∈ [0,1] with a single abstention symbol. Generalizing to multiple treatment arms or to settings where deferral could be to several different experts or defaults is an open direction.

• Characterizing the optimal abstention bonus. The regret bounds scale inversely in the bonus p, and the experiments use a fixed grid P = {0, 0.01, 0.05, 0.10, 0.20}. How to choose p — or adapt it to the data — is left open, and the boundary case p = 0 provably cannot deliver fast rates.

• Connecting to adjacent areas the paper explicitly sets aside. The authors list policy evaluation, offline reinforcement learning, policy learning under partial identification, and inference on the value of optimal treatment policies as complementary areas not directly considered, each of which could potentially absorb the abstention viewpoint.

Target Audience

This paper is aimed at researchers and graduate students working in causal inference, offline policy learning, and statistical learning theory, particularly those interested in the theory of individualized treatment rules and in semiparametric methods such as doubly robust estimation. It will also be valuable to applied machine learning practitioners in safety-critical domains — clinical decision support, ad delivery, public policy — who need policies that can decline to act, and to methodologists working on safe policy improvement who want a new lens on certification against a baseline. A reader should be comfortable with regret analysis, VC-dimension arguments, and propensity-score-based estimators to get full value from the theorems, though the framing and intuition are accessible to a broader audience.

Authors’ abstract

Policy learning algorithms are widely used in areas such as personalized medicine and advertising to develop individualized treatment regimes. However, most methods force a decision even when predictions are uncertain, which is risky in high-stakes settings. We study policy learning with abstention, where a policy may defer to a safe default or an expert. When a policy abstains, it receives a small additive reward on top of the value of a random guess. We propose a two-stage learner that first identifies a set of near-optimal policies and then constructs an abstention rule from their disagreements. We establish fast O(1/n)-type regret guarantees when propensities are known, and extend these guarantees to the unknown-propensity case via a doubly robust (DR) objective. We further show that abstention is a versatile tool with direct applications to other core problems in policy learning: it yields improved guarantees under margin conditions without the common realizability assumption, connects to distributionally robust policy learning by hedging against small data shifts, and supports safe policy improvement by ensuring improvement over a baseline policy with high probability.