General Bayesian Policy Learning

General Bayesian Policy Learning

Introduction

Policy learning aims to train a policy function $\delta(x)$ that maps context features $x\in{\mathcal{X}}$ to an action $a$ in a constrained action set, so as to maximize an expected outcome or, equivalently, to minimize an expected task loss. When the action set is discrete, many objectives can be written as a cost-sensitive $0$-$1$ loss, where the loss penalizes choosing a suboptimal action with a weight given by an outcome gap. This perspective suggests using a learning rule that targets a task loss directly rather than fitting a potentially misspecified generative model.

General Bayes is a framework for updating beliefs using a loss function rather than a likelihood function Bissiri2016generalframework. Given a prior distribution and an empirical loss, General Bayes yields a generalized posterior that is coherent from a decision-theoretic viewpoint. This framework is attractive when a probabilistic model is misspecified, inconvenient, or unnecessary for the downstream task. This study develops a General Bayes framework for policy learning.

A practical obstacle is that welfare objectives are typically linear in the policy, so they do not directly correspond to a convenient likelihood. We address this by rewriting welfare maximization as a squared-loss minimization problem. In the binary action case, the key surrogate loss has the form $\left(\frac{1}{\sqrt{\zeta}}\left(y\left(1\right)-y\left(0\right)\right))-y\left(0\right)}-\sqrt{\zeta}f\left(x\right)\right)right)}-\sqrt{\zeta}f\left(x\right)}^2$, where $f\left(x\right)\in\left[-1,1\right]$ encodes a policy and $\zeta>0$ is a tuning parameter. This surrogate admits a Gaussian pseudo-likelihood interpretation, so it enables a computationally convenient General Bayes posterior and standard approximation methods.

General Bayes allows any loss, so one could update directly using the negative welfare contribution as a loss. This direct update yields objectives that are linear in the policy. It does not induce the quadratic regularization in Theorem (ref), and it does not provide a Gaussian pseudo-likelihood that can be exploited by standard Bayesian computation. The squared-loss surrogate yields a regression-style objective with stable gradients, an explicit regularization path indexed by $\zeta$, and a unified extension to $K$ actions and missing outcome settings.

For $K$ actions, a baseline-gap formulation in terms of outcome differences is natural, but it introduces a dependence on the baseline action through the regularization term. We therefore develop a baseline-free symmetric surrogate that operates on the full feedback vector. In missing outcome settings, for example, in observational studies and logged bandit feedback, outcome differences are not observed. We show that inverse propensity weighting (IPW) and doubly robust (DR) pseudo-outcomes yield explicit empirical losses that can be used for General Bayes updating, and we provide population-level target characterizations.

\paragraph{Contributions.} We list our contributions as follows:

We introduce related works in Appendix (ref). We refer to our framework as General Bayesian Policy Learning (GBPL) and our proposed implementation with neural networks as GBPLNet. Our proposed method can be applied to various tasks, including treatment choice Kitagawa2018whoshould,Athey2021policylearning,Zhou2023offlinemultiaction and portfolio optimization Tallman2023bayesianpredictive,Tallman2024predictivedecision,Kato2024bayesianportfolio,Kato2024generalbayesian. Note that we use empirical welfare maximization in a general sense, not restricted to counterfactual risk minimization Swaminathan2015batchlearning,Kitagawa2018whoshould.

Setup

Assume that there are $K$ actions, indexed by $1,2,\dots,K$. Let $X\in{\mathcal{X}}\subseteq {\mathbb{R}}^k$ be a regressor and $Y\left(a\right)\in{\mathcal{Y}}\subseteq {\mathbb{R}}$ be an outcome for each action $a\in\left\{1,\dots,K\right\}$. Let $Z=\left(X,Y\left(1\right),\dots,Y\left(K\right)\right)s,Y\left(K\right)}$ be a random vector following a distribution $P$ on ${\mathcal{Z}}={\mathcal{X}}\times{\mathcal{Y}}^K$. Given i.i.d. observations ${\mathcal{D}}=\left\{z_i\right\}_{i=1}^n$ with $z_i=\left(x_i,y_i\left(1\right),\dots,y_i\left(K\right)\right)y_i\left(K\right)}$, a decision-maker trains a decision rule $\delta\colon {\mathcal{X}}\to\left\{1,\dots,K\right\}$ to maximize the expected welfare

The Bayes optimal rule satisfies

for almost every $x$.

In many applications, outcomes are not fully observed, for example, under the potential outcome framework in causal inference and in bandit feedback in online learning. The missing outcome setting is introduced in Section (ref). Sections (ref) to (ref) focus on the full feedback setting to isolate the surrogate construction.

Background: Generalized Bayesian Updating

Let $\Theta$ be a parameter space and let $\theta\in\Theta$ index a model for a decision rule or a score function. Let $\Pi$ be a prior distribution on $\Theta$. Given a loss function $\ell\left(\theta;z\right)$ and data ${\mathcal{D}}=\left\{z_i\right\}_{i=1}^n$, generalized Bayesian updating defines a generalized posterior

where $\eta>0$ is a temperature parameter.

\paragraph{Decision-theoretic derivation.} A key result of Bissiri2016generalframework is that ((ref)) is the optimizer of a variational problem. Let $Q$ be any probability distribution on $\Theta$ such that $Q$ is absolutely continuous with respect to $\Pi$. Define

The proof is provided in Appendix (ref).

\paragraph{Relation to ordinary Bayes.} If the loss is chosen as the negative log-likelihood, $\ell\left(\theta;z\right)=-\log p\left(z\mid \theta\right)$, and $\eta=1$, then ((ref)) reduces to ordinary Bayes' rule. For other losses, ((ref)) yields a coherent update even when no global likelihood model is posited Bissiri2016generalframework.

GBPL with Binary Actions

For simplicity, we first consider the binary action case, $K=2$. We reindex actions $1$ and $2$ as $1$ and $0$. We consider a possibly randomized policy $\delta\colon{\mathcal{X}}\to\left[0,1\right]$. The expected welfare can be written as

Empirical Welfare Maximization

Let ${\mathcal{H}}_{\mathrm{pol}}$ be a class of measurable maps $\delta\colon{\mathcal{X}}\to\left[0,1\right]$. The empirical welfare is

A natural estimator is

From a Bayesian perspective, ((ref)) does not specify a likelihood, so it is unclear how to define a posterior distribution over $\delta$ without further modeling choices.

Squared-loss surrogate

Define the score function $f\left(x\right)=2\delta\left(x\right)-1$ and the induced class ${\mathcal{F}}_{{\mathcal{H}}_{\mathrm{pol}}} = \left\{f\colon f\left(x\right)=2\delta\left(x\right)-1,\delta\in{\mathcal{H}}_{\mathrm{pol}}\right\}$. Let $\zeta>0$ be a constant. We consider the squared-loss surrogate

To state the equivalence, define the penalized welfare maximizer

The proof is provided in Appendix (ref). As $\zeta\to 0$, the penalty in ((ref)) vanishes. If the maximizer in ((ref)) is unique, then $\widetilde{\delta}\left(\lambda\right)\to \widehat{\delta}$ as $\lambda\to 0$. Without uniqueness, any limit point of $\widetilde{\delta}\left(\lambda\right)$ belongs to the set of maximizers in ((ref)). The penalty term in ((ref)) is minimized at $\delta\left(x\right)=1/2$, so for finite $\zeta$ the surrogate induces shrinkage toward randomized policies.

Loss and Generalized Posterior

Let $f_{\theta}$ be a parametric score model indexed by $\theta\in\Theta$. In the binary action case, we define the loss as

The generalized posterior is then

where $z_i=\left(x_i,y_i\left(1\right),y_i\left(0\right)\right)y_i\left(0\right)}$.

Pseudo-Likelihood Interpretation

Let $u =y\left(1\right)-y\left(0\right)$ denote the outcome difference. The exponential weight in ((ref)) can be rewritten as \[ \exp\left(-\eta\ell\left(\theta;\left(x,y\left(1\right),y\left(0\right)\right)),y\left(0\right)}\right)}\right)),y\left(0\right)}}\right)}\right)}\right)),y\left(0\right)}}} = \exp\left(-\frac{\eta}{2\zeta}\left(u-\zeta f_{\theta}\left(x\right)\right)x\right)}^2\right)right)x\right)}^2}. \] Therefore, ((ref)) coincides with ordinary Bayesian updating under the working Gaussian model

We emphasize that ((ref)) is a working model introduced for computation and interpretation, not a claim about the true data-generating mechanism.

General Bayes Interpretation

The posterior ((ref)) supports two complementary Bayesian interpretations. First, under the working model ((ref)), it is an ordinary Bayesian posterior. Second, without committing to any likelihood, it is the unique optimizer of the variational problem in Proposition (ref), so it is a coherent loss-based belief update Bissiri2016generalframework.

The parameters $\zeta$ and $\eta$ play different roles. The parameter $\zeta$ changes the learning objective itself, because it determines the strength of the quadratic regularization in Theorem (ref). The parameter $\eta$ changes posterior concentration, and it can be viewed as a calibration parameter. The parameter $\eta$ also affects point estimators derived from the generalized posterior, for example, the MAP, because it rescales the empirical loss relative to the prior through ((ref)). Under the working model ((ref)), $\zeta$ sets the scale of the conditional mean and $\zeta/\eta$ sets the working conditional variance, so $\zeta$ and $\eta$ jointly determine the pseudo-likelihood scale. Selecting $\eta$ is therefore a calibration problem, and practical strategies include information matching procedures Lyddon2019generalbayesian and adaptive learning rate procedures Grunwald2012thesafe.

GBPL with Multiple Actions

We now extend the squared-loss surrogate and the posterior construction to $K$ actions. We present two surrogate constructions. The first is based on outcome gaps relative to a baseline action. The second is baseline-free and symmetric, and it uses the full feedback vector.

Baseline-Gap Surrogate

Fix a baseline action $K$. Define outcome gaps \[\overline{U}\left(a\right)=Y\left(a\right)-Y\left(K\right)\] for $a\in\left\{1,\dots,K-1\right\}$. For a possibly randomized policy $\delta\colon{\mathcal{X}}\to\Delta_K$, where $\Delta_K=\left\{p\in{\mathbb{R}}^K\colon p_a\ge 0,\sum_{a=1}^K p_a=1\right\}$, welfare can be written as

The baseline term ${\mathbb{E}}\left[Y\left(K\right)\right]$ does not depend on $\delta$.

Let ${\mathcal{H}}_{\mathrm{pol}}^{\left(K\right)}$ be a class of policies $\delta\colon{\mathcal{X}}\to\Delta_K$. For $a\in\left\{1,\dots,K-1\right\}$, define the componentwise score transform $f_a\left(x\right)=2\delta_a\left(x\right)-1$. Let ${\mathcal{F}}_{{\mathcal{H}}_{\mathrm{pol}}^{\left(K\right)}}$ denote the induced class of score vectors $f\left(x\right)=\left(f_1\left(x\right),\dots,f_{K-1}\left(x\right)\right)-1}\left(x\right)}$. For $\zeta>0$, define the empirical surrogate

Define the empirical welfare $\widehat{V}\left(\delta\right) = \frac{1}{n}\sum^n_{i=1}\sum_{a=1}^{K}\delta_a\left(X_i\right)Y_i\left(a\right)$ and the penalized empirical welfare maximizer

The proof is provided in Appendix (ref).

Baseline-Free Symmetric Full-Vector Surrogate

The baseline-gap surrogate is convenient, but its regularization depends on the chosen baseline through the term $\sum_{a=1}^{K-1}\left(2\delta_a\left(X\right)-1\right)right)-1}^2$. To remove baseline dependence, we use a full-vector surrogate that treats all actions symmetrically.

Let ${\mathcal{H}}_{\mathrm{pol}}^{\left(K\right)}$ be a class of policies $\delta\colon{\mathcal{X}}\to\Delta_K$. For $\zeta>0$, define the empirical full-vector surrogate

Define the penalized empirical welfare maximizer

The penalty term in ((ref)) is minimized at $\delta_a\left(x\right)=1/K$ for all $a$, so for finite $\zeta$ the full-vector surrogate induces shrinkage toward uniform randomization.

The proof is provided in Appendix (ref).

Baseline Dependence and Symmetric Alternatives

Theorems (ref) and (ref) show that both surrogate constructions correspond to penalized welfare maximization, but with different regularizers. In particular, the baseline-gap penalty depends on which action is chosen as baseline, while the full-vector penalty is baseline-free and symmetric across actions.

If the interest is the unregularized welfare maximizer, one may take $\zeta$ small so that the regularization is negligible. In that regime, baseline choice does not materially affect the maximizer, but it can affect numerical stability, because it changes the scaling and the variance of estimated gaps in missing outcome settings. If one prefers an invariant formulation for finite $\zeta$, the full-vector surrogate provides a direct baseline-free alternative.

Another symmetric alternative is to average baseline-gap surrogates over baselines. This produces a pairwise gap objective that depends on all differences $Y\left(a\right)-Y\left(b\right)$. We do not pursue this direction further, because it introduces $K\left(K-1\right)/2$ pairwise terms.

Shift Invariance and Centering

Welfare is invariant to adding an action-common shift that depends on $X$, because $\sum_{a=1}^K \delta_a\left(X\right)=1$. That is, for any measurable $c\left(X\right)$, we have ${\mathbb{E}}\left[\sum_{a=1}^K \delta_a\left(X\right)\left(Y\left(a\right)+c\left(X\right)\right))+c\left(X\right)\right]} = V\left(\delta\right)+{\mathbb{E}}\left[c\left(X\right)\right]$. The argmax over $\delta$ is therefore unchanged. The full-vector surrogate inherits the same invariance, because replacing $Y_i\left(a\right)$ by $Y_i\left(a\right)+c_i$ for any $c_i$ that does not depend on $a$ changes ((ref)) only through additive constants, using again that $\sum_{a=1}^K \delta_a\left(X_i\right)=1$. This observation provides a simple baseline-free variance reduction idea. One can center outcomes by an action-common baseline without changing the minimizer.

In the full feedback setting, a symmetric choice is $c_i=-\frac{1}{K}\sum_{a=1}^K Y_i\left(a\right)$, which centers the outcomes across actions for each unit. In missing outcome settings, one may use a regression estimate $c\left(x\right)$ of an action-common component of the outcome and replace pseudo-outcomes by centered pseudo-outcomes, which can reduce the variance of IPW and DR constructions without changing the optimal policy.

General Bayes Posterior

Let $\delta_{\theta}\left(x\right)\in\Delta_K$ be a parametric policy model indexed by $\theta\in\Theta$. For the full-vector surrogate, define the loss

The generalized posterior is \[ {\mathrm{d}} \Pi_{\eta}\left(\theta\mid {\mathcal{D}}\right) \propto {\mathrm{d}} \Pi\left(\theta\right)\exp\left(-\eta\sum^n_{i=1}\ell^{\mathrm{Full}}_K\left(\theta;z_i\right)\right)i\right)}. \] As in the binary case, ((ref)) admits a Gaussian pseudo-likelihood interpretation, for example, $Y\left(a\right)\mid X=x,\theta\sim\mathcal{N}\left(\zeta\delta_{\theta,a}\left(x\right),\zeta/\eta\right)eta/\eta}$ as a working model.

GBPLNet: An Implementation Example of Neural Networks

The generalized posterior is typically intractable for flexible models. Standard approximation approaches include MAP, Gaussian approximations, and stochastic gradient Langevin dynamics Welling2011bayesianlearning. As one implementation example, we parameterize the bounded score $f_{\theta}\left(x\right)\in\left[-1,1\right]$ by a neural network with a tanh-squashed output, $f_{w}\left(x\right)=\tanh\left(g_{w}\left(x\right)\right)x\right)}$. We refer to this implementation as the GBPLNet.

Given a fitted score, we form a policy by $\delta_{w}\left(x\right)=\left(f_{w}\left(x\right)+1\right)right)+1}/2$. For evaluation in the full feedback setting, it is also convenient to use the deterministic decision rule $a_{w}\left(x\right)=\mathbbm{1}\left[f_{w}\left(x\right)\ge 0\right]$. This deterministic rule selects the more likely action under the randomized policy $\delta_{w}\left(x\right)$. For $K$ actions, a deterministic decision rule can be formed by $a_{\theta}\left(x\right)\in\arg\max_{a\in\left\{1,\dots,K\right\}} \delta_{\theta,a}\left(x\right)$. MAP computation corresponds to minimizing the negative log generalized posterior, which is the empirical surrogate loss plus the negative log prior. Gaussian approximations and SGLD provide complementary ways to represent posterior uncertainty. Neural networks and tanh are not essential to the framework, and any model class that produces bounded scores can be used. Appendix (ref) provides details.

\paragraph{Implementation summary} The GBPL update is determined by a choice of empirical loss, a prior, and tuning parameters.

GBPL with Missing Outcomes

We now consider a missing outcome setting where only the outcome under the chosen action is observed. We focus on a standard observational study setting with bandit feedback.

Observation Model and Assumptions

We observe i.i.d. data $\left\{\left(X_i,A_i,Y_i\right)\right\}_{i=1}^n$, where $A_i\in\left\{1,\dots,K\right\}$ is the chosen action and $Y_i=Y_i\left(A_i\right)$ is the observed outcome. Let $e_a\left(x\right)={\mathbb{P}}\left(A=a\mid X=x\right)$ be the propensity score. We impose standard conditions Rosenbaum1983centralrole.

Under these assumptions, the welfare of a policy $\delta\colon{\mathcal{X}}\to\Delta_K$ satisfies $V\left(\delta\right)={\mathbb{E}}\left[\sum_{a=1}^K \delta_a\left(X\right){\mathbb{E}}\left[Y\left(a\right)\mid X\right]\right]right]}$.

IPW Pseudo-Outcomes

If the propensity score $e_a\left(x\right)$ is known, or replaced by an estimator $\widehat e_a\left(x\right)$, IPW constructs pseudo-outcomes Horvitz1952generalization

Using the full-vector surrogate, we define the IPW empirical loss for a parametric policy $\delta_{\theta}$ as $\widehat{L}^{\mathrm{IPW}}_n\left(\theta\right) = \sum^n_{i=1}\frac{1}{2}\sum_{a=1}^{K}\left(\frac{1}{\sqrt{\zeta}}\widetilde{Y}^{\mathrm{IPW}}_i\left(a\right)-\sqrt{\zeta}\delta_{\theta,a}\left(X_i\right)\right)}\left(X_i\right)}^2$. The corresponding generalized posterior is ${\mathrm{d}} \Pi^{\mathrm{IPW}}_{\eta}\left(\theta\mid {\mathcal{D}}\right) \propto {\mathrm{d}} \Pi\left(\theta\right)\exp\left(-\eta \widehat{L}^{\mathrm{IPW}}_n\left(\theta\right)\right)a\right)}$.

The proof is provided in Appendix (ref).

DR Pseudo-Outcomes

Let $\gamma_a\left(x\right)={\mathbb{E}}\left[Y\left(a\right)\mid X=x\right]$. Let $\widehat\gamma_a\left(x\right)$ be an estimator of $\gamma_a\left(x\right)$. The DR pseudo-outcome is given as follows Bang2005doublyrobust:

Using the full-vector surrogate, define the DR empirical loss $\widehat{L}^{\mathrm{DR}}_n\left(\theta\right) = \sum^n_{i=1}\frac{1}{2}\sum_{a=1}^{K}\left(\frac{1}{\sqrt{\zeta}}\widetilde{Y}^{\mathrm{DR}}_i\left(a\right)-\sqrt{\zeta}\delta_{\theta,a}\left(X_i\right)\right)}\left(X_i\right)}^2$. The generalized posterior is ${\mathrm{d}} \Pi^{\mathrm{DR}}_{\eta}\left(\theta\mid {\mathcal{D}}\right) \propto {\mathrm{d}} \Pi\left(\theta\right)\exp\left(-\eta \widehat{L}^{\mathrm{DR}}_n\left(\theta\right)\right)a\right)}$.

In practice, $\widehat\gamma_a$ and $\widehat e_a$ may be estimated from the same data used to fit $\delta_{\theta}$. To reduce bias from the empirical process, we usually assume Donsker conditions for the nuisance estimators, or apply sample splitting, also called cross-fitting Klaassen1987consistentestimation,VanderVaart2002semiparametricstatistics. See Chernozhukov2018doubledebiased, Zhou2023offlinemultiaction, and Schuler2024introductionmodern for details.

Binary Action Case

When $K=2$, the surrogate in Section (ref) is expressed in terms of the outcome difference $U=Y\left(1\right)-Y\left(0\right)$. In the missing outcome setting, $U$ is not observed, but it can be replaced by an IPW or DR pseudo-difference.

For IPW, define \[\widetilde{U}^{\mathrm{IPW}}_i = \frac{\mathbbm{1}\left[A_i=1\right]Y_i}{\widehat e_1\left(X_i\right)} - \frac{\mathbbm{1}\left[A_i=0\right]Y_i}{\widehat e_0\left(X_i\right)}.\] For DR, define \[\widetilde{U}^{\mathrm{DR}}_i = \Big(\widehat\gamma_1\left(X_i\right)-\widehat\gamma_0\left(X_i\right)\Big) + \frac{\mathbbm{1}\left[A_i=1\right]\left(Y_i-\widehat\gamma_1\left(X_i\right)\right)i\right)}}{\widehat e_1\left(X_i\right)} - \frac{\mathbbm{1}\left[A_i=0\right]\left(Y_i-\widehat\gamma_0\left(X_i\right)\right)i\right)}}{\widehat e_0\left(X_i\right)}.\] Let $\widetilde{U}_i$ denote a generic pseudo-difference, with $\widetilde{U}_i=\widetilde{U}^{\mathrm{IPW}}_i$ for IPW and $\widetilde{U}_i=\widetilde{U}^{\mathrm{DR}}_i$ for DR. Let $f_{\theta}\left(x\right)$ be a score model with $f_{\theta}\left(x\right)\in\left[-1,1\right]$ and define the binary loss by replacing $y\left(1\right)-y\left(0\right)$ in ((ref)) with $\widetilde{U}_i$. For example, the DR empirical loss is $\widehat{L}^{\mathrm{DR}}_{n,\mathrm{bin}}\left(\theta\right) = \sum^n_{i=1}\frac{1}{2}\left(\frac{1}{\sqrt{\zeta}}\widetilde{U}^{\mathrm{DR}}_i-\sqrt{\zeta}f_{\theta}\left(X_i\right)\right)i\right)}^2$, and the corresponding generalized posterior is defined as in ((ref)) with $\widehat{L}^{\mathrm{DR}}_{n,\mathrm{bin}}$. The objective identity in Theorem (ref) continues to hold after replacing $U_i$ by $\widetilde{U}_i$, see Appendix (ref).

Population Targets

The following proposition summarizes the population interpretation of IPW and DR pseudo-outcomes. Its proof is provided in Appendix (ref).

Theoretical Analysis

This section summarizes theoretical guarantees for the surrogate loss and their implications for welfare. Proofs are provided in the Appendix.

Population Surrogate Minimizer

\paragraph{Binary case. } In the binary action case, define $U=Y\left(1\right)-Y\left(0\right)$. For a measurable $f\colon{\mathcal{X}}\to\left[-1,1\right]$, define the population surrogate risk \[R_{\zeta}\left(f\right) = {\mathbb{E}}\left[ \frac{1}{2}\left(\frac{1}{\sqrt{\zeta}}U-\sqrt{\zeta}f\left(X\right)\right)X\right)\right]^2 }.\] Let $m\left(x\right)={\mathbb{E}}\left[U\mid X=x\right]$. Then the unconstrained minimizer satisfies $f^\star\left(x\right)=m\left(x\right)/\zeta$. Under the constraint $f\left(x\right)\in\left[-1,1\right]$, the minimizer is the clipped conditional mean $f^\star\left(x\right)=\max\left(-1,\min\left(1,m\left(x\right)/\zeta\right)t)/\zeta}\right)a\right)t)/\zeta}}$. A formal statement and proof are given in Appendix (ref).

\paragraph{$K \geq 3$ case.} For the full-vector surrogate, define $\gamma\left(x\right)=\left(\gamma_1\left(x\right),\dots,\gamma_K\left(x\right)\right)a_K\left(x\right)}$ where $\gamma_a\left(x\right)={\mathbb{E}}\left[Y\left(a\right)\mid X=x\right]$. Consider the population risk associated with ((ref)). Conditionally on $X=x$, minimizing the conditional risk over $\delta\in\Delta_K$ yields the Euclidean projection of $\gamma\left(x\right)/\zeta$ onto $\Delta_K$. A formal statement and proof are given in Appendix (ref).

PAC-Bayes Bound under a Moment Condition

We state a PAC-Bayes bound for the surrogate loss under a sub-exponential moment condition, following unbounded-loss PAC-Bayes analyses Catoni2008pacbayesian,Haddouche2021pacbayes,Alquier2024userfriendly. Let $\widehat{R}_{\zeta}\left(\theta\right)=\frac{1}{n}\sum^n_{i=1}\ell\left(\theta;z_i\right)$ and $R_{\zeta}\left(\theta\right)={\mathbb{E}}\left[\ell\left(\theta;Z\right)\right]$.

Assumption (ref) is a moment condition on the loss. In the full feedback binary setting with bounded scores, it holds under standard tail conditions on the outcome difference $U=Y\left(1\right)-Y\left(0\right)$. In missing outcome settings, the IPW and DR pseudo-outcomes in ((ref)) and ((ref)) involve inverse propensity weights, so the constants $v$ and $b$ can deteriorate as the overlap constant $\epsilon$ in Assumption (ref) decreases. When nuisance estimators are learned from the same sample, the empirical losses are no longer i.i.d., and PAC-Bayes analyses typically rely on sample splitting or cross-fitting to control this dependence. In heavy-tailed regimes, one can replace the squared surrogate by a robust loss, for example, Catoni type or Huber type losses, and apply generalized Bayesian updating with the modified loss. Also see Catoni2008pacbayesian.

A proof sketch is provided in Appendix (ref).

A proof is provided in Appendix (ref).

Welfare Corollaries

The squared-loss surrogate and penalized welfare are linked by an identity. The proofs in this subsection are provided in Appendices (ref)--(ref).

\paragraph{Binary case.} Define the population penalized welfare \[W_{\lambda}\left(\delta\right)=V\left(\delta\right)-\lambda{\mathbb{E}}\left[\left(2\delta\left(X\right)-1\right)right)-1\right]^2}.\] Let $\lambda=\zeta/4$ and let $f=2\delta-1$. Then $W_{\zeta/4}\left(\delta\right)$ differs from $R_{\zeta}\left(f\right)$ only through an additive constant that does not depend on $\delta$. The next corollary makes the link explicit.

\paragraph{$K \geq 3$ case.} For the full-vector surrogate, define the penalized welfare $W^{\mathrm{Full}}_{\lambda}\left(\delta\right) = V\left(\delta\right) - \lambda{\mathbb{E}}\left[\sum_{a=1}^K \delta_a\left(X\right)^2\right]$. The next corollary parallels Corollary (ref).

PAC-Bayes to Welfare

Combining Theorem (ref) with Corollaries (ref) and (ref) yields welfare bounds. For example, in the binary case, we have \[W_{\lambda}\left(\delta^\star_{\lambda}\right)-W_{\lambda}\left(\delta\right) = \frac{1}{2}\left(R_{\zeta}\left(f\right)-R_{\zeta}\left(f^\star_{\zeta}\right)\right)ar_{\zeta}\right)}.\] Thus, any upper bound on $R_{\zeta}\left(f\right)$ translates into a lower bound on penalized welfare, relative to the best-in-class penalized welfare. The same translation applies for missing outcome settings after replacing $Y\left(a\right)$ by IPW or DR pseudo-outcomes, because the optimization identities in Theorems (ref) and Proposition (ref) determine the population targets.

Experiments

This section evaluates GBPLNet as an implementation example of the proposed General Bayes policy learning framework. Throughout, we focus on the full feedback setting, in which each observation contains the entire outcome vector $\left(Y\left(1\right),\dots,Y\left(K\right)\right)s,Y\left(K\right)}$. We report welfare and regret computed on held-out test data, and we aggregate results over $100$ independent trials. In this section, we only show the results with binary actions using simulation datasets. The other detailed explanations and additional results are shown in Appendices (ref)--(ref).

\paragraph{Evaluation protocol} For each trial, we generate or load a dataset and split it into training, validation, and test sets. We train each method on the training set and select tuning parameters using the validation set. We then evaluate welfare on the test set using the empirical welfare formulas in ((ref)) and their analogs for $K$ actions. For binary actions, we report the realized test welfare of the deterministic policy induced by the learned score function $f$, $\widehat{\delta}\left(x\right)=\mathbbm{1}\left[f\left(x\right)\ge 0\right]$. For $K$ actions, we evaluate the welfare of the learned policy $\widehat{\delta}\left(x\right)\in\Delta_K$ by plugging $\widehat{\delta}\left(X_i\right)$ into ((ref)) on the test set.

We define test regret relative to the oracle policy that maximizes the realized outcome for each test observation. In the binary case, the oracle action is $\delta^\mathrm{orc}\left(X_i\right)\in\arg\max_{a\in\left\{0,1\right\}} Y_i\left(a\right)$. We report the empirical regret on the test set, $\widehat{V}_\mathrm{test}\left(\delta^\mathrm{orc}\right)-\widehat{V}_\mathrm{test}\left(\widehat{\delta}\right)$. For each scenario and method, the tables report the mean, variance, and standard error of welfare across $100$ trials, and the mean and standard error of regret across $100$ trials.

We compare four values of the surrogate tuning parameter $\zeta$ for the proposed method. Three are fixed at $\zeta\in\left\{1,0.1,0.01\right\}$. One is selected by validation from the grid $\left\{1,0.1,0.01,0.001\right\}$, and is reported as GBPL (CV).

For GBPL (CV), we select $\zeta$ by maximizing validation welfare rather than minimizing the validation surrogate loss. In the binary case, expanding the squared loss in ((ref)) yields the term $\frac{1}{2\zeta}\frac{1}{n}\sum^n_{i=1}U_i^2$, where $U_i=Y_i\left(1\right)-Y_i\left(0\right)$, which does not depend on the fitted score, see Proposition (ref) in Appendix (ref). If the surrogate loss value itself is used for validation, this term can bias selection toward larger $\zeta$ values.

\paragraph{Synthetic Experiments} We consider three synthetic data generating processes, denoted by DGP1, DGP2, and DGP3. Each trial uses $n=5000$ observations and $d=10$ features. Table (ref) reports welfare and regret, and Figure (ref) visualizes the distribution of welfare. We consider cases with $K = 2$. For the results with $K = 5$, see Appendix (ref).

The comparison methods are as follows. DiffReg fits a regression model for the outcome difference $Y\left(1\right)-Y\left(0\right)$ and chooses actions by its sign. PluginReg fits separate regression models for $Y\left(1\right)$ and $Y\left(0\right)$ and chooses the action with the larger predicted outcome. WeightedLogistic fits a weighted logistic classifier for the best action, with weights given by the observed outcome gap magnitude. GBPLNet trains the tanh transformed neural score model in Section (ref) using the squared-loss surrogate.

Table (ref) shows that GBPLNet is competitive across DGP1, DGP2, and DGP3. In DGP2, GBPLNet yields substantial welfare gains relative to DiffReg and PluginReg, while WeightedLogistic underperforms. The sensitivity to $\zeta$ is scenario dependent.

Conclusion

This paper proposes a General Bayes framework for policy learning (GBPL). The key technical device is a squared-loss surrogate that rewrites welfare maximization as a regression-style objective, with a Gaussian pseudo-likelihood interpretation. We demonstrate its implementation in both full-feedback and bandit-feedback settings. In our theoretical analysis, we provide PAC-Bayes style generalization bounds and explicit corollaries for welfare.

\onecolumn