PAC-Bayesian Reward-Certified Outcome Weighted Learning

Yuya Ishikawa; Shu Tamano

cs.LGApr 2, 2026

PAC-Bayesian Reward-Certified Outcome Weighted Learning

Authors: Yuya Ishikawa, Shu Tamano

Organizations: Department of Clinical Biostatistics, Graduate School of Medical and Dental Sciences, Institute of Science Tokyo, 1-5-45 Yushima, Bunkyo-ku, Tokyo, 113-8510, Japan · Biostatistics Division, Center for Research Administration and Support, National Cancer Center, 5-1-1 Tsukiji, Chuo-ku, Tokyo 104-0045, Japan · Department Multidisciplinary Sciences, Graduate School of Arts and Sciences, The University of Tokyo, 3-8-1 Komaba, Meguro-Ku, Tokyo 153-8902, Japan · Department of Epidemiology, National Institute of Infectious Diseases, Japan Institute for Health Security, 1-23-1 Toyama, Shinjuku-Ku, Tokyo 162-0052, Japan

Abstract

Estimating optimal individualized treatment rules (ITRs) via outcome weighted learning (OWL) often relies on observed rewards that are noisy or optimistic proxies for the true latent utility. Ignoring this reward uncertainty leads to the selection of policies with inflated apparent performance, yet existing OWL frameworks lack the finite-sample guarantees required to systematically embed such uncertainty into the learning objective. To address this issue, we propose PAC-Bayesian Reward-Certified Outcome Weighted Learning (PROWL). Given a one-sided uncertainty certificate, PROWL constructs a conservative reward and a strictly policy-dependent lower bound on the true expected value. Theoretically, we prove an exact certified reduction that transforms robust policy learning into a unified, split-free cost-sensitive classification task. This formulation enables the derivation of a nonasymptotic PAC-Bayes lower bound for randomized ITRs, where we establish that the optimal posterior maximizing this bound is exactly characterized by a general Bayes update. To overcome the learning-rate selection problem inherent in generalized Bayesian inference, we introduce a fully automated, bounds-based calibration procedure, coupled with a Fisher-consistent certified hinge surrogate for efficient optimization. Our experiments demonstrate that PROWL achieves improvements in estimating robust, high-value treatment regimes under severe reward uncertainty compared to standard methods for ITR estimation.

Figures & tables

Figure 1: Performance comparison across varying uncertainty levels (

\rho

) for Scenario 1. The left panel reports target regret against the

R

-family baselines, and the right panel reports robust regret against the

\underline{R}

-family baselines.

Figure 2: Performance comparison across varying uncertainty levels (

\rho

) for Scenario 2. The left panel reports target regret against the

R

-family baselines, and the right panel reports robust regret against the

\underline{R}

-family baselines.

Figure 3: Performance comparison across varying sample sizes (

N

) for Scenario 1. The left panel reports target regret against the

R

-family baselines, and the right panel reports robust regret against the

\underline{R}

-family baselines.

Figure 4: Performance comparison across varying sample sizes (

N

) for Scenario 1. The left panel reports target regret against the

R

-family baselines, and the right panel reports robust regret against the

\underline{R}

-family baselines.

Methods	Certified value	Nominal value	Composite-free value	Mortality risk	Alert rate
Never alert	$0.905\,(0.001)$	$0.918\,(0.001)$	$0.797\,(0.001)$	$0.086\,(0.001)$	$0.000\,(0.000)$
Always alert	$0.900\,(0.001)$	$0.914\,(0.001)$	$0.783\,(0.002)$	$0.088\,(0.001)$	$1.000\,(0.000)$
OWL ( $R$ )	$0.900\,(0.001)$	$0.914\,(0.001)$	$0.786\,(0.002)$	$0.089\,(0.001)$	$0.672\,(0.024)$
Q-learning ( $R$ )	$0.903\,(0.001)$	$0.917\,(0.001)$	$0.793\,(0.002)$	$0.086\,(0.001)$	$0.500\,(0.005)$
RWL ( $R$ )	$0.903\,(0.001)$	$0.916\,(0.001)$	$0.793\,(0.002)$	$0.088\,(0.001)$	$0.450\,(0.030)$
OWL ( $\underline{R}$ )	$0.901\,(0.001)$	$0.915\,(0.001)$	$0.787\,(0.002)$	$0.087\,(0.001)$	$0.715\,(0.024)$

Table 1: Repeated sample-split AIPW comparison on the ELAIA-1 AKI-alert trial under the three-component uncertain hard-clinical utility with

\rho=1

. Entries are means with standard errors over

30

repeated

70/30

hospital-by-alert stratified splits. Certified value is the estimate of

V_{\underline{R}}

, nominal value is the corresponding estimate of

V_{R}

, and composite-free value uses

R^{\mathrm{comp}}=1-\texttt{composite\_outcome}

. Among adaptive learned policies, PROWL attains the largest certified value and the largest composite-free value.

Figure 5: Point-range summary of the primary actual-data comparison on ELAIA-1. Left panel reports the estimated certified value

\hat{V}_{\underline{R}}

with

95\%

confidence intervals across the

30

repeated sample splits. Right panel reports the estimated composite-free value

\hat{V}_{\mathrm{comp}}

with the same uncertainty summary. The figure complements Table 1 by emphasizing method ordering on the certified objective and on the trial’s primary hard-outcome anchor.

Figure 6: Learned alert allocation by baseline-risk quintile and hospital type. Baseline-risk quintiles are obtained within each training fold by fitting a control-arm risk model for the primary composite outcome and then evaluating the resulting risk scores on the held-out fold. Each panel displays the mean recommended alert rate across the repeated sample splits for one policy. Teaching hospitals are shown in blue and non-teaching hospitals in red. PROWL is compared against the strongest raw-reward baseline, Q-learning

(R)

, and the strongest lower-reward plug-in baseline, Q-learning

(\underline{R})

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Complementary performance comparison across varying uncertainty levels (

\rho

) for Scenario 1. The left panel reports target regret against the

\underline{R}

-family baselines, and the right panel reports robust regret against the

R

-family baselines.

Figure 8: Complementary performance comparison across varying uncertainty levels (

\rho

) for Scenario 2. The left panel reports target regret against the

\underline{R}

-family baselines, and the right panel reports robust regret against the

R

-family baselines.

Figure 9: Complementary performance comparison across varying sample sizes (

N

) for Scenario 1. The left panel reports target regret against the

\underline{R}

-family baselines, and the right panel reports robust regret against the

R

-family baselines.

Figure 10: Complementary performance comparison across varying sample sizes (

N

) for Scenario 2. The left panel reports target regret against the

\underline{R}

-family baselines, and the right panel reports robust regret against the

R

-family baselines.

Figure 11: Certificate-diagnostic plots. Rows report the proxy-target gap

V(\hat{d})-V^{\ast}(\hat{d})

and the target-certified gap

V^{\ast}(\hat{d})-V_{\underline{R}}(\hat{d})

, while columns correspond to Scenarios 1 and 2. Methods shown are PROWL, PROWL (

U=0

), and the classical

\underline{R}

-family baselines.

	Scenario 1			Scenario 2
$\rho$	$\mathbb{E}[U]$	Clip	Valid	$\mathbb{E}[U]$	Clip	Valid
$0$	$0.000$	$0.000$	$1.000$	$0.000$	$0.000$	$1.000$
$0.25$	$0.008$	$0.000$	$1.000$	$0.023$	$0.000$	$1.000$
$0.5$	$0.015$	$0.000$	$1.000$	$0.047$	$0.000$	$1.000$
$0.75$	$0.023$	$0.000$	$1.000$	$0.070$	$0.000$	$1.000$
$1$	$0.030$	$0.000$	$1.000$	$0.093$	$0.000$	$1.000$

Appendix

Table 2: Dataset-level certificate diagnostics across uncertainty levels

\rho

. Here

\mathbb{E}[U]

is the logged-sample mean certificate magnitude, Clip is the logged-sample rate

\mathbb{P}(U>R)

, and Valid is the independent-test validity rate that the oracle envelope upper-bounds the true proxy bias in both arms.

Figure 12: Split-free ablation at

\rho=1.5

. Rows report target regret and robust regret, while columns correspond to Scenarios 1 and 2. PROWL (split-free) fits nuisance functions and the policy on the same sample, whereas PROWL (sample-split) uses a stratified honest half-sample split.

Methods	death14 risk	dialysis14 risk	aki_progression14 risk	composite_outcome risk	discharge_to_home rate
Never alert	$0.086\,(0.001)$	$0.031\,(0.000)$	$0.151\,(0.001)$	$0.203\,(0.001)$	$0.507\,(0.002)$
Always alert	$0.088\,(0.001)$	$0.034\,(0.001)$	$0.162\,(0.002)$	$0.217\,(0.002)$	$0.490\,(0.003)$
OWL ( $R$ )	$0.089\,(0.001)$	$0.034\,(0.001)$	$0.160\,(0.002)$	$0.214\,(0.002)$	$0.493\,(0.003)$
Q-learning ( $R$ )	$0.086\,(0.001)$	$0.032\,(0.001)$	$0.156\,(0.002)$	$0.207\,(0.002)$	$0.494\,(0.002)$
RWL ( $R$ )	$0.088\,(0.001)$	$0.031\,(0.001)$	$0.154\,(0.002)$	$0.207\,(0.002)$	$0.500\,(0.002)$
OWL ( $\underline{R}$ )	$0.087\,(0.001)$	$0.033\,(0.001)$	$0.160\,(0.002)$	$0.213\,(0.002)$	$0.493\,(0.003)$

Appendix

Table 3: Component-wise decomposition of the repeated sample-split AIPW evaluation on ELAIA-1. Entries are means with standard errors over

30

repeated

70/30

hospital-by-alert stratified splits. This table complements Table 1 by demonstrating that the improvements in certified value are not driven by deterioration in any single hard-outcome component.

Figure 13: Hospital-level intention-to-treat heterogeneity in the ELAIA-1 trial. The plotted quantity is the risk difference of alert versus control on the 14-day composite outcome, accompanied by Wald-type

95\%

confidence intervals. Positive values favor usual care. The two non-teaching hospitals exhibit the clearest adverse alert signal, strongly motivating the use of site-aware individualized alerting rules.

Figure 14: Procedural sensitivity to the preference-uncertainty radius

\rho

in the actual-data experiment. The panels report the estimated certified value, composite-free value, and alert rate for PROWL, PROWL

(U=0)

, Q-learning

(R)

, and Q-learning

(\underline{R})

across varying levels of

\rho

. While a larger

\rho

structurally induces a more conservative lower utility, PROWL strategically responds by modestly reducing its alerting frequency, consistently preserving the clinically anchored composite-free value.

Figure 15: Ablation sensitivity analysis excluding hospital covariates. For each method, the points and intervals compare the baseline analysis against a refitted model that completely removes hospital indicators from the policy covariates. Although absolute performance experiences a slight drop across all methods, the relative ordering remains robustly unchanged. This confirms that the observed gains of PROWL are not driven merely by the memorization of site-specific labels.

Figure 16: Deployment derandomization sensitivity for PROWL. This figure compares the deterministic mean-rule (utilized in the main text analysis) against MAP and Gibbs deployment strategies derived from the identically fitted posterior. Across all three configurations, the certified value, composite-free value, and alert rates are virtually indistinguishable, underscoring the robustness of our empirical findings.

Jun 17, 2026

Wasserstein Policy Learning for Distributional Outcomes

Offline policy learning has received growing attention in causal inference. The primary objective is to learn a policy (individualized treatment rule) as a mapping from covariates to treatment that maximizes the empirical welfare defined as the mean of scalar-valued potential outcomes. In this paper, we study offline policy learning with distribution-valued outcomes, where each potential outcome is a probability measure on

\mathbb{R}

and the reward is defined through a utility functional applied to the Wasserstein barycenter of induced outcome distributions. We establish statistical guarantees for the policy learning framework based on both Inverse Probability Weighting (IPW) and Doubly Robust (DR) estimators. By handling the challenging uniform deviation over the product of the combinatorial policy class and the infinite-dimensional quantile domain, we prove that the finite-sample regret has leading dependence

\widetilde{\mathcal{O}}(\sqrt{\mathrm{N\text{-}dim}(Π)/N})

. In the one-dimensional Wasserstein setting and under the stated regularity conditions, the leading regret rate is still governed by the policy-class complexity. Moreover, we provide a minimax lower bound establishing the sharpness of the leading dependence on

N

and

\mathrm{N\text{-}dim}(Π)

Yiyan Huang, Cheuk Hang Leung, Qi Wu +1

School of Statistics and Data Science & Institute of Big Data Research, Shanghai University of Finance and Economics, Shanghai, China

Offline RL Wasserstein Barycenters Causal Inference

Ask PDF

Feb 27, 2026

General Bayesian Policy Learning

This study proposes a General Bayes framework for policy learning. We consider decision problems in which a decision-maker chooses an action from a given set to maximize expected welfare. Typical examples include treatment choice and portfolio optimization. In such problems, the statistical target is a decision rule, and predicting each potential outcome is not necessarily of primary interest. We formulate this policy-learning problem through loss-based Bayesian updating. Our main technical device is a squared-loss surrogate for welfare maximization. We show that maximizing empirical welfare over a policy class with a quadratic penalty controlled by a tuning parameter

ζ>0

is equivalent to minimizing a scaled squared error in the outcome difference. The resulting General Bayes posterior over decision rules admits two equivalent characterizations: a Gaussian pseudo-likelihood representation and a decision-theoretic loss-based characterization. As one implementation, we introduce GBPLNet, a neural network implementation with a tanh-squashed output. Finally, we establish PAC-Bayes-type guarantees for the surrogate risk and the corresponding penalized welfare.

Masahiro Kato

1Mizuho-DL Financial Technology Co., Ltd., Tokyo, Japan

Generalized Bayesian Inference Policy Learning Bayesian Decision Theory Bayesian Inference

Ask PDF

Jul 24, 2023

A Differentially Private Weighted Empirical Risk Minimization Procedure and its Application to Outcome Weighted Learning

Data used to train predictive models via empirical risk minimization (ERM) often contain sensitive personal information. While differential privacy (DP) provides mathematically provable bounds to protect such data, previous work has focused almost exclusively on unweighted ERM. We consider weighted ERM (wERM) -- an important generalization where individual contributions to the objective function vary. We propose the first DP algorithm for general wERM with formal privacy guarantees and derive both its empirical and population excess risk bounds. Crucially, this general wERM framework provides a pathway for deriving privacy-preserving learning methods for individualized treatment rules, including the popular outcome-weighted learning (OWL) approach. We evaluate DP-wERM applied to OWL in simulated and real data experiments. Our empirical results demonstrate that training OWL models via wERM provides strong DP guarantees while maintaining robust performance, proving the method is practical for sensitive, real-world data.

Spencer Giddens, Yiwang Zhou, Kevin R. Krull +3

University of Notre Dame, Notre Dame, IN 46556, USA · St. Jude Children’s Research Hospital, Memphis, TN 38138, USA · University of Michigan, Ann Arbor, MI 48109, USA

Privacy-Preserving ML Differential Privacy Statistical Learning Theory Empirical Risk Minimization

Ask PDF

Abstract

Explore similar work

Wasserstein Policy Learning for Distributional Outcomes

General Bayesian Policy Learning

A Differentially Private Weighted Empirical Risk Minimization Procedure and its Application to Outcome Weighted Learning