High Probability Risk Control for Online Policy Learning

AI summary

Overview

• Research area: Machine learning — sequential decision-making, online policy learning, distribution shift, and statistical risk control.
• Technical level: Advanced. The paper relies on conformal risk control, importance sampling, feedback covariate shift, optimal transport / Sinkhorn matrix scaling, and Gaussian multiplier bootstrap theory.
• Scope (one sentence): The paper develops an online policy-learning method that keeps the risk of each deployed policy below a user-specified threshold with high probability, using a permutation-based risk estimator approximated by the Sinkhorn algorithm and a bootstrap-based upper confidence bound for policy selection.
• Affiliation and venue metadata: Yihong Guo, Drew Prinster, Suchi Saria, and Anqi Liu (Johns Hopkins University); arXiv:2611.13163v1 [cs.LG], 09 Nov 2026.

What This Paper Is About

In online policy learning, an agent repeatedly updates a policy from data it collected itself, deploys the policy, and then collects more data with that policy. Because the data-collection process depends on past policy choices, the data are not independent and identically distributed — a phenomenon the paper calls feedback covariate shift (FCS). The authors ask how to pick, at every round, a policy from a candidate family whose risk (such as the false-negative rate or the error rate) stays under a tolerance threshold with high probability, not merely in expectation. Achieving this requires solving two subproblems: estimating each candidate policy's risk efficiently and validly despite FCS, and selecting a policy in a way that selection bias does not silently break the guarantee.

Key Contributions

1. A permutation-based risk evaluator that accounts for feedback covariate shift, made tractable. The authors extend the permutation-weighting idea from conformal prediction under FCS to risk evaluation. Because exact permutation marginal weights have factorial complexity, they reformulate weight computation as an entropy-relaxed KL projection onto doubly stochastic matrices (Theorem 3.1) and solve it efficiently with the Sinkhorn algorithm.

2. A variance guarantee relative to standard importance weighting. Proposition 4.1 and Theorem 4.3 show that, under stated conditions (uniformly positive action probabilities, and the log range of the weight matrix not increasing after permutation), the permutation weights have a no-larger maximum-weight upper bound than importance weights, and the permutation risk estimator has a no-larger variance upper bound than the importance-weighted estimator.

3. A high-probability policy-selection procedure (RCL-HP). Instead of controlling each candidate separately, the method builds a simultaneous upper confidence bound over the whole candidate set by estimating the (1−δ)-quantile of the worst-case underestimation error T = sup_λ {ℓ(λ) − ℓ̂(λ)}. A Gaussian multiplier bootstrap is used for this quantile because ordinary resampling is inappropriate for adaptively collected data.

4. An expectation-level variant (RCL-E) and theory for both. RCL-E recovers a conformal-risk-control / conformal-policy-control style guarantee in expectation when the risk is monotonic, while RCL-HP yields an asymptotically valid simultaneous risk bound and asymptotic high-probability risk control (Theorem 4.5, Corollary 4.7).

Main Findings

• Variance reduction is provable (in upper bound form): Theorem 4.3 gives Var(ℓ̂_perm) ≤ U_perm = Γ_n · a_perm² and Var(ℓ̂_IPS) ≤ U_IPS = Γ_n · a_IPS², with U_perm ≤ U_IPS, where a_perm = M_ε/(M_ε + n − 1) and a_IPS = M_ε/n, and M_ε = [1 − (m−1)ε]/ε with m = |𝒜| and n = t + 1.
• Permutation weights generalize known estimators: Proposition 3.3 shows that under a fixed behavior policy the exact permutation weight reduces to the Self-Normalized Importance Sampling (SNIPS) weight; Proposition 3.4 expresses it as a global, normalized importance weight depending on the entire data bag and all feasible assignments.
• The Sinkhorn relaxation is empirically accurate: The authors report in Appendix G.6.1 that the approximation error of the Sinkhorn approximation is small, and Appendix G.6.3 provides empirical support for the log-range condition Δ_Q^(r) ≤ Δ_J.
• Bootstrap consistency holds asymptotically: Theorem 4.5 states that the sup-norm distance between the distribution of T_{t+1} and the bootstrap distribution of T*_{t+1} converges to zero in probability, and that the simultaneous coverage probability of the bootstrap upper bound converges to 1 − δ.
• High-probability safety follows for the selected policy: Corollary 4.7 gives lim sup_{t→∞} P(ℓ_{t+1}(λ̂_{t+1}) > α) ≤ δ when the selected λ̂ is restricted to the bootstrap-screened feasible set or falls back to π_safe; it also bounds the family-wise error of selecting any unsafe candidate.
• Empirical claims across three task families: The method reportedly maintains the target risk level, achieves lower risk than baselines, has lower variance in risk estimation, and yields reliable policy selection with high-probability risk control, with ablation studies supporting the theoretical claims. Specific numerical values for these experimental results are not reported in the provided (truncated) text, which cuts off mid-sentence in Section 5.1.
• Experimental setup details: Tasks are online contextual bandit problems on breast cancer diagnosis and stock investment, plus a SafetyBench question-answering task; data budgets are (m₀, m) = (50, 20) for breast cancer and (300, 50) for stock investment, and m₀ = 512 with m = 128 per round for safe question answering. Breast cancer diagnosis is a monotonic-risk setting; stock investment and SafetyBench are non-monotonic.
• Baselines compared: Online CRC (O-CRC), RollingRC (built on adaptive conformal inference, ACI), Conformal Policy Control (CPC), and the original permutation-weight estimation method of Prinster et al. Learn-then-Test (LTT) is not compared on these tasks because it requires i.i.d. data, though it is discussed in a later section.
• Two distinct shift regimes are studied: The paper is motivated by FCS from online policy updates, but also evaluates a non-adaptive covariate-shift setting (SafetyBench with fixed SFT checkpoints used as a fixed policy sequence), explicitly noting this is not feedback covariate shift.

Methodology in Plain English

• Set up a candidate family instead of a single policy. At each round the agent trains a base policy, then chooses a tuning parameter λ (for example a temperature) that indexes a family of policies {π_{θ,λ}} with different reward–risk trade-offs. A known fallback policy π_safe, whose risk is at most α, is always available. The task is to pick λ̂ whose risk ℓ(λ̂) is below tolerance α.
• Correct for the fact that the data were gathered by changing policies. Rather than weighting each sample by a simple ratio of target policy to behavior policy (classic importance weighting), the authors treat the observed losses as a "bag" of values and ask: if we permuted the assignment of positions to samples according to the historical policies, how much probability mass would land on each sample being the next-round loss? Those permutation marginals are the weights.
• Make the weights computable. Summing over all permutations is intractable, so the authors replace the exact problem with a KL projection of the policy matrix J (where J_{k,i} = π_k(A_i | X_i)) onto the set of doubly stochastic matrices, and read approximate weights off the last row of the solution. This is solved by iterative matrix scaling (Sinkhorn/IPFP).
• Handle the unknown next-round loss conservatively. Because the risk is estimated before the next action is drawn, a hypothetical test point is assumed to take the worst-case loss B, giving ℓ̂(λ) = Σ_{i=1}^{t} l_i p_i(λ) + B · p_{t+1}(λ).
• Screen candidates with a uniform safety margin. For high-probability control, one needs the worst underestimation across all candidates, not per-candidate error. The authors estimate the (1−δ)-quantile of that worst-case error by repeatedly drawing Gaussian multipliers ζ_i ~ N(0,1) shared across candidates, forming Σ ζ_i φ̂_i(λ) with φ̂_i(λ) = p_i(λ)(L_i − ℓ̂(λ)), taking the supremum over λ, and reading off the quantile. Any λ whose risk estimate plus this margin is ≤ α goes into the feasible set.
• Test empirically on three settings. Online bandit experiments (breast cancer false-negative rate; stock investment worst-action probability) and a SafetyBench refusal-based error-rate experiment, compared against O-CRC, RollingRC, and CPC, plus ablation and diagnostic studies in the appendix.

Why This Matters

• For research: The paper supplies the missing statistical bridge in the conformal-policy-control line of work. It gives a theoretically characterized, computationally practical surrogate for exact permutation weights (which CPC's guarantee required but its implementation did not analyze), and it upgrades the guarantee from "in expectation" to "asymptotically with high probability" — a substantially stronger and more deployment-relevant promise.
• Real-world applications:
  • Clinical decision support: controlling a diagnostic policy's false-negative rate in breast cancer screening, where a threshold on missed cases is a hard operational requirement.
  • Financial decision policies: bounding the probability of choosing the worst available action in portfolio or trading decisions.
  • Safety-filtered LLM deployment: controlling the error rate of a model that may decline to answer, using SafetyBench-style evaluation.
  • Medical claim or record filtering: the abstract lists medical claim filtering among the evaluated tasks, i.e., gating which items pass an automated screen under a bounded error budget.
• Industry relevance: Any team that updates a model on its own logged data — recommender systems, ad allocation, risk scoring, autonomous systems — inherits feedback covariate shift. A policy-selection rule that certifies risk levels per round, rather than only on average across retraining runs, is directly usable as a deployment gate with an auditable δ.

Future Directions

• Finite-sample rather than asymptotic guarantees. Corollary 4.7 is asymptotic; the authors explicitly note that bootstrap consistency yields asymptotic, not finite-sample, high-probability control. Closing this gap is the obvious next theoretical step.
• Quantifying how Sinkhorn approximation error propagates into the guarantee. The paper defers a theoretical analysis of how the entropy relaxation's error affects downstream risk estimation and guarantees to Appendix G.6.2, leaving room for sharper end-to-end bounds.
• Scaling to larger candidate families and longer histories. The theory assumes a finite candidate set Λ and uniformly positive action probabilities (π(a|x) ≥ ε), and the weight analysis assumes the log range of the weight matrix does not grow under permutation (Δ_p ≤ Δ_J), which is supported empirically rather than proven in general.
• Extending beyond the i.i.d.-context and mixing assumptions. The estimator assumes an i.i.d. context distribution so that context terms cancel in the permutation weights, and the variance bound assumes geometrically decaying covariances; relaxing either would broaden applicability to genuinely non-stationary deployments.

Target Audience

Researchers and practitioners working on safe or risk-constrained sequential decision-making: conformal prediction and conformal risk control specialists, bandit and reinforcement learning theorists interested in adaptive data collection and off-policy evaluation, and applied ML engineers in healthcare, finance, or LLM safety who need per-deployment statistical guarantees rather than average-case performance. Readers should be comfortable with importance sampling, exchangeability, and bootstrap arguments; the paper is not an introductory treatment.

Authors’ abstract

We study risk control for online policy learning, where an agent adaptively collects data and aims to learn a new policy whose risk stays below a tolerance threshold at every round with high probability. Such policy updates and data collection induce feedback covariate shift (FCS): observed data depend on history through evolving policies and the data-collection process. FCS brings two substantive problems for risk control: 1) efficient risk evaluation for each candidate…