diff --git a/docs/src/tutorials/cmab.ipynb b/docs/src/tutorials/cmab.ipynb index a709ec46..c7a9b4de 100644 --- a/docs/src/tutorials/cmab.ipynb +++ b/docs/src/tutorials/cmab.ipynb @@ -2,11 +2,7 @@ "cells": [ { "cell_type": "markdown", - "metadata": { - "pycharm": { - "name": "#%% md\n" - } - }, + "metadata": {}, "source": [ "# Contextual Multi-Armed Bandit" ] @@ -29,22 +25,42 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "The following notebook contains an example of usage of the class Cmab, which implements the algorithm above." + "The following notebook contains an example of usage of the classes `CmabBernoulli` (binary rewards) and `CmabGaussian` (continuous, real-valued rewards), which implement the algorithm above." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "pybandits supports three kinds of reward for contextual bandits:\n", + "\n", + "| reward | type | model | bandit |\n", + "|---|---|---|---|\n", + "| binary 0 / 1 | `BinaryReward` | `BayesianNeuralNetwork` (Bernoulli) | `CmabBernoulli` (and the CC / MO / DP variants) |\n", + "| soft label in $[0, 1]$ | `SoftReward` (`use_soft_rewards=True`) | `BayesianNeuralNetwork` (Bernoulli on fractional targets) | the same |\n", + "| real-valued | `ContinuousReward` | `GaussianBayesianNeuralNetwork` | `CmabGaussian` |\n", + "\n", + "The bandit entry points accept any of these reward types, and the actions manager checks each reward against what its models support." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Binary rewards: `CmabBernoulli`\n", + "\n", + "Use `CmabBernoulli` when the reward is binary (or a soft reward in $[0, 1]$), e.g. a click or a purchase." ] }, { "cell_type": "code", "execution_count": null, - "metadata": { - "pycharm": { - "is_executing": false - } - }, + "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "\n", - "from pybandits.cmab import CmabBernoulli\n", + "from pybandits.cmab import CmabBernoulli, CmabGaussian\n", "from pybandits.model import BayesianNeuralNetwork, BnnLayerParams, BnnParams, FeaturesConfig, StudentTArray" ] }, @@ -179,6 +195,198 @@ "# update model\n", "cmab.update(context=X, actions=pred_actions, rewards=simulated_rewards)" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Continuous rewards: `CmabGaussian`\n", + "\n", + "Use `CmabGaussian` when the reward is real-valued, e.g. revenue or an uplift $R - E_0(x)$ against a control baseline computed outside the bandit. Each action is a `GaussianBayesianNeuralNetwork`: a BNN with a single mean head $\\mu_k(x)$ and a learned scalar noise std $\\sigma_k$, trained with a Gaussian likelihood $r \\sim \\mathcal{N}(\\mu_k(x), \\sigma_k)$. At prediction time the weights are sampled from the posterior and the bandit selects $a_t = argmax_k \\mu_k(x_t)$.\n", + "\n", + "Unlike `CmabBernoulli`, `predict` returns the sampled reward means in place of the probabilities, and the noise std $\\sigma_k$ (for monitoring only) in place of the weighted sums. The adaptive window (`delta`) is not supported, since its change-point test assumes binary rewards." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### The Gaussian model\n", + "\n", + "- **$\\sigma_k$ is a single noise std per action,** not a function of the context. It is a latent ($\\log \\sigma_k \\sim \\mathcal{N}$) whose posterior is stored in the model state (`noise_log_sigma`) and carried across updates, like the weights. Its initial prior is centered on `noise_sigma` (in reward units) if given, otherwise on the std of the first update batch.\n", + "- **Training is on the targets** $(r - \\text{reward\\_loc}) / \\text{reward\\_scale}$. They are set in exactly one of two ways. With `fit_reward_standardization=False` (the default), both must be given as fixed constants and are never re-fitted: choose them so the default $O(1)$ weight priors fit (e.g. `reward_loc=0, reward_scale=20` for a centered uplift in dollars), or use `reward_loc=0, reward_scale=1` to train on raw rewards. With `fit_reward_standardization=True`, leave both unset: they are fitted (mean / std) on the first batch, then frozen until `reset()`.\n", + "- **Thompson sampling** draws the weights from the posterior and ranks actions on $\\mu_k(x)$.\n", + "- **`reset()`** returns the model to its cold-start state, including the standardization: fixed values are restored, fitted ones are re-fitted on the next update.\n", + "\n", + "#### Why one $\\sigma$ and not $\\sigma(x)$\n", + "\n", + "With a single noise level, the fit of $\\mu$ is a (Bayesian) least-squares fit, whose target is $E[r \\mid x]$ whatever the shape of the noise. A context-dependent $\\sigma(x)$ weights each residual by $1/\\sigma(x)^2$. On zero-inflated or skewed rewards (e.g. revenue), the model can then explain large rewards as noise instead of raising $\\mu$, which biases $\\mu$ low exactly where the rewards are most variable. That bias is what Thompson sampling would rank on." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We simulate a two-action environment where the expected reward of each action depends linearly on the first context feature, plus Gaussian noise. Action `a1` is best when $x_0 > 0.25$ and `a2` otherwise." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "rng = np.random.default_rng(0)\n", + "noise_std = 0.5\n", + "\n", + "# expected reward of each action as a function of the context\n", + "true_mean = {\n", + " \"a1\": lambda X: 1.0 + 2.0 * X[:, 0],\n", + " \"a2\": lambda X: 2.0 - 2.0 * X[:, 0],\n", + "}\n", + "\n", + "X_gauss = 2 * rng.random((n_samples, n_features)) - 1 # random float in the interval (-1, 1)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The bandit is built with `cold_start`, which creates one `GaussianBayesianNeuralNetwork` per action id. We use a narrow prior on the weights and a larger optimizer step size, which make the unbounded Gaussian likelihood fit much faster. The rewards here are already $O(1)$, so we train on the raw rewards with the fixed standardization `reward_loc=0.0, reward_scale=1.0`. For rewards far from $O(1)$ (e.g. revenue in dollars), pass fixed values that bring them to $O(1)$, or set `fit_reward_standardization=True` (and no `reward_loc` / `reward_scale`) to fit both on the first batch. `noise_sigma` (in reward units) sets an initial guess of the noise std." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "cmab_gauss = CmabGaussian.cold_start(\n", + " action_ids={\"a1\", \"a2\"},\n", + " n_features=n_features,\n", + " hidden_dim_list=[16],\n", + " reward_loc=0.0,\n", + " reward_scale=1.0,\n", + " dist_type=\"normal\",\n", + " dist_params_init={\"mu\": 0, \"sigma\": 0.1},\n", + " update_kwargs={\"num_steps\": 1000, \"optimizer_kwargs\": {\"step_size\": 3e-3}},\n", + " random_seed=0,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Before any update, the actions are selected from the prior. `predict` returns, for each sample, the selected action, the sampled mean reward of every action and its noise std." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "pred_actions_g, mus, sigmas = cmab_gauss.predict(X_gauss)\n", + "print(\"Recommended action: {}\".format(pred_actions_g[:10]))\n", + "print(\"Sampled mean rewards: {}\".format(mus[0]))\n", + "print(\"Noise std: {}\".format(sigmas[0]))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now observe the (continuous) rewards of the selected actions. Here they are simulated from the true model above." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "simulated_rewards_g = [\n", + " float(true_mean[a](X_gauss[i : i + 1])[0] + rng.normal(scale=noise_std)) for i, a in enumerate(pred_actions_g)\n", + "]\n", + "print(\"Simulated rewards: {}\".format(np.round(simulated_rewards_g[:10], 2)))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finally, we update the model providing per each sample the context $x_t$, the selected action $a_t$ and the corresponding reward $r_t$." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "cmab_gauss.update(context=X_gauss, actions=pred_actions_g, rewards=simulated_rewards_g)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "After the update, the bandit selects the best action for each context far more often. We compare its recommendations on fresh contexts with the truly best action." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "X_test = 2 * rng.random((500, n_features)) - 1\n", + "pred_actions_test, mus_test, sigmas_test = cmab_gauss.predict(X_test)\n", + "best_actions = np.where(true_mean[\"a1\"](X_test) > true_mean[\"a2\"](X_test), \"a1\", \"a2\")\n", + "print(\"Share of best-action recommendations: {:.2f}\".format(np.mean(np.array(pred_actions_test) == best_actions)))\n", + "print(\"Sampled mean rewards (first context): {}\".format(mus_test[0]))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Recommended use: uplift against a control baseline\n", + "\n", + "For rewards with a lot of player-to-player variance (e.g. revenue), train on the **uplift** $r = R - E_0(x)$, where $E_0(x)$ is the expected reward under a control policy, from a separate baseline model. $E_0$ doesn't depend on the action, so the ranking of actions is unchanged, and the shared variance is removed (a control variate). The posterior then contracts much faster with per-action data volumes. Note that $r$ can be negative.\n", + "\n", + "```python\n", + "mab = CmabGaussian.cold_start(\n", + " action_ids={\"a1\", \"a2\", \"a3\"},\n", + " n_features=n_features,\n", + " hidden_dim_list=[16],\n", + " dist_type=\"normal\",\n", + " dist_params_init={\"mu\": 0, \"sigma\": 0.1}, # a narrow prior fits much faster than the default sigma=1\n", + " update_kwargs={\"num_steps\": 400, \"optimizer_kwargs\": {\"step_size\": 3e-3}},\n", + " reward_loc=0.0, # fixed values, e.g. from historical data;\n", + " reward_scale=..., # or fit_reward_standardization=True instead of both, to fit them on the first batch\n", + " # noise_sigma=..., # optional initial guess of the noise std, in reward units\n", + " # decay_factor=..., # if the reward process drifts\n", + ")\n", + "actions, mu, sigma = mab.predict(context=X)\n", + "mab.update(actions=actions, rewards=(R - E0).tolist(), context=X)\n", + "```\n", + "\n", + "A shared backbone (`backbone_hidden_dims=...`) works as for `CmabBernoulli`; each action keeps its own $\\sigma_k$ and standardization.\n", + "\n", + "### Alternatives that were evaluated and rejected\n", + "\n", + "- **$\\sigma(x)$ noise head:** biases $\\mu$ on skewed rewards (see above).\n", + "- **Bernoulli on a $[0, 1]$ rescaling of a continuous reward:** the mean is consistent, but the Bernoulli variance $p(1 - p)$ can overstate the real noise by orders of magnitude. Each observation then carries far too little weight, and the model learns little beyond an intercept.\n", + "- **Hurdle $\\times$ log-normal amount:** a good fit of the distribution's shape, but the back-transform $\\exp(m + \\sigma^2/2)$ can bias the level, and Thompson draws of it have extreme tails. If a split of the effect into conversion vs. amount is needed, a Bernoulli $\\times$ Gaussian amount (on the raw scale) is the better hurdle.\n", + "\n", + "### Not supported (yet)\n", + "\n", + "- Gaussian variants of the cost-control, multi-objective, dynamic-pricing and quantitative bandits.\n", + "- The adaptive window (`delta`) with continuous rewards.\n", + "- Offline policy evaluation and the simulators with continuous rewards." + ] } ], "metadata": { diff --git a/pybandits/actions_manager.py b/pybandits/actions_manager.py index 4fc18a71..7d26e699 100644 --- a/pybandits/actions_manager.py +++ b/pybandits/actions_manager.py @@ -52,6 +52,7 @@ from pybandits.base import ( ActionId, + AnyReward, PositiveProbability, PyBanditsBaseModel, Reward, @@ -78,6 +79,7 @@ BetaDP, BetaMO, BetaMOCC, + GaussianBayesianNeuralNetwork, Model, ModelMO, ) @@ -125,7 +127,8 @@ class ActionsManager(PyBanditsBaseModel, ABC): The confidence level for the adaptive window. None for skipping the change point detection. use_soft_rewards : bool If True, rewards may be soft labels in [0, 1]; otherwise they must be binary (0/1). - Not supported together with the adaptive window (delta). + Not supported together with the adaptive window (delta). Ignored when the models take + continuous rewards (``supports_continuous_rewards``), whose rewards may be any finite float. """ meta_model: BaseMetaModel @@ -217,8 +220,46 @@ def _check_soft_rewards_delta(self) -> "ActionsManager": """The adaptive window's change-point test and memory checks assume binary rewards.""" if self.delta is not None and self.use_soft_rewards: raise ValueError("Soft rewards are not supported with the adaptive window (delta).") + if self.delta is not None and self.continuous_rewards: + raise ValueError("Continuous rewards are not supported with the adaptive window (delta).") return self + @model_validator(mode="after") + def _check_uniform_reward_support(self) -> "ActionsManager": + """All action models must take the same kind of reward: either all continuous or none.""" + supports = {model.supports_continuous_rewards for model in self.actions.values()} + if len(supports) > 1: + raise ValueError("Action models must either all take continuous rewards or none of them.") + return self + + @property + def continuous_rewards(self) -> bool: + """Whether the action models take real-valued rewards rather than binary / soft ones in [0, 1].""" + return all(model.supports_continuous_rewards for model in self.actions.values()) + + def _check_rewards(self, *rewards: Union[List[AnyReward], List[List[AnyReward]]]) -> None: + """Check the rewards against what the action models accept. + + Parameters + ---------- + *rewards : Union[List[AnyReward], List[List[AnyReward]]] + Reward lists to check (e.g. the batch and the memory). + + Raises + ------ + ValueError + If a reward is not finite (continuous models), outside [0, 1] (soft rewards), or not 0/1. + """ + arrays = [np.asarray(r, dtype=float) for r in rewards] + if self.continuous_rewards: + if not all(np.isfinite(a).all() for a in arrays): + raise ValueError("Rewards must be finite.") + elif self.use_soft_rewards: + if not all(((a >= 0) & (a <= 1)).all() for a in arrays): + raise ValueError("Soft rewards must be in [0, 1].") + elif not all(np.isin(a, (0, 1)).all() for a in arrays): + raise ValueError("Rewards must be binary (0/1); set use_soft_rewards=True to pass soft rewards in [0, 1].") + def _validate_update_params( self, actions: List[ActionId], @@ -293,10 +334,10 @@ def sample_proba( def update( self, actions: List[ActionId], - rewards: Union[List[Reward], List[List[Reward]]], + rewards: Union[List[AnyReward], List[List[AnyReward]]], quantities: Optional[List[Union[float, List[float], None]]] = None, actions_memory: Optional[List[ActionId]] = None, - rewards_memory: Optional[Union[List[Reward], List[List[Reward]]]] = None, + rewards_memory: Optional[Union[List[AnyReward], List[List[AnyReward]]]] = None, **kwargs, ): """ @@ -317,8 +358,7 @@ def update( List of previously collected rewards. """ self.actions_with_change.clear() - if not self.use_soft_rewards and not all(np.isin(r, (0, 1)).all() for r in (rewards, rewards_memory or [])): - raise ValueError("Rewards must be binary (0/1); set use_soft_rewards=True to pass soft rewards in [0, 1].") + self._check_rewards(rewards, rewards_memory or []) if self.delta is None and (actions_memory or rewards_memory): raise AttributeError("Adaptive window size is not set, so memory should not be provided.") if self.delta is not None and (actions_memory is None or rewards_memory is None): @@ -851,11 +891,11 @@ def _check_backbone_delta(self) -> "CmabActionsManager": def update( self, actions: List[ActionId], - rewards: Union[List[Reward], List[List[Reward]]], + rewards: Union[List[AnyReward], List[List[AnyReward]]], quantities: Optional[List[Union[float, List[float], None]]], context: np.ndarray, actions_memory: Optional[List[ActionId]] = None, - rewards_memory: Optional[Union[List[Reward], List[List[Reward]]]] = None, + rewards_memory: Optional[Union[List[AnyReward], List[List[AnyReward]]]] = None, context_memory: Optional[np.ndarray] = None, ): """ @@ -967,3 +1007,4 @@ def _update_actions( CmabActionsManagerDP = CmabActionsManager[Union[BayesianNeuralNetworkDP, QuantitativeBayesianNeuralNetworkDP]] CmabActionsManagerMO = CmabActionsManager[BayesianNeuralNetworkMO] CmabActionsManagerMOCC = CmabActionsManager[BayesianNeuralNetworkMOCC] +CmabActionsManagerGaussian = CmabActionsManager[GaussianBayesianNeuralNetwork] diff --git a/pybandits/base.py b/pybandits/base.py index cf16b819..74762fca 100644 --- a/pybandits/base.py +++ b/pybandits/base.py @@ -41,6 +41,7 @@ BaseModel, ConfigDict, Field, + FiniteFloat, confloat, conint, constr, @@ -55,6 +56,8 @@ Probability = NewType("Probability", Float01) PositiveProbability = NewType("PositiveProbability", confloat(gt=0, le=1)) ProbabilityWeight = Tuple[Probability, float] +# Gaussian-likelihood sample: (mean, noise std) of the reward, in the ProbabilityWeight slot of a Bernoulli model. +GaussianSample = Tuple[float, float] MOProbability = List[Probability] MOProbabilityWeight = List[ProbabilityWeight] # QuantitativeProbability generalizes probability to include both action quantities and their associated probability @@ -108,6 +111,11 @@ SoftReward = NewType("SoftReward", Float01) # Binary or soft label; left_to_right keeps 0/1 (incl. numpy ints) as int. Reward = NewType("Reward", Annotated[Union[BinaryReward, SoftReward], Field(union_mode="left_to_right")]) +ContinuousReward = NewType("ContinuousReward", FiniteFloat) +# Any reward a bandit entry point accepts; the actions manager narrows it to what its models support. +AnyReward = NewType( + "AnyReward", Annotated[Union[BinaryReward, SoftReward, ContinuousReward], Field(union_mode="left_to_right")] +) ActionRewardLikelihood = NewType( "ActionRewardLikelihood", Union[ diff --git a/pybandits/base_model.py b/pybandits/base_model.py index a890bfab..3de9cac4 100644 --- a/pybandits/base_model.py +++ b/pybandits/base_model.py @@ -52,6 +52,9 @@ class BaseModel(PyBanditsBaseModel, ABC): Class to model the prior distributions of standard actions and quantitative actions. """ + # Whether the model accepts real-valued rewards (otherwise binary or soft rewards in [0, 1]). + supports_continuous_rewards: ClassVar[bool] = False + @abstractmethod def sample_proba( self, rng: Generator, **kwargs diff --git a/pybandits/cmab.py b/pybandits/cmab.py index aa486763..1dd15d2f 100644 --- a/pybandits/cmab.py +++ b/pybandits/cmab.py @@ -30,6 +30,7 @@ CmabActionsManager, CmabActionsManagerCC, CmabActionsManagerDP, + CmabActionsManagerGaussian, CmabActionsManagerMO, CmabActionsManagerMOCC, CmabActionsManagerSO, @@ -37,11 +38,11 @@ ) from pybandits.base import ( ActionId, + AnyReward, CmabPredictions, ForbiddenActions, MOProbabilityWeight, ProbabilityWeight, - Reward, Serializable, ) from pybandits.mab import BaseMab @@ -162,11 +163,11 @@ def predict( def update( self, actions: List[ActionId], - rewards: Union[List[Reward], List[List[Reward]]], + rewards: Union[List[AnyReward], List[List[AnyReward]]], context: np.ndarray, quantities: Optional[List[Union[float, List[float], None]]] = None, actions_memory: Optional[List[ActionId]] = None, - rewards_memory: Optional[Union[List[Reward], List[List[Reward]]]] = None, + rewards_memory: Optional[Union[List[AnyReward], List[List[AnyReward]]]] = None, context_memory: Optional[np.ndarray] = None, ): """ @@ -177,8 +178,8 @@ def update( ---------- actions : List[ActionId] of shape (n_samples,), e.g. ['a1', 'a2', 'a3', 'a4', 'a5'] The selected action for each sample. - rewards : List[Union[Reward, List[Reward]]] of shape (n_samples, n_objectives) - The binary reward for each sample. + rewards : List[Union[AnyReward, List[AnyReward]]] of shape (n_samples, n_objectives) + The reward for each sample: binary, soft in [0, 1], or real-valued for continuous-reward models. If strategy is not MultiObjectiveBandit, rewards should be a list, e.g. rewards = [1, 0, 1, 1, 1, ...] If strategy is MultiObjectiveBandit, rewards should be a list of list, e.g. (with n_objectives=2): @@ -189,7 +190,7 @@ def update( The value associated with each action. If none, the value is not used, i.e. non-quantitative action. actions_memory : Optional[List[ActionId]] List of previously selected actions. - rewards_memory : Optional[Union[List[Reward], List[List[Reward]]]] + rewards_memory : Optional[Union[List[AnyReward], List[List[AnyReward]]]] List of previously collected rewards. context_memory : Optional[ArrayLike] of shape (n_samples, n_features) Matrix of contextual features. @@ -306,6 +307,43 @@ class CmabBernoulli(BaseCmabBernoulli): _predict_with_proba: bool = False +class CmabGaussian(BaseCmabBernoulli): + """ + Contextual Multi-Armed Bandit with Thompson Sampling for continuous (real-valued) rewards. + + Each action is a :class:`~pybandits.model.GaussianBayesianNeuralNetwork`: a BNN with a single mean head + ``mu(x)`` and a learned scalar noise std per action, trained with a Gaussian likelihood. At prediction time the + weights are sampled from the posterior and the action with the highest sampled reward mean is selected. + + ``predict`` returns the sampled reward means in place of the probabilities and the noise std (for monitoring + only) in place of the weighted sums. + + A typical reward is an uplift ``R - E0(x)`` against a control baseline computed outside the bandit: it removes + the noise shared by all actions without changing their ranking. + + References + ---------- + Thompson Sampling for Contextual Bandits with Linear Payoffs (Agrawal and Goyal, 2014) + https://arxiv.org/pdf/1209.3352.pdf + + Parameters + ---------- + actions_manager: CmabActionsManagerGaussian + The manager for actions and their associated models. + strategy: ClassicBandit + The strategy used to select actions. + + Notes + ----- + A shared backbone is requested at ``cold_start`` exactly as for :class:`CmabBernoulli`. The adaptive + window (``delta``) is not supported, since its change-point test assumes binary rewards. + """ + + actions_manager: CmabActionsManagerGaussian + strategy: ClassicBandit + _predict_with_proba: bool = True + + class CmabBernoulliBAI(BaseCmabBernoulli): """ Contextual Bernoulli Multi-Armed Bandit with Thompson Sampling, and Best Action Identification strategy. diff --git a/pybandits/mab.py b/pybandits/mab.py index dd6c4946..b438f5b5 100644 --- a/pybandits/mab.py +++ b/pybandits/mab.py @@ -41,6 +41,7 @@ from pybandits.base import ( ActionId, ActionRewardLikelihood, + AnyReward, Float01, ForbiddenActions, MOProbability, @@ -54,7 +55,6 @@ QuantitativeMOProbabilityWeight, QuantitativeProbability, QuantitativeProbabilityWeight, - Reward, Serializable, UnifiedActionId, ) @@ -343,10 +343,10 @@ def actions(self) -> Dict[ActionId, BaseModel]: def update( self, actions: List[ActionId], - rewards: Union[List[Reward], List[List[Reward]]], + rewards: Union[List[AnyReward], List[List[AnyReward]]], quantities: Optional[List[Union[float, List[float], None]]] = None, actions_memory: Optional[List[ActionId]] = None, - rewards_memory: Optional[Union[List[Reward], List[List[Reward]]]] = None, + rewards_memory: Optional[Union[List[AnyReward], List[List[AnyReward]]]] = None, **kwargs, ): """ @@ -356,8 +356,8 @@ def update( ---------- actions: List[ActionId] The selected action for each sample. - rewards : Union[List[Reward], List[List[Reward]]] of shape (n_samples, n_objectives) - The binary reward for each sample. + rewards : Union[List[AnyReward], List[List[AnyReward]]] of shape (n_samples, n_objectives) + The reward for each sample: binary, soft in [0, 1], or real-valued for continuous-reward models. If strategy is not MultiObjectiveBandit, rewards should be a list, e.g. rewards = [1, 0, 1, 1, 1, ...] If strategy is MultiObjectiveBandit, rewards should be a list of list, e.g. (with n_objectives=2): @@ -366,7 +366,7 @@ def update( The value associated with each action. If none, the value is not used, i.e. non-quantitative action. actions_memory : Optional[List[ActionId]] List of previously selected actions. - rewards_memory : Optional[Union[List[Reward], List[List[Reward]]]] + rewards_memory : Optional[Union[List[AnyReward], List[List[AnyReward]]]] List of previously collected rewards. """ self.actions_manager.update( diff --git a/pybandits/meta_model/__init__.py b/pybandits/meta_model/__init__.py index 2b71fe78..d1b4f7cf 100644 --- a/pybandits/meta_model/__init__.py +++ b/pybandits/meta_model/__init__.py @@ -29,6 +29,7 @@ CmabMetaModel, CmabMetaModelCC, CmabMetaModelDP, + CmabMetaModelGaussian, CmabMetaModelMO, CmabMetaModelMOCC, CmabMetaModelSO, @@ -58,4 +59,5 @@ "CmabMetaModelDP", "CmabMetaModelMO", "CmabMetaModelMOCC", + "CmabMetaModelGaussian", ] diff --git a/pybandits/meta_model/cmab_meta_model.py b/pybandits/meta_model/cmab_meta_model.py index bec60290..9ce221c5 100644 --- a/pybandits/meta_model/cmab_meta_model.py +++ b/pybandits/meta_model/cmab_meta_model.py @@ -43,7 +43,6 @@ import numpyro import optax from loguru import logger -from numpyro.distributions import Bernoulli as NumpyroBernoulli from numpyro.infer import TraceMeanField_ELBO from pydantic import ConfigDict, NonNegativeInt, PrivateAttr, model_validator @@ -58,6 +57,7 @@ BayesianNeuralNetworkDP, BayesianNeuralNetworkMO, BayesianNeuralNetworkMOCC, + GaussianBayesianNeuralNetwork, ) from pybandits.model.bnn._guide import ParameterizedScaleAutoNormal from pybandits.model.bnn._svi import forward_layers, per_sample_linear, run_svi @@ -237,6 +237,10 @@ def _check_models(self) -> "CmabMetaModel": raise AttributeError("All actions should have the same input size.") if first_bnn.update_kwargs != self._rep_bnn(head).update_kwargs: raise AttributeError("All actions should have the same update kwargs.") + # The joint pass observes every row under the representative head's likelihood. + head_bnn = self._rep_bnn(head) + if type(head_bnn).output_distribution is not type(first_bnn).output_distribution: + raise AttributeError("All actions should have the same reward likelihood.") # What must match the embedding is the width the head's first layer consumes, which its own # categorical expansion (if any) widens beyond the context dim. feature_config = first_bnn.feature_config @@ -449,11 +453,12 @@ def _joint_svi_update( declare_backbone = functools.partial(self._declare_backbone, backbone_w, backbone_b, n_bb_layers) batch_size = resolved_kwargs.get("batch_size") + targets = self._prepare_targets(arm_to_rows, rewards_arr) if self._heads_support_minibatching: # One shared data plate for every arm's rows pooled together; _so_model subsamples it # (the arm-indexing trick) when batch_size < n_samples, else it's a full-batch nullcontext. action_index = jnp.asarray(self._build_action_index(arm_to_rows, batch_arms), dtype=jnp.int32) - y_all = jnp.asarray(rewards_arr, dtype=BaseBayesianNeuralNetwork._reward_dtype) + y_all = jnp.asarray(targets, dtype=BaseBayesianNeuralNetwork.reward_dtype) model = self._make_so_model(batch_arms, representative_bnn, batch_size, n_samples, declare_backbone) model_args: Tuple[Any, ...] = (x, action_index, y_all) else: @@ -465,7 +470,7 @@ def _joint_svi_update( arm_data = { arm: ( jnp.asarray(rows, dtype=jnp.int32), - jnp.asarray(rewards_arr[rows], dtype=BaseBayesianNeuralNetwork._reward_dtype), + jnp.asarray(targets[rows], dtype=BaseBayesianNeuralNetwork.reward_dtype), ) for arm, rows in arm_to_rows.items() } @@ -565,11 +570,11 @@ def _full_batch_model( to MO's multiple weight sets per arm or quantitative's per-row quantity concatenation). """ z = self.backbone.forward_jax(x_, declare_backbone()) if self.backbone is not None else x_ - for arm, (row_idx, arm_rewards) in arm_data_.items(): + for arm, (row_idx, arm_targets) in arm_data_.items(): z_arm = z[row_idx] for bnn, obj_index, is_quantitative in self._arm_units(self.actions[arm]): x_unit = jnp.concatenate([arm_q[arm], z_arm], axis=1) if is_quantitative else z_arm - y_unit = arm_rewards if obj_index is None else arm_rewards[:, obj_index] + y_unit = arm_targets if obj_index is None else arm_targets[:, obj_index] # scope() prepends "{arm}/" (and "/objN" for MO) to every site this unit emits, so # the otherwise-identical per-head site names stay unique across arms. with numpyro.handlers.scope(prefix=self._unit_scope(arm, obj_index), divider=self._scope_divider): @@ -677,12 +682,15 @@ def _so_model( ) -> None: """The single-BNN-per-arm joint NumPyro model; built and bound by :meth:`_make_so_model`.""" backbone_wb = declare_backbone() if self.backbone is not None else None - per_arm = [] + per_arm, per_arm_extra = [], [] for arm in batch_arms: ((bnn, obj_index, _q),) = self._arm_units(self.actions[arm]) with numpyro.handlers.scope(prefix=self._unit_scope(arm, obj_index), divider=self._scope_divider): per_arm.append(bnn.sample_head_sites(cast(Any, kl_annealing_factor))) + per_arm_extra.append(bnn.sample_extra_sites(cast(Any, kl_annealing_factor))) stacked_wb, stacked_emb = self._stack_head_sites(per_arm) + # Head-level likelihood latents (e.g. the Gaussian noise std), stacked to (num_arms, *site_shape). + stacked_extra = {name: jnp.stack([extra[name] for extra in per_arm_extra]) for name in per_arm_extra[0]} use_minibatch = batch_size is not None and batch_size < n_samples plate_ctx = numpyro.plate("data", size=n_samples, subsample_size=batch_size) if use_minibatch else nullcontext() @@ -695,15 +703,16 @@ def _so_model( else: head_input = self._per_sample_head_input(ctx, arm_of_row, stacked_emb, representative_bnn) per_sample_wb = [(w[arm_of_row], b[arm_of_row]) for w, b in stacked_wb] - logit = forward_layers( + linear_out = forward_layers( next_layer_input=head_input, weights_biases=per_sample_wb, activation_fn=representative_bnn._jax_activation_fn, linear_fn=per_sample_linear, backend=jnp, use_residual_connections=representative_bnn.use_residual_connections, - ).squeeze(-1) - numpyro.sample("out", NumpyroBernoulli(logits=logit), obs=y_batch) + ) + row_extra = {name: site[arm_of_row] for name, site in stacked_extra.items()} + numpyro.sample("out", representative_bnn.output_distribution(linear_out, row_extra), obs=y_batch) @staticmethod def _stack_head_sites(per_arm: List[Tuple[list, list]]) -> Tuple[list, list]: @@ -866,6 +875,7 @@ def _store_head(self, head: BaseModel, site_mu: dict, site_sigma: dict, arm: Act unit_sigma = {name[len(scope) :]: v for name, v in site_sigma.items() if name.startswith(scope)} bnn.model_params.bnn_layer_params = bnn.layer_params_from_posterior(unit_mu, unit_sigma) bnn.update_embedding_params_from_vi(unit_mu, unit_sigma) + bnn.update_extra_params_from_vi(unit_mu, unit_sigma) def _store_backbone(self, params: dict, n_bb_layers: int) -> None: """Write back the trained backbone point estimates (no-op when no backbone). @@ -884,6 +894,31 @@ def _store_backbone(self, params: dict, n_bb_layers: int) -> None: new_b = [np.asarray(params[bn]) for _, bn in layer_names] self.backbone = self.backbone.with_weights_and_biases(new_w, new_b) + def _prepare_targets(self, arm_to_rows: Dict[ActionId, List[int]], rewards_arr: np.ndarray) -> np.ndarray: + """Training targets for the joint pass: each single-BNN head maps its own rows' rewards. + + Identity for Bernoulli heads; Gaussian heads map with their own location and scale (fixed, or fitted on + their first batch), so arms on different reward scales share one likelihood in the joint ELBO. + + Parameters + ---------- + arm_to_rows : Dict[ActionId, List[int]] + Row indices of the batch belonging to each arm. + rewards_arr : np.ndarray + Rewards for the batch (1-D single-objective, 2-D ``(n_samples, n_objectives)`` for MO). + + Returns + ------- + np.ndarray + Targets, row-aligned with ``rewards_arr`` (the raw rewards are left untouched for the counters). + """ + targets = np.array(rewards_arr, dtype=float) + for arm, rows in arm_to_rows.items(): + head = self.actions[arm] + if isinstance(head, BaseBayesianNeuralNetwork): + targets[rows] = head.prepare_rewards(rewards_arr[rows].tolist()) + return targets + def _increment_counters(self, arm_to_rows: Dict[ActionId, List[int]], rewards_arr: np.ndarray) -> None: """Keep each head's success/failure counters in sync (the joint engine bypasses per-head ``update``). @@ -915,3 +950,4 @@ def _increment_counters(self, arm_to_rows: Dict[ActionId, List[int]], rewards_ar CmabMetaModelDP = CmabMetaModel[Union[BayesianNeuralNetworkDP, QuantitativeBayesianNeuralNetworkDP]] CmabMetaModelMO = CmabMetaModel[BayesianNeuralNetworkMO] CmabMetaModelMOCC = CmabMetaModel[BayesianNeuralNetworkMOCC] +CmabMetaModelGaussian = CmabMetaModel[GaussianBayesianNeuralNetwork] diff --git a/pybandits/model/__init__.py b/pybandits/model/__init__.py index 34e0301e..cf1da8b0 100644 --- a/pybandits/model/__init__.py +++ b/pybandits/model/__init__.py @@ -60,6 +60,7 @@ BayesianNeuralNetworkDP, BayesianNeuralNetworkMO, BayesianNeuralNetworkMOCC, + GaussianBayesianNeuralNetwork, ) from pybandits.model.bnn.priors import ( BaseLocationScaleArray, @@ -100,4 +101,5 @@ "BaseBayesianNeuralNetworkMO", "BayesianNeuralNetworkMO", "BayesianNeuralNetworkMOCC", + "GaussianBayesianNeuralNetwork", ] diff --git a/pybandits/model/bnn/network.py b/pybandits/model/bnn/network.py index d77accbc..f69d0d90 100644 --- a/pybandits/model/bnn/network.py +++ b/pybandits/model/bnn/network.py @@ -34,6 +34,8 @@ import numpyro.optim as noptim import optax from numpyro.distributions import Bernoulli as NumpyroBernoulli +from numpyro.distributions import Distribution as NumpyroDistribution +from numpyro.distributions import Normal as NumpyroNormal from numpyro.infer import Trace_ELBO, TraceMeanField_ELBO from numpyro.infer.autoguide import AutoMultivariateNormal from numpyro.infer.initialization import init_to_median, init_to_value @@ -51,6 +53,8 @@ from typing_extensions import Self from pybandits.base import ( + ContinuousReward, + GaussianSample, PositiveFloat01, ProbabilityWeight, Reward, @@ -140,7 +144,9 @@ class BaseBayesianNeuralNetwork(Model, DNNMixin, ABC): model_params: BnnParams - _reward_dtype: ClassVar[Any] = jnp.float32 # float, not int: soft rewards are fractional labels + reward_dtype: ClassVar[Any] = jnp.float32 # float, not int: soft rewards are fractional labels + # Width of the output layer: one logit for the Bernoulli likelihood. + default_output_dim: ClassVar[PositiveInt] = 1 _logit_var_name: ClassVar[str] = "logit" _prob_var_name: ClassVar[str] = "prob" weight_var_name: ClassVar[str] = "weight" @@ -367,6 +373,7 @@ def create_model_params( use_layerwise_scaling: bool = False, dist_class: type[BaseLocationScaleArray] = StudentTArray, bias_std: Optional[PositiveFloat] = None, + output_dim: Optional[PositiveInt] = None, **dist_params_init, ) -> BnnParams: """ @@ -390,6 +397,8 @@ def create_model_params( If provided, overrides ``sigma`` from ``dist_params_init`` for all layers' bias priors. Applied to every layer's bias (including the output layer's logit bias), leaving weight priors unchanged. Default is None (use ``sigma`` from ``dist_params_init``). + output_dim : Optional[PositiveInt] + Width of the output layer. None uses the class attribute ``default_output_dim``. **dist_params_init : dict, optional Additional parameters for initializing the distribution of weights and biases. @@ -405,7 +414,7 @@ def create_model_params( else: _dim_list = [effective_n_features] + hidden_dim_list - _dim_list.append(1) + _dim_list.append(output_dim or cls.default_output_dim) layer_params_init = [] for layer_ind in range(len(_dim_list) - 1): @@ -607,11 +616,11 @@ def _create_update_model(self) -> Callable: (current posteriors used as new priors). 2. Sample embedding matrices for categorical features (if any). 3. Apply linear transformations and activations through the layers. - 4. Apply sigmoid activation at the output. - 5. Use Bernoulli likelihood for binary classification + 4. Observe the rewards under :meth:`output_distribution` on the raw output layer (Bernoulli on the + logit for binary rewards; a subclass may override it, e.g. Gaussian for continuous rewards). Steps 1-2 happen inside ``numpyro.handlers.scale(scale=kl_annealing_factor)`` so that - the KL portion of the ELBO can be scheduled across training steps. Step 5 stays outside + the KL portion of the ELBO can be scheduled across training steps. Step 4 stays outside that context so the likelihood term is not scaled. """ @@ -631,8 +640,9 @@ def emit_submodel( ) -> None: """Emit this BNN's NumPyro sites on the supplied input array. - Samples the layer weights/biases and categorical embeddings, runs the forward pass on ``x``, - and registers the ``logit`` deterministic + ``out`` Bernoulli likelihood. Factored out of + Samples the layer weights/biases, categorical embeddings and any :meth:`extra_site_params` latents + (e.g. the Gaussian head's noise std), runs the forward pass on ``x``, and registers the ``logit`` + deterministic + ``out`` likelihood (:meth:`output_distribution`). Factored out of :meth:`_create_update_model` so a meta-model can compose several arms' sub-models into one joint model on a shared backbone embedding *or* on raw context — the meta-model wraps each call in ``numpyro.handlers.scope(prefix=arm)`` to keep the (otherwise identical) site names @@ -643,14 +653,15 @@ def emit_submodel( x : jax.Array Network input — raw context (standalone / no-backbone head) or a backbone embedding. y : jax.Array - Binary rewards for these rows. + Training targets for these rows (from :meth:`prepare_rewards`). kl_annealing_factor : Union[PositiveFloat01, jax.Array] Scales the prior-site log-probabilities (KL term); ``1.0`` is a no-op. batch_size : Optional[PositiveInt] Minibatch size for the data plate; ``None`` (or ≥ n_samples) means full batch. """ weights_biases, embedding_matrices = self.sample_head_sites(kl_annealing_factor) - self._observe(x, y, weights_biases, embedding_matrices, batch_size) + extra_sites = self.sample_extra_sites(kl_annealing_factor) + self._observe(x, y, weights_biases, embedding_matrices, batch_size, extra_sites=extra_sites) def sample_head_sites( self, kl_annealing_factor: Union[PositiveFloat01, jax.Array] = 1.0 @@ -698,12 +709,13 @@ def _observe( weights_biases: List[Tuple[jax.Array, jax.Array]], embedding_matrices: List[jax.Array], batch_size: Optional[PositiveInt] = None, + extra_sites: Optional[Dict[str, jax.Array]] = None, ) -> None: """Forward the (optionally minibatched) input through the sampled sites and observe ``out``. Builds the network input (numerical columns + embedded categoricals), runs the forward pass, - registers the ``logit`` deterministic, and observes the ``out`` Bernoulli likelihood — inside a - subsampling ``data`` plate when ``batch_size`` is set. Kept separate from + registers the ``logit`` deterministic, and observes the ``out`` likelihood (:meth:`output_distribution`) + — inside a subsampling ``data`` plate when ``batch_size`` is set. Kept separate from :meth:`sample_head_sites` so the latent sites can be sampled once and reused. Parameters @@ -711,13 +723,15 @@ def _observe( x : jax.Array Network input — raw context (standalone / no-backbone head) or a backbone embedding. y : jax.Array - Binary rewards for these rows. + Training targets for these rows (from :meth:`prepare_rewards`). weights_biases : List[Tuple[jax.Array, jax.Array]] Per-layer ``(weight, bias)`` from :meth:`sample_head_sites`. embedding_matrices : List[jax.Array] Categorical embedding matrices from :meth:`sample_head_sites` (empty if none). batch_size : Optional[PositiveInt] Minibatch size for the data plate; ``None`` (or ≥ n_samples) means full batch. + extra_sites : Optional[Dict[str, jax.Array]] + Head-level likelihood latents from :meth:`sample_extra_sites` (empty for the Bernoulli head). """ numerical_indices = self.feature_config.numerical_indices cat_configs = self.feature_config.categorical_features_configs @@ -755,14 +769,110 @@ def _observe( backend=jnp, ) - # Final output processing - logit = numpyro.deterministic( - self._logit_var_name, - linear_transform.squeeze(-1), - ) - numpyro.sample( - "out", NumpyroBernoulli(logits=logit), obs=y_batch - ) # "The observed reward follows a Bernoulli distribution given the network output" + self._observe_output(linear_transform, y_batch, extra_sites) + + def extra_site_params(self) -> Dict[str, BaseLocationScaleArray]: + """Head-level latents of the likelihood (beyond weights / embeddings), keyed by site name. + + They are sampled once per head, outside the data plate, and carried across updates like the weights + (current posterior → next prior). None for the Bernoulli head; e.g. the Gaussian head's + log noise std. + """ + return {} + + def sample_extra_sites(self, kl_annealing_factor: Union[PositiveFloat01, jax.Array] = 1.0) -> Dict[str, jax.Array]: + """Sample the :meth:`extra_site_params` latents as NumPyro sites (KL-annealed like the weights).""" + sites = {} + with numpyro.handlers.scale(scale=kl_annealing_factor): + for name, dist in self.extra_site_params().items(): + sites[name] = numpyro.sample(name, dist.to_numpyro_distribution()) + return sites + + def sample_extra_numpy(self, n_samples: PositiveInt, rng: np.random.Generator) -> Dict[str, np.ndarray]: + """Draw the :meth:`extra_site_params` latents for ``n_samples`` rows, shape ``(n_samples, *site_shape)``.""" + return { + name: dist.sample_rvs(size=(n_samples, *dist.shape), rng=rng) + for name, dist in self.extra_site_params().items() + } + + def update_extra_params_from_vi(self, site_mu: dict, site_sigma: dict) -> None: + """Store the VI posterior of the :meth:`extra_site_params` latents (no-op without any).""" + + def _inflate_extra_params(self, inflation: float) -> None: + """Widen the stored :meth:`extra_site_params` posteriors by ``inflation`` (forgetting; no-op without any).""" + + def output_distribution( + self, linear_out: jax.Array, extra_sites: Optional[Dict[str, jax.Array]] = None + ) -> NumpyroDistribution: + """The reward likelihood given the raw output layer. + + Shared by this BNN's own model and the joint cMAB meta-model, so a subclass only overrides this + (plus :meth:`_postprocess_output`) to change the likelihood. + + Parameters + ---------- + linear_out : jax.Array of shape (batch, default_output_dim) + Raw (pre-activation) output of the final layer. + extra_sites : Optional[Dict[str, jax.Array]] + Sampled :meth:`extra_site_params` latents, broadcastable to the batch (unused here). + + Returns + ------- + NumpyroDistribution + Bernoulli on the logit. + """ + return NumpyroBernoulli(logits=linear_out.squeeze(-1)) + + def _observe_output( + self, linear_out: jax.Array, y: jax.Array, extra_sites: Optional[Dict[str, jax.Array]] = None + ) -> None: + """Register the ``logit`` deterministic and observe ``y`` under :meth:`output_distribution`.""" + numpyro.deterministic(self._logit_var_name, linear_out.squeeze(-1)) + # The observed reward follows output_distribution (Bernoulli by default) given the network output + numpyro.sample("out", self.output_distribution(linear_out, extra_sites), obs=y) + + def _postprocess_output( + self, linear_out: np.ndarray, extra_samples: Optional[Dict[str, np.ndarray]] = None + ) -> List[ProbabilityWeight]: + """Map the sampled raw output layer to the per-sample values returned by :meth:`sample_proba`. + + Parameters + ---------- + linear_out : np.ndarray of shape (n_samples, default_output_dim) + Raw output of the final layer, one row per sampled network. + extra_samples : Optional[Dict[str, np.ndarray]] + Per-row draws of the :meth:`extra_site_params` latents (unused here). + + Returns + ------- + List[ProbabilityWeight] + ``(sigmoid(logit), logit)`` per sample. + """ + weighted_sum = linear_out.squeeze(-1) + prob = _numpy_sigmoid(weighted_sum) + return list(zip(prob, weighted_sum)) + + def prepare_rewards(self, rewards: List[Reward]) -> np.ndarray: + """Rewards as the training targets of the likelihood (identity for the Bernoulli head). + + Also called by the joint cMAB meta-model, which trains every head in one pass. + + Parameters + ---------- + rewards : List[Reward] + Rewards observed for this model. + + Returns + ------- + np.ndarray + Training targets, row-aligned with ``rewards``. + """ + return np.asarray(rewards, dtype=float) + + @property + def _is_first_fit(self) -> bool: + """Whether no reward has been recorded yet (the counters still hold the prior pseudo-counts).""" + return self.n_successes == self._prior_pseudo_count and self.n_failures == self._prior_pseudo_count @validate_call(config=dict(arbitrary_types_allowed=True)) def sample_weights(self, n_samples: PositiveInt, rng: np.random.Generator) -> List[Tuple[np.ndarray, np.ndarray]]: @@ -890,6 +1000,7 @@ def forward_pass( sampled_weights: List[Tuple[np.ndarray, np.ndarray]], context: np.ndarray, sampled_embeddings: Optional[List[np.ndarray]] = None, + sampled_extras: Optional[Dict[str, np.ndarray]] = None, ) -> List[ProbabilityWeight]: """ Apply the neural network forward pass using pre-sampled weights, biases, and embeddings. @@ -909,6 +1020,8 @@ def forward_pass( Pre-sampled embedding vectors from ``sample_embeddings``, one array per categorical feature, each of shape ``(n_samples, emb_dim)``. ``None`` when the model has no categorical features. + sampled_extras : Optional[Dict[str, np.ndarray]] + Per-row draws of the :meth:`extra_site_params` latents from :meth:`sample_extra_numpy`. Returns ------- @@ -925,9 +1038,7 @@ def forward_pass( backend=np, ) - weighted_sum = linear_transform.squeeze(-1) - prob = _numpy_sigmoid(weighted_sum) - return list(zip(prob, weighted_sum)) + return self._postprocess_output(linear_transform, sampled_extras) @validate_call(config=dict(arbitrary_types_allowed=True)) def sample_proba(self, context: np.ndarray, rng: np.random.Generator) -> List[ProbabilityWeight]: @@ -956,8 +1067,13 @@ def sample_proba(self, context: np.ndarray, rng: np.random.Generator) -> List[Pr sampled_weights = self.sample_weights(n_samples, rng=rng) sampled_embeddings = self.sample_embeddings(_context, rng=rng) + sampled_extras = self.sample_extra_numpy(n_samples, rng=rng) + return self.forward_pass( - sampled_weights=sampled_weights, context=_context, sampled_embeddings=sampled_embeddings + sampled_weights=sampled_weights, + context=_context, + sampled_embeddings=sampled_embeddings, + sampled_extras=sampled_extras, ) def _create_updated_layer_params( @@ -1068,6 +1184,12 @@ def collect_guide_init_arrays(self) -> Tuple[dict, dict, list, list]: all_sigmas.append(emb.params["sigma"].ravel()) site_sigmas[emb_name] = emb.params["sigma"] + for name, dist in self.extra_site_params().items(): + values[name] = jnp.array(dist.params["mu"]) + all_mus.append(dist.params["mu"].ravel()) + all_sigmas.append(dist.params["sigma"].ravel()) + site_sigmas[name] = dist.params["sigma"] + return values, site_sigmas, all_mus, all_sigmas def build_guide_init_fns(self, values: dict, site_sigmas: dict, all_mus: list, all_sigmas: list) -> tuple: @@ -1264,6 +1386,7 @@ def _extract_fullrank_advi_params(self, guide, params: dict) -> tuple: if self.model_params.embedding_params is not None: for _i in range(len(self.model_params.embedding_params.embeddings)): sampled_site_names.add(self.get_embedding_var_name(_i)) + sampled_site_names.update(self.extra_site_params()) site_mu = {k: v for k, v in site_mu_all.items() if k in sampled_site_names} offset = 0 site_sigma = {} @@ -1306,6 +1429,7 @@ def _extract_vi_params(self, x_jnp: jnp.ndarray, y_jnp: jnp.ndarray, n_samples: # Update layer params + embeddings from per-site posterior updated_layer_params_list = self.layer_params_from_posterior(site_mu, site_sigma) self.update_embedding_params_from_vi(site_mu, site_sigma) + self.update_extra_params_from_vi(site_mu, site_sigma) return updated_layer_params_list def layer_params_from_posterior(self, site_mu: dict, site_sigma: dict) -> List: @@ -1348,6 +1472,14 @@ def _update(self, context: np.ndarray, rewards: List[Reward]): This method updates the model's parameters by fitting the posterior distribution with Variational Inference (VI). """ + self._fit(context=context, rewards=rewards) + + def _fit(self, context: np.ndarray, rewards: List[Any]) -> None: + """Fit the posterior on ``(context, rewards)``; the body of :meth:`_update` once rewards are validated. + + Kept separate so a subclass with a different reward type (e.g. the Gaussian head) can validate + its own rewards in ``_update`` and reuse the fit unchanged. + """ self.check_context_matrix(context=context) if len(context) != len(rewards): @@ -1355,7 +1487,7 @@ def _update(self, context: np.ndarray, rewards: List[Reward]): _context = np.atleast_2d(context) x_jnp = jnp.array(_context, dtype=jnp.float32) - y_jnp = jnp.asarray(rewards, dtype=self._reward_dtype) + y_jnp = jnp.asarray(self.prepare_rewards(rewards), dtype=self.reward_dtype) n_samples = _context.shape[0] if self.calibrate_output_bias: @@ -1364,8 +1496,7 @@ def _update(self, context: np.ndarray, rewards: List[Reward]): # Forgetting: widen the current posterior (used as the prior for this fit) before re-fitting, # so fresh data dominates old evidence. No-op on the first fit (no posterior to inflate yet), # and no-op when decay_factor is None or 1. - is_first_fit = self.n_successes == self._prior_pseudo_count and self.n_failures == self._prior_pseudo_count - if not is_first_fit: + if not self._is_first_fit: self._inflate_prior_variance() self.model_params.bnn_layer_params = self._extract_vi_params(x_jnp, y_jnp, n_samples) @@ -1393,6 +1524,7 @@ def _inflate_prior_variance(self): emb.with_dist_parameters(sigma=(np.asarray(emb.params["sigma"]) * inflation).tolist()) for emb in self.model_params.embedding_params.embeddings ] + self._inflate_extra_params(inflation) @classmethod @validate_call @@ -1608,6 +1740,352 @@ class BayesianNeuralNetworkDP(BaseBayesianNeuralNetwork, ModelDP): """ +class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): + """Bayesian Neural Network for continuous (real-valued) rewards with a Gaussian likelihood. + + Rewards are modelled as ``Normal(mu(x), sigma)``: the network has a single output unit ``mu(x)``, and the noise + std ``sigma`` is one latent of the model (not a function of the context), ``log sigma ~ Normal``, whose posterior + is stored in ``noise_log_sigma`` and carried across updates like the weights. Thompson sampling draws the weights + (and ``log sigma``) from the posterior and returns ``mu(x)``, so arm selection is driven by the (epistemic) + uncertainty on the mean; ``sigma`` is the aleatoric reward noise and sets how much each observation is trusted + during training, i.e. how fast the posterior of ``mu`` contracts. + + Why a single ``sigma``: with one noise level the fit of ``mu`` is a (Bayesian) least-squares fit, whose target is + ``E[reward | x]`` whatever the shape of the noise (zero-inflated, skewed, heavy-tailed, e.g. revenue). A + context-dependent ``sigma(x)`` would let the model explain large rewards as noise instead of raising ``mu(x)``, + which biases ``mu`` low exactly where the rewards are most variable. + + The network is trained on the targets ``(reward - reward_loc) / reward_scale``, and ``sample_proba`` maps ``mu`` + and ``sigma`` back to the reward scale. The standardization is set in exactly one of two ways: either both + ``reward_loc`` and ``reward_scale`` are given as fixed constants (``fit_reward_standardization=False``, the + default), or both are fitted on the first update batch (``fit_reward_standardization=True``, leaving them unset). + For rewards far from ``O(1)``, choose them so the default ``O(1)`` weight priors fit; ``reward_loc=0`` and + ``reward_scale=1`` train on the raw rewards. + + Parameters + ---------- + fit_reward_standardization : bool + Whether to fit ``reward_loc`` / ``reward_scale`` (mean / std) on the first update batch, kept fixed afterwards + and re-fitted after ``reset()``. When True, leave ``reward_loc`` / ``reward_scale`` unset: values given + together are taken as an already fitted standardization (e.g. a reloaded model), used until ``reset()``. + When False, both must be given, and are used as fixed constants. Default is False. + reward_loc : Optional[float] + Location subtracted from the rewards. Required when ``fit_reward_standardization`` is False; otherwise None + until fitted on the first update. + reward_scale : Optional[PositiveFloat] + Scale the rewards are divided by. Required when ``fit_reward_standardization`` is False; otherwise None + until fitted on the first update. + n_observations : NonNegativeInt + Number of rewards observed. + reward_sum : float + Sum of the rewards observed. + noise_sigma : Optional[PositiveFloat] + Initial guess of the noise std, **in reward units**. It sets the mean of the ``log sigma`` prior on the first + update (once ``reward_scale`` is known). None to use the std of the first update batch. Default is None. + noise_log_sigma_prior_std : PositiveFloat + Std of the initial ``log sigma`` prior (0.5 allows about a factor 1.65 at one std). Default is 0.5. + noise_log_sigma : Optional[NormalArray] + The current posterior of ``log sigma`` (in target units, i.e. divided by ``reward_scale``), shape (1,). None until the first update + initializes it from ``noise_sigma`` or the batch. + + Notes + ----- + - ``n_successes`` / ``n_failures`` are not meaningful for continuous rewards and stay at their prior + pseudo-counts; ``count`` and ``mean`` are computed from ``n_observations`` / ``reward_sum`` instead. + - The unbounded likelihood is more sensitive than the Bernoulli one to a wide prior: the guide starts + at the prior's width, so with the default ``sigma=1`` the posterior may not tighten within the + default number of steps. A narrower prior (e.g. ``dist_params_init={"mu": 0, "sigma": 0.1, ...}``) + and a larger step size (e.g. ``3e-3``) fit much faster. + - ``reset()`` returns the model to its cold-start state: the weights to their initial prior, ``noise_log_sigma`` + to None, the counters to 0, and, with ``fit_reward_standardization``, ``reward_loc`` / ``reward_scale`` to None + (re-fitted on the next update); fixed values are kept. + + Examples + -------- + >>> bnn = GaussianBayesianNeuralNetwork.cold_start( + ... n_features=2, + ... reward_loc=0.0, + ... reward_scale=1.0, + ... hidden_dim_list=[16], + ... dist_type="normal", + ... dist_params_init={"mu": 0, "sigma": 0.1}, + ... update_kwargs={"optimizer_kwargs": {"step_size": 3e-3}}, + ... ) + """ + + fit_reward_standardization: bool = False + reward_loc: Optional[float] = None + reward_scale: Optional[PositiveFloat] = None + n_observations: NonNegativeInt = 0 + reward_sum: float = 0.0 + noise_sigma: Optional[PositiveFloat] = None + noise_log_sigma_prior_std: PositiveFloat = 0.5 + noise_log_sigma: Optional[NormalArray] = None + + supports_continuous_rewards: ClassVar[bool] = True + # [mu]; the noise std is the scalar latent ``noise_log_sigma`` + default_output_dim: ClassVar[PositiveInt] = 1 + _noise_var_name: ClassVar[str] = "noise_log_sigma" + + _transfer_learned_keys: ClassVar[Tuple[str, ...]] = ( + "reward_loc", + "reward_scale", + "n_observations", + "reward_sum", + "noise_log_sigma", + ) + # The weights are trained relative to the standardization, so its mode can't change across a transfer. + _transfer_structural_keys: ClassVar[Tuple[str, ...]] = ("fit_reward_standardization",) + + @model_validator(mode="after") + def validate_reward_standardization(self) -> "GaussianBayesianNeuralNetwork": + if self.fit_reward_standardization: + if (self.reward_loc is None) != (self.reward_scale is None): + raise ValueError( + "With fit_reward_standardization=True, reward_loc and reward_scale are either both unset " + "(not fitted yet) or both set (fitted)." + ) + elif self.reward_loc is None or self.reward_scale is None: + raise ValueError( + "reward_loc and reward_scale must both be given when fit_reward_standardization=False " + "(use reward_loc=0, reward_scale=1 to train on raw rewards - with caution, as the weight priors are in reward units)." + ) + return self + + @model_validator(mode="after") + def validate_output_layer(self) -> "GaussianBayesianNeuralNetwork": + output_dim = self.model_params.bnn_layer_params[-1].weight.shape[-1] + if output_dim != self.default_output_dim: + raise ValueError(f"The output layer has {output_dim} units, expected {self.default_output_dim} (mu).") + if self.noise_log_sigma is not None and self.noise_log_sigma.shape != (1,): + raise ValueError("noise_log_sigma must have shape (1,).") + return self + + @property + def _loc_scale(self) -> Tuple[float, float]: + """``(loc, scale)`` mapping targets back to rewards (0 / 1 while a fitted standardization is not yet fitted).""" + loc = 0.0 if self.reward_loc is None else self.reward_loc + scale = 1.0 if self.reward_scale is None else self.reward_scale + return loc, scale + + def _initial_noise_log_sigma(self, targets: Optional[np.ndarray] = None) -> NormalArray: + """The cold-start ``log sigma`` prior, in target units. + + Centered on ``log(noise_sigma / reward_scale)`` when ``noise_sigma`` is given, otherwise on the log std of + ``targets`` (the first batch of targets; 0, i.e. sigma = 1, when there is no batch or no spread). + """ + _, scale = self._loc_scale + if self.noise_sigma is not None: + sigma = self.noise_sigma / scale + else: + std = float(np.std(targets)) if targets is not None and len(targets) > 1 else 0.0 + sigma = std if std > self._numerical_eps else 1.0 + return NormalArray.cold_start(shape=1, mu=float(np.log(sigma)), sigma=self.noise_log_sigma_prior_std) + + def extra_site_params(self) -> Dict[str, BaseLocationScaleArray]: + """The noise latent ``log sigma`` (its current posterior, or the cold-start prior before the first update).""" + return {self._noise_var_name: self.noise_log_sigma or self._initial_noise_log_sigma()} + + def update_extra_params_from_vi(self, site_mu: dict, site_sigma: dict) -> None: + mu = np.asarray(site_mu[self._noise_var_name], dtype=float).reshape(1) + sigma = np.asarray(site_sigma[self._noise_var_name], dtype=float).reshape(1) + self.noise_log_sigma = NormalArray(mu=mu.tolist(), sigma=sigma.tolist()) + + def _inflate_extra_params(self, inflation: float) -> None: + if self.noise_log_sigma is not None: + self.noise_log_sigma = self.noise_log_sigma.with_dist_parameters( + sigma=(self.noise_log_sigma.params["sigma"] * inflation).tolist() + ) + + @staticmethod + def _noise_log_sigma_of(extra: Dict[str, Any], name: str) -> Any: + """The ``log sigma`` draw(s) without the trailing site axis: () or (batch,).""" + log_sigma = extra[name] + return log_sigma.reshape(log_sigma.shape[:-1]) + + def output_distribution( + self, linear_out: jax.Array, extra_sites: Optional[Dict[str, jax.Array]] = None + ) -> NumpyroDistribution: + """``Normal(mu, exp(log sigma))`` on the targets ``(reward - reward_loc) / reward_scale``. + + Parameters + ---------- + linear_out : jax.Array of shape (batch, 1) + Raw output of the final layer: ``[mu]``. + extra_sites : Optional[Dict[str, jax.Array]] + The sampled ``log sigma`` site, shape (1,) or (batch, 1) (per-row in the joint meta-model, where rows of + several arms share one batch). + + Returns + ------- + NumpyroDistribution + The Gaussian reward likelihood. + """ + sigma = jnp.exp(self._noise_log_sigma_of(extra_sites, self._noise_var_name)) + return NumpyroNormal(loc=linear_out[..., 0], scale=sigma) + + def _observe_output( + self, linear_out: jax.Array, y: jax.Array, extra_sites: Optional[Dict[str, jax.Array]] = None + ) -> None: + """Observe ``y`` under the Gaussian likelihood (no ``logit`` site).""" + numpyro.sample("out", self.output_distribution(linear_out, extra_sites), obs=y) + + def _postprocess_output( + self, linear_out: np.ndarray, extra_samples: Optional[Dict[str, np.ndarray]] = None + ) -> List[GaussianSample]: + """Map the sampled raw output layer to ``(mu, sigma)`` on the reward scale. + + Parameters + ---------- + linear_out : np.ndarray of shape (n_samples, 1) + Raw output of the final layer, one row per sampled network. + extra_samples : Optional[Dict[str, np.ndarray]] + Per-row draws of ``log sigma``, shape (n_samples, 1). None uses its posterior mean. + + Returns + ------- + List[GaussianSample] + ``(mu, sigma)`` per sample: the sampled reward mean and the noise std (the latter is for monitoring; arm + selection uses ``mu`` only). + """ + loc, scale = self._loc_scale + mu = linear_out[..., 0] * scale + loc + if extra_samples is None: # no per-row draws given: use the posterior mean of log sigma + log_sigma_mu = self.extra_site_params()[self._noise_var_name].params["mu"] + extra_samples = {self._noise_var_name: np.broadcast_to(log_sigma_mu, (len(mu), 1))} + sigma = np.exp(self._noise_log_sigma_of(extra_samples, self._noise_var_name)) * scale + return list(zip(mu, sigma)) + + def sample_proba(self, context: np.ndarray, rng: np.random.Generator) -> List[GaussianSample]: + """ + Sample the reward mean from the posterior (Thompson sampling), with the noise std. + + Parameters + ---------- + context : np.ndarray + The context matrix for which the rewards are to be sampled. + rng : np.random.Generator + Numpy random generator for weight/embedding/noise sampling. + + Returns + ------- + List[GaussianSample] + ``(mu, sigma)`` per context row. + """ + return super().sample_proba(context=context, rng=rng) + + def prepare_rewards(self, rewards: List[ContinuousReward]) -> np.ndarray: + """Map rewards to training targets, fitting the standardization and the noise prior on first use. + + With ``fit_reward_standardization``, ``reward_loc`` / ``reward_scale`` are fitted (mean / std) on the first + batch if not set; the scale falls back to ``max(|loc|, 1)`` when that batch has (near) zero spread, e.g. a + single observation. The ``log sigma`` prior is initialized here too, on first use (see ``noise_sigma``). + + Parameters + ---------- + rewards : List[ContinuousReward] + Rewards observed for this model. + + Returns + ------- + np.ndarray + ``(rewards - loc) / scale``. + """ + y = np.asarray(rewards, dtype=float) + if not np.all(np.isfinite(y)): + raise ValueError("Rewards must be finite.") + if self.fit_reward_standardization and self.reward_loc is None and len(y) > 0: + loc = float(y.mean()) + std = float(y.std()) + self.reward_loc = loc + self.reward_scale = std if std > self._numerical_eps else max(abs(loc), 1.0) + loc, scale = self._loc_scale + targets = (y - loc) / scale + if self.noise_log_sigma is None and len(y) > 0: + self.noise_log_sigma = self._initial_noise_log_sigma(targets) + return targets + + @validate_call(config=dict(arbitrary_types_allowed=True)) + def update(self, rewards: List[ContinuousReward], **kwargs): + """ + Update the model parameters. + + Parameters + ---------- + rewards : List[ContinuousReward] + The real-valued reward for each sample. + """ + self._update(rewards=rewards, **kwargs) + self.record_rewards(rewards) + + @validate_call(config=dict(arbitrary_types_allowed=True)) + def _update(self, context: np.ndarray, rewards: List[ContinuousReward]): + """ + Update the model_params with new context and real-valued rewards (see :meth:`_fit`). + + Parameters + ---------- + context : np.ndarray + The context matrix where each row represents a context vector. + rewards : List[ContinuousReward] + The real-valued reward for each context vector. + """ + self._fit(context=context, rewards=rewards) + + def record_rewards(self, rewards: List[ContinuousReward]) -> None: + """Accumulate the number and the sum of the observed rewards. + + Parameters + ---------- + rewards : List[ContinuousReward] + The real-valued reward for each sample. + """ + self.n_observations += len(rewards) + self.reward_sum += float(np.sum(rewards)) + + @property + def _is_first_fit(self) -> bool: + return self.n_observations == 0 + + @property + def count(self) -> NonNegativeInt: + """The number of rewards observed.""" + return self.n_observations + + @property + def mean(self) -> float: + """The average observed reward (0 before any observation).""" + return self.reward_sum / self.n_observations if self.n_observations else 0.0 + + def _calibrate_output_bias(self, rewards: List[ContinuousReward]) -> None: + """Set the output-layer bias to the first batch's target mean on the first update call. + + With ``fit_reward_standardization`` the standardization is fitted on that same batch, so this is 0. The noise prior is initialized by + :meth:`prepare_rewards` instead. + + Parameters + ---------- + rewards : List[ContinuousReward] + Rewards observed in the current update batch. + """ + if self.bias_calibrated or len(rewards) == 0: + return + y = self.prepare_rewards(rewards) + output_layer = self.model_params.bnn_layer_params[-1] + new_bias = output_layer.bias.with_dist_parameters(mu=[float(y.mean())]) + self.model_params.bnn_layer_params[-1] = BnnLayerParams(weight=output_layer.weight, bias=new_bias) + self.bias_calibrated = True + + def _reset(self): + """Return to the cold-start state: weights, noise latent, counters and standardization.""" + super()._reset() + self.n_observations = 0 + self.reward_sum = 0.0 + self.noise_log_sigma = None + if self.fit_reward_standardization: + self.reward_loc = self.reward_scale = None + + class BaseBayesianNeuralNetworkMO(ModelMO, ABC): """ Base class for Bayesian Neural Network with multi-objective. diff --git a/pybandits/transfer.py b/pybandits/transfer.py index 1cfa2946..c82d2195 100644 --- a/pybandits/transfer.py +++ b/pybandits/transfer.py @@ -98,26 +98,29 @@ def _reconstruct_mab(mab_class: Type[BaseMab], state_dict: Dict[str, Any]) -> Ba ("CMAB", "SMAB", "context type (CMAB vs SMAB)"), ("MO", "SO", "objective type (MO vs SO)"), ("quantitative", "standard", "action type (quantitative vs standard)"), + ("continuous", "binary", "reward type (continuous vs binary)"), ) -def _mab_variant(mab: BaseMab) -> Tuple[bool, bool, bool]: - """Return (is_cmab, is_mo, is_quantitative) for a MAB based on its type and action models.""" +def _mab_variant(mab: BaseMab) -> Tuple[bool, bool, bool, bool]: + """Return (is_cmab, is_mo, is_quantitative, is_continuous) for a MAB based on its type and action models.""" first_model = next(iter(mab.actions.values()), None) return ( isinstance(mab, BaseCmabBernoulli), isinstance(first_model, BaseModelMO), isinstance(first_model, QuantitativeModel), + getattr(first_model, "supports_continuous_rewards", False), ) def _validate_mab_compatibility(mab1: BaseMab, mab2: BaseMab) -> None: """Validate that two MABs can be merged. - Checks that both MABs share the same three variant dimensions: + Checks that both MABs share the same four variant dimensions: - Context type: CMAB (contextual) vs SMAB (stochastic) - Objective type: MO (multi-objective) vs SO (single-objective) - Action type: quantitative (zooming) vs standard (Bernoulli/Beta) + - Reward type: continuous (Gaussian likelihood) vs binary / soft (Bernoulli likelihood) Parameters ---------- @@ -762,6 +765,10 @@ def edit_model_on_the_fly( **STRUCTURAL_IMMUTABLE (must match - raises ValueError):** - activation: Activation function (tanh, relu, sigmoid, gelu) - use_residual_connections: Whether skip connections are used + - fit_reward_standardization (Gaussian models): Whether reward_loc / reward_scale are fitted or fixed. + The source's loc / scale are carried over with the weights, which were trained relative to them. + + Bernoulli and Gaussian (continuous-reward) bandits cannot be merged (raises TypeError). **CONFIGURABLE_MUTABLE (safe to change):** - update_kwargs: Training hyperparameters (epochs, learning rate, etc.) diff --git a/pyproject.toml b/pyproject.toml index 46e50293..6542ea17 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "pybandits" -version = "8.3.0" +version = "8.4.0" description = "Python Multi-Armed Bandit Library" authors = [ "Dario d'Andrea ", diff --git a/tests/test_actions_manager.py b/tests/test_actions_manager.py index 19e25d5c..d1f34149 100644 --- a/tests/test_actions_manager.py +++ b/tests/test_actions_manager.py @@ -39,7 +39,7 @@ ) from pybandits.base import ActionId, BinaryReward from pybandits.meta_model import BaseMetaModel, CmabMetaModel, SmabMetaModel -from pybandits.model import BayesianNeuralNetwork, Beta, BetaMO +from pybandits.model import BayesianNeuralNetwork, Beta, BetaMO, GaussianBayesianNeuralNetwork from pybandits.quantitative_model import QuantitativeBayesianNeuralNetwork, Zooming from tests.utils import make_action_ids, make_binary_rewards, make_context_matrix, mock_update @@ -217,6 +217,13 @@ def test_cmab_mixed_action_types_error(n_features=1): CmabActionsManager[Union[BayesianNeuralNetwork, QuantitativeBayesianNeuralNetwork]](actions=actions) +def test_mixed_continuous_and_binary_reward_models_error(n_features: int = 1) -> None: + """A manager whose type admits both kinds of model still rejects mixing continuous and binary-reward arms.""" + gaussian = GaussianBayesianNeuralNetwork.cold_start(n_features=n_features, reward_loc=0.0, reward_scale=1.0) + with pytest.raises(ValidationError, match="continuous rewards"): + SmabActionsManager[Union[Beta, GaussianBayesianNeuralNetwork]](actions={"a1": Beta(), "a2": gaussian}) + + @given( n_successes=st.just(100), n_failures=st.just(1), diff --git a/tests/test_cmab.py b/tests/test_cmab.py index ec3b4313..ef2465f8 100644 --- a/tests/test_cmab.py +++ b/tests/test_cmab.py @@ -23,7 +23,7 @@ import json from copy import deepcopy from functools import partial -from typing import Any, Dict, List, Optional, Tuple, Type, Union, get_args +from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union, get_args from unittest.mock import patch import numpy as np @@ -48,7 +48,9 @@ CmabBernoulliDP, CmabBernoulliMO, CmabBernoulliMOCC, + CmabGaussian, ) +from pybandits.meta_model import CmabMetaModel from pybandits.model import ( BaseBayesianNeuralNetwork, BaseBayesianNeuralNetworkMO, @@ -57,6 +59,7 @@ BayesianNeuralNetworkDP, BayesianNeuralNetworkMO, BayesianNeuralNetworkMOCC, + GaussianBayesianNeuralNetwork, StudentTArray, ) from pybandits.model.bnn._typing import ActivationFunctions @@ -1636,3 +1639,162 @@ def test_cold_start_rejects_delta_with_backbone(delta: float) -> None: backbone_hidden_dims=_BACKBONE_HIDDEN_DIMS, delta=delta, ) + + +######################################################################################################################## +# --- Gaussian (continuous-reward) CmabGaussian constants --- +_GAUSSIAN_N_FEATURES = 3 +_GAUSSIAN_HIDDEN_DIMS = [8] +_GAUSSIAN_RANDOM_SEED = 0 +# Narrow prior + larger step size: the setting the class docstring recommends for a fast fit. +_GAUSSIAN_FAST_FIT_KWARGS: Dict[str, Any] = { + "dist_type": "normal", + "dist_params_init": {"mu": 0, "sigma": 0.1}, + "update_kwargs": {"num_steps": 1000, "optimizer_kwargs": {"step_size": 3e-3}}, +} +_GAUSSIAN_SHORT_FIT_KWARGS: Dict[str, Any] = {"update_kwargs": {"num_steps": 5}} +_GAUSSIAN_ACTION_IDS = {"a", "b", "c"} +_GAUSSIAN_ARM_MEANS = {"a": -1.0, "b": 0.0, "c": 2.0} +_GAUSSIAN_BEST_ARM = "c" +_GAUSSIAN_N_ROUNDS = 4 +_GAUSSIAN_N_PER_ROUND = 200 +_GAUSSIAN_MIN_BEST_ARM_SHARE = 0.8 +_GAUSSIAN_BACKBONE_KWARGS: Dict[str, Any] = {"backbone_hidden_dims": [8], "backbone_embedding_dim": 3} +_GAUSSIAN_MINIBATCH_SIZE = 64 +_GAUSSIAN_BAD_REWARDS = [float("nan"), float("inf"), -float("inf")] +_GAUSSIAN_OUT_OF_RANGE_SOFT_REWARD = 2.5 +_GAUSSIAN_DELTA = 0.1 +_GAUSSIAN_EPSILON = 0.1 +# Fixed identity standardization: train on the raw rewards. +_GAUSSIAN_RAW_STANDARDIZATION: Dict[str, float] = {"reward_loc": 0.0, "reward_scale": 1.0} + + +@pytest.fixture(scope="module") +def make_cmab_gaussian() -> Callable[..., CmabGaussian]: + """Factory: cold-starts a CmabGaussian over ACTION_IDS with the fast-fit prior and extra kwargs.""" + + def _factory(**kwargs: Any) -> CmabGaussian: + return CmabGaussian.cold_start( + action_ids=_GAUSSIAN_ACTION_IDS, + n_features=_GAUSSIAN_N_FEATURES, + hidden_dim_list=_GAUSSIAN_HIDDEN_DIMS, + random_seed=_GAUSSIAN_RANDOM_SEED, + **{**_GAUSSIAN_RAW_STANDARDIZATION, **_GAUSSIAN_FAST_FIT_KWARGS, **kwargs}, + ) + + return _factory + + +class TestGaussianMetaModel: + """Tests for Gaussian heads in CmabMetaModel.""" + + def test_rejects_mixed_likelihoods( + self, + n_features: int = _GAUSSIAN_N_FEATURES, + hidden_dims: List[int] = _GAUSSIAN_HIDDEN_DIMS, + raw_standardization: Dict[str, float] = _GAUSSIAN_RAW_STANDARDIZATION, + ) -> None: + """Bernoulli and Gaussian heads cannot share one meta-model.""" + gaussian = GaussianBayesianNeuralNetwork.cold_start( + n_features=n_features, hidden_dim_list=hidden_dims, **raw_standardization + ) + bernoulli = BayesianNeuralNetwork.cold_start(n_features=n_features, hidden_dim_list=hidden_dims) + with pytest.raises(AttributeError): + CmabMetaModel(actions={"a": gaussian, "b": bernoulli}) + + +class TestCmabGaussian: + """Tests for the CmabGaussian bandit.""" + + @pytest.mark.parametrize( + "extra_kwargs", + [ + {}, + { + **_GAUSSIAN_BACKBONE_KWARGS, + "update_kwargs": {**_GAUSSIAN_FAST_FIT_KWARGS["update_kwargs"], "batch_size": _GAUSSIAN_MINIBATCH_SIZE}, + }, + ], + ids=["no_backbone", "backbone_minibatch"], + ) + def test_learns_best_arm_and_per_arm_noise( + self, + make_cmab_gaussian: Callable[..., CmabGaussian], + extra_kwargs: Dict[str, Any], + rng: np.random.Generator, + arm_means: Dict[str, float] = _GAUSSIAN_ARM_MEANS, + best_arm: str = _GAUSSIAN_BEST_ARM, + n_rounds: int = _GAUSSIAN_N_ROUNDS, + n_per_round: int = _GAUSSIAN_N_PER_ROUND, + n_features: int = _GAUSSIAN_N_FEATURES, + min_share: float = _GAUSSIAN_MIN_BEST_ARM_SHARE, + ) -> None: + """After a few rounds of real-valued feedback (incl. negative rewards) the best-mean arm dominates.""" + mab = make_cmab_gaussian(**extra_kwargs) + for _ in range(n_rounds): + context = rng.normal(size=(n_per_round, n_features)) + actions, _, _ = mab.predict(context=context) + mab.update(actions=actions, rewards=[arm_means[a] + rng.normal() for a in actions], context=context) + actions, _, sigmas = mab.predict(context=rng.normal(size=(n_per_round, n_features))) + assert np.mean(np.array(actions) == best_arm) >= min_share + assert all(sigma > 0 for row in sigmas for sigma in row.values()) + assert mab.actions[best_arm].noise_log_sigma is not None + + def test_state_round_trip( + self, + make_cmab_gaussian: Callable[..., CmabGaussian], + rng: np.random.Generator, + n_per_round: int = _GAUSSIAN_N_PER_ROUND, + n_features: int = _GAUSSIAN_N_FEATURES, + short_fit_kwargs: Dict[str, Any] = _GAUSSIAN_SHORT_FIT_KWARGS, + ) -> None: + """The bandit state round-trips through get_state / from_state.""" + mab = make_cmab_gaussian(**short_fit_kwargs) + context = rng.normal(size=(n_per_round, n_features)) + actions, _, _ = mab.predict(context=context) + mab.update(actions=actions, rewards=rng.normal(size=n_per_round).tolist(), context=context) + restored = CmabGaussian.from_state(mab.get_state()[1]) + assert restored.actions == mab.actions + + @pytest.mark.parametrize("bad_reward", _GAUSSIAN_BAD_REWARDS) + def test_rejects_non_finite_rewards( + self, + make_cmab_gaussian: Callable[..., CmabGaussian], + bad_reward: float, + n_features: int = _GAUSSIAN_N_FEATURES, + short_fit_kwargs: Dict[str, Any] = _GAUSSIAN_SHORT_FIT_KWARGS, + ) -> None: + """NaN / inf rewards are rejected at the bandit entry point.""" + mab = make_cmab_gaussian(**short_fit_kwargs) + with pytest.raises(ValueError): + mab.update(actions=[next(iter(mab.actions))], rewards=[bad_reward], context=np.zeros((1, n_features))) + + def test_rejects_adaptive_window( + self, + action_ids: set = _GAUSSIAN_ACTION_IDS, + n_features: int = _GAUSSIAN_N_FEATURES, + delta: float = _GAUSSIAN_DELTA, + epsilon: float = _GAUSSIAN_EPSILON, + raw_standardization: Dict[str, float] = _GAUSSIAN_RAW_STANDARDIZATION, + ) -> None: + """The adaptive window assumes binary rewards, so delta is refused.""" + with pytest.raises(ValueError): + CmabGaussian.cold_start( + action_ids=action_ids, n_features=n_features, delta=delta, epsilon=epsilon, **raw_standardization + ) + + +@pytest.mark.parametrize("use_soft_rewards", [False, True]) +def test_bernoulli_bandit_still_rejects_out_of_range_rewards( + use_soft_rewards: bool, + action_ids: set = _GAUSSIAN_ACTION_IDS, + n_features: int = _GAUSSIAN_N_FEATURES, + reward: float = _GAUSSIAN_OUT_OF_RANGE_SOFT_REWARD, + short_fit_kwargs: Dict[str, Any] = _GAUSSIAN_SHORT_FIT_KWARGS, +) -> None: + """Widening the entry-point reward type must not let a Bernoulli bandit accept rewards outside [0, 1].""" + mab = CmabBernoulli.cold_start( + action_ids=action_ids, n_features=n_features, use_soft_rewards=use_soft_rewards, **short_fit_kwargs + ) + with pytest.raises(ValueError): + mab.update(actions=[next(iter(mab.actions))], rewards=[reward], context=np.zeros((1, n_features))) diff --git a/tests/test_model.py b/tests/test_model.py index 05b1a352..cc171cae 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -20,7 +20,7 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -from typing import List, Literal, Optional, Tuple, get_args +from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, get_args from unittest.mock import patch import jax @@ -54,6 +54,7 @@ EarlyStopping, EmbeddingParams, FeaturesConfig, + GaussianBayesianNeuralNetwork, NormalArray, OptaxKind, StudentTArray, @@ -2576,3 +2577,471 @@ def test_advi_zero_lr_posterior_equals_prior( np.testing.assert_allclose(post_b.params["mu"], prior_b.params["mu"], atol=atol_mu) np.testing.assert_allclose(post_w.params["sigma"], prior_w.params["sigma"], rtol=rtol_sigma) np.testing.assert_allclose(post_b.params["sigma"], prior_b.params["sigma"], rtol=rtol_sigma) + + +######################################################################################################################## +# --- Gaussian (continuous-reward) BNN constants --- +_GAUSSIAN_N_FEATURES = 3 +_GAUSSIAN_HIDDEN_DIMS = [8] +_GAUSSIAN_OUTPUT_DIM = 1 +_GAUSSIAN_RANDOM_SEED = 0 +_GAUSSIAN_N_TRAIN = 2000 +_GAUSSIAN_N_POSTERIOR_SAMPLES = 200 +# Narrow prior + larger step size: the setting the class docstring recommends for a fast fit. +_GAUSSIAN_FAST_FIT_KWARGS: Dict[str, Any] = { + "dist_type": "normal", + "dist_params_init": {"mu": 0, "sigma": 0.1}, + "update_kwargs": {"num_steps": 1000, "optimizer_kwargs": {"step_size": 3e-3}}, +} +_GAUSSIAN_SHORT_FIT_KWARGS: Dict[str, Any] = {"update_kwargs": {"num_steps": 5}} +# Linear synthetic reward: REWARD_OFFSET + REWARD_SLOPE * x0 + Normal(0, NOISE_STD). +_GAUSSIAN_REWARD_OFFSET = 100.0 +_GAUSSIAN_REWARD_SLOPE = 30.0 +_GAUSSIAN_NOISE_STD = 10.0 +_GAUSSIAN_MEAN_REL_TOL = 0.1 +_GAUSSIAN_SIGMA_REL_TOL = 0.3 +# Probe contexts on both sides of x0. +_GAUSSIAN_PROBE_CONTEXTS = [[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0], [0.5, 1.0, -1.0]] +# Zero-inflated, skewed reward (as revenue): a buy with probability BUY_PROB, then a log-normal amount, so the mean +# is SKEW_OFFSET + SKEW_SLOPE * x0. +_GAUSSIAN_BUY_PROB = 0.15 +_GAUSSIAN_SKEW_OFFSET = 2.0 +_GAUSSIAN_SKEW_SLOPE = 1.5 +_GAUSSIAN_SKEW_LOG_SD = 1.0 +_GAUSSIAN_LEVEL_REL_TOL = 0.15 +_GAUSSIAN_ROUND_TRIP_TOL = 1e-6 +_GAUSSIAN_MAX_ABS_REWARD = 1e6 +_GAUSSIAN_MAX_REWARDS = 50 +_GAUSSIAN_NOISE_SIGMA = 3.0 +_GAUSSIAN_DECAY_FACTOR = 0.5 +_GAUSSIAN_USER_LOC = 5.0 +_GAUSSIAN_USER_SCALE = 2.0 +# Fixed identity standardization: train on the raw rewards. +_GAUSSIAN_RAW_STANDARDIZATION: Dict[str, float] = {"reward_loc": 0.0, "reward_scale": 1.0} +_GAUSSIAN_SECOND_BATCH_SHIFT = 50.0 +_GAUSSIAN_BAD_REWARDS = [float("nan"), float("inf"), -float("inf")] + + +def _linear_rewards(x: np.ndarray, rng: np.random.Generator) -> np.ndarray: + """Linear-in-x0 rewards with constant Gaussian noise.""" + return _GAUSSIAN_REWARD_OFFSET + _GAUSSIAN_REWARD_SLOPE * x[:, 0] + rng.normal(size=len(x)) * _GAUSSIAN_NOISE_STD + + +def _skewed_rewards(x: np.ndarray, rng: np.random.Generator) -> np.ndarray: + """Zero-inflated log-normal rewards with mean SKEW_OFFSET + SKEW_SLOPE * x0 (x0 clipped to keep it positive).""" + mean = _GAUSSIAN_SKEW_OFFSET + _GAUSSIAN_SKEW_SLOPE * np.clip(x[:, 0], -1.0, 2.0) + amount = rng.lognormal(np.log(mean / _GAUSSIAN_BUY_PROB) - _GAUSSIAN_SKEW_LOG_SD**2 / 2, _GAUSSIAN_SKEW_LOG_SD) + return (rng.uniform(size=len(x)) < _GAUSSIAN_BUY_PROB) * amount + + +def _noise_sigma_in_reward_units(bnn: GaussianBayesianNeuralNetwork) -> float: + """Posterior mean of the noise std, mapped back to reward units.""" + _, scale = bnn._loc_scale + return float(np.exp(bnn.noise_log_sigma.params["mu"][0]) * scale) + + +@pytest.fixture(scope="module") +def make_gaussian_bnn() -> Callable[..., GaussianBayesianNeuralNetwork]: + """Factory: cold-starts a GaussianBayesianNeuralNetwork with the module's shape and extra kwargs. + + Unless fit_reward_standardization=True, the fixed standardization defaults to the identity (raw rewards). + """ + + def _factory(**kwargs: Any) -> GaussianBayesianNeuralNetwork: + standardization = {} if kwargs.get("fit_reward_standardization") else _GAUSSIAN_RAW_STANDARDIZATION + return GaussianBayesianNeuralNetwork.cold_start( + n_features=_GAUSSIAN_N_FEATURES, + hidden_dim_list=_GAUSSIAN_HIDDEN_DIMS, + random_seed=_GAUSSIAN_RANDOM_SEED, + **{**standardization, **kwargs}, + ) + + return _factory + + +@pytest.fixture(scope="module") +def fitted_gaussian_bnn( + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], +) -> GaussianBayesianNeuralNetwork: + """A Gaussian BNN fitted on linear synthetic rewards with constant noise.""" + data_rng = np.random.default_rng(_GAUSSIAN_RANDOM_SEED) + bnn = make_gaussian_bnn(fit_reward_standardization=True, **_GAUSSIAN_FAST_FIT_KWARGS) + x = data_rng.normal(size=(_GAUSSIAN_N_TRAIN, _GAUSSIAN_N_FEATURES)) + bnn.update(context=x, rewards=_linear_rewards(x, data_rng).tolist()) + return bnn + + +class TestGaussianBayesianNeuralNetwork: + """Tests for the standalone Gaussian BNN (one mean unit, a learned scalar noise latent).""" + + def test_output_layer_has_only_the_mean_unit( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + output_dim: int = _GAUSSIAN_OUTPUT_DIM, + ) -> None: + """The output layer has a single unit and the noise latent is unset before the first update.""" + bnn = make_gaussian_bnn() + output_layer = bnn.model_params.bnn_layer_params[-1] + assert output_layer.weight.shape[-1] == output_dim + assert output_layer.bias.shape == (output_dim,) + assert bnn.noise_log_sigma is None + + def test_rejects_legacy_homoscedastic_argument( + self, make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork] + ) -> None: + """The removed noise-model switch is an unknown argument.""" + with pytest.raises(ValidationError): + make_gaussian_bnn(homoscedastic=True) + + def test_sample_proba_returns_mean_and_positive_noise( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = _GAUSSIAN_N_TRAIN, + n_features: int = _GAUSSIAN_N_FEATURES, + ) -> None: + """Each sample is a (mu, sigma) pair with a positive sigma, also before any update.""" + bnn = make_gaussian_bnn() + samples = bnn.sample_proba(context=rng.normal(size=(n_samples, n_features)), rng=rng) + assert len(samples) == n_samples + assert all(sigma > 0 for _, sigma in samples) + + def test_fit_recovers_mean_trend_and_noise_level( + self, + fitted_gaussian_bnn: GaussianBayesianNeuralNetwork, + rng: np.random.Generator, + n_posterior_samples: int = _GAUSSIAN_N_POSTERIOR_SAMPLES, + offset: float = _GAUSSIAN_REWARD_OFFSET, + slope: float = _GAUSSIAN_REWARD_SLOPE, + noise_std: float = _GAUSSIAN_NOISE_STD, + mean_rel_tol: float = _GAUSSIAN_MEAN_REL_TOL, + sigma_rel_tol: float = _GAUSSIAN_SIGMA_REL_TOL, + probe_contexts: List[List[float]] = _GAUSSIAN_PROBE_CONTEXTS, + ) -> None: + """The posterior mean tracks the linear trend and the learned sigma matches the noise std.""" + probes = np.array(probe_contexts) + draws = np.array( + [fitted_gaussian_bnn.sample_proba(context=probes, rng=rng) for _ in range(n_posterior_samples)] + ) + np.testing.assert_allclose(draws[..., 0].mean(axis=0), offset + slope * probes[:, 0], rtol=mean_rel_tol) + assert _noise_sigma_in_reward_units(fitted_gaussian_bnn) == pytest.approx(noise_std, rel=sigma_rel_tol) + + def test_counters_track_observations( + self, + fitted_gaussian_bnn: GaussianBayesianNeuralNetwork, + n_train: int = _GAUSSIAN_N_TRAIN, + ) -> None: + """count / mean come from the observed rewards; the Beta pseudo-counts are untouched.""" + assert fitted_gaussian_bnn.count == n_train + assert fitted_gaussian_bnn.mean == pytest.approx(fitted_gaussian_bnn.reward_sum / n_train) + assert fitted_gaussian_bnn.n_successes == fitted_gaussian_bnn._prior_pseudo_count + assert fitted_gaussian_bnn.n_failures == fitted_gaussian_bnn._prior_pseudo_count + + @given( + rewards=st.lists( + st.floats(min_value=-_GAUSSIAN_MAX_ABS_REWARD, max_value=_GAUSSIAN_MAX_ABS_REWARD, allow_nan=False), + min_size=2, + max_size=_GAUSSIAN_MAX_REWARDS, + unique=True, + ), + n_features=st.just(_GAUSSIAN_N_FEATURES), + tol=st.just(_GAUSSIAN_ROUND_TRIP_TOL), + ) + def test_standardization_is_fitted_once(self, rewards: List[float], n_features: int, tol: float) -> None: + """The first batch fixes reward_loc / reward_scale; later batches are mapped with the same values.""" + bnn = GaussianBayesianNeuralNetwork.cold_start(n_features=n_features, fit_reward_standardization=True) + targets = bnn.prepare_rewards(rewards) + loc, scale = bnn.reward_loc, bnn.reward_scale + assert loc == pytest.approx(np.mean(rewards)) + np.testing.assert_allclose(targets * scale + loc, rewards, rtol=tol, atol=tol * scale) + bnn.prepare_rewards([value + scale for value in rewards]) + assert (bnn.reward_loc, bnn.reward_scale) == (loc, scale) + + @given( + reward=st.floats(min_value=-_GAUSSIAN_MAX_ABS_REWARD, max_value=_GAUSSIAN_MAX_ABS_REWARD, allow_nan=False), + n_features=st.just(_GAUSSIAN_N_FEATURES), + ) + def test_single_reward_standardization_has_unit_or_larger_scale(self, reward: float, n_features: int) -> None: + """A zero-spread first batch falls back to max(|reward|, 1) instead of dividing by ~0.""" + bnn = GaussianBayesianNeuralNetwork.cold_start(n_features=n_features, fit_reward_standardization=True) + bnn.prepare_rewards([reward]) + assert bnn.reward_scale == max(abs(reward), 1.0) + + @given( + rewards=st.lists( + st.floats(min_value=-_GAUSSIAN_MAX_ABS_REWARD, max_value=_GAUSSIAN_MAX_ABS_REWARD, allow_nan=False), + min_size=2, + max_size=_GAUSSIAN_MAX_REWARDS, + unique=True, + ), + n_features=st.just(_GAUSSIAN_N_FEATURES), + tol=st.just(_GAUSSIAN_ROUND_TRIP_TOL), + ) + def test_noise_initialized_from_first_batch(self, rewards: List[float], n_features: int, tol: float) -> None: + """Without noise_sigma, the log-sigma prior is centered on the first batch's std (1 if it has no spread).""" + bnn = GaussianBayesianNeuralNetwork.cold_start(n_features=n_features, fit_reward_standardization=True) + targets = bnn.prepare_rewards(rewards) + expected = targets.std() if targets.std() > bnn._numerical_eps else 1.0 + assert np.exp(bnn.noise_log_sigma.params["mu"][0]) == pytest.approx(expected, rel=tol) + + def test_noise_initialized_from_noise_sigma( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + noise_sigma: float = _GAUSSIAN_NOISE_SIGMA, + n_samples: int = _GAUSSIAN_N_TRAIN, + tol: float = _GAUSSIAN_ROUND_TRIP_TOL, + ) -> None: + """noise_sigma (reward units) sets the prior once the reward scale is known.""" + bnn = make_gaussian_bnn(noise_sigma=noise_sigma) + bnn.prepare_rewards(rng.normal(size=n_samples).tolist()) + assert _noise_sigma_in_reward_units(bnn) == pytest.approx(noise_sigma, rel=tol) + assert bnn.noise_log_sigma.params["sigma"][0] == pytest.approx(bnn.noise_log_sigma_prior_std, rel=tol) + + def test_noise_posterior_is_carried_across_updates( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = _GAUSSIAN_N_TRAIN, + n_features: int = _GAUSSIAN_N_FEATURES, + short_fit_kwargs: Dict[str, Any] = _GAUSSIAN_SHORT_FIT_KWARGS, + ) -> None: + """Each update replaces the stored noise posterior, and the next update starts from it.""" + bnn = make_gaussian_bnn(**short_fit_kwargs) + bnn.update(context=rng.normal(size=(n_samples, n_features)), rewards=rng.normal(size=n_samples).tolist()) + first = bnn.noise_log_sigma + assert first is not None + bnn.update(context=rng.normal(size=(n_samples, n_features)), rewards=rng.normal(size=n_samples).tolist()) + assert bnn.noise_log_sigma != first + _, site_sigmas, _, _ = bnn.collect_guide_init_arrays() + np.testing.assert_array_equal(site_sigmas[bnn._noise_var_name], bnn.noise_log_sigma.params["sigma"]) + + def test_sample_proba_noise_is_shared_across_contexts( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = _GAUSSIAN_N_TRAIN, + n_features: int = _GAUSSIAN_N_FEATURES, + short_fit_kwargs: Dict[str, Any] = _GAUSSIAN_SHORT_FIT_KWARGS, + rel_tol: float = _GAUSSIAN_SIGMA_REL_TOL, + ) -> None: + """The noise std doesn't depend on the context: its spread comes only from the log-sigma posterior.""" + bnn = make_gaussian_bnn(**short_fit_kwargs) + bnn.update(context=rng.normal(size=(n_samples, n_features)), rewards=rng.normal(size=n_samples).tolist()) + sigmas = np.array([s for _, s in bnn.sample_proba(context=rng.normal(size=(n_samples, n_features)), rng=rng)]) + assert np.all(sigmas > 0) + assert np.log(sigmas).std() == pytest.approx(bnn.noise_log_sigma.params["sigma"][0], rel=rel_tol) + + def test_mean_is_level_calibrated_on_skewed_rewards( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = _GAUSSIAN_N_TRAIN, + n_features: int = _GAUSSIAN_N_FEATURES, + n_posterior_samples: int = _GAUSSIAN_N_POSTERIOR_SAMPLES, + rel_tol: float = _GAUSSIAN_LEVEL_REL_TOL, + fast_fit_kwargs: Dict[str, Any] = _GAUSSIAN_FAST_FIT_KWARGS, + ) -> None: + """On zero-inflated skewed rewards, the average posterior mean matches the average reward (least squares).""" + data_rng = np.random.default_rng(_GAUSSIAN_RANDOM_SEED) + x = data_rng.normal(size=(n_samples, n_features)) + y = _skewed_rewards(x, data_rng) + bnn = make_gaussian_bnn(fit_reward_standardization=True, **fast_fit_kwargs) + bnn.update(context=x, rewards=y.tolist()) + mu = np.mean([[m for m, _ in bnn.sample_proba(context=x, rng=rng)] for _ in range(n_posterior_samples)], axis=0) + assert mu.mean() == pytest.approx(y.mean(), rel=rel_tol) + + @pytest.mark.parametrize("bad_reward", _GAUSSIAN_BAD_REWARDS) + def test_update_rejects_non_finite_rewards( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + bad_reward: float, + n_features: int = _GAUSSIAN_N_FEATURES, + ) -> None: + """NaN / inf rewards are rejected.""" + bnn = make_gaussian_bnn() + with pytest.raises(ValidationError): + bnn.update(context=np.zeros((1, n_features)), rewards=[bad_reward]) + + @pytest.mark.parametrize( + "reward_loc, reward_scale", + [(_GAUSSIAN_USER_LOC, None), (None, _GAUSSIAN_USER_SCALE), (None, None)], + ids=["loc_only", "scale_only", "neither"], + ) + def test_rejects_missing_fixed_values_without_fitted_standardization( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + reward_loc: Optional[float], + reward_scale: Optional[float], + ) -> None: + """fit_reward_standardization=False (the default) uses fixed loc / scale, so both must be given.""" + with pytest.raises(ValidationError): + make_gaussian_bnn(reward_loc=reward_loc, reward_scale=reward_scale) + + def test_fixed_standardization_maps_rewards( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + reward_loc: float = _GAUSSIAN_USER_LOC, + reward_scale: float = _GAUSSIAN_USER_SCALE, + n_samples: int = _GAUSSIAN_N_TRAIN, + tol: float = _GAUSSIAN_ROUND_TRIP_TOL, + ) -> None: + """Fixed loc / scale map the rewards to (reward - loc) / scale and are never re-fitted.""" + bnn = make_gaussian_bnn(reward_loc=reward_loc, reward_scale=reward_scale) + rewards = rng.normal(size=n_samples) + np.testing.assert_allclose( + bnn.prepare_rewards(rewards.tolist()), (rewards - reward_loc) / reward_scale, rtol=tol + ) + assert (bnn.reward_loc, bnn.reward_scale) == (reward_loc, reward_scale) + + @pytest.mark.parametrize( + "reward_loc, reward_scale", + [(_GAUSSIAN_USER_LOC, None), (None, _GAUSSIAN_USER_SCALE)], + ids=["loc_only", "scale_only"], + ) + def test_rejects_partial_fitted_standardization( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + reward_loc: Optional[float], + reward_scale: Optional[float], + ) -> None: + """With fit_reward_standardization=True, loc / scale are fitted together, so only one of them is invalid.""" + with pytest.raises(ValidationError): + make_gaussian_bnn(fit_reward_standardization=True, reward_loc=reward_loc, reward_scale=reward_scale) + + def test_given_values_with_fitted_standardization_are_taken_as_fitted( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + reward_loc: float = _GAUSSIAN_USER_LOC, + reward_scale: float = _GAUSSIAN_USER_SCALE, + n_samples: int = _GAUSSIAN_N_TRAIN, + tol: float = _GAUSSIAN_ROUND_TRIP_TOL, + ) -> None: + """Both values with fit_reward_standardization=True are an already fitted state: used, not re-fitted.""" + bnn = make_gaussian_bnn(fit_reward_standardization=True, reward_loc=reward_loc, reward_scale=reward_scale) + rewards = rng.normal(size=n_samples) + np.testing.assert_allclose( + bnn.prepare_rewards(rewards.tolist()), (rewards - reward_loc) / reward_scale, rtol=tol + ) + assert (bnn.reward_loc, bnn.reward_scale) == (reward_loc, reward_scale) + + def test_identity_standardization_trains_on_raw_rewards( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = _GAUSSIAN_N_TRAIN, + noise_sigma: float = _GAUSSIAN_NOISE_SIGMA, + short_fit_kwargs: Dict[str, Any] = _GAUSSIAN_SHORT_FIT_KWARGS, + raw_standardization: Dict[str, float] = _GAUSSIAN_RAW_STANDARDIZATION, + tol: float = _GAUSSIAN_ROUND_TRIP_TOL, + ) -> None: + """reward_loc=0, reward_scale=1: targets are the raw rewards and noise_sigma is in raw units.""" + bnn = make_gaussian_bnn(noise_sigma=noise_sigma, **raw_standardization, **short_fit_kwargs) + rewards = rng.normal(size=n_samples) + np.testing.assert_array_equal(bnn.prepare_rewards(rewards.tolist()), rewards) + assert (bnn.reward_loc, bnn.reward_scale) == ( + raw_standardization["reward_loc"], + raw_standardization["reward_scale"], + ) + assert np.exp(bnn.noise_log_sigma.params["mu"][0]) == pytest.approx(noise_sigma, rel=tol) + + def test_calibrate_output_bias_sets_mean_bias( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = _GAUSSIAN_N_TRAIN, + tol: float = _GAUSSIAN_ROUND_TRIP_TOL, + ) -> None: + """Calibration puts the output bias at the standardized target mean.""" + bnn = make_gaussian_bnn(calibrate_output_bias=True) + rewards = rng.normal(size=n_samples).tolist() + bnn._calibrate_output_bias(rewards) + (mu_bias,) = bnn.model_params.bnn_layer_params[-1].bias.params["mu"] + assert bnn.bias_calibrated + assert mu_bias == pytest.approx(bnn.prepare_rewards(rewards).mean(), abs=tol) + + def test_reset_refits_fitted_standardization( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = _GAUSSIAN_N_TRAIN, + n_features: int = _GAUSSIAN_N_FEATURES, + shift: float = _GAUSSIAN_SECOND_BATCH_SHIFT, + short_fit_kwargs: Dict[str, Any] = _GAUSSIAN_SHORT_FIT_KWARGS, + ) -> None: + """Fitted loc / scale are cleared by reset() and re-fitted on the next batch, with the noise prior.""" + bnn = make_gaussian_bnn(fit_reward_standardization=True, **short_fit_kwargs) + bnn.update(context=rng.normal(size=(n_samples, n_features)), rewards=rng.normal(size=n_samples).tolist()) + first_loc = bnn.reward_loc + bnn.reset() + assert bnn.reward_loc is None and bnn.reward_scale is None and bnn.noise_log_sigma is None + assert bnn.n_observations == bnn.reward_sum == 0 + assert bnn.model_params.bnn_layer_params == bnn.model_params.bnn_layer_params_init + second = (rng.normal(size=n_samples) + shift).tolist() + bnn.update(context=rng.normal(size=(n_samples, n_features)), rewards=second) + assert bnn.reward_loc == pytest.approx(np.mean(second)) + assert bnn.reward_loc != first_loc + assert bnn.noise_log_sigma is not None + + def test_reset_restores_user_supplied_standardization( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + reward_loc: float = _GAUSSIAN_USER_LOC, + reward_scale: float = _GAUSSIAN_USER_SCALE, + n_samples: int = _GAUSSIAN_N_TRAIN, + n_features: int = _GAUSSIAN_N_FEATURES, + short_fit_kwargs: Dict[str, Any] = _GAUSSIAN_SHORT_FIT_KWARGS, + ) -> None: + """loc / scale given at cold start are never re-fitted, and survive reset() and a JSON round trip.""" + bnn = make_gaussian_bnn(reward_loc=reward_loc, reward_scale=reward_scale, **short_fit_kwargs) + bnn.update(context=rng.normal(size=(n_samples, n_features)), rewards=rng.normal(size=n_samples).tolist()) + assert (bnn.reward_loc, bnn.reward_scale) == (reward_loc, reward_scale) + restored = GaussianBayesianNeuralNetwork.model_validate_json(bnn.model_dump_json()) + for model in (bnn, restored): + model.reset() + assert (model.reward_loc, model.reward_scale) == (reward_loc, reward_scale) + + def test_reset_after_round_trip_clears_fitted_standardization( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = _GAUSSIAN_N_TRAIN, + n_features: int = _GAUSSIAN_N_FEATURES, + short_fit_kwargs: Dict[str, Any] = _GAUSSIAN_SHORT_FIT_KWARGS, + ) -> None: + """A fitted standardization stays 'fitted' through serialization: reset() still clears it.""" + bnn = make_gaussian_bnn(fit_reward_standardization=True, **short_fit_kwargs) + bnn.update(context=rng.normal(size=(n_samples, n_features)), rewards=rng.normal(size=n_samples).tolist()) + restored = GaussianBayesianNeuralNetwork.model_validate_json(bnn.model_dump_json()) + assert restored.reward_loc == bnn.reward_loc + restored.reset() + assert restored.reward_loc is None and restored.reward_scale is None + + def test_decay_inflates_noise_posterior( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + decay_factor: float = _GAUSSIAN_DECAY_FACTOR, + n_samples: int = _GAUSSIAN_N_TRAIN, + n_features: int = _GAUSSIAN_N_FEATURES, + short_fit_kwargs: Dict[str, Any] = _GAUSSIAN_SHORT_FIT_KWARGS, + tol: float = _GAUSSIAN_ROUND_TRIP_TOL, + ) -> None: + """Forgetting widens the noise posterior by 1 / decay_factor, like the weights.""" + bnn = make_gaussian_bnn(decay_factor=decay_factor, **short_fit_kwargs) + bnn.update(context=rng.normal(size=(n_samples, n_features)), rewards=rng.normal(size=n_samples).tolist()) + before = bnn.noise_log_sigma.params["sigma"][0] + bnn._inflate_prior_variance() + assert bnn.noise_log_sigma.params["sigma"][0] == pytest.approx(before / decay_factor, rel=tol) + + def test_serialization_round_trip(self, fitted_gaussian_bnn: GaussianBayesianNeuralNetwork) -> None: + """The JSON state round-trips, including the standardization, the noise posterior and the counters.""" + restored = GaussianBayesianNeuralNetwork.model_validate_json(fitted_gaussian_bnn.model_dump_json()) + assert restored.reward_loc == fitted_gaussian_bnn.reward_loc + assert restored.reward_scale == fitted_gaussian_bnn.reward_scale + assert restored.noise_log_sigma == fitted_gaussian_bnn.noise_log_sigma + assert restored.n_observations == fitted_gaussian_bnn.n_observations + assert restored.model_params == fitted_gaussian_bnn.model_params diff --git a/tests/test_transfer.py b/tests/test_transfer.py index 7e6eedd6..5da8d225 100644 --- a/tests/test_transfer.py +++ b/tests/test_transfer.py @@ -24,7 +24,7 @@ import copy import json -from typing import Any, Dict, List, Tuple +from typing import Any, Callable, Dict, List, Tuple import numpy as np import pytest @@ -33,7 +33,7 @@ from pytest import MonkeyPatch import pybandits -from pybandits.cmab import CmabBernoulli, CmabBernoulliMO +from pybandits.cmab import CmabBernoulli, CmabBernoulliMO, CmabGaussian from pybandits.model import BaseBayesianNeuralNetwork from pybandits.smab import SmabBernoulli, SmabBernoulliCC, SmabBernoulliMO from pybandits.strategy import ( @@ -1423,3 +1423,94 @@ def test_non_overlapping_action_no_embedding_expansion( ) result = edit_model_on_the_fly(current, template) assert set(result.actions.keys()) == {"a2"} + + +# --------------------------------------------------------------------------- +# Gaussian (continuous-reward) cMAB +# --------------------------------------------------------------------------- +_GAUSSIAN_ACTION_IDS = {"a1", "a2"} +_GAUSSIAN_N_SAMPLES = 100 +_GAUSSIAN_REWARD_OFFSET = 3.0 +_GAUSSIAN_UPDATE_KWARGS: Dict[str, Any] = {"num_steps": 20} +_GAUSSIAN_STANDARDIZATIONS: Dict[str, Dict[str, Any]] = { + "fixed": {"reward_loc": 0.0, "reward_scale": 1.0}, + "fitted": {"fit_reward_standardization": True}, +} + + +@pytest.fixture(scope="module") +def make_cmab_gaussian() -> Callable[..., CmabGaussian]: + """Factory: cold-starts a CmabGaussian with the given standardization mode ("fixed" or "fitted").""" + + def _factory(standardization: str) -> CmabGaussian: + return CmabGaussian.cold_start( + action_ids=_GAUSSIAN_ACTION_IDS, + n_features=_N_FEATURES, + update_kwargs=_GAUSSIAN_UPDATE_KWARGS, + **_GAUSSIAN_STANDARDIZATIONS[standardization], + ) + + return _factory + + +def _update_with_continuous_rewards( + mab: CmabGaussian, rng: np.random.Generator, n_samples: int = _GAUSSIAN_N_SAMPLES +) -> CmabGaussian: + """One predict / update round with real-valued rewards away from 0 (so a fitted loc differs from 0).""" + context = rng.normal(size=(n_samples, _N_FEATURES)) + actions, _, _ = mab.predict(context=context) + rewards = (rng.normal(size=n_samples) + _GAUSSIAN_REWARD_OFFSET).tolist() + mab.update(context=context, actions=actions, rewards=rewards) + return mab + + +class TestGaussianTransfer: + """Transfer learning for CmabGaussian: learned state, standardization mode and likelihood compatibility.""" + + @pytest.mark.parametrize("standardization", list(_GAUSSIAN_STANDARDIZATIONS)) + def test_transfer_carries_learned_state( + self, + make_cmab_gaussian: Callable[..., CmabGaussian], + rng: np.random.Generator, + standardization: str, + ) -> None: + """Weights, standardization, noise posterior and counters are carried over to the template.""" + current = _update_with_continuous_rewards(make_cmab_gaussian(standardization), rng) + result = edit_model_on_the_fly(current, make_cmab_gaussian(standardization)) + assert isinstance(result, CmabGaussian) + for action_id, source in current.actions.items(): + transferred = result.actions[action_id] + assert transferred.model_params == source.model_params + assert (transferred.reward_loc, transferred.reward_scale) == (source.reward_loc, source.reward_scale) + assert transferred.noise_log_sigma == source.noise_log_sigma + assert transferred.n_observations == source.n_observations + + @pytest.mark.parametrize("source, template", [("fixed", "fitted"), ("fitted", "fixed")]) + def test_standardization_mode_change_raises( + self, + make_cmab_gaussian: Callable[..., CmabGaussian], + rng: np.random.Generator, + source: str, + template: str, + ) -> None: + """Weights only mean something relative to their loc / scale, so the mode can't change across a transfer.""" + current = _update_with_continuous_rewards(make_cmab_gaussian(source), rng) + with pytest.raises(ValueError, match="fit_reward_standardization"): + edit_model_on_the_fly(current, make_cmab_gaussian(template)) + + @pytest.mark.parametrize( + "gaussian_is_source", [True, False], ids=["gaussian_to_bernoulli", "bernoulli_to_gaussian"] + ) + def test_bernoulli_and_gaussian_cannot_be_merged( + self, + make_cmab_gaussian: Callable[..., CmabGaussian], + gaussian_is_source: bool, + action_ids: set = _GAUSSIAN_ACTION_IDS, + n_features: int = _N_FEATURES, + ) -> None: + """Logit weights and mean-reward weights are not interchangeable: mixing likelihoods raises TypeError.""" + gaussian = make_cmab_gaussian("fixed") + bernoulli = CmabBernoulli.cold_start(action_ids=action_ids, n_features=n_features, strategy=ClassicBandit()) + current, template = (gaussian, bernoulli) if gaussian_is_source else (bernoulli, gaussian) + with pytest.raises(TypeError, match="reward type"): + edit_model_on_the_fly(current, template)