From 1042258772ad63f5fe76cb58591129cdf3f09b99 Mon Sep 17 00:00:00 2001 From: ronshiff1 Date: Fri, 9 Oct 2026 10:36:27 +0000 Subject: [PATCH 01/14] added continous cmab --- docs/continuous_reward_cmab.md | 73 +++ pybandits/actions_manager.py | 47 +- pybandits/base.py | 8 + pybandits/base_model.py | 3 + pybandits/cmab.py | 47 +- pybandits/mab.py | 12 +- pybandits/meta_model/__init__.py | 2 + pybandits/meta_model/cmab_meta_model.py | 54 +- pybandits/model/__init__.py | 2 + pybandits/model/bnn/network.py | 538 +++++++++++++++++++- tests/test_gaussian_bnn.py | 622 ++++++++++++++++++++++++ 11 files changed, 1365 insertions(+), 43 deletions(-) create mode 100644 docs/continuous_reward_cmab.md create mode 100644 tests/test_gaussian_bnn.py diff --git a/docs/continuous_reward_cmab.md b/docs/continuous_reward_cmab.md new file mode 100644 index 00000000..fd002723 --- /dev/null +++ b/docs/continuous_reward_cmab.md @@ -0,0 +1,73 @@ +# Continuous rewards in [-1, 1] for cMAB + +## Current state + +The library does not take -1/1 rewards. It takes binary {0, 1} rewards: + +- `BinaryReward = conint(ge=0, le=1)` in `pybandits/base.py`. +- Every cMAB model uses a Bernoulli likelihood on a logit: + - `pybandits/model/bnn/network.py` (`NumpyroBernoulli(logits=logit)`) + - `pybandits/meta_model/cmab_meta_model.py` (`NumpyroBernoulli(logits=logit)`) + +So any -1/1 reward is already being mapped to 0/1 before it reaches the library. + +## Option A: rescale to [0, 1] and keep the Bernoulli model (easy) + +Map the reward with `y = (r + 1) / 2`, which puts it in [0, 1]. + +### A1. Binarization trick (no library change) + +For each observation, draw `b ~ Bernoulli((r + 1) / 2)` and pass `b` as the reward. + +- This is Agrawal & Goyal's standard way to extend Thompson Sampling to bounded rewards, and the regret guarantees still hold. +- Cost: extra noise, so the model learns somewhat slower. +- It can be used today with no code change. + +### A2. Pass the fractional target directly + +This fits a logistic regression with targets that are not 0/1. The Bernoulli log-likelihood with logits is BCE-with-logits, which is valid for any `y` in [0, 1], and numpyro does not check the support by default. + +Changes needed: + +- Relax `BinaryReward` to a float in [0, 1]. The type appears in many signatures, but the change is mechanical. +- Check `_calibrate_output_bias`. It should still work because it uses the mean reward. +- Update the simulators and tests. + +Behaviour: + +- The mean is estimated correctly: `E[y | x] = sigmoid(f)`. +- The uncertainty is not. The model assumes variance `p(1 - p)`, which is the largest possible variance for a variable in [0, 1]. The posterior is therefore too wide and the bandit **over-explores**. This errs on the safe side. + +Estimated effort: 1–2 days, including tests. + +## Option B: a proper continuous likelihood (harder) + +Use a Gaussian, Beta or truncated-normal likelihood, with an identity or tanh output and a learned noise `σ`. + +The likelihood itself is simple. The work is in what is built around it: + +- `sample_proba` returns a `Probability` (`Float01`) everywhere, and the strategies and meta-model all rely on that. +- `_calibrate_output_bias` uses `logit(rate)`. +- The CostControl `subsidy_factor` is a **relative** threshold ("within X of the best"). Relative thresholds break once the expected reward can be negative, so they must be rescaled or redefined. +- The OPE estimators, the simulators and the whole test suite need changes. + +Estimated effort: 1–2+ weeks. It touches every cMAB / MO / CC variant. + +## Do we need to understand the reward distribution first? + +**For Option A: mostly no.** Arm selection depends only on the mean, and the Bernoulli model's uncertainty is conservative whatever the shape. A histogram per arm is still worth plotting, because some shapes change the recommendation: + +| Reward shape | Implication | +|---|---| +| Mostly at ±1 with a few values in between | The Bernoulli model is almost exact. Use Option A. | +| Spike at 0 (e.g. "no response" coded as 0) | Rewards are zero-inflated. Consider modelling "any response" and "sign/size given a response" separately. | +| Concentrated in the middle (e.g. 0 ± 0.2) | Option A's posterior is much too wide and exploration is slow. Option B with a learned `σ` pays off here. | +| Heteroscedastic (noise varies with context) | Only Option B can model this. | + +**For Option B: yes.** The likelihood family and a sensible prior on `σ` have to be chosen from the data. + +## Recommendation + +1. Start with **A2** (fractional targets), or **A1** for zero code changes. +2. Back-test it on logged data with the existing offline policy evaluator. +3. Build **Option B** only if the histogram shows rewards bunched in the middle and A2 explores too slowly. diff --git a/pybandits/actions_manager.py b/pybandits/actions_manager.py index 4fc18a71..13d61a70 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,38 @@ 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 + @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 +326,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 +350,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 +883,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 +999,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..2604fada 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,40 @@ 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 mean head and + a (heteroscedastic) noise head, 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 predicted reward + noise std in place of the weighted sums. + + 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..c8b45987 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,14 @@ 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) + same_likelihood = type(head_bnn).output_distribution is type(first_bnn).output_distribution and all( + getattr(head_bnn, attr, None) == getattr(first_bnn, attr, None) + for attr in ("sigma_min", "homoscedastic") + ) + if not same_likelihood: + 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 +457,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 +474,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() } @@ -677,12 +686,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. a homoscedastic 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 +707,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 +879,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 +898,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 standardize with their own (first-batch fitted) + location and scale, 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 +954,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..278a6b40 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, @@ -141,6 +145,8 @@ class BaseBayesianNeuralNetwork(Model, DNNMixin, ABC): model_params: BnnParams _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. + _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 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._output_dim) layer_params_init = [] for layer_ind in range(len(_dim_list) - 1): @@ -650,7 +659,8 @@ def emit_submodel( 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,6 +708,7 @@ 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``. @@ -718,6 +729,8 @@ def _observe( 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 +768,115 @@ 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 homoscedastic 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, _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 a Bernoulli distribution 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, _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) + + @classmethod + def _resolve_output_dim(cls, **kwargs) -> PositiveInt: + """Output-layer width for a cold start with these extra constructor kwargs (the class default here).""" + return cls._output_dim + + @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 +1004,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 +1024,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 +1042,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 +1071,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 +1188,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 +1390,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 +1433,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 +1476,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 +1491,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 +1500,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 +1528,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 @@ -1494,6 +1630,7 @@ def cold_start( use_layerwise_scaling=use_layerwise_scaling, dist_class=dist_class, bias_std=bias_std, + output_dim=cls._resolve_output_dim(**kwargs), **dist_params_init, ) return cls( @@ -1608,6 +1745,373 @@ 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)``. Thompson sampling draws the network weights 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 only sets how much each observation is trusted during training. + Two noise models: + + - **heteroscedastic** (default): the output layer has two units that share the hidden layers, a mean head + ``mu(x)`` and a noise head ``sigma(x) = softplus(.) + sigma_min``. + - **homoscedastic** (``homoscedastic=True``): the output layer has one unit ``mu(x)``, and ``sigma`` is a + single latent of the model, ``log sigma ~ Normal``, whose posterior is stored in ``noise_log_sigma`` and + carried across updates like the weights. + + Use the homoscedastic model when the reward noise is skewed or zero-inflated (e.g. revenue). The Gaussian + likelihood weights each residual by ``1 / sigma(x)^2``, so a learned ``sigma(x)`` lets the model explain + the large rewards as noise instead of moving ``mu(x)``: ``mu`` then sits below ``E[reward | x]`` where the + rewards are most variable. With a single ``sigma`` the fit of ``mu`` is a (Bayesian) least-squares fit, + whose target is ``E[reward | x]`` whatever the shape of the noise. + + The network is trained on standardized targets ``(reward - reward_loc) / reward_scale`` so the default + ``O(1)`` weight priors fit rewards on any scale (e.g. revenue); ``sample_proba`` maps the outputs + back to the reward scale. + + Parameters + ---------- + sigma_min : PositiveFloat + Floor on the noise std, in standardized units. Keeps the likelihood (and its gradients) bounded + while the early, wide posterior samples push the noise head down. Default is 0.05. + standardize_rewards : bool + Whether to train on standardized targets. When ``reward_loc`` / ``reward_scale`` are not given they + are fitted (mean / std) on the first update batch and kept fixed afterwards, including across + ``reset()`` (the prior lives in standardized space). Default is True. + reward_loc : Optional[float] + Location used to standardize rewards. None to fit it on the first update. + reward_scale : Optional[PositiveFloat] + Scale used to standardize rewards. None to fit it on the first update. + n_observations : NonNegativeInt + Number of rewards observed. + reward_sum : float + Sum of the rewards observed. + homoscedastic : bool + Whether the noise std is one latent of the model rather than a function of the context. Default is False. + noise_sigma : Optional[PositiveFloat] + Homoscedastic model only: the 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 + Homoscedastic model only: 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] + Homoscedastic model only: the current posterior of ``log sigma`` (in standardized units), 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. + - ``sigma_min`` only applies to the heteroscedastic noise head. + - ``reset()`` clears ``noise_log_sigma``: it is re-initialized on the next update, like the weights return to + their cold-start prior. The reward standardization is kept. + + Examples + -------- + >>> bnn = GaussianBayesianNeuralNetwork.cold_start( + ... n_features=2, + ... hidden_dim_list=[16], + ... dist_type="normal", + ... dist_params_init={"mu": 0, "sigma": 0.1}, + ... update_kwargs={"optimizer_kwargs": {"step_size": 3e-3}}, + ... ) + """ + + sigma_min: PositiveFloat = 0.05 + standardize_rewards: bool = True + reward_loc: Optional[float] = None + reward_scale: Optional[PositiveFloat] = None + n_observations: NonNegativeInt = 0 + reward_sum: float = 0.0 + homoscedastic: bool = False + 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, raw sigma]; the homoscedastic model has [mu] only + _output_dim: ClassVar[PositiveInt] = 2 + _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", + ) + _transfer_structural_keys: ClassVar[Tuple[str, ...]] = ("homoscedastic",) + + @model_validator(mode="after") + def validate_reward_standardization(self) -> "GaussianBayesianNeuralNetwork": + if (self.reward_loc is None) != (self.reward_scale is None): + raise ValueError("reward_loc and reward_scale must be either both set or both None.") + return self + + @model_validator(mode="after") + def validate_noise_model(self) -> "GaussianBayesianNeuralNetwork": + expected = self._resolve_output_dim(homoscedastic=self.homoscedastic) + output_dim = self.model_params.bnn_layer_params[-1].weight.shape[-1] + if output_dim != expected: + raise ValueError( + f"The output layer has {output_dim} units, but homoscedastic={self.homoscedastic} needs {expected}." + ) + if not self.homoscedastic and (self.noise_log_sigma is not None or self.noise_sigma is not None): + raise ValueError("noise_sigma / noise_log_sigma require homoscedastic=True.") + 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 + + @classmethod + def _resolve_output_dim(cls, **kwargs) -> PositiveInt: + return 1 if kwargs.get("homoscedastic", False) else cls._output_dim + + def _initial_noise_log_sigma(self, targets: Optional[np.ndarray] = None) -> NormalArray: + """The cold-start ``log sigma`` prior, in standardized units. + + Centered on ``log(noise_sigma / reward_scale)`` when ``noise_sigma`` is given, otherwise on the log std of + ``targets`` (the first standardized batch; 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 homoscedastic noise latent ``log sigma`` (none for the heteroscedastic model).""" + if not self.homoscedastic: + return {} + 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: + if not self.homoscedastic: + return + 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=(np.asarray(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]) + + @property + def _loc_scale(self) -> Tuple[float, float]: + """``(loc, scale)`` mapping standardized targets back to rewards (identity until fitted).""" + if not self.standardize_rewards or self.reward_loc is None: + return 0.0, 1.0 + return self.reward_loc, self.reward_scale + + def output_distribution( + self, linear_out: jax.Array, extra_sites: Optional[Dict[str, jax.Array]] = None + ) -> NumpyroDistribution: + """``Normal(mu, sigma)`` on the standardized targets. + + Parameters + ---------- + linear_out : jax.Array of shape (batch, 2) or (batch, 1) + Raw output of the final layer: ``[mu, raw_sigma]`` (heteroscedastic) or ``[mu]`` (homoscedastic). + extra_sites : Optional[Dict[str, jax.Array]] + Homoscedastic model: 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, with ``sigma = exp(log sigma)`` (homoscedastic) or + ``softplus(raw_sigma) + sigma_min`` (heteroscedastic). + """ + mu = linear_out[..., 0] + if self.homoscedastic: + sigma = jnp.exp(self._noise_log_sigma_of(extra_sites, self._noise_var_name)) + else: + sigma = jax.nn.softplus(linear_out[..., 1]) + self.sigma_min + return NumpyroNormal(loc=mu, 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, 2) or (n_samples, 1) + Raw output of the final layer, one row per sampled network. + extra_samples : Optional[Dict[str, np.ndarray]] + Homoscedastic model: 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. + """ + loc, scale = self._loc_scale + mu = linear_out[..., 0] * scale + loc + if self.homoscedastic: + 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 + else: + sigma = (np.logaddexp(0.0, linear_out[..., 1]) + self.sigma_min) * 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 predicted 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 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: + """Standardize rewards into training targets, fitting ``reward_loc`` / ``reward_scale`` on first use. + + The scale falls back to ``max(|loc|, 1)`` when the first batch has (near) zero spread, e.g. a + single observation. The homoscedastic model also initializes its ``log sigma`` prior here, 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.standardize_rewards 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.homoscedastic and 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 biases to the first batch's target mean and std on the first update call. + + The mean head's bias gets the mean of the (standardized) targets and the noise head's bias gets + ``softplus^-1(std - sigma_min)``, so the cold-start predictive matches the observed reward spread. + With ``standardize_rewards`` (fitted on that same batch) this is ``(0, softplus^-1(1 - sigma_min))``. + + 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] + bias_mu = [float(y.mean())] + if not self.homoscedastic: # the homoscedastic noise is initialized by prepare_rewards instead + std = float(y.std()) + std = std if std > self._numerical_eps else 1.0 # no spread to calibrate on (e.g. a single reward) + raw_sigma = max(std - self.sigma_min, self._numerical_eps) + bias_mu.append(float(np.log(np.expm1(raw_sigma)))) + new_bias = output_layer.bias.with_dist_parameters(mu=bias_mu) + self.model_params.bnn_layer_params[-1] = BnnLayerParams(weight=output_layer.weight, bias=new_bias) + self.bias_calibrated = True + + def _reset(self): + """Reset the network, the noise latent and the reward counters; the reward standardization is kept.""" + super()._reset() + self.n_observations = 0 + self.reward_sum = 0.0 + self.noise_log_sigma = None + + class BaseBayesianNeuralNetworkMO(ModelMO, ABC): """ Base class for Bayesian Neural Network with multi-objective. diff --git a/tests/test_gaussian_bnn.py b/tests/test_gaussian_bnn.py new file mode 100644 index 00000000..bc9d7336 --- /dev/null +++ b/tests/test_gaussian_bnn.py @@ -0,0 +1,622 @@ +# MIT License +# +# Copyright (c) 2023 Playtika Ltd. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""Tests for the Gaussian (continuous-reward) BNN, its meta-model wiring and the CmabGaussian bandit.""" + +from typing import Any, Callable, Dict, List, Optional + +import numpy as np +import pytest +from hypothesis import given +from hypothesis import strategies as st +from pydantic import ValidationError + +from pybandits.cmab import CmabBernoulli, CmabGaussian +from pybandits.meta_model import CmabMetaModel +from pybandits.model import BayesianNeuralNetwork, GaussianBayesianNeuralNetwork + +# --- test configuration constants --- +N_FEATURES = 3 +HIDDEN_DIMS = [8] +OUTPUT_DIM = 2 +RANDOM_SEED = 0 +N_TRAIN = 2000 +N_POSTERIOR_SAMPLES = 200 +# Narrow prior + larger step size: the setting the class docstring recommends for a fast fit. +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}}, +} +SHORT_FIT_KWARGS: Dict[str, Any] = {"update_kwargs": {"num_steps": 5}} +# Synthetic reward: REWARD_OFFSET + REWARD_SLOPE * x0 + noise, noise std LOW / HIGH for x1 <= 0 / x1 > 0. +REWARD_OFFSET = 100.0 +REWARD_SLOPE = 30.0 +NOISE_STD_LOW = 5.0 +NOISE_STD_HIGH = 15.0 +# Relative tolerance on the recovered mean / noise std at the probe points. +MEAN_REL_TOL = 0.1 +SIGMA_REL_TOL = 0.5 +ACTION_IDS = {"a", "b", "c"} +ARM_MEANS = {"a": -1.0, "b": 0.0, "c": 2.0} +BEST_ARM = "c" +N_ROUNDS = 4 +N_PER_ROUND = 200 +MIN_BEST_ARM_SHARE = 0.8 +# Probe contexts: (x0, x1) in {-1, 1}^2 hit both mean levels and both noise regions. +PROBE_CONTEXTS = [[1.0, -1.0, 0.0], [1.0, 1.0, 0.0], [-1.0, -1.0, 0.0], [-1.0, 1.0, 0.0]] +ROUND_TRIP_TOL = 1e-6 +MAX_ABS_REWARD = 1e6 +MAX_REWARDS = 50 +BACKBONE_KWARGS: Dict[str, Any] = {"backbone_hidden_dims": [8], "backbone_embedding_dim": 3} +MINIBATCH_SIZE = 64 +BAD_REWARDS = [float("nan"), float("inf"), -float("inf")] +OUT_OF_RANGE_SOFT_REWARD = 2.5 +DELTA = 0.1 +OTHER_SIGMA_MIN = 0.2 +EPSILON = 0.1 +HOMOSCEDASTIC_OUTPUT_DIM = 1 +NOISE_SIGMA = 3.0 +DECAY_FACTOR = 0.5 +# 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. The noise grows with the mean, which a learned sigma(x) uses to shrink mu. +BUY_PROB = 0.15 +SKEW_OFFSET = 2.0 +SKEW_SLOPE = 1.5 +SKEW_LOG_SD = 1.0 +LEVEL_REL_TOL = 0.15 + + +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 = SKEW_OFFSET + SKEW_SLOPE * np.clip(x[:, 0], -1.0, 2.0) + amount = rng.lognormal(np.log(mean / BUY_PROB) - SKEW_LOG_SD**2 / 2, SKEW_LOG_SD) + return (rng.uniform(size=len(x)) < BUY_PROB) * amount + + +def _heteroscedastic_rewards(x: np.ndarray, rng: np.random.Generator) -> np.ndarray: + """Linear-in-x0 rewards whose noise std switches with the sign of x1.""" + noise_std = np.where(x[:, 1] > 0, NOISE_STD_HIGH, NOISE_STD_LOW) + return REWARD_OFFSET + REWARD_SLOPE * x[:, 0] + rng.normal(size=len(x)) * noise_std + + +@pytest.fixture(scope="module") +def make_gaussian_bnn() -> Callable[..., GaussianBayesianNeuralNetwork]: + """Factory: cold-starts a GaussianBayesianNeuralNetwork with the module's shape and extra kwargs.""" + + def _factory(**kwargs: Any) -> GaussianBayesianNeuralNetwork: + return GaussianBayesianNeuralNetwork.cold_start( + n_features=N_FEATURES, hidden_dim_list=HIDDEN_DIMS, random_seed=RANDOM_SEED, **kwargs + ) + + return _factory + + +@pytest.fixture(scope="module") +def fitted_gaussian_bnn( + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], +) -> GaussianBayesianNeuralNetwork: + """A Gaussian BNN fitted on heteroscedastic synthetic rewards.""" + data_rng = np.random.default_rng(RANDOM_SEED) + bnn = make_gaussian_bnn(**FAST_FIT_KWARGS) + x = data_rng.normal(size=(N_TRAIN, N_FEATURES)) + bnn.update(context=x, rewards=_heteroscedastic_rewards(x, data_rng).tolist()) + return bnn + + +@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=ACTION_IDS, + n_features=N_FEATURES, + hidden_dim_list=HIDDEN_DIMS, + random_seed=RANDOM_SEED, + **{**FAST_FIT_KWARGS, **kwargs}, + ) + + return _factory + + +class TestGaussianBayesianNeuralNetwork: + """Tests for the standalone Gaussian BNN.""" + + def test_output_layer_has_mean_and_noise_units( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + output_dim: int = OUTPUT_DIM, + ) -> None: + """The output layer has one unit for the mean and one for the noise std.""" + 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,) + + def test_sample_proba_returns_mean_and_positive_noise( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = N_TRAIN, + n_features: int = N_FEATURES, + ) -> None: + """Each sample is a (mu, sigma) pair with sigma above the configured floor.""" + bnn = make_gaussian_bnn() + samples = bnn.sample_proba(context=rng.normal(size=(n_samples, n_features)), rng=rng) + assert len(samples) == n_samples + sigmas = np.array([sigma for _, sigma in samples]) + assert np.all(sigmas >= bnn.sigma_min) + + def test_fit_recovers_mean_and_heteroscedastic_noise( + self, + fitted_gaussian_bnn: GaussianBayesianNeuralNetwork, + rng: np.random.Generator, + n_posterior_samples: int = N_POSTERIOR_SAMPLES, + offset: float = REWARD_OFFSET, + slope: float = REWARD_SLOPE, + noise_low: float = NOISE_STD_LOW, + noise_high: float = NOISE_STD_HIGH, + mean_rel_tol: float = MEAN_REL_TOL, + sigma_rel_tol: float = SIGMA_REL_TOL, + probe_contexts: List[List[float]] = PROBE_CONTEXTS, + ) -> None: + """The posterior mean tracks the linear trend and the noise head separates low / high noise regions.""" + probes = np.array(probe_contexts) + true_mu = offset + slope * probes[:, 0] + true_sigma = np.where(probes[:, 1] > 0, noise_high, noise_low) + draws = np.array( + [fitted_gaussian_bnn.sample_proba(context=probes, rng=rng) for _ in range(n_posterior_samples)] + ) + mu, sigma = draws[..., 0].mean(axis=0), draws[..., 1].mean(axis=0) + np.testing.assert_allclose(mu, true_mu, rtol=mean_rel_tol) + high_noise = probes[:, 1] > 0 + np.testing.assert_allclose(sigma[high_noise], true_sigma[high_noise], rtol=sigma_rel_tol) + # VI pulls the noise head towards the pooled noise level, so the low-noise region is only required to + # sit clearly below it (a homoscedastic fit would put every probe at the pooled std). + pooled_std = np.sqrt((noise_low**2 + noise_high**2) / 2) + assert np.all(sigma[~high_noise] < pooled_std) + + def test_counters_track_observations( + self, + fitted_gaussian_bnn: GaussianBayesianNeuralNetwork, + n_train: int = 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=-MAX_ABS_REWARD, max_value=MAX_ABS_REWARD, allow_nan=False), + min_size=2, + max_size=MAX_REWARDS, + unique=True, + ), + n_features=st.just(N_FEATURES), + tol=st.just(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) + 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=-MAX_ABS_REWARD, max_value=MAX_ABS_REWARD, allow_nan=False), + n_features=st.just(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) + bnn.prepare_rewards([reward]) + assert bnn.reward_scale == max(abs(reward), 1.0) + + @pytest.mark.parametrize("bad_reward", BAD_REWARDS) + def test_update_rejects_non_finite_rewards( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + bad_reward: float, + n_features: int = 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", [(1.0, None), (None, 1.0)]) + def test_rejects_half_specified_standardization( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + reward_loc: Optional[float], + reward_scale: Optional[float], + ) -> None: + """reward_loc and reward_scale must be given together.""" + with pytest.raises(ValidationError): + make_gaussian_bnn(reward_loc=reward_loc, reward_scale=reward_scale) + + def test_calibrate_output_bias_sets_mean_and_noise_bias( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = N_TRAIN, + tol: float = ROUND_TRIP_TOL, + ) -> None: + """Calibration puts the output biases at the standardized target mean and inverse-softplus std.""" + bnn = make_gaussian_bnn(calibrate_output_bias=True) + rewards = rng.normal(size=n_samples).tolist() + bnn._calibrate_output_bias(rewards) + targets = bnn.prepare_rewards(rewards) + mu_bias, raw_sigma_bias = bnn.model_params.bnn_layer_params[-1].bias.params["mu"] + assert bnn.bias_calibrated + assert mu_bias == pytest.approx(targets.mean(), abs=tol) + assert np.logaddexp(0, raw_sigma_bias) + bnn.sigma_min == pytest.approx(targets.std(), rel=tol) + + def test_reset_clears_counters_and_keeps_standardization( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = N_TRAIN, + n_features: int = N_FEATURES, + short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, + ) -> None: + """reset() zeroes the reward counters and the network, but keeps the fitted standardization.""" + bnn = make_gaussian_bnn(**short_fit_kwargs) + bnn.update(context=rng.normal(size=(n_samples, n_features)), rewards=rng.normal(size=n_samples).tolist()) + loc, scale = bnn.reward_loc, bnn.reward_scale + bnn.reset() + assert bnn.n_observations == bnn.reward_sum == 0 + assert (bnn.reward_loc, bnn.reward_scale) == (loc, scale) + assert bnn.model_params.bnn_layer_params == bnn.model_params.bnn_layer_params_init + + def test_serialization_round_trip(self, fitted_gaussian_bnn: GaussianBayesianNeuralNetwork) -> None: + """The JSON state round-trips, including the standardization and 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.n_observations == fitted_gaussian_bnn.n_observations + assert restored.model_params == fitted_gaussian_bnn.model_params + + +class TestGaussianMetaModel: + """Tests for Gaussian heads in CmabMetaModel.""" + + def test_rejects_mixed_likelihoods( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + n_features: int = N_FEATURES, + hidden_dims: List[int] = HIDDEN_DIMS, + ) -> None: + """Bernoulli and Gaussian heads cannot share one meta-model.""" + bernoulli = BayesianNeuralNetwork.cold_start(n_features=n_features, hidden_dim_list=hidden_dims) + with pytest.raises(AttributeError): + CmabMetaModel(actions={"a": make_gaussian_bnn(), "b": bernoulli}) + + def test_rejects_different_sigma_min( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + sigma_min: float = OTHER_SIGMA_MIN, + ) -> None: + """Gaussian heads must share sigma_min, which the joint likelihood takes from one head.""" + with pytest.raises(AttributeError): + CmabMetaModel(actions={"a": make_gaussian_bnn(), "b": make_gaussian_bnn(sigma_min=sigma_min)}) + + +class TestCmabGaussian: + """Tests for the CmabGaussian bandit.""" + + @pytest.mark.parametrize( + "extra_kwargs", + [ + {}, + {**BACKBONE_KWARGS, "update_kwargs": {**FAST_FIT_KWARGS["update_kwargs"], "batch_size": MINIBATCH_SIZE}}, + ], + ids=["no_backbone", "backbone_minibatch"], + ) + def test_learns_best_arm( + self, + make_cmab_gaussian: Callable[..., CmabGaussian], + extra_kwargs: Dict[str, Any], + rng: np.random.Generator, + arm_means: Dict[str, float] = ARM_MEANS, + best_arm: str = BEST_ARM, + n_rounds: int = N_ROUNDS, + n_per_round: int = N_PER_ROUND, + n_features: int = N_FEATURES, + min_share: float = 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, mus, sigmas = mab.predict(context=context) + rewards = [arm_means[a] + rng.normal() for a in actions] + mab.update(actions=actions, rewards=rewards, context=context) + actions, mus, 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()) + + def test_state_round_trip( + self, + make_cmab_gaussian: Callable[..., CmabGaussian], + rng: np.random.Generator, + n_per_round: int = N_PER_ROUND, + n_features: int = N_FEATURES, + short_fit_kwargs: Dict[str, Any] = 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", BAD_REWARDS) + def test_rejects_non_finite_rewards( + self, + make_cmab_gaussian: Callable[..., CmabGaussian], + bad_reward: float, + n_features: int = N_FEATURES, + short_fit_kwargs: Dict[str, Any] = 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 = ACTION_IDS, + n_features: int = N_FEATURES, + delta: float = DELTA, + epsilon: float = EPSILON, + ) -> 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) + + +@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 = ACTION_IDS, + n_features: int = N_FEATURES, + reward: float = OUT_OF_RANGE_SOFT_REWARD, + short_fit_kwargs: Dict[str, Any] = 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))) + + +class TestHomoscedasticGaussian: + """Tests for the homoscedastic 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 = HOMOSCEDASTIC_OUTPUT_DIM, + ) -> None: + """The output layer has a single unit and the noise latent is unset before the first update.""" + bnn = make_gaussian_bnn(homoscedastic=True) + assert bnn.model_params.bnn_layer_params[-1].weight.shape[-1] == output_dim + assert bnn.noise_log_sigma is None + + @given( + rewards=st.lists( + st.floats(min_value=-MAX_ABS_REWARD, max_value=MAX_ABS_REWARD, allow_nan=False), + min_size=2, + max_size=MAX_REWARDS, + unique=True, + ), + n_features=st.just(N_FEATURES), + tol=st.just(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, homoscedastic=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 = NOISE_SIGMA, + n_samples: int = N_TRAIN, + tol: float = ROUND_TRIP_TOL, + ) -> None: + """noise_sigma (reward units) sets the prior once the reward scale is known.""" + bnn = make_gaussian_bnn(homoscedastic=True, noise_sigma=noise_sigma) + bnn.prepare_rewards(rng.normal(size=n_samples).tolist()) + prior_sigma = np.exp(bnn.noise_log_sigma.params["mu"][0]) * bnn.reward_scale + assert prior_sigma == 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 = N_TRAIN, + n_features: int = N_FEATURES, + short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, + ) -> None: + """Each update replaces the stored noise posterior, and the next update starts from it.""" + bnn = make_gaussian_bnn(homoscedastic=True, **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 = N_TRAIN, + n_features: int = N_FEATURES, + short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, + ) -> None: + """The predicted noise std doesn't depend on the context: its spread comes only from the log-sigma posterior.""" + bnn = make_gaussian_bnn(homoscedastic=True, **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)]) + log_sigma_sd = bnn.noise_log_sigma.params["sigma"][0] + assert np.all(sigmas > 0) + assert np.log(sigmas).std() == pytest.approx(log_sigma_sd, rel=SIGMA_REL_TOL) + + def test_mean_is_level_calibrated_on_skewed_rewards( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = N_TRAIN, + n_features: int = N_FEATURES, + n_posterior_samples: int = N_POSTERIOR_SAMPLES, + rel_tol: float = LEVEL_REL_TOL, + fast_fit_kwargs: Dict[str, Any] = 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(RANDOM_SEED) + x = data_rng.normal(size=(n_samples, n_features)) + y = _skewed_rewards(x, data_rng) + bnn = make_gaussian_bnn(homoscedastic=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) + + def test_reset_clears_noise_and_keeps_standardization( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = N_TRAIN, + n_features: int = N_FEATURES, + short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, + ) -> None: + """reset() clears the noise latent (re-initialized on the next update) but keeps loc / scale.""" + bnn = make_gaussian_bnn(homoscedastic=True, **short_fit_kwargs) + bnn.update(context=rng.normal(size=(n_samples, n_features)), rewards=rng.normal(size=n_samples).tolist()) + loc, scale = bnn.reward_loc, bnn.reward_scale + bnn.reset() + assert bnn.noise_log_sigma is None + assert (bnn.reward_loc, bnn.reward_scale) == (loc, scale) + + def test_decay_inflates_noise_posterior( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + decay_factor: float = DECAY_FACTOR, + n_samples: int = N_TRAIN, + n_features: int = N_FEATURES, + short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, + tol: float = ROUND_TRIP_TOL, + ) -> None: + """Forgetting widens the noise posterior by 1 / decay_factor, like the weights.""" + bnn = make_gaussian_bnn(homoscedastic=True, 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, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = N_TRAIN, + n_features: int = N_FEATURES, + short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, + ) -> None: + """The noise posterior survives the JSON round trip.""" + bnn = make_gaussian_bnn(homoscedastic=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.noise_log_sigma == bnn.noise_log_sigma + assert restored.homoscedastic + + @pytest.mark.parametrize( + "kwargs", + [{"noise_sigma": NOISE_SIGMA}, {"noise_log_sigma": {"mu": [0.0], "sigma": [1.0]}}], + ids=["noise_sigma", "noise_log_sigma"], + ) + def test_heteroscedastic_rejects_noise_settings( + self, make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], kwargs: Dict[str, Any] + ) -> None: + """noise_sigma / noise_log_sigma only make sense for the homoscedastic model.""" + with pytest.raises(ValidationError): + make_gaussian_bnn(**kwargs) + + def test_rejects_output_width_mismatch( + self, make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork] + ) -> None: + """A heteroscedastic network (2 output units) can't be flagged homoscedastic.""" + bnn = make_gaussian_bnn() + with pytest.raises(ValidationError): + GaussianBayesianNeuralNetwork(**{**bnn.model_dump(), "homoscedastic": True}) + + def test_meta_model_rejects_mixed_noise_models( + self, make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork] + ) -> None: + """Homoscedastic and heteroscedastic heads can't share one meta-model.""" + with pytest.raises(AttributeError): + CmabMetaModel(actions={"a": make_gaussian_bnn(), "b": make_gaussian_bnn(homoscedastic=True)}) + + @pytest.mark.parametrize( + "extra_kwargs", + [ + {}, + {**BACKBONE_KWARGS, "update_kwargs": {**FAST_FIT_KWARGS["update_kwargs"], "batch_size": MINIBATCH_SIZE}}, + ], + ids=["no_backbone", "backbone_minibatch"], + ) + def test_cmab_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] = ARM_MEANS, + best_arm: str = BEST_ARM, + n_rounds: int = N_ROUNDS, + n_per_round: int = N_PER_ROUND, + n_features: int = N_FEATURES, + min_share: float = MIN_BEST_ARM_SHARE, + ) -> None: + """The homoscedastic bandit finds the best arm, and every played arm carries its own noise posterior.""" + mab = make_cmab_gaussian(homoscedastic=True, **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, _, _ = mab.predict(context=rng.normal(size=(n_per_round, n_features))) + assert np.mean(np.array(actions) == best_arm) >= min_share + assert mab.actions[best_arm].noise_log_sigma is not None + assert CmabGaussian.from_state(mab.get_state()[1]).actions == mab.actions From 8201a70ce45934198f42c4452261e4c6b9a23b5a Mon Sep 17 00:00:00 2001 From: ronshiff1 Date: Fri, 9 Oct 2026 10:57:24 +0000 Subject: [PATCH 02/14] updates continous cmab --- .claude/skills/gaussian-cmab/SKILL.md | 134 ++++++ docs/continuous_reward_cmab.md | 151 +++---- pybandits/cmab.py | 13 +- pybandits/meta_model/cmab_meta_model.py | 8 +- pybandits/model/bnn/network.py | 205 +++++----- pyproject.toml | 2 +- tests/test_gaussian_bnn.py | 514 +++++++++++------------- 7 files changed, 549 insertions(+), 478 deletions(-) create mode 100644 .claude/skills/gaussian-cmab/SKILL.md diff --git a/.claude/skills/gaussian-cmab/SKILL.md b/.claude/skills/gaussian-cmab/SKILL.md new file mode 100644 index 00000000..b8ef9b0d --- /dev/null +++ b/.claude/skills/gaussian-cmab/SKILL.md @@ -0,0 +1,134 @@ +--- +name: gaussian-cmab +description: | + Background and conventions for continuous (real-valued) rewards in pybandits: GaussianBayesianNeuralNetwork, + CmabGaussian, the uplift reward R − E0(x), and why the model has one μ head with a learned scalar σ. + Use when: working on GaussianBayesianNeuralNetwork / CmabGaussian, continuous or revenue-like rewards, uplift + rewards against a control baseline, the extra-site likelihood hook, or extending the Gaussian model (offsets, + hurdle models, new likelihoods). +license: MIT +metadata: + author: ronshiff1 + version: "1.0.0" +--- + +# Gaussian cMAB: continuous rewards + +## What exists (pybandits ≥ 8.4.0) + +- **`GaussianBayesianNeuralNetwork`** (`pybandits/model/bnn/network.py`): + - a Bayesian MLP with **one output unit, μ(x)**, and the likelihood `Normal(μ(x), σ)`; + - **σ is one latent per model** (`log σ ~ Normal`, site `noise_log_sigma`), with its posterior stored in the state + and carried across updates like the weights; + - **standardized targets** `(r − reward_loc) / reward_scale`, using the user's values or else fitted on the first + batch, then frozen. `standardize_rewards=False` trains on raw rewards (and then rejects loc / scale). +- **`CmabGaussian`** (`pybandits/cmab.py`): `ClassicBandit` strategy, `_predict_with_proba = True`. Actions are + ranked on μ; `predict` returns `(actions, mu, sigma)`, where σ is for monitoring only. +- **Reward types** (`pybandits/base.py`): `BinaryReward`, `SoftReward`, `ContinuousReward`, `AnyReward`. + - Entry points (`BaseMab.update`, `BaseCmabBernoulli.update`, `ActionsManager.update`) accept `AnyReward`. + - `ActionsManager._check_rewards` enforces what the models support (`BaseModel.supports_continuous_rewards`). + - `delta` + continuous rewards is rejected. + +## Core decisions and why + +1. **One σ, not σ(x).** The Gaussian NLL weights each residual by 1/σ(x)². On zero-inflated or skewed rewards + (revenue), a learned σ(x) lets the model call large rewards noise instead of raising μ, so μ is biased low where + rewards are most variable. In evaluations this under-predicted the level ~3×, and the epistemic sd collapsed + (confidently wrong). With one σ, the μ fit is least squares, whose target is E[r | x] for any noise shape. The + σ(x) head was therefore removed entirely. +2. **σ learned, not fixed.** A fixed σ mis-sets the posterior width, i.e. the exploration. The learned σ converged to + the true residual sd and gave calibrated predictive intervals. Changing σ by hand between updates would make the + stored posterior inconsistent: never do it. Use `noise_sigma` for the initial prior and `decay_factor` for + forgetting. +3. **Standardize.** The weight priors are O(1). Raw revenue-scale targets need weights about 100× outside the + prior, and the model barely moves (with tanh or ReLU alike). Putting the prior on the data scale instead is + equivalent in theory but badly conditioned for Adam. `loc` / `scale` are frozen because the weights only mean + something relative to them. +4. **`reset()` = back to the cold-start state:** weights, noise latent, counters, and `loc` / `scale` back to their + cold-start values. `reward_loc_init` / `reward_scale_init` hold the user-supplied values (None if fitted), so fitted + values are re-fitted on the next update and user values survive. +5. **Reward = uplift `R − E0(x)`** against a control baseline computed outside pybandits. + - `E0` doesn't depend on the action, so the ranking is unchanged and the shared variance is removed (a control + variate). + - Per-action data volumes are too small to learn revenue from scratch; models without a baseline lost about + a third of the explained variance. + - The uplift can be negative, and that is expected. + +## Rejected alternatives (don't re-propose without new evidence) + +- **σ(x) noise head:** see decision 1. +- **A Bernoulli likelihood on a continuous reward rescaled to [0, 1]** (e.g. `(min(Y, c) − k·baseline + a) / b`): + - The mean is consistent, but the Bernoulli variance p(1 − p) overstated the real noise ~128×, so each observation + carried ~1/128 of its information and the model learned only an intercept. + - Weighting the likelihood by ≈ p(1 − p) / var(z) matches a Gaussian only to second order (same minimum and + curvature, but a different loss shape, link and per-row weighting). Use the Gaussian directly. +- **Hurdle × log-normal amount** (`p0 · exp(m + σ²/2)`): + - The best distributional fit (flat PIT), but the back-transform around a baseline biased the level by 10–15%. + - Averaging or drawing `exp(m + σ²/2)` over posterior samples has extreme tails (10⁸× outliers). +- **The hurdle that works, if a conversion vs. amount split is needed:** a Bernoulli with a `logit p0_control` + offset × a Gaussian on the raw amount residual `R − E_buy_control`, for buyers only (called "H3n"). It is as + accurate and calibrated as the single-Gaussian uplift model. It needs per-row offset support (see below). + +## Code map (internals) + +**Base-class likelihood hooks in `BaseBayesianNeuralNetwork`** (the Bernoulli defaults reproduce the original +behavior): +- `_output_dim` (ClassVar), `create_model_params(..., output_dim=None)`; +- `output_distribution(linear_out, extra_sites)`, `_observe_output`, `_postprocess_output(linear_out, extra_samples)`; +- `prepare_rewards(rewards)` → training targets; `_fit(context, rewards)` (the `_update` body); `_is_first_fit`. + +**The extra-site hook, for head-level latents beyond the weights and embeddings:** +- `extra_site_params()` → `{name: BaseLocationScaleArray}`; +- `sample_extra_sites()` (NumPyro, KL-annealed), `sample_extra_numpy(n, rng)` (Thompson); +- `update_extra_params_from_vi()`, `_inflate_extra_params()`. + +It is wired into guide init (`collect_guide_init_arrays`), VI readback, the full-rank ADVI site filter, decay and +sampling. A new likelihood latent (e.g. a Student-t ν) only needs these methods. + +**Meta-model** (`pybandits/meta_model/cmab_meta_model.py`): +- `_check_models` requires one likelihood class across heads; +- `_so_model` stacks each arm's extra sites and gathers them per row; +- `_store_head` reads them back; +- `_prepare_targets` applies each head's `prepare_rewards` to its own rows (each arm keeps its own standardization). + +**Aliases:** `CmabMetaModelGaussian`, `CmabActionsManagerGaussian`. + +## Usage + +```python +mab = CmabGaussian.cold_start( + action_ids={...}, n_features=d, hidden_dim_list=[16], + dist_type="normal", dist_params_init={"mu": 0, "sigma": 0.1}, # narrow prior: default sigma=1 is too slow + update_kwargs={"num_steps": 400, "optimizer_kwargs": {"step_size": 3e-3}}, + # production: reward_loc / reward_scale / noise_sigma from history; decay_factor if drifting +) +actions, mu, sigma = mab.predict(context=X) +mab.update(actions=actions, rewards=(R - E0).tolist(), context=X) +``` + +## Gotchas + +- **Evaluate the mean with a plug-in** (posterior mean of μ) or by averaging draws of μ. Never average + `exp(...)` of log-scale draws. +- **A wide prior** (`sigma=1`) on an unbounded likelihood converges slowly. Narrow it (≈0.1) and raise the step size. +- **Additive residual models can predict E0 + μ < 0** on low-reward rows. That's irrelevant for ranking an uplift, + but floor it when reading the prediction back as revenue. +- **Whale-type under-prediction** inherited from the baseline isn't fixed by the BNN. Fix the baseline (more history). + +## Possible next work + +- **Per-row offset support in `BaseBayesianNeuralNetwork`**, for the H3n hurdle: `logit = offset + f(x)`. + - A working prototype treats the context's last column as the offset: it strips it in `check_context_matrix`, + `emit_submodel` and `forward_pass`, and adds it in `_observe_output` / `_postprocess_output`. It is full-batch + only. + - A proper version must carry the offset through the data plate (minibatching) and the joint meta-model. +- **Gaussian variants** of CC / MO / DP / quantitative; `delta` with continuous rewards; OPE and simulators. + +## Tests + +- `tests/test_gaussian_bnn.py` (the project style: see the `pytest-writer` skill). +- Run with `PYTHONPATH=$PWD python -m pytest tests/test_gaussian_bnn.py -q -n 8`. Force the local checkout: some + environments have another checkout installed as editable. +- Pre-existing flaky failures, unrelated: + - `test_model.py::{test_bnn_svi_nan_loss_raises_error, test_vi_training_options, test_bnn_vi_update_with_categorical_features_updates_embeddings}`; + - `test_quantitative_model.py::test_sample_proba_reproducible_with_same_rng`. diff --git a/docs/continuous_reward_cmab.md b/docs/continuous_reward_cmab.md index fd002723..f74f3c5d 100644 --- a/docs/continuous_reward_cmab.md +++ b/docs/continuous_reward_cmab.md @@ -1,73 +1,78 @@ -# Continuous rewards in [-1, 1] for cMAB - -## Current state - -The library does not take -1/1 rewards. It takes binary {0, 1} rewards: - -- `BinaryReward = conint(ge=0, le=1)` in `pybandits/base.py`. -- Every cMAB model uses a Bernoulli likelihood on a logit: - - `pybandits/model/bnn/network.py` (`NumpyroBernoulli(logits=logit)`) - - `pybandits/meta_model/cmab_meta_model.py` (`NumpyroBernoulli(logits=logit)`) - -So any -1/1 reward is already being mapped to 0/1 before it reaches the library. - -## Option A: rescale to [0, 1] and keep the Bernoulli model (easy) - -Map the reward with `y = (r + 1) / 2`, which puts it in [0, 1]. - -### A1. Binarization trick (no library change) - -For each observation, draw `b ~ Bernoulli((r + 1) / 2)` and pass `b` as the reward. - -- This is Agrawal & Goyal's standard way to extend Thompson Sampling to bounded rewards, and the regret guarantees still hold. -- Cost: extra noise, so the model learns somewhat slower. -- It can be used today with no code change. - -### A2. Pass the fractional target directly - -This fits a logistic regression with targets that are not 0/1. The Bernoulli log-likelihood with logits is BCE-with-logits, which is valid for any `y` in [0, 1], and numpyro does not check the support by default. - -Changes needed: - -- Relax `BinaryReward` to a float in [0, 1]. The type appears in many signatures, but the change is mechanical. -- Check `_calibrate_output_bias`. It should still work because it uses the mean reward. -- Update the simulators and tests. - -Behaviour: - -- The mean is estimated correctly: `E[y | x] = sigmoid(f)`. -- The uncertainty is not. The model assumes variance `p(1 - p)`, which is the largest possible variance for a variable in [0, 1]. The posterior is therefore too wide and the bandit **over-explores**. This errs on the safe side. - -Estimated effort: 1–2 days, including tests. - -## Option B: a proper continuous likelihood (harder) - -Use a Gaussian, Beta or truncated-normal likelihood, with an identity or tanh output and a learned noise `σ`. - -The likelihood itself is simple. The work is in what is built around it: - -- `sample_proba` returns a `Probability` (`Float01`) everywhere, and the strategies and meta-model all rely on that. -- `_calibrate_output_bias` uses `logit(rate)`. -- The CostControl `subsidy_factor` is a **relative** threshold ("within X of the best"). Relative thresholds break once the expected reward can be negative, so they must be rescaled or redefined. -- The OPE estimators, the simulators and the whole test suite need changes. - -Estimated effort: 1–2+ weeks. It touches every cMAB / MO / CC variant. - -## Do we need to understand the reward distribution first? - -**For Option A: mostly no.** Arm selection depends only on the mean, and the Bernoulli model's uncertainty is conservative whatever the shape. A histogram per arm is still worth plotting, because some shapes change the recommendation: - -| Reward shape | Implication | -|---|---| -| Mostly at ±1 with a few values in between | The Bernoulli model is almost exact. Use Option A. | -| Spike at 0 (e.g. "no response" coded as 0) | Rewards are zero-inflated. Consider modelling "any response" and "sign/size given a response" separately. | -| Concentrated in the middle (e.g. 0 ± 0.2) | Option A's posterior is much too wide and exploration is slow. Option B with a learned `σ` pays off here. | -| Heteroscedastic (noise varies with context) | Only Option B can model this. | - -**For Option B: yes.** The likelihood family and a sensible prior on `σ` have to be chosen from the data. - -## Recommendation - -1. Start with **A2** (fractional targets), or **A1** for zero code changes. -2. Back-test it on logged data with the existing offline policy evaluator. -3. Build **Option B** only if the histogram shows rewards bunched in the middle and A2 explores too slowly. +# Continuous rewards for cMAB + +pybandits supports three kinds of reward for contextual bandits: + +| reward | type | model | bandit | +|---|---|---|---| +| binary 0 / 1 | `BinaryReward` | `BayesianNeuralNetwork` (Bernoulli) | `CmabBernoulli` (and the CC / MO / DP variants) | +| soft label in [0, 1] | `SoftReward` (`use_soft_rewards=True`) | `BayesianNeuralNetwork` (Bernoulli on fractional targets) | the same | +| real-valued | `ContinuousReward` | `GaussianBayesianNeuralNetwork` | `CmabGaussian` | + +The bandit entry points accept `AnyReward`, and the actions manager checks each reward against what its models +support (`ActionsManager._check_rewards`). + +## The Gaussian model + +Each action is a Bayesian MLP with **one output unit, μ(x)**, and the reward is modelled as `Normal(μ(x), σ)`: + +- **σ is a single noise std per action,** not a function of the context. It is a latent (`log σ ~ Normal`) 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` (reward units) if given, otherwise on the std of the first update batch. +- **Training is on standardized targets** `(r − reward_loc) / reward_scale`, so the default O(1) weight priors fit + rewards on any scale. `reward_loc` / `reward_scale` are the values given at cold start, otherwise fitted on the first + batch and then frozen. `standardize_rewards=False` trains on raw rewards. +- **Thompson sampling** draws the weights from the posterior and ranks actions on μ. `predict` returns μ (in the + "probabilities" slot) and σ (for monitoring) per action and row. +- **`reset()`** returns the model to its cold-start state, including the standardization (re-fitted on the next + update unless it was given at cold start). + +### Why one σ and not σ(x) + +With a single noise level the fit of μ is a (Bayesian) least-squares fit, whose target is `E[reward | x]` whatever +the shape of the noise. A context-dependent σ(x) weights each residual by 1/σ(x)². On zero-inflated or skewed rewards +(e.g. revenue), the model can then explain large rewards as noise instead of raising μ, which biases μ low exactly +where the rewards are most variable. That bias is what Thompson sampling would rank on. + +### Recommended use: uplift against a control baseline + +For rewards with a lot of player-to-player variance (e.g. revenue), train on the **uplift** `r = R − E0(x)`, where +`E0(x)` is the expected reward under a control policy, from a separate baseline model. `E0` 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. `r` can be negative. + +```python +from pybandits.cmab import CmabGaussian + +mab = CmabGaussian.cold_start( + action_ids={"a1", "a2", "a3"}, + n_features=n_features, + hidden_dim_list=[16], + dist_type="normal", + dist_params_init={"mu": 0, "sigma": 0.1}, # a narrow prior fits much faster than the default sigma=1 + update_kwargs={"num_steps": 400, "optimizer_kwargs": {"step_size": 3e-3}}, + # optionally, from historical data rather than the first batch: + # reward_loc=..., reward_scale=..., noise_sigma=..., + # decay_factor=..., # if the reward process drifts +) +actions, mu, sigma = mab.predict(context=X) +mab.update(actions=actions, rewards=(R - E0).tolist(), context=X) +``` + +A shared backbone (`backbone_hidden_dims=...`) works as for `CmabBernoulli`; each action keeps its own σ and +standardization. + +### Alternatives that were evaluated and rejected + +- **σ(x) noise head:** biases μ on skewed rewards (see above). +- **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. +- **Hurdle × log-normal amount:** a good fit of the distribution's shape, but the back-transform `exp(m + σ²/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 × Gaussian amount (on the raw scale) is the better hurdle. + +### Not supported (yet) + +- Gaussian variants of the cost-control, multi-objective, dynamic-pricing and quantitative bandits. +- The adaptive window (`delta`) with continuous rewards. +- Offline policy evaluation and the simulators with continuous rewards. diff --git a/pybandits/cmab.py b/pybandits/cmab.py index 2604fada..1dd15d2f 100644 --- a/pybandits/cmab.py +++ b/pybandits/cmab.py @@ -311,12 +311,15 @@ 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 mean head and - a (heteroscedastic) noise head, 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. + 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 predicted reward - noise std in place of the weighted sums. + ``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 ---------- diff --git a/pybandits/meta_model/cmab_meta_model.py b/pybandits/meta_model/cmab_meta_model.py index c8b45987..30eaa318 100644 --- a/pybandits/meta_model/cmab_meta_model.py +++ b/pybandits/meta_model/cmab_meta_model.py @@ -239,11 +239,7 @@ def _check_models(self) -> "CmabMetaModel": 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) - same_likelihood = type(head_bnn).output_distribution is type(first_bnn).output_distribution and all( - getattr(head_bnn, attr, None) == getattr(first_bnn, attr, None) - for attr in ("sigma_min", "homoscedastic") - ) - if not same_likelihood: + 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. @@ -693,7 +689,7 @@ def _so_model( 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. a homoscedastic noise std), stacked to (num_arms, *site_shape). + # 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 diff --git a/pybandits/model/bnn/network.py b/pybandits/model/bnn/network.py index 278a6b40..74704909 100644 --- a/pybandits/model/bnn/network.py +++ b/pybandits/model/bnn/network.py @@ -774,7 +774,7 @@ 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 homoscedastic Gaussian head's + (current posterior → next prior). None for the Bernoulli head; e.g. the Gaussian head's log noise std. """ return {} @@ -868,11 +868,6 @@ def prepare_rewards(self, rewards: List[Reward]) -> np.ndarray: """ return np.asarray(rewards, dtype=float) - @classmethod - def _resolve_output_dim(cls, **kwargs) -> PositiveInt: - """Output-layer width for a cold start with these extra constructor kwargs (the class default here).""" - return cls._output_dim - @property def _is_first_fit(self) -> bool: """Whether no reward has been recorded yet (the counters still hold the prior pseudo-counts).""" @@ -1630,7 +1625,6 @@ def cold_start( use_layerwise_scaling=use_layerwise_scaling, dist_class=dist_class, bias_std=bias_std, - output_dim=cls._resolve_output_dim(**kwargs), **dist_params_init, ) return cls( @@ -1748,56 +1742,50 @@ 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)``. Thompson sampling draws the network weights 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 only sets how much each observation is trusted during training. - Two noise models: + 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. - - **heteroscedastic** (default): the output layer has two units that share the hidden layers, a mean head - ``mu(x)`` and a noise head ``sigma(x) = softplus(.) + sigma_min``. - - **homoscedastic** (``homoscedastic=True``): the output layer has one unit ``mu(x)``, and ``sigma`` is a - single latent of the model, ``log sigma ~ Normal``, whose posterior is stored in ``noise_log_sigma`` and - carried across updates like the weights. - - Use the homoscedastic model when the reward noise is skewed or zero-inflated (e.g. revenue). The Gaussian - likelihood weights each residual by ``1 / sigma(x)^2``, so a learned ``sigma(x)`` lets the model explain - the large rewards as noise instead of moving ``mu(x)``: ``mu`` then sits below ``E[reward | x]`` where the - rewards are most variable. With a single ``sigma`` the fit of ``mu`` is a (Bayesian) least-squares fit, - whose target is ``E[reward | x]`` whatever the shape of the noise. + 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 standardized targets ``(reward - reward_loc) / reward_scale`` so the default - ``O(1)`` weight priors fit rewards on any scale (e.g. revenue); ``sample_proba`` maps the outputs - back to the reward scale. + ``O(1)`` weight priors fit rewards on any scale; ``sample_proba`` maps ``mu`` and ``sigma`` back to the reward + scale. Parameters ---------- - sigma_min : PositiveFloat - Floor on the noise std, in standardized units. Keeps the likelihood (and its gradients) bounded - while the early, wide posterior samples push the noise head down. Default is 0.05. standardize_rewards : bool - Whether to train on standardized targets. When ``reward_loc`` / ``reward_scale`` are not given they - are fitted (mean / std) on the first update batch and kept fixed afterwards, including across - ``reset()`` (the prior lives in standardized space). Default is True. + Whether to train on standardized targets. When True, ``reward_loc`` / ``reward_scale`` are used: the values + given at cold start, or else fitted (mean / std) on the first update batch and kept fixed afterwards. When + False, they are neither fitted nor used (the model trains on raw rewards). Default is True. reward_loc : Optional[float] Location used to standardize rewards. None to fit it on the first update. reward_scale : Optional[PositiveFloat] Scale used to standardize rewards. None to fit it on the first update. + reward_loc_init : Optional[float] + The cold-start value of ``reward_loc`` (None if it is fitted), restored by ``reset()``. Set automatically + from ``reward_loc`` at construction. + reward_scale_init : Optional[PositiveFloat] + The cold-start value of ``reward_scale`` (None if it is fitted), restored by ``reset()``. Set automatically + from ``reward_scale`` at construction. n_observations : NonNegativeInt Number of rewards observed. reward_sum : float Sum of the rewards observed. - homoscedastic : bool - Whether the noise std is one latent of the model rather than a function of the context. Default is False. noise_sigma : Optional[PositiveFloat] - Homoscedastic model only: the 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. + 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 - Homoscedastic model only: std of the initial ``log sigma`` prior (0.5 allows about a factor 1.65 at one - std). Default is 0.5. + 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] - Homoscedastic model only: the current posterior of ``log sigma`` (in standardized units), shape (1,). None - until the first update initializes it from ``noise_sigma`` or the batch. + The current posterior of ``log sigma`` (in standardized units), shape (1,). None until the first update + initializes it from ``noise_sigma`` or the batch. Notes ----- @@ -1807,9 +1795,9 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): 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. - - ``sigma_min`` only applies to the heteroscedastic noise head. - - ``reset()`` clears ``noise_log_sigma``: it is re-initialized on the next update, like the weights return to - their cold-start prior. The reward standardization is kept. + - ``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 ``reward_loc`` / ``reward_scale`` to their cold-start values (re-fitted on the + next update if they were fitted). Examples -------- @@ -1822,54 +1810,66 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): ... ) """ - sigma_min: PositiveFloat = 0.05 standardize_rewards: bool = True reward_loc: Optional[float] = None reward_scale: Optional[PositiveFloat] = None + reward_loc_init: Optional[float] = None + reward_scale_init: Optional[PositiveFloat] = None n_observations: NonNegativeInt = 0 reward_sum: float = 0.0 - homoscedastic: bool = False 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, raw sigma]; the homoscedastic model has [mu] only - _output_dim: ClassVar[PositiveInt] = 2 + # [mu]; the noise std is the scalar latent ``noise_log_sigma`` + _output_dim: ClassVar[PositiveInt] = 1 _noise_var_name: ClassVar[str] = "noise_log_sigma" _transfer_learned_keys: ClassVar[Tuple[str, ...]] = ( "reward_loc", "reward_scale", + "reward_loc_init", + "reward_scale_init", "n_observations", "reward_sum", "noise_log_sigma", ) - _transfer_structural_keys: ClassVar[Tuple[str, ...]] = ("homoscedastic",) + + @model_validator(mode="before") + @classmethod + def set_standardization_init(cls, data: Any) -> Any: + """Record user-supplied ``reward_loc`` / ``reward_scale`` as their cold-start values (restored by reset).""" + if isinstance(data, dict): + for name in ("reward_loc", "reward_scale"): + if f"{name}_init" not in data and data.get(name) is not None: + data = {**data, f"{name}_init": data[name]} + return data @model_validator(mode="after") def validate_reward_standardization(self) -> "GaussianBayesianNeuralNetwork": - if (self.reward_loc is None) != (self.reward_scale is None): - raise ValueError("reward_loc and reward_scale must be either both set or both None.") + for loc, scale in [(self.reward_loc, self.reward_scale), (self.reward_loc_init, self.reward_scale_init)]: + if (loc is None) != (scale is None): + raise ValueError("reward_loc and reward_scale must be either both set or both None.") + if not self.standardize_rewards and loc is not None: + raise ValueError("reward_loc / reward_scale require standardize_rewards=True.") return self @model_validator(mode="after") - def validate_noise_model(self) -> "GaussianBayesianNeuralNetwork": - expected = self._resolve_output_dim(homoscedastic=self.homoscedastic) + def validate_output_layer(self) -> "GaussianBayesianNeuralNetwork": output_dim = self.model_params.bnn_layer_params[-1].weight.shape[-1] - if output_dim != expected: - raise ValueError( - f"The output layer has {output_dim} units, but homoscedastic={self.homoscedastic} needs {expected}." - ) - if not self.homoscedastic and (self.noise_log_sigma is not None or self.noise_sigma is not None): - raise ValueError("noise_sigma / noise_log_sigma require homoscedastic=True.") + if output_dim != self._output_dim: + raise ValueError(f"The output layer has {output_dim} units, expected {self._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 - @classmethod - def _resolve_output_dim(cls, **kwargs) -> PositiveInt: - return 1 if kwargs.get("homoscedastic", False) else cls._output_dim + @property + def _loc_scale(self) -> Tuple[float, float]: + """``(loc, scale)`` mapping standardized targets back to rewards (identity until fitted, or if disabled).""" + if not self.standardize_rewards or self.reward_loc is None: + return 0.0, 1.0 + return self.reward_loc, self.reward_scale def _initial_noise_log_sigma(self, targets: Optional[np.ndarray] = None) -> NormalArray: """The cold-start ``log sigma`` prior, in standardized units. @@ -1886,14 +1886,10 @@ def _initial_noise_log_sigma(self, targets: Optional[np.ndarray] = None) -> Norm 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 homoscedastic noise latent ``log sigma`` (none for the heteroscedastic model).""" - if not self.homoscedastic: - return {} + """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: - if not self.homoscedastic: - return 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()) @@ -1910,38 +1906,26 @@ def _noise_log_sigma_of(extra: Dict[str, Any], name: str) -> Any: log_sigma = extra[name] return log_sigma.reshape(log_sigma.shape[:-1]) - @property - def _loc_scale(self) -> Tuple[float, float]: - """``(loc, scale)`` mapping standardized targets back to rewards (identity until fitted).""" - if not self.standardize_rewards or self.reward_loc is None: - return 0.0, 1.0 - return self.reward_loc, self.reward_scale - def output_distribution( self, linear_out: jax.Array, extra_sites: Optional[Dict[str, jax.Array]] = None ) -> NumpyroDistribution: - """``Normal(mu, sigma)`` on the standardized targets. + """``Normal(mu, exp(log sigma))`` on the standardized targets. Parameters ---------- - linear_out : jax.Array of shape (batch, 2) or (batch, 1) - Raw output of the final layer: ``[mu, raw_sigma]`` (heteroscedastic) or ``[mu]`` (homoscedastic). + linear_out : jax.Array of shape (batch, 1) + Raw output of the final layer: ``[mu]``. extra_sites : Optional[Dict[str, jax.Array]] - Homoscedastic model: 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). + 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, with ``sigma = exp(log sigma)`` (homoscedastic) or - ``softplus(raw_sigma) + sigma_min`` (heteroscedastic). + The Gaussian reward likelihood. """ - mu = linear_out[..., 0] - if self.homoscedastic: - sigma = jnp.exp(self._noise_log_sigma_of(extra_sites, self._noise_var_name)) - else: - sigma = jax.nn.softplus(linear_out[..., 1]) + self.sigma_min - return NumpyroNormal(loc=mu, scale=sigma) + 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 @@ -1956,37 +1940,35 @@ def _postprocess_output( Parameters ---------- - linear_out : np.ndarray of shape (n_samples, 2) or (n_samples, 1) + 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]] - Homoscedastic model: per-row draws of ``log sigma``, shape (n_samples, 1). None uses its posterior mean. + 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. + ``(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 self.homoscedastic: - 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 - else: - sigma = (np.logaddexp(0.0, linear_out[..., 1]) + self.sigma_min) * scale + 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 predicted noise std. + 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 sampling. + Numpy random generator for weight/embedding/noise sampling. Returns ------- @@ -1996,11 +1978,11 @@ def sample_proba(self, context: np.ndarray, rng: np.random.Generator) -> List[Ga return super().sample_proba(context=context, rng=rng) def prepare_rewards(self, rewards: List[ContinuousReward]) -> np.ndarray: - """Standardize rewards into training targets, fitting ``reward_loc`` / ``reward_scale`` on first use. + """Standardize rewards into training targets, fitting the standardization and the noise prior on first use. - The scale falls back to ``max(|loc|, 1)`` when the first batch has (near) zero spread, e.g. a - single observation. The homoscedastic model also initializes its ``log sigma`` prior here, on first - use (see ``noise_sigma``). + With ``standardize_rewards``, ``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 ---------- @@ -2022,7 +2004,7 @@ def prepare_rewards(self, rewards: List[ContinuousReward]) -> np.ndarray: 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.homoscedastic and self.noise_log_sigma is None and len(y) > 0: + if self.noise_log_sigma is None and len(y) > 0: self.noise_log_sigma = self._initial_noise_log_sigma(targets) return targets @@ -2079,11 +2061,10 @@ def mean(self) -> float: 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 biases to the first batch's target mean and std on the first update call. + """Set the output-layer bias to the first batch's (standardized) target mean on the first update call. - The mean head's bias gets the mean of the (standardized) targets and the noise head's bias gets - ``softplus^-1(std - sigma_min)``, so the cold-start predictive matches the observed reward spread. - With ``standardize_rewards`` (fitted on that same batch) this is ``(0, softplus^-1(1 - sigma_min))``. + With ``standardize_rewards`` fitted on that same batch this is 0. The noise prior is initialized by + :meth:`prepare_rewards` instead. Parameters ---------- @@ -2094,22 +2075,18 @@ def _calibrate_output_bias(self, rewards: List[ContinuousReward]) -> None: return y = self.prepare_rewards(rewards) output_layer = self.model_params.bnn_layer_params[-1] - bias_mu = [float(y.mean())] - if not self.homoscedastic: # the homoscedastic noise is initialized by prepare_rewards instead - std = float(y.std()) - std = std if std > self._numerical_eps else 1.0 # no spread to calibrate on (e.g. a single reward) - raw_sigma = max(std - self.sigma_min, self._numerical_eps) - bias_mu.append(float(np.log(np.expm1(raw_sigma)))) - new_bias = output_layer.bias.with_dist_parameters(mu=bias_mu) + 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): - """Reset the network, the noise latent and the reward counters; the reward standardization is kept.""" + """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 + self.reward_loc = self.reward_loc_init + self.reward_scale = self.reward_scale_init class BaseBayesianNeuralNetworkMO(ModelMO, ABC): 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_gaussian_bnn.py b/tests/test_gaussian_bnn.py index bc9d7336..f7c0673a 100644 --- a/tests/test_gaussian_bnn.py +++ b/tests/test_gaussian_bnn.py @@ -36,7 +36,7 @@ # --- test configuration constants --- N_FEATURES = 3 HIDDEN_DIMS = [8] -OUTPUT_DIM = 2 +OUTPUT_DIM = 1 RANDOM_SEED = 0 N_TRAIN = 2000 N_POSTERIOR_SAMPLES = 200 @@ -47,42 +47,46 @@ "update_kwargs": {"num_steps": 1000, "optimizer_kwargs": {"step_size": 3e-3}}, } SHORT_FIT_KWARGS: Dict[str, Any] = {"update_kwargs": {"num_steps": 5}} -# Synthetic reward: REWARD_OFFSET + REWARD_SLOPE * x0 + noise, noise std LOW / HIGH for x1 <= 0 / x1 > 0. +# Linear synthetic reward: REWARD_OFFSET + REWARD_SLOPE * x0 + Normal(0, NOISE_STD). REWARD_OFFSET = 100.0 REWARD_SLOPE = 30.0 -NOISE_STD_LOW = 5.0 -NOISE_STD_HIGH = 15.0 -# Relative tolerance on the recovered mean / noise std at the probe points. +NOISE_STD = 10.0 MEAN_REL_TOL = 0.1 -SIGMA_REL_TOL = 0.5 +SIGMA_REL_TOL = 0.3 +# Probe contexts on both sides of x0. +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. +BUY_PROB = 0.15 +SKEW_OFFSET = 2.0 +SKEW_SLOPE = 1.5 +SKEW_LOG_SD = 1.0 +LEVEL_REL_TOL = 0.15 +ROUND_TRIP_TOL = 1e-6 +MAX_ABS_REWARD = 1e6 +MAX_REWARDS = 50 +NOISE_SIGMA = 3.0 +DECAY_FACTOR = 0.5 +USER_LOC = 5.0 +USER_SCALE = 2.0 +SECOND_BATCH_SHIFT = 50.0 ACTION_IDS = {"a", "b", "c"} ARM_MEANS = {"a": -1.0, "b": 0.0, "c": 2.0} BEST_ARM = "c" N_ROUNDS = 4 N_PER_ROUND = 200 MIN_BEST_ARM_SHARE = 0.8 -# Probe contexts: (x0, x1) in {-1, 1}^2 hit both mean levels and both noise regions. -PROBE_CONTEXTS = [[1.0, -1.0, 0.0], [1.0, 1.0, 0.0], [-1.0, -1.0, 0.0], [-1.0, 1.0, 0.0]] -ROUND_TRIP_TOL = 1e-6 -MAX_ABS_REWARD = 1e6 -MAX_REWARDS = 50 BACKBONE_KWARGS: Dict[str, Any] = {"backbone_hidden_dims": [8], "backbone_embedding_dim": 3} MINIBATCH_SIZE = 64 BAD_REWARDS = [float("nan"), float("inf"), -float("inf")] OUT_OF_RANGE_SOFT_REWARD = 2.5 DELTA = 0.1 -OTHER_SIGMA_MIN = 0.2 EPSILON = 0.1 -HOMOSCEDASTIC_OUTPUT_DIM = 1 -NOISE_SIGMA = 3.0 -DECAY_FACTOR = 0.5 -# 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. The noise grows with the mean, which a learned sigma(x) uses to shrink mu. -BUY_PROB = 0.15 -SKEW_OFFSET = 2.0 -SKEW_SLOPE = 1.5 -SKEW_LOG_SD = 1.0 -LEVEL_REL_TOL = 0.15 + + +def _linear_rewards(x: np.ndarray, rng: np.random.Generator) -> np.ndarray: + """Linear-in-x0 rewards with constant Gaussian noise.""" + return REWARD_OFFSET + REWARD_SLOPE * x[:, 0] + rng.normal(size=len(x)) * NOISE_STD def _skewed_rewards(x: np.ndarray, rng: np.random.Generator) -> np.ndarray: @@ -92,10 +96,10 @@ def _skewed_rewards(x: np.ndarray, rng: np.random.Generator) -> np.ndarray: return (rng.uniform(size=len(x)) < BUY_PROB) * amount -def _heteroscedastic_rewards(x: np.ndarray, rng: np.random.Generator) -> np.ndarray: - """Linear-in-x0 rewards whose noise std switches with the sign of x1.""" - noise_std = np.where(x[:, 1] > 0, NOISE_STD_HIGH, NOISE_STD_LOW) - return REWARD_OFFSET + REWARD_SLOPE * x[:, 0] + rng.normal(size=len(x)) * noise_std +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") @@ -114,11 +118,11 @@ def _factory(**kwargs: Any) -> GaussianBayesianNeuralNetwork: def fitted_gaussian_bnn( make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], ) -> GaussianBayesianNeuralNetwork: - """A Gaussian BNN fitted on heteroscedastic synthetic rewards.""" + """A Gaussian BNN fitted on linear synthetic rewards with constant noise.""" data_rng = np.random.default_rng(RANDOM_SEED) bnn = make_gaussian_bnn(**FAST_FIT_KWARGS) x = data_rng.normal(size=(N_TRAIN, N_FEATURES)) - bnn.update(context=x, rewards=_heteroscedastic_rewards(x, data_rng).tolist()) + bnn.update(context=x, rewards=_linear_rewards(x, data_rng).tolist()) return bnn @@ -139,18 +143,26 @@ def _factory(**kwargs: Any) -> CmabGaussian: class TestGaussianBayesianNeuralNetwork: - """Tests for the standalone Gaussian BNN.""" + """Tests for the standalone Gaussian BNN (one mean unit, a learned scalar noise latent).""" - def test_output_layer_has_mean_and_noise_units( + def test_output_layer_has_only_the_mean_unit( self, make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], output_dim: int = OUTPUT_DIM, ) -> None: - """The output layer has one unit for the mean and one for the noise std.""" + """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, @@ -159,41 +171,31 @@ def test_sample_proba_returns_mean_and_positive_noise( n_samples: int = N_TRAIN, n_features: int = N_FEATURES, ) -> None: - """Each sample is a (mu, sigma) pair with sigma above the configured floor.""" + """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 - sigmas = np.array([sigma for _, sigma in samples]) - assert np.all(sigmas >= bnn.sigma_min) + assert all(sigma > 0 for _, sigma in samples) - def test_fit_recovers_mean_and_heteroscedastic_noise( + def test_fit_recovers_mean_trend_and_noise_level( self, fitted_gaussian_bnn: GaussianBayesianNeuralNetwork, rng: np.random.Generator, n_posterior_samples: int = N_POSTERIOR_SAMPLES, offset: float = REWARD_OFFSET, slope: float = REWARD_SLOPE, - noise_low: float = NOISE_STD_LOW, - noise_high: float = NOISE_STD_HIGH, + noise_std: float = NOISE_STD, mean_rel_tol: float = MEAN_REL_TOL, sigma_rel_tol: float = SIGMA_REL_TOL, probe_contexts: List[List[float]] = PROBE_CONTEXTS, ) -> None: - """The posterior mean tracks the linear trend and the noise head separates low / high noise regions.""" + """The posterior mean tracks the linear trend and the learned sigma matches the noise std.""" probes = np.array(probe_contexts) - true_mu = offset + slope * probes[:, 0] - true_sigma = np.where(probes[:, 1] > 0, noise_high, noise_low) draws = np.array( [fitted_gaussian_bnn.sample_proba(context=probes, rng=rng) for _ in range(n_posterior_samples)] ) - mu, sigma = draws[..., 0].mean(axis=0), draws[..., 1].mean(axis=0) - np.testing.assert_allclose(mu, true_mu, rtol=mean_rel_tol) - high_noise = probes[:, 1] > 0 - np.testing.assert_allclose(sigma[high_noise], true_sigma[high_noise], rtol=sigma_rel_tol) - # VI pulls the noise head towards the pooled noise level, so the low-noise region is only required to - # sit clearly below it (a homoscedastic fit would put every probe at the pooled std). - pooled_std = np.sqrt((noise_low**2 + noise_high**2) / 2) - assert np.all(sigma[~high_noise] < pooled_std) + 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, @@ -236,6 +238,90 @@ def test_single_reward_standardization_has_unit_or_larger_scale(self, reward: fl bnn.prepare_rewards([reward]) assert bnn.reward_scale == max(abs(reward), 1.0) + @given( + rewards=st.lists( + st.floats(min_value=-MAX_ABS_REWARD, max_value=MAX_ABS_REWARD, allow_nan=False), + min_size=2, + max_size=MAX_REWARDS, + unique=True, + ), + n_features=st.just(N_FEATURES), + tol=st.just(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) + 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 = NOISE_SIGMA, + n_samples: int = N_TRAIN, + tol: float = 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 = N_TRAIN, + n_features: int = N_FEATURES, + short_fit_kwargs: Dict[str, Any] = 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 = N_TRAIN, + n_features: int = N_FEATURES, + short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, + rel_tol: float = 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 = N_TRAIN, + n_features: int = N_FEATURES, + n_posterior_samples: int = N_POSTERIOR_SAMPLES, + rel_tol: float = LEVEL_REL_TOL, + fast_fit_kwargs: Dict[str, Any] = 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(RANDOM_SEED) + x = data_rng.normal(size=(n_samples, n_features)) + y = _skewed_rewards(x, data_rng) + bnn = make_gaussian_bnn(**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", BAD_REWARDS) def test_update_rejects_non_finite_rewards( self, @@ -248,7 +334,7 @@ def test_update_rejects_non_finite_rewards( with pytest.raises(ValidationError): bnn.update(context=np.zeros((1, n_features)), rewards=[bad_reward]) - @pytest.mark.parametrize("reward_loc, reward_scale", [(1.0, None), (None, 1.0)]) + @pytest.mark.parametrize("reward_loc, reward_scale", [(USER_LOC, None), (None, USER_SCALE)]) def test_rejects_half_specified_standardization( self, make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], @@ -259,45 +345,129 @@ def test_rejects_half_specified_standardization( with pytest.raises(ValidationError): make_gaussian_bnn(reward_loc=reward_loc, reward_scale=reward_scale) - def test_calibrate_output_bias_sets_mean_and_noise_bias( + def test_rejects_standardization_values_when_disabled( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + reward_loc: float = USER_LOC, + reward_scale: float = USER_SCALE, + ) -> None: + """With standardize_rewards=False, reward_loc / reward_scale would be ignored, so they are rejected.""" + with pytest.raises(ValidationError): + make_gaussian_bnn(standardize_rewards=False, reward_loc=reward_loc, reward_scale=reward_scale) + + def test_unstandardized_model_trains_and_predicts_on_raw_rewards( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = N_TRAIN, + n_features: int = N_FEATURES, + noise_sigma: float = NOISE_SIGMA, + short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, + tol: float = ROUND_TRIP_TOL, + ) -> None: + """standardize_rewards=False: no loc / scale are fitted, targets are raw and noise_sigma is in raw units.""" + bnn = make_gaussian_bnn(standardize_rewards=False, noise_sigma=noise_sigma, **short_fit_kwargs) + rewards = rng.normal(size=n_samples) + np.testing.assert_array_equal(bnn.prepare_rewards(rewards.tolist()), rewards) + assert bnn.reward_loc is None and bnn.reward_scale is None + 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 = N_TRAIN, tol: float = ROUND_TRIP_TOL, ) -> None: - """Calibration puts the output biases at the standardized target mean and inverse-softplus std.""" + """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) - targets = bnn.prepare_rewards(rewards) - mu_bias, raw_sigma_bias = bnn.model_params.bnn_layer_params[-1].bias.params["mu"] + (mu_bias,) = bnn.model_params.bnn_layer_params[-1].bias.params["mu"] assert bnn.bias_calibrated - assert mu_bias == pytest.approx(targets.mean(), abs=tol) - assert np.logaddexp(0, raw_sigma_bias) + bnn.sigma_min == pytest.approx(targets.std(), rel=tol) + assert mu_bias == pytest.approx(bnn.prepare_rewards(rewards).mean(), abs=tol) - def test_reset_clears_counters_and_keeps_standardization( + def test_reset_refits_fitted_standardization( self, make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], rng: np.random.Generator, n_samples: int = N_TRAIN, n_features: int = N_FEATURES, + shift: float = SECOND_BATCH_SHIFT, short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, ) -> None: - """reset() zeroes the reward counters and the network, but keeps the fitted standardization.""" + """Fitted loc / scale are cleared by reset() and re-fitted on the next batch, with the noise prior.""" bnn = make_gaussian_bnn(**short_fit_kwargs) bnn.update(context=rng.normal(size=(n_samples, n_features)), rewards=rng.normal(size=n_samples).tolist()) - loc, scale = bnn.reward_loc, bnn.reward_scale + 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.reward_loc, bnn.reward_scale) == (loc, scale) 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 = USER_LOC, + reward_scale: float = USER_SCALE, + n_samples: int = N_TRAIN, + n_features: int = N_FEATURES, + short_fit_kwargs: Dict[str, Any] = 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 = N_TRAIN, + n_features: int = N_FEATURES, + short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, + ) -> None: + """A fitted standardization stays 'fitted' through serialization: reset() still clears 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()) + restored = GaussianBayesianNeuralNetwork.model_validate_json(bnn.model_dump_json()) + assert restored.reward_loc == bnn.reward_loc and restored.reward_loc_init is None + 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 = DECAY_FACTOR, + n_samples: int = N_TRAIN, + n_features: int = N_FEATURES, + short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, + tol: float = 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 and counters.""" + """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 @@ -316,15 +486,6 @@ def test_rejects_mixed_likelihoods( with pytest.raises(AttributeError): CmabMetaModel(actions={"a": make_gaussian_bnn(), "b": bernoulli}) - def test_rejects_different_sigma_min( - self, - make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], - sigma_min: float = OTHER_SIGMA_MIN, - ) -> None: - """Gaussian heads must share sigma_min, which the joint likelihood takes from one head.""" - with pytest.raises(AttributeError): - CmabMetaModel(actions={"a": make_gaussian_bnn(), "b": make_gaussian_bnn(sigma_min=sigma_min)}) - class TestCmabGaussian: """Tests for the CmabGaussian bandit.""" @@ -337,7 +498,7 @@ class TestCmabGaussian: ], ids=["no_backbone", "backbone_minibatch"], ) - def test_learns_best_arm( + def test_learns_best_arm_and_per_arm_noise( self, make_cmab_gaussian: Callable[..., CmabGaussian], extra_kwargs: Dict[str, Any], @@ -353,12 +514,12 @@ def test_learns_best_arm( mab = make_cmab_gaussian(**extra_kwargs) for _ in range(n_rounds): context = rng.normal(size=(n_per_round, n_features)) - actions, mus, sigmas = mab.predict(context=context) - rewards = [arm_means[a] + rng.normal() for a in actions] - mab.update(actions=actions, rewards=rewards, context=context) - actions, mus, sigmas = mab.predict(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, @@ -415,208 +576,3 @@ def test_bernoulli_bandit_still_rejects_out_of_range_rewards( ) with pytest.raises(ValueError): mab.update(actions=[next(iter(mab.actions))], rewards=[reward], context=np.zeros((1, n_features))) - - -class TestHomoscedasticGaussian: - """Tests for the homoscedastic 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 = HOMOSCEDASTIC_OUTPUT_DIM, - ) -> None: - """The output layer has a single unit and the noise latent is unset before the first update.""" - bnn = make_gaussian_bnn(homoscedastic=True) - assert bnn.model_params.bnn_layer_params[-1].weight.shape[-1] == output_dim - assert bnn.noise_log_sigma is None - - @given( - rewards=st.lists( - st.floats(min_value=-MAX_ABS_REWARD, max_value=MAX_ABS_REWARD, allow_nan=False), - min_size=2, - max_size=MAX_REWARDS, - unique=True, - ), - n_features=st.just(N_FEATURES), - tol=st.just(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, homoscedastic=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 = NOISE_SIGMA, - n_samples: int = N_TRAIN, - tol: float = ROUND_TRIP_TOL, - ) -> None: - """noise_sigma (reward units) sets the prior once the reward scale is known.""" - bnn = make_gaussian_bnn(homoscedastic=True, noise_sigma=noise_sigma) - bnn.prepare_rewards(rng.normal(size=n_samples).tolist()) - prior_sigma = np.exp(bnn.noise_log_sigma.params["mu"][0]) * bnn.reward_scale - assert prior_sigma == 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 = N_TRAIN, - n_features: int = N_FEATURES, - short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, - ) -> None: - """Each update replaces the stored noise posterior, and the next update starts from it.""" - bnn = make_gaussian_bnn(homoscedastic=True, **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 = N_TRAIN, - n_features: int = N_FEATURES, - short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, - ) -> None: - """The predicted noise std doesn't depend on the context: its spread comes only from the log-sigma posterior.""" - bnn = make_gaussian_bnn(homoscedastic=True, **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)]) - log_sigma_sd = bnn.noise_log_sigma.params["sigma"][0] - assert np.all(sigmas > 0) - assert np.log(sigmas).std() == pytest.approx(log_sigma_sd, rel=SIGMA_REL_TOL) - - def test_mean_is_level_calibrated_on_skewed_rewards( - self, - make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], - rng: np.random.Generator, - n_samples: int = N_TRAIN, - n_features: int = N_FEATURES, - n_posterior_samples: int = N_POSTERIOR_SAMPLES, - rel_tol: float = LEVEL_REL_TOL, - fast_fit_kwargs: Dict[str, Any] = 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(RANDOM_SEED) - x = data_rng.normal(size=(n_samples, n_features)) - y = _skewed_rewards(x, data_rng) - bnn = make_gaussian_bnn(homoscedastic=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) - - def test_reset_clears_noise_and_keeps_standardization( - self, - make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], - rng: np.random.Generator, - n_samples: int = N_TRAIN, - n_features: int = N_FEATURES, - short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, - ) -> None: - """reset() clears the noise latent (re-initialized on the next update) but keeps loc / scale.""" - bnn = make_gaussian_bnn(homoscedastic=True, **short_fit_kwargs) - bnn.update(context=rng.normal(size=(n_samples, n_features)), rewards=rng.normal(size=n_samples).tolist()) - loc, scale = bnn.reward_loc, bnn.reward_scale - bnn.reset() - assert bnn.noise_log_sigma is None - assert (bnn.reward_loc, bnn.reward_scale) == (loc, scale) - - def test_decay_inflates_noise_posterior( - self, - make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], - rng: np.random.Generator, - decay_factor: float = DECAY_FACTOR, - n_samples: int = N_TRAIN, - n_features: int = N_FEATURES, - short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, - tol: float = ROUND_TRIP_TOL, - ) -> None: - """Forgetting widens the noise posterior by 1 / decay_factor, like the weights.""" - bnn = make_gaussian_bnn(homoscedastic=True, 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, - make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], - rng: np.random.Generator, - n_samples: int = N_TRAIN, - n_features: int = N_FEATURES, - short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, - ) -> None: - """The noise posterior survives the JSON round trip.""" - bnn = make_gaussian_bnn(homoscedastic=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.noise_log_sigma == bnn.noise_log_sigma - assert restored.homoscedastic - - @pytest.mark.parametrize( - "kwargs", - [{"noise_sigma": NOISE_SIGMA}, {"noise_log_sigma": {"mu": [0.0], "sigma": [1.0]}}], - ids=["noise_sigma", "noise_log_sigma"], - ) - def test_heteroscedastic_rejects_noise_settings( - self, make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], kwargs: Dict[str, Any] - ) -> None: - """noise_sigma / noise_log_sigma only make sense for the homoscedastic model.""" - with pytest.raises(ValidationError): - make_gaussian_bnn(**kwargs) - - def test_rejects_output_width_mismatch( - self, make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork] - ) -> None: - """A heteroscedastic network (2 output units) can't be flagged homoscedastic.""" - bnn = make_gaussian_bnn() - with pytest.raises(ValidationError): - GaussianBayesianNeuralNetwork(**{**bnn.model_dump(), "homoscedastic": True}) - - def test_meta_model_rejects_mixed_noise_models( - self, make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork] - ) -> None: - """Homoscedastic and heteroscedastic heads can't share one meta-model.""" - with pytest.raises(AttributeError): - CmabMetaModel(actions={"a": make_gaussian_bnn(), "b": make_gaussian_bnn(homoscedastic=True)}) - - @pytest.mark.parametrize( - "extra_kwargs", - [ - {}, - {**BACKBONE_KWARGS, "update_kwargs": {**FAST_FIT_KWARGS["update_kwargs"], "batch_size": MINIBATCH_SIZE}}, - ], - ids=["no_backbone", "backbone_minibatch"], - ) - def test_cmab_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] = ARM_MEANS, - best_arm: str = BEST_ARM, - n_rounds: int = N_ROUNDS, - n_per_round: int = N_PER_ROUND, - n_features: int = N_FEATURES, - min_share: float = MIN_BEST_ARM_SHARE, - ) -> None: - """The homoscedastic bandit finds the best arm, and every played arm carries its own noise posterior.""" - mab = make_cmab_gaussian(homoscedastic=True, **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, _, _ = mab.predict(context=rng.normal(size=(n_per_round, n_features))) - assert np.mean(np.array(actions) == best_arm) >= min_share - assert mab.actions[best_arm].noise_log_sigma is not None - assert CmabGaussian.from_state(mab.get_state()[1]).actions == mab.actions From cc2e3f90872a200d6abd4c27ae1130b8850ee970 Mon Sep 17 00:00:00 2001 From: ronshiff1 Date: Fri, 9 Oct 2026 11:24:50 +0000 Subject: [PATCH 03/14] changed normlizing to false in default + updated tutorial --- .claude/skills/gaussian-cmab/SKILL.md | 5 +- docs/continuous_reward_cmab.md | 6 +- docs/src/tutorials/cmab.ipynb | 159 ++++++++++++++++++++++++-- pybandits/model/bnn/network.py | 4 +- tests/test_gaussian_bnn.py | 18 +-- 5 files changed, 165 insertions(+), 27 deletions(-) diff --git a/.claude/skills/gaussian-cmab/SKILL.md b/.claude/skills/gaussian-cmab/SKILL.md index b8ef9b0d..f171b6b6 100644 --- a/.claude/skills/gaussian-cmab/SKILL.md +++ b/.claude/skills/gaussian-cmab/SKILL.md @@ -20,8 +20,9 @@ metadata: - a Bayesian MLP with **one output unit, μ(x)**, and the likelihood `Normal(μ(x), σ)`; - **σ is one latent per model** (`log σ ~ Normal`, site `noise_log_sigma`), with its posterior stored in the state and carried across updates like the weights; - - **standardized targets** `(r − reward_loc) / reward_scale`, using the user's values or else fitted on the first - batch, then frozen. `standardize_rewards=False` trains on raw rewards (and then rejects loc / scale). + - **raw rewards by default** (`standardize_rewards=False`, which rejects loc / scale); with `standardize_rewards=True`, + **standardized targets** `(r − reward_loc) / reward_scale`, using the user's values or else fitted on the first + batch, then frozen. - **`CmabGaussian`** (`pybandits/cmab.py`): `ClassicBandit` strategy, `_predict_with_proba = True`. Actions are ranked on μ; `predict` returns `(actions, mu, sigma)`, where σ is for monitoring only. - **Reward types** (`pybandits/base.py`): `BinaryReward`, `SoftReward`, `ContinuousReward`, `AnyReward`. diff --git a/docs/continuous_reward_cmab.md b/docs/continuous_reward_cmab.md index f74f3c5d..9186b99c 100644 --- a/docs/continuous_reward_cmab.md +++ b/docs/continuous_reward_cmab.md @@ -18,9 +18,9 @@ Each action is a Bayesian MLP with **one output unit, μ(x)**, and the reward is - **σ is a single noise std per action,** not a function of the context. It is a latent (`log σ ~ Normal`) 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` (reward units) if given, otherwise on the std of the first update batch. -- **Training is on standardized targets** `(r − reward_loc) / reward_scale`, so the default O(1) weight priors fit - rewards on any scale. `reward_loc` / `reward_scale` are the values given at cold start, otherwise fitted on the first - batch and then frozen. `standardize_rewards=False` trains on raw rewards. +- **Training is on raw rewards by default** (`standardize_rewards=False`). With `standardize_rewards=True` it is on + standardized targets `(r − reward_loc) / reward_scale`, so the default O(1) weight priors fit rewards on any scale; + `reward_loc` / `reward_scale` are the values given at cold start, otherwise fitted on the first batch and then frozen. - **Thompson sampling** draws the weights from the posterior and ranks actions on μ. `predict` returns μ (in the "probabilities" slot) and σ (for monitoring) per action and row. - **`reset()`** returns the model to its cold-start state, including the standardization (re-fitted on the next diff --git a/docs/src/tutorials/cmab.ipynb b/docs/src/tutorials/cmab.ipynb index a709ec46..322ca6fc 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,27 @@ "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": [ + "## 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 +180,140 @@ "# 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": [ + "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. By default the model is trained on raw rewards (`standardize_rewards=False`); pass `standardize_rewards=True` to train on standardized rewards, and `noise_sigma` (in reward units) to set 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", + " 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]))" + ] } ], "metadata": { diff --git a/pybandits/model/bnn/network.py b/pybandits/model/bnn/network.py index 74704909..03b501cd 100644 --- a/pybandits/model/bnn/network.py +++ b/pybandits/model/bnn/network.py @@ -1763,7 +1763,7 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): standardize_rewards : bool Whether to train on standardized targets. When True, ``reward_loc`` / ``reward_scale`` are used: the values given at cold start, or else fitted (mean / std) on the first update batch and kept fixed afterwards. When - False, they are neither fitted nor used (the model trains on raw rewards). Default is True. + False, they are neither fitted nor used (the model trains on raw rewards). Default is False. reward_loc : Optional[float] Location used to standardize rewards. None to fit it on the first update. reward_scale : Optional[PositiveFloat] @@ -1810,7 +1810,7 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): ... ) """ - standardize_rewards: bool = True + standardize_rewards: bool = False reward_loc: Optional[float] = None reward_scale: Optional[PositiveFloat] = None reward_loc_init: Optional[float] = None diff --git a/tests/test_gaussian_bnn.py b/tests/test_gaussian_bnn.py index f7c0673a..a0a07322 100644 --- a/tests/test_gaussian_bnn.py +++ b/tests/test_gaussian_bnn.py @@ -120,7 +120,7 @@ def fitted_gaussian_bnn( ) -> GaussianBayesianNeuralNetwork: """A Gaussian BNN fitted on linear synthetic rewards with constant noise.""" data_rng = np.random.default_rng(RANDOM_SEED) - bnn = make_gaussian_bnn(**FAST_FIT_KWARGS) + bnn = make_gaussian_bnn(standardize_rewards=True, **FAST_FIT_KWARGS) x = data_rng.normal(size=(N_TRAIN, N_FEATURES)) bnn.update(context=x, rewards=_linear_rewards(x, data_rng).tolist()) return bnn @@ -220,7 +220,7 @@ def test_counters_track_observations( ) 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) + bnn = GaussianBayesianNeuralNetwork.cold_start(n_features=n_features, standardize_rewards=True) targets = bnn.prepare_rewards(rewards) loc, scale = bnn.reward_loc, bnn.reward_scale assert loc == pytest.approx(np.mean(rewards)) @@ -234,7 +234,7 @@ def test_standardization_is_fitted_once(self, rewards: List[float], 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) + bnn = GaussianBayesianNeuralNetwork.cold_start(n_features=n_features, standardize_rewards=True) bnn.prepare_rewards([reward]) assert bnn.reward_scale == max(abs(reward), 1.0) @@ -250,7 +250,7 @@ def test_single_reward_standardization_has_unit_or_larger_scale(self, reward: fl ) 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) + bnn = GaussianBayesianNeuralNetwork.cold_start(n_features=n_features, standardize_rewards=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) @@ -317,7 +317,7 @@ def test_mean_is_level_calibrated_on_skewed_rewards( data_rng = np.random.default_rng(RANDOM_SEED) x = data_rng.normal(size=(n_samples, n_features)) y = _skewed_rewards(x, data_rng) - bnn = make_gaussian_bnn(**fast_fit_kwargs) + bnn = make_gaussian_bnn(standardize_rewards=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) @@ -397,7 +397,7 @@ def test_reset_refits_fitted_standardization( short_fit_kwargs: Dict[str, Any] = 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(**short_fit_kwargs) + bnn = make_gaussian_bnn(standardize_rewards=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() @@ -421,7 +421,9 @@ def test_reset_restores_user_supplied_standardization( short_fit_kwargs: Dict[str, Any] = 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 = make_gaussian_bnn( + standardize_rewards=True, 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()) @@ -438,7 +440,7 @@ def test_reset_after_round_trip_clears_fitted_standardization( short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, ) -> None: """A fitted standardization stays 'fitted' through serialization: reset() still clears it.""" - bnn = make_gaussian_bnn(**short_fit_kwargs) + bnn = make_gaussian_bnn(standardize_rewards=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 and restored.reward_loc_init is None From 1c4dd2f762075ab5df8c0af75244d3336d835484 Mon Sep 17 00:00:00 2001 From: ronshiff1 Date: Fri, 9 Oct 2026 13:19:18 +0000 Subject: [PATCH 04/14] fix docstring --- pybandits/model/bnn/network.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/pybandits/model/bnn/network.py b/pybandits/model/bnn/network.py index 03b501cd..57c35af1 100644 --- a/pybandits/model/bnn/network.py +++ b/pybandits/model/bnn/network.py @@ -616,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. """ @@ -640,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 @@ -652,7 +653,7 @@ 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] @@ -713,8 +714,8 @@ def _observe( """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 @@ -722,7 +723,7 @@ 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] @@ -827,7 +828,7 @@ def _observe_output( ) -> 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 a Bernoulli distribution given the network output" + # 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( From b3da4be2aa5d9c388096e28620ecc01a02981611 Mon Sep 17 00:00:00 2001 From: ronshiff1 Date: Fri, 9 Oct 2026 14:34:31 +0000 Subject: [PATCH 05/14] change in loc/scalre mechansim --- docs/continuous_reward_cmab.md | 14 ++--- docs/src/tutorials/cmab.ipynb | 2 +- pybandits/meta_model/cmab_meta_model.py | 4 +- pybandits/model/bnn/network.py | 71 +++++++++++++------------ tests/test_gaussian_bnn.py | 56 +++++++++++-------- 5 files changed, 81 insertions(+), 66 deletions(-) diff --git a/docs/continuous_reward_cmab.md b/docs/continuous_reward_cmab.md index 9186b99c..452c1e13 100644 --- a/docs/continuous_reward_cmab.md +++ b/docs/continuous_reward_cmab.md @@ -18,13 +18,15 @@ Each action is a Bayesian MLP with **one output unit, μ(x)**, and the reward is - **σ is a single noise std per action,** not a function of the context. It is a latent (`log σ ~ Normal`) 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` (reward units) if given, otherwise on the std of the first update batch. -- **Training is on raw rewards by default** (`standardize_rewards=False`). With `standardize_rewards=True` it is on - standardized targets `(r − reward_loc) / reward_scale`, so the default O(1) weight priors fit rewards on any scale; - `reward_loc` / `reward_scale` are the values given at cold start, otherwise fitted on the first batch and then frozen. +- **Training is on the targets** `(r − reward_loc) / reward_scale`. By default both are None, i.e. 0 / 1, so the + model trains on raw rewards and the weight priors are in reward units. Fixed constants can be given (alone or + together, e.g. only `reward_scale=20` for a centered uplift in dollars) so the default O(1) weight priors fit; they + are never re-fitted. With `fit_reward_standardization=True` they are instead fitted on the first batch and then + frozen (and must not be given). - **Thompson sampling** draws the weights from the posterior and ranks actions on μ. `predict` returns μ (in the "probabilities" slot) and σ (for monitoring) per action and row. -- **`reset()`** returns the model to its cold-start state, including the standardization (re-fitted on the next - update unless it was given at cold start). +- **`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). ### Why one σ and not σ(x) @@ -50,7 +52,7 @@ mab = CmabGaussian.cold_start( dist_type="normal", dist_params_init={"mu": 0, "sigma": 0.1}, # a narrow prior fits much faster than the default sigma=1 update_kwargs={"num_steps": 400, "optimizer_kwargs": {"step_size": 3e-3}}, - # optionally, from historical data rather than the first batch: + # optionally, fixed values from historical data (or fit_reward_standardization=True to fit on the first batch): # reward_loc=..., reward_scale=..., noise_sigma=..., # decay_factor=..., # if the reward process drifts ) diff --git a/docs/src/tutorials/cmab.ipynb b/docs/src/tutorials/cmab.ipynb index 322ca6fc..21b67e0d 100644 --- a/docs/src/tutorials/cmab.ipynb +++ b/docs/src/tutorials/cmab.ipynb @@ -221,7 +221,7 @@ "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. By default the model is trained on raw rewards (`standardize_rewards=False`); pass `standardize_rewards=True` to train on standardized rewards, and `noise_sigma` (in reward units) to set an initial guess of the noise std." + "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. By default the model is trained on raw rewards. For rewards far from $O(1)$ (e.g. revenue in dollars), pass a fixed `reward_scale` (and optionally `reward_loc`) so the default $O(1)$ weight priors fit, or `fit_reward_standardization=True` to fit both on the first batch. `noise_sigma` (in reward units) sets an initial guess of the noise std." ] }, { diff --git a/pybandits/meta_model/cmab_meta_model.py b/pybandits/meta_model/cmab_meta_model.py index 30eaa318..650098d5 100644 --- a/pybandits/meta_model/cmab_meta_model.py +++ b/pybandits/meta_model/cmab_meta_model.py @@ -897,8 +897,8 @@ def _store_backbone(self, params: dict, n_bb_layers: int) -> None: 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 standardize with their own (first-batch fitted) - location and scale, so arms on different reward scales share one likelihood in the joint ELBO. + 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 ---------- diff --git a/pybandits/model/bnn/network.py b/pybandits/model/bnn/network.py index 57c35af1..69b15379 100644 --- a/pybandits/model/bnn/network.py +++ b/pybandits/model/bnn/network.py @@ -1755,26 +1755,27 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): 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 standardized targets ``(reward - reward_loc) / reward_scale`` so the default - ``O(1)`` weight priors fit rewards on any scale; ``sample_proba`` maps ``mu`` and ``sigma`` back to the reward - scale. + The network is trained on the targets ``(reward - reward_loc) / reward_scale``, and ``sample_proba`` maps ``mu`` + and ``sigma`` back to the reward scale. By default (both None) the targets are the raw rewards, so the weight + priors are in reward units: for rewards far from ``O(1)``, set a fixed ``reward_scale`` (or fit it from data with + ``fit_reward_standardization``) so the default ``O(1)`` weight priors fit. Parameters ---------- - standardize_rewards : bool - Whether to train on standardized targets. When True, ``reward_loc`` / ``reward_scale`` are used: the values - given at cold start, or else fitted (mean / std) on the first update batch and kept fixed afterwards. When - False, they are neither fitted nor used (the model trains on raw rewards). Default is False. + 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, ``reward_loc`` / ``reward_scale`` must not be given. When False, + the values given at cold start are used as fixed constants. Default is False. reward_loc : Optional[float] - Location used to standardize rewards. None to fit it on the first update. + Location subtracted from the rewards. None means 0 (or fitted, with ``fit_reward_standardization``). reward_scale : Optional[PositiveFloat] - Scale used to standardize rewards. None to fit it on the first update. + Scale the rewards are divided by. None means 1 (or fitted, with ``fit_reward_standardization``). reward_loc_init : Optional[float] - The cold-start value of ``reward_loc`` (None if it is fitted), restored by ``reset()``. Set automatically - from ``reward_loc`` at construction. + The cold-start value of ``reward_loc``, restored by ``reset()``. Set automatically from ``reward_loc`` at + construction. reward_scale_init : Optional[PositiveFloat] - The cold-start value of ``reward_scale`` (None if it is fitted), restored by ``reset()``. Set automatically - from ``reward_scale`` at construction. + The cold-start value of ``reward_scale``, restored by ``reset()``. Set automatically from ``reward_scale`` at + construction. n_observations : NonNegativeInt Number of rewards observed. reward_sum : float @@ -1785,7 +1786,7 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): 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 standardized units), shape (1,). None until the first update + 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 @@ -1798,7 +1799,7 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): 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 ``reward_loc`` / ``reward_scale`` to their cold-start values (re-fitted on the - next update if they were fitted). + next update with ``fit_reward_standardization``). Examples -------- @@ -1811,7 +1812,7 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): ... ) """ - standardize_rewards: bool = False + fit_reward_standardization: bool = False reward_loc: Optional[float] = None reward_scale: Optional[PositiveFloat] = None reward_loc_init: Optional[float] = None @@ -1849,11 +1850,11 @@ def set_standardization_init(cls, data: Any) -> Any: @model_validator(mode="after") def validate_reward_standardization(self) -> "GaussianBayesianNeuralNetwork": - for loc, scale in [(self.reward_loc, self.reward_scale), (self.reward_loc_init, self.reward_scale_init)]: - if (loc is None) != (scale is None): - raise ValueError("reward_loc and reward_scale must be either both set or both None.") - if not self.standardize_rewards and loc is not None: - raise ValueError("reward_loc / reward_scale require standardize_rewards=True.") + if self.fit_reward_standardization: + if self.reward_loc_init is not None or self.reward_scale_init is not None: + raise ValueError("reward_loc / reward_scale cannot be given with fit_reward_standardization=True.") + if (self.reward_loc is None) != (self.reward_scale is None): + raise ValueError("A fitted reward_loc and reward_scale must be either both set or both None.") return self @model_validator(mode="after") @@ -1867,16 +1868,16 @@ def validate_output_layer(self) -> "GaussianBayesianNeuralNetwork": @property def _loc_scale(self) -> Tuple[float, float]: - """``(loc, scale)`` mapping standardized targets back to rewards (identity until fitted, or if disabled).""" - if not self.standardize_rewards or self.reward_loc is None: - return 0.0, 1.0 - return self.reward_loc, self.reward_scale + """``(loc, scale)`` mapping targets back to rewards (0 / 1 for an unset or not yet fitted value).""" + 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 standardized units. + """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 standardized batch; 0, i.e. sigma = 1, when there is no batch or no spread). + ``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: @@ -1910,7 +1911,7 @@ def _noise_log_sigma_of(extra: Dict[str, Any], name: str) -> Any: def output_distribution( self, linear_out: jax.Array, extra_sites: Optional[Dict[str, jax.Array]] = None ) -> NumpyroDistribution: - """``Normal(mu, exp(log sigma))`` on the standardized targets. + """``Normal(mu, exp(log sigma))`` on the targets ``(reward - reward_loc) / reward_scale``. Parameters ---------- @@ -1979,11 +1980,11 @@ def sample_proba(self, context: np.ndarray, rng: np.random.Generator) -> List[Ga return super().sample_proba(context=context, rng=rng) def prepare_rewards(self, rewards: List[ContinuousReward]) -> np.ndarray: - """Standardize rewards into training targets, fitting the standardization and the noise prior on first use. + """Map rewards to training targets, fitting the standardization and the noise prior on first use. - With ``standardize_rewards``, ``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``). + 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 ---------- @@ -1998,7 +1999,7 @@ def prepare_rewards(self, rewards: List[ContinuousReward]) -> np.ndarray: y = np.asarray(rewards, dtype=float) if not np.all(np.isfinite(y)): raise ValueError("Rewards must be finite.") - if self.standardize_rewards and self.reward_loc is None and len(y) > 0: + 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 @@ -2062,9 +2063,9 @@ def mean(self) -> float: 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 (standardized) target mean on the first update call. + """Set the output-layer bias to the first batch's target mean on the first update call. - With ``standardize_rewards`` fitted on that same batch this is 0. The noise prior is initialized by + 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 diff --git a/tests/test_gaussian_bnn.py b/tests/test_gaussian_bnn.py index a0a07322..2126787e 100644 --- a/tests/test_gaussian_bnn.py +++ b/tests/test_gaussian_bnn.py @@ -120,7 +120,7 @@ def fitted_gaussian_bnn( ) -> GaussianBayesianNeuralNetwork: """A Gaussian BNN fitted on linear synthetic rewards with constant noise.""" data_rng = np.random.default_rng(RANDOM_SEED) - bnn = make_gaussian_bnn(standardize_rewards=True, **FAST_FIT_KWARGS) + bnn = make_gaussian_bnn(fit_reward_standardization=True, **FAST_FIT_KWARGS) x = data_rng.normal(size=(N_TRAIN, N_FEATURES)) bnn.update(context=x, rewards=_linear_rewards(x, data_rng).tolist()) return bnn @@ -220,7 +220,7 @@ def test_counters_track_observations( ) 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, standardize_rewards=True) + 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)) @@ -234,7 +234,7 @@ def test_standardization_is_fitted_once(self, rewards: List[float], 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, standardize_rewards=True) + 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) @@ -250,7 +250,7 @@ def test_single_reward_standardization_has_unit_or_larger_scale(self, reward: fl ) 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, standardize_rewards=True) + 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) @@ -317,7 +317,7 @@ def test_mean_is_level_calibrated_on_skewed_rewards( data_rng = np.random.default_rng(RANDOM_SEED) x = data_rng.normal(size=(n_samples, n_features)) y = _skewed_rewards(x, data_rng) - bnn = make_gaussian_bnn(standardize_rewards=True, **fast_fit_kwargs) + 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) @@ -334,26 +334,40 @@ def test_update_rejects_non_finite_rewards( with pytest.raises(ValidationError): bnn.update(context=np.zeros((1, n_features)), rewards=[bad_reward]) - @pytest.mark.parametrize("reward_loc, reward_scale", [(USER_LOC, None), (None, USER_SCALE)]) - def test_rejects_half_specified_standardization( + @pytest.mark.parametrize( + "reward_loc, reward_scale", [(USER_LOC, None), (None, USER_SCALE)], ids=["loc_only", "scale_only"] + ) + def test_fixed_standardization_can_be_partial( self, make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, reward_loc: Optional[float], reward_scale: Optional[float], + n_samples: int = N_TRAIN, + tol: float = ROUND_TRIP_TOL, ) -> None: - """reward_loc and reward_scale must be given together.""" - with pytest.raises(ValidationError): - make_gaussian_bnn(reward_loc=reward_loc, reward_scale=reward_scale) + """A fixed reward_loc or reward_scale can be given alone: the other one defaults to 0 / 1, never fitted.""" + bnn = make_gaussian_bnn(reward_loc=reward_loc, reward_scale=reward_scale) + rewards = rng.normal(size=n_samples) + loc = 0.0 if reward_loc is None else reward_loc + scale = 1.0 if reward_scale is None else reward_scale + np.testing.assert_allclose(bnn.prepare_rewards(rewards.tolist()), (rewards - loc) / scale, rtol=tol) + assert (bnn.reward_loc, bnn.reward_scale) == (reward_loc, reward_scale) - def test_rejects_standardization_values_when_disabled( + @pytest.mark.parametrize( + "reward_loc, reward_scale", + [(USER_LOC, None), (None, USER_SCALE), (USER_LOC, USER_SCALE)], + ids=["loc_only", "scale_only", "both"], + ) + def test_rejects_fixed_values_with_fitted_standardization( self, make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], - reward_loc: float = USER_LOC, - reward_scale: float = USER_SCALE, + reward_loc: Optional[float], + reward_scale: Optional[float], ) -> None: - """With standardize_rewards=False, reward_loc / reward_scale would be ignored, so they are rejected.""" + """fit_reward_standardization=True fits loc / scale from data, so given values would conflict.""" with pytest.raises(ValidationError): - make_gaussian_bnn(standardize_rewards=False, reward_loc=reward_loc, reward_scale=reward_scale) + make_gaussian_bnn(fit_reward_standardization=True, reward_loc=reward_loc, reward_scale=reward_scale) def test_unstandardized_model_trains_and_predicts_on_raw_rewards( self, @@ -365,8 +379,8 @@ def test_unstandardized_model_trains_and_predicts_on_raw_rewards( short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, tol: float = ROUND_TRIP_TOL, ) -> None: - """standardize_rewards=False: no loc / scale are fitted, targets are raw and noise_sigma is in raw units.""" - bnn = make_gaussian_bnn(standardize_rewards=False, noise_sigma=noise_sigma, **short_fit_kwargs) + """By default no loc / scale are set or fitted: targets are raw and noise_sigma is in raw units.""" + bnn = make_gaussian_bnn(noise_sigma=noise_sigma, **short_fit_kwargs) rewards = rng.normal(size=n_samples) np.testing.assert_array_equal(bnn.prepare_rewards(rewards.tolist()), rewards) assert bnn.reward_loc is None and bnn.reward_scale is None @@ -397,7 +411,7 @@ def test_reset_refits_fitted_standardization( short_fit_kwargs: Dict[str, Any] = 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(standardize_rewards=True, **short_fit_kwargs) + 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() @@ -421,9 +435,7 @@ def test_reset_restores_user_supplied_standardization( short_fit_kwargs: Dict[str, Any] = 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( - standardize_rewards=True, reward_loc=reward_loc, reward_scale=reward_scale, **short_fit_kwargs - ) + 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()) @@ -440,7 +452,7 @@ def test_reset_after_round_trip_clears_fitted_standardization( short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, ) -> None: """A fitted standardization stays 'fitted' through serialization: reset() still clears it.""" - bnn = make_gaussian_bnn(standardize_rewards=True, **short_fit_kwargs) + 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 and restored.reward_loc_init is None From 0e4fe94948ff05601d08534906bfcf5cd3de7b38 Mon Sep 17 00:00:00 2001 From: ronshiff1 Date: Mon, 12 Oct 2026 06:21:47 +0000 Subject: [PATCH 06/14] skill updated --- .claude/skills/gaussian-cmab/SKILL.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.claude/skills/gaussian-cmab/SKILL.md b/.claude/skills/gaussian-cmab/SKILL.md index f171b6b6..64a640d7 100644 --- a/.claude/skills/gaussian-cmab/SKILL.md +++ b/.claude/skills/gaussian-cmab/SKILL.md @@ -20,9 +20,9 @@ metadata: - a Bayesian MLP with **one output unit, μ(x)**, and the likelihood `Normal(μ(x), σ)`; - **σ is one latent per model** (`log σ ~ Normal`, site `noise_log_sigma`), with its posterior stored in the state and carried across updates like the weights; - - **raw rewards by default** (`standardize_rewards=False`, which rejects loc / scale); with `standardize_rewards=True`, - **standardized targets** `(r − reward_loc) / reward_scale`, using the user's values or else fitted on the first - batch, then frozen. + - targets `(r − reward_loc) / reward_scale`: **raw rewards by default** (both None, i.e. 0 / 1); user-given values + are fixed constants (loc or scale may be given alone); `fit_reward_standardization=True` fits both on the first + batch, then freezes them (and rejects user-given values). - **`CmabGaussian`** (`pybandits/cmab.py`): `ClassicBandit` strategy, `_predict_with_proba = True`. Actions are ranked on μ; `predict` returns `(actions, mu, sigma)`, where σ is for monitoring only. - **Reward types** (`pybandits/base.py`): `BinaryReward`, `SoftReward`, `ContinuousReward`, `AnyReward`. From 7a415f67df15635ba8a2764a8d123499ef2c9e09 Mon Sep 17 00:00:00 2001 From: ronshiff1 Date: Mon, 12 Oct 2026 06:22:44 +0000 Subject: [PATCH 07/14] remove skill --- .claude/skills/gaussian-cmab/SKILL.md | 135 -------------------------- 1 file changed, 135 deletions(-) delete mode 100644 .claude/skills/gaussian-cmab/SKILL.md diff --git a/.claude/skills/gaussian-cmab/SKILL.md b/.claude/skills/gaussian-cmab/SKILL.md deleted file mode 100644 index 64a640d7..00000000 --- a/.claude/skills/gaussian-cmab/SKILL.md +++ /dev/null @@ -1,135 +0,0 @@ ---- -name: gaussian-cmab -description: | - Background and conventions for continuous (real-valued) rewards in pybandits: GaussianBayesianNeuralNetwork, - CmabGaussian, the uplift reward R − E0(x), and why the model has one μ head with a learned scalar σ. - Use when: working on GaussianBayesianNeuralNetwork / CmabGaussian, continuous or revenue-like rewards, uplift - rewards against a control baseline, the extra-site likelihood hook, or extending the Gaussian model (offsets, - hurdle models, new likelihoods). -license: MIT -metadata: - author: ronshiff1 - version: "1.0.0" ---- - -# Gaussian cMAB: continuous rewards - -## What exists (pybandits ≥ 8.4.0) - -- **`GaussianBayesianNeuralNetwork`** (`pybandits/model/bnn/network.py`): - - a Bayesian MLP with **one output unit, μ(x)**, and the likelihood `Normal(μ(x), σ)`; - - **σ is one latent per model** (`log σ ~ Normal`, site `noise_log_sigma`), with its posterior stored in the state - and carried across updates like the weights; - - targets `(r − reward_loc) / reward_scale`: **raw rewards by default** (both None, i.e. 0 / 1); user-given values - are fixed constants (loc or scale may be given alone); `fit_reward_standardization=True` fits both on the first - batch, then freezes them (and rejects user-given values). -- **`CmabGaussian`** (`pybandits/cmab.py`): `ClassicBandit` strategy, `_predict_with_proba = True`. Actions are - ranked on μ; `predict` returns `(actions, mu, sigma)`, where σ is for monitoring only. -- **Reward types** (`pybandits/base.py`): `BinaryReward`, `SoftReward`, `ContinuousReward`, `AnyReward`. - - Entry points (`BaseMab.update`, `BaseCmabBernoulli.update`, `ActionsManager.update`) accept `AnyReward`. - - `ActionsManager._check_rewards` enforces what the models support (`BaseModel.supports_continuous_rewards`). - - `delta` + continuous rewards is rejected. - -## Core decisions and why - -1. **One σ, not σ(x).** The Gaussian NLL weights each residual by 1/σ(x)². On zero-inflated or skewed rewards - (revenue), a learned σ(x) lets the model call large rewards noise instead of raising μ, so μ is biased low where - rewards are most variable. In evaluations this under-predicted the level ~3×, and the epistemic sd collapsed - (confidently wrong). With one σ, the μ fit is least squares, whose target is E[r | x] for any noise shape. The - σ(x) head was therefore removed entirely. -2. **σ learned, not fixed.** A fixed σ mis-sets the posterior width, i.e. the exploration. The learned σ converged to - the true residual sd and gave calibrated predictive intervals. Changing σ by hand between updates would make the - stored posterior inconsistent: never do it. Use `noise_sigma` for the initial prior and `decay_factor` for - forgetting. -3. **Standardize.** The weight priors are O(1). Raw revenue-scale targets need weights about 100× outside the - prior, and the model barely moves (with tanh or ReLU alike). Putting the prior on the data scale instead is - equivalent in theory but badly conditioned for Adam. `loc` / `scale` are frozen because the weights only mean - something relative to them. -4. **`reset()` = back to the cold-start state:** weights, noise latent, counters, and `loc` / `scale` back to their - cold-start values. `reward_loc_init` / `reward_scale_init` hold the user-supplied values (None if fitted), so fitted - values are re-fitted on the next update and user values survive. -5. **Reward = uplift `R − E0(x)`** against a control baseline computed outside pybandits. - - `E0` doesn't depend on the action, so the ranking is unchanged and the shared variance is removed (a control - variate). - - Per-action data volumes are too small to learn revenue from scratch; models without a baseline lost about - a third of the explained variance. - - The uplift can be negative, and that is expected. - -## Rejected alternatives (don't re-propose without new evidence) - -- **σ(x) noise head:** see decision 1. -- **A Bernoulli likelihood on a continuous reward rescaled to [0, 1]** (e.g. `(min(Y, c) − k·baseline + a) / b`): - - The mean is consistent, but the Bernoulli variance p(1 − p) overstated the real noise ~128×, so each observation - carried ~1/128 of its information and the model learned only an intercept. - - Weighting the likelihood by ≈ p(1 − p) / var(z) matches a Gaussian only to second order (same minimum and - curvature, but a different loss shape, link and per-row weighting). Use the Gaussian directly. -- **Hurdle × log-normal amount** (`p0 · exp(m + σ²/2)`): - - The best distributional fit (flat PIT), but the back-transform around a baseline biased the level by 10–15%. - - Averaging or drawing `exp(m + σ²/2)` over posterior samples has extreme tails (10⁸× outliers). -- **The hurdle that works, if a conversion vs. amount split is needed:** a Bernoulli with a `logit p0_control` - offset × a Gaussian on the raw amount residual `R − E_buy_control`, for buyers only (called "H3n"). It is as - accurate and calibrated as the single-Gaussian uplift model. It needs per-row offset support (see below). - -## Code map (internals) - -**Base-class likelihood hooks in `BaseBayesianNeuralNetwork`** (the Bernoulli defaults reproduce the original -behavior): -- `_output_dim` (ClassVar), `create_model_params(..., output_dim=None)`; -- `output_distribution(linear_out, extra_sites)`, `_observe_output`, `_postprocess_output(linear_out, extra_samples)`; -- `prepare_rewards(rewards)` → training targets; `_fit(context, rewards)` (the `_update` body); `_is_first_fit`. - -**The extra-site hook, for head-level latents beyond the weights and embeddings:** -- `extra_site_params()` → `{name: BaseLocationScaleArray}`; -- `sample_extra_sites()` (NumPyro, KL-annealed), `sample_extra_numpy(n, rng)` (Thompson); -- `update_extra_params_from_vi()`, `_inflate_extra_params()`. - -It is wired into guide init (`collect_guide_init_arrays`), VI readback, the full-rank ADVI site filter, decay and -sampling. A new likelihood latent (e.g. a Student-t ν) only needs these methods. - -**Meta-model** (`pybandits/meta_model/cmab_meta_model.py`): -- `_check_models` requires one likelihood class across heads; -- `_so_model` stacks each arm's extra sites and gathers them per row; -- `_store_head` reads them back; -- `_prepare_targets` applies each head's `prepare_rewards` to its own rows (each arm keeps its own standardization). - -**Aliases:** `CmabMetaModelGaussian`, `CmabActionsManagerGaussian`. - -## Usage - -```python -mab = CmabGaussian.cold_start( - action_ids={...}, n_features=d, hidden_dim_list=[16], - dist_type="normal", dist_params_init={"mu": 0, "sigma": 0.1}, # narrow prior: default sigma=1 is too slow - update_kwargs={"num_steps": 400, "optimizer_kwargs": {"step_size": 3e-3}}, - # production: reward_loc / reward_scale / noise_sigma from history; decay_factor if drifting -) -actions, mu, sigma = mab.predict(context=X) -mab.update(actions=actions, rewards=(R - E0).tolist(), context=X) -``` - -## Gotchas - -- **Evaluate the mean with a plug-in** (posterior mean of μ) or by averaging draws of μ. Never average - `exp(...)` of log-scale draws. -- **A wide prior** (`sigma=1`) on an unbounded likelihood converges slowly. Narrow it (≈0.1) and raise the step size. -- **Additive residual models can predict E0 + μ < 0** on low-reward rows. That's irrelevant for ranking an uplift, - but floor it when reading the prediction back as revenue. -- **Whale-type under-prediction** inherited from the baseline isn't fixed by the BNN. Fix the baseline (more history). - -## Possible next work - -- **Per-row offset support in `BaseBayesianNeuralNetwork`**, for the H3n hurdle: `logit = offset + f(x)`. - - A working prototype treats the context's last column as the offset: it strips it in `check_context_matrix`, - `emit_submodel` and `forward_pass`, and adds it in `_observe_output` / `_postprocess_output`. It is full-batch - only. - - A proper version must carry the offset through the data plate (minibatching) and the joint meta-model. -- **Gaussian variants** of CC / MO / DP / quantitative; `delta` with continuous rewards; OPE and simulators. - -## Tests - -- `tests/test_gaussian_bnn.py` (the project style: see the `pytest-writer` skill). -- Run with `PYTHONPATH=$PWD python -m pytest tests/test_gaussian_bnn.py -q -n 8`. Force the local checkout: some - environments have another checkout installed as editable. -- Pre-existing flaky failures, unrelated: - - `test_model.py::{test_bnn_svi_nan_loss_raises_error, test_vi_training_options, test_bnn_vi_update_with_categorical_features_updates_embeddings}`; - - `test_quantitative_model.py::test_sample_proba_reproducible_with_same_rng`. From 6e380250d3beb4d0137d91dd0d444086ec00e0db Mon Sep 17 00:00:00 2001 From: ronshiff1 Date: Mon, 12 Oct 2026 06:25:35 +0000 Subject: [PATCH 08/14] tutorial moved --- docs/{ => src}/continuous_reward_cmab.md | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename docs/{ => src}/continuous_reward_cmab.md (100%) diff --git a/docs/continuous_reward_cmab.md b/docs/src/continuous_reward_cmab.md similarity index 100% rename from docs/continuous_reward_cmab.md rename to docs/src/continuous_reward_cmab.md From 6cf277f560a6703825048e1a1d475c61b6e60ac9 Mon Sep 17 00:00:00 2001 From: ronshiff1 Date: Mon, 12 Oct 2026 06:50:15 +0000 Subject: [PATCH 09/14] tests mobed to old files --- tests/test_cmab.py | 156 +++++++++- tests/test_gaussian_bnn.py | 592 ------------------------------------- tests/test_model.py | 435 ++++++++++++++++++++++++++- 3 files changed, 589 insertions(+), 594 deletions(-) delete mode 100644 tests/test_gaussian_bnn.py diff --git a/tests/test_cmab.py b/tests/test_cmab.py index ec3b4313..f601c919 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,154 @@ 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 + + +@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_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, + ) -> None: + """Bernoulli and Gaussian heads cannot share one meta-model.""" + gaussian = GaussianBayesianNeuralNetwork.cold_start(n_features=n_features, hidden_dim_list=hidden_dims) + 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, + ) -> 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) + + +@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_gaussian_bnn.py b/tests/test_gaussian_bnn.py deleted file mode 100644 index 2126787e..00000000 --- a/tests/test_gaussian_bnn.py +++ /dev/null @@ -1,592 +0,0 @@ -# MIT License -# -# Copyright (c) 2023 Playtika Ltd. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -"""Tests for the Gaussian (continuous-reward) BNN, its meta-model wiring and the CmabGaussian bandit.""" - -from typing import Any, Callable, Dict, List, Optional - -import numpy as np -import pytest -from hypothesis import given -from hypothesis import strategies as st -from pydantic import ValidationError - -from pybandits.cmab import CmabBernoulli, CmabGaussian -from pybandits.meta_model import CmabMetaModel -from pybandits.model import BayesianNeuralNetwork, GaussianBayesianNeuralNetwork - -# --- test configuration constants --- -N_FEATURES = 3 -HIDDEN_DIMS = [8] -OUTPUT_DIM = 1 -RANDOM_SEED = 0 -N_TRAIN = 2000 -N_POSTERIOR_SAMPLES = 200 -# Narrow prior + larger step size: the setting the class docstring recommends for a fast fit. -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}}, -} -SHORT_FIT_KWARGS: Dict[str, Any] = {"update_kwargs": {"num_steps": 5}} -# Linear synthetic reward: REWARD_OFFSET + REWARD_SLOPE * x0 + Normal(0, NOISE_STD). -REWARD_OFFSET = 100.0 -REWARD_SLOPE = 30.0 -NOISE_STD = 10.0 -MEAN_REL_TOL = 0.1 -SIGMA_REL_TOL = 0.3 -# Probe contexts on both sides of x0. -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. -BUY_PROB = 0.15 -SKEW_OFFSET = 2.0 -SKEW_SLOPE = 1.5 -SKEW_LOG_SD = 1.0 -LEVEL_REL_TOL = 0.15 -ROUND_TRIP_TOL = 1e-6 -MAX_ABS_REWARD = 1e6 -MAX_REWARDS = 50 -NOISE_SIGMA = 3.0 -DECAY_FACTOR = 0.5 -USER_LOC = 5.0 -USER_SCALE = 2.0 -SECOND_BATCH_SHIFT = 50.0 -ACTION_IDS = {"a", "b", "c"} -ARM_MEANS = {"a": -1.0, "b": 0.0, "c": 2.0} -BEST_ARM = "c" -N_ROUNDS = 4 -N_PER_ROUND = 200 -MIN_BEST_ARM_SHARE = 0.8 -BACKBONE_KWARGS: Dict[str, Any] = {"backbone_hidden_dims": [8], "backbone_embedding_dim": 3} -MINIBATCH_SIZE = 64 -BAD_REWARDS = [float("nan"), float("inf"), -float("inf")] -OUT_OF_RANGE_SOFT_REWARD = 2.5 -DELTA = 0.1 -EPSILON = 0.1 - - -def _linear_rewards(x: np.ndarray, rng: np.random.Generator) -> np.ndarray: - """Linear-in-x0 rewards with constant Gaussian noise.""" - return REWARD_OFFSET + REWARD_SLOPE * x[:, 0] + rng.normal(size=len(x)) * 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 = SKEW_OFFSET + SKEW_SLOPE * np.clip(x[:, 0], -1.0, 2.0) - amount = rng.lognormal(np.log(mean / BUY_PROB) - SKEW_LOG_SD**2 / 2, SKEW_LOG_SD) - return (rng.uniform(size=len(x)) < 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.""" - - def _factory(**kwargs: Any) -> GaussianBayesianNeuralNetwork: - return GaussianBayesianNeuralNetwork.cold_start( - n_features=N_FEATURES, hidden_dim_list=HIDDEN_DIMS, random_seed=RANDOM_SEED, **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(RANDOM_SEED) - bnn = make_gaussian_bnn(fit_reward_standardization=True, **FAST_FIT_KWARGS) - x = data_rng.normal(size=(N_TRAIN, N_FEATURES)) - bnn.update(context=x, rewards=_linear_rewards(x, data_rng).tolist()) - return bnn - - -@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=ACTION_IDS, - n_features=N_FEATURES, - hidden_dim_list=HIDDEN_DIMS, - random_seed=RANDOM_SEED, - **{**FAST_FIT_KWARGS, **kwargs}, - ) - - return _factory - - -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 = 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 = N_TRAIN, - n_features: int = 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 = N_POSTERIOR_SAMPLES, - offset: float = REWARD_OFFSET, - slope: float = REWARD_SLOPE, - noise_std: float = NOISE_STD, - mean_rel_tol: float = MEAN_REL_TOL, - sigma_rel_tol: float = SIGMA_REL_TOL, - probe_contexts: List[List[float]] = 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 = 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=-MAX_ABS_REWARD, max_value=MAX_ABS_REWARD, allow_nan=False), - min_size=2, - max_size=MAX_REWARDS, - unique=True, - ), - n_features=st.just(N_FEATURES), - tol=st.just(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=-MAX_ABS_REWARD, max_value=MAX_ABS_REWARD, allow_nan=False), - n_features=st.just(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=-MAX_ABS_REWARD, max_value=MAX_ABS_REWARD, allow_nan=False), - min_size=2, - max_size=MAX_REWARDS, - unique=True, - ), - n_features=st.just(N_FEATURES), - tol=st.just(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 = NOISE_SIGMA, - n_samples: int = N_TRAIN, - tol: float = 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 = N_TRAIN, - n_features: int = N_FEATURES, - short_fit_kwargs: Dict[str, Any] = 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 = N_TRAIN, - n_features: int = N_FEATURES, - short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, - rel_tol: float = 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 = N_TRAIN, - n_features: int = N_FEATURES, - n_posterior_samples: int = N_POSTERIOR_SAMPLES, - rel_tol: float = LEVEL_REL_TOL, - fast_fit_kwargs: Dict[str, Any] = 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(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", BAD_REWARDS) - def test_update_rejects_non_finite_rewards( - self, - make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], - bad_reward: float, - n_features: int = 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", [(USER_LOC, None), (None, USER_SCALE)], ids=["loc_only", "scale_only"] - ) - def test_fixed_standardization_can_be_partial( - self, - make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], - rng: np.random.Generator, - reward_loc: Optional[float], - reward_scale: Optional[float], - n_samples: int = N_TRAIN, - tol: float = ROUND_TRIP_TOL, - ) -> None: - """A fixed reward_loc or reward_scale can be given alone: the other one defaults to 0 / 1, never fitted.""" - bnn = make_gaussian_bnn(reward_loc=reward_loc, reward_scale=reward_scale) - rewards = rng.normal(size=n_samples) - loc = 0.0 if reward_loc is None else reward_loc - scale = 1.0 if reward_scale is None else reward_scale - np.testing.assert_allclose(bnn.prepare_rewards(rewards.tolist()), (rewards - loc) / scale, rtol=tol) - assert (bnn.reward_loc, bnn.reward_scale) == (reward_loc, reward_scale) - - @pytest.mark.parametrize( - "reward_loc, reward_scale", - [(USER_LOC, None), (None, USER_SCALE), (USER_LOC, USER_SCALE)], - ids=["loc_only", "scale_only", "both"], - ) - def test_rejects_fixed_values_with_fitted_standardization( - self, - make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], - reward_loc: Optional[float], - reward_scale: Optional[float], - ) -> None: - """fit_reward_standardization=True fits loc / scale from data, so given values would conflict.""" - with pytest.raises(ValidationError): - make_gaussian_bnn(fit_reward_standardization=True, reward_loc=reward_loc, reward_scale=reward_scale) - - def test_unstandardized_model_trains_and_predicts_on_raw_rewards( - self, - make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], - rng: np.random.Generator, - n_samples: int = N_TRAIN, - n_features: int = N_FEATURES, - noise_sigma: float = NOISE_SIGMA, - short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, - tol: float = ROUND_TRIP_TOL, - ) -> None: - """By default no loc / scale are set or fitted: targets are raw and noise_sigma is in raw units.""" - bnn = make_gaussian_bnn(noise_sigma=noise_sigma, **short_fit_kwargs) - rewards = rng.normal(size=n_samples) - np.testing.assert_array_equal(bnn.prepare_rewards(rewards.tolist()), rewards) - assert bnn.reward_loc is None and bnn.reward_scale is None - 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 = N_TRAIN, - tol: float = 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 = N_TRAIN, - n_features: int = N_FEATURES, - shift: float = SECOND_BATCH_SHIFT, - short_fit_kwargs: Dict[str, Any] = 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 = USER_LOC, - reward_scale: float = USER_SCALE, - n_samples: int = N_TRAIN, - n_features: int = N_FEATURES, - short_fit_kwargs: Dict[str, Any] = 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 = N_TRAIN, - n_features: int = N_FEATURES, - short_fit_kwargs: Dict[str, Any] = 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 and restored.reward_loc_init is None - 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 = DECAY_FACTOR, - n_samples: int = N_TRAIN, - n_features: int = N_FEATURES, - short_fit_kwargs: Dict[str, Any] = SHORT_FIT_KWARGS, - tol: float = 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 - - -class TestGaussianMetaModel: - """Tests for Gaussian heads in CmabMetaModel.""" - - def test_rejects_mixed_likelihoods( - self, - make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], - n_features: int = N_FEATURES, - hidden_dims: List[int] = HIDDEN_DIMS, - ) -> None: - """Bernoulli and Gaussian heads cannot share one meta-model.""" - bernoulli = BayesianNeuralNetwork.cold_start(n_features=n_features, hidden_dim_list=hidden_dims) - with pytest.raises(AttributeError): - CmabMetaModel(actions={"a": make_gaussian_bnn(), "b": bernoulli}) - - -class TestCmabGaussian: - """Tests for the CmabGaussian bandit.""" - - @pytest.mark.parametrize( - "extra_kwargs", - [ - {}, - {**BACKBONE_KWARGS, "update_kwargs": {**FAST_FIT_KWARGS["update_kwargs"], "batch_size": 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] = ARM_MEANS, - best_arm: str = BEST_ARM, - n_rounds: int = N_ROUNDS, - n_per_round: int = N_PER_ROUND, - n_features: int = N_FEATURES, - min_share: float = 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 = N_PER_ROUND, - n_features: int = N_FEATURES, - short_fit_kwargs: Dict[str, Any] = 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", BAD_REWARDS) - def test_rejects_non_finite_rewards( - self, - make_cmab_gaussian: Callable[..., CmabGaussian], - bad_reward: float, - n_features: int = N_FEATURES, - short_fit_kwargs: Dict[str, Any] = 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 = ACTION_IDS, - n_features: int = N_FEATURES, - delta: float = DELTA, - epsilon: float = EPSILON, - ) -> 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) - - -@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 = ACTION_IDS, - n_features: int = N_FEATURES, - reward: float = OUT_OF_RANGE_SOFT_REWARD, - short_fit_kwargs: Dict[str, Any] = 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..5a7bd8f4 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,435 @@ 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 +_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.""" + + def _factory(**kwargs: Any) -> GaussianBayesianNeuralNetwork: + return GaussianBayesianNeuralNetwork.cold_start( + n_features=_GAUSSIAN_N_FEATURES, + hidden_dim_list=_GAUSSIAN_HIDDEN_DIMS, + random_seed=_GAUSSIAN_RANDOM_SEED, + **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)], + ids=["loc_only", "scale_only"], + ) + def test_fixed_standardization_can_be_partial( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + reward_loc: Optional[float], + reward_scale: Optional[float], + n_samples: int = _GAUSSIAN_N_TRAIN, + tol: float = _GAUSSIAN_ROUND_TRIP_TOL, + ) -> None: + """A fixed reward_loc or reward_scale can be given alone: the other one defaults to 0 / 1, never fitted.""" + bnn = make_gaussian_bnn(reward_loc=reward_loc, reward_scale=reward_scale) + rewards = rng.normal(size=n_samples) + loc = 0.0 if reward_loc is None else reward_loc + scale = 1.0 if reward_scale is None else reward_scale + np.testing.assert_allclose(bnn.prepare_rewards(rewards.tolist()), (rewards - loc) / 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), (_GAUSSIAN_USER_LOC, _GAUSSIAN_USER_SCALE)], + ids=["loc_only", "scale_only", "both"], + ) + def test_rejects_fixed_values_with_fitted_standardization( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + reward_loc: Optional[float], + reward_scale: Optional[float], + ) -> None: + """fit_reward_standardization=True fits loc / scale from data, so given values would conflict.""" + with pytest.raises(ValidationError): + make_gaussian_bnn(fit_reward_standardization=True, reward_loc=reward_loc, reward_scale=reward_scale) + + def test_unstandardized_model_trains_and_predicts_on_raw_rewards( + self, + make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], + rng: np.random.Generator, + n_samples: int = _GAUSSIAN_N_TRAIN, + n_features: int = _GAUSSIAN_N_FEATURES, + noise_sigma: float = _GAUSSIAN_NOISE_SIGMA, + short_fit_kwargs: Dict[str, Any] = _GAUSSIAN_SHORT_FIT_KWARGS, + tol: float = _GAUSSIAN_ROUND_TRIP_TOL, + ) -> None: + """By default no loc / scale are set or fitted: targets are raw and noise_sigma is in raw units.""" + bnn = make_gaussian_bnn(noise_sigma=noise_sigma, **short_fit_kwargs) + rewards = rng.normal(size=n_samples) + np.testing.assert_array_equal(bnn.prepare_rewards(rewards.tolist()), rewards) + assert bnn.reward_loc is None and bnn.reward_scale is None + 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 and restored.reward_loc_init is None + 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 From 452c7ce23b2b967eb11a1ea6e12c45a1e08c1822 Mon Sep 17 00:00:00 2001 From: ronshiff1 Date: Mon, 12 Oct 2026 12:02:22 +0000 Subject: [PATCH 10/14] some PR fixes --- docs/src/continuous_reward_cmab.md | 80 ------------------------- docs/src/tutorials/cmab.ipynb | 70 ++++++++++++++++++++++ pybandits/meta_model/cmab_meta_model.py | 4 +- pybandits/model/bnn/network.py | 4 +- 4 files changed, 74 insertions(+), 84 deletions(-) delete mode 100644 docs/src/continuous_reward_cmab.md diff --git a/docs/src/continuous_reward_cmab.md b/docs/src/continuous_reward_cmab.md deleted file mode 100644 index 452c1e13..00000000 --- a/docs/src/continuous_reward_cmab.md +++ /dev/null @@ -1,80 +0,0 @@ -# Continuous rewards for cMAB - -pybandits supports three kinds of reward for contextual bandits: - -| reward | type | model | bandit | -|---|---|---|---| -| binary 0 / 1 | `BinaryReward` | `BayesianNeuralNetwork` (Bernoulli) | `CmabBernoulli` (and the CC / MO / DP variants) | -| soft label in [0, 1] | `SoftReward` (`use_soft_rewards=True`) | `BayesianNeuralNetwork` (Bernoulli on fractional targets) | the same | -| real-valued | `ContinuousReward` | `GaussianBayesianNeuralNetwork` | `CmabGaussian` | - -The bandit entry points accept `AnyReward`, and the actions manager checks each reward against what its models -support (`ActionsManager._check_rewards`). - -## The Gaussian model - -Each action is a Bayesian MLP with **one output unit, μ(x)**, and the reward is modelled as `Normal(μ(x), σ)`: - -- **σ is a single noise std per action,** not a function of the context. It is a latent (`log σ ~ Normal`) 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` (reward units) if given, otherwise on the std of the first update batch. -- **Training is on the targets** `(r − reward_loc) / reward_scale`. By default both are None, i.e. 0 / 1, so the - model trains on raw rewards and the weight priors are in reward units. Fixed constants can be given (alone or - together, e.g. only `reward_scale=20` for a centered uplift in dollars) so the default O(1) weight priors fit; they - are never re-fitted. With `fit_reward_standardization=True` they are instead fitted on the first batch and then - frozen (and must not be given). -- **Thompson sampling** draws the weights from the posterior and ranks actions on μ. `predict` returns μ (in the - "probabilities" slot) and σ (for monitoring) per action and row. -- **`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). - -### Why one σ and not σ(x) - -With a single noise level the fit of μ is a (Bayesian) least-squares fit, whose target is `E[reward | x]` whatever -the shape of the noise. A context-dependent σ(x) weights each residual by 1/σ(x)². On zero-inflated or skewed rewards -(e.g. revenue), the model can then explain large rewards as noise instead of raising μ, which biases μ low exactly -where the rewards are most variable. That bias is what Thompson sampling would rank on. - -### Recommended use: uplift against a control baseline - -For rewards with a lot of player-to-player variance (e.g. revenue), train on the **uplift** `r = R − E0(x)`, where -`E0(x)` is the expected reward under a control policy, from a separate baseline model. `E0` 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. `r` can be negative. - -```python -from pybandits.cmab import CmabGaussian - -mab = CmabGaussian.cold_start( - action_ids={"a1", "a2", "a3"}, - n_features=n_features, - hidden_dim_list=[16], - dist_type="normal", - dist_params_init={"mu": 0, "sigma": 0.1}, # a narrow prior fits much faster than the default sigma=1 - update_kwargs={"num_steps": 400, "optimizer_kwargs": {"step_size": 3e-3}}, - # optionally, fixed values from historical data (or fit_reward_standardization=True to fit on the first batch): - # reward_loc=..., reward_scale=..., noise_sigma=..., - # decay_factor=..., # if the reward process drifts -) -actions, mu, sigma = mab.predict(context=X) -mab.update(actions=actions, rewards=(R - E0).tolist(), context=X) -``` - -A shared backbone (`backbone_hidden_dims=...`) works as for `CmabBernoulli`; each action keeps its own σ and -standardization. - -### Alternatives that were evaluated and rejected - -- **σ(x) noise head:** biases μ on skewed rewards (see above). -- **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. -- **Hurdle × log-normal amount:** a good fit of the distribution's shape, but the back-transform `exp(m + σ²/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 × Gaussian amount (on the raw scale) is the better hurdle. - -### Not supported (yet) - -- Gaussian variants of the cost-control, multi-objective, dynamic-pricing and quantitative bandits. -- The adaptive window (`delta`) with continuous rewards. -- Offline policy evaluation and the simulators with continuous rewards. diff --git a/docs/src/tutorials/cmab.ipynb b/docs/src/tutorials/cmab.ipynb index 21b67e0d..0f552600 100644 --- a/docs/src/tutorials/cmab.ipynb +++ b/docs/src/tutorials/cmab.ipynb @@ -28,6 +28,21 @@ "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": {}, @@ -192,6 +207,22 @@ "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}$. By default both are `None`, i.e. 0 and 1, so the model trains on raw rewards and the weight priors are in reward units. Fixed constants can be given, alone or together (e.g. only `reward_scale=20` for a centered uplift in dollars), so the default $O(1)$ weight priors fit; they are never re-fitted. With `fit_reward_standardization=True` they are instead fitted on the first batch and then frozen (and must not be given).\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": {}, @@ -314,6 +345,45 @@ "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", + " # optionally, fixed values from historical data (or fit_reward_standardization=True to fit on the first batch):\n", + " # reward_loc=..., reward_scale=..., noise_sigma=...,\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/meta_model/cmab_meta_model.py b/pybandits/meta_model/cmab_meta_model.py index 650098d5..4b8f8dc1 100644 --- a/pybandits/meta_model/cmab_meta_model.py +++ b/pybandits/meta_model/cmab_meta_model.py @@ -458,7 +458,7 @@ def _joint_svi_update( # 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(targets, 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: @@ -470,7 +470,7 @@ def _joint_svi_update( arm_data = { arm: ( jnp.asarray(rows, dtype=jnp.int32), - jnp.asarray(targets[rows], dtype=BaseBayesianNeuralNetwork._reward_dtype), + jnp.asarray(targets[rows], dtype=BaseBayesianNeuralNetwork.reward_dtype), ) for arm, rows in arm_to_rows.items() } diff --git a/pybandits/model/bnn/network.py b/pybandits/model/bnn/network.py index 69b15379..c5bd1101 100644 --- a/pybandits/model/bnn/network.py +++ b/pybandits/model/bnn/network.py @@ -144,7 +144,7 @@ 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. _output_dim: ClassVar[PositiveInt] = 1 _logit_var_name: ClassVar[str] = "logit" @@ -1487,7 +1487,7 @@ def _fit(self, context: np.ndarray, rewards: List[Any]) -> None: _context = np.atleast_2d(context) x_jnp = jnp.array(_context, dtype=jnp.float32) - y_jnp = jnp.asarray(self.prepare_rewards(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: From d5b843a554c7b89230e9d898ed81987ec0b4a508 Mon Sep 17 00:00:00 2001 From: ronshiff1 Date: Mon, 12 Oct 2026 13:34:45 +0000 Subject: [PATCH 11/14] fix scaling + deafult output size --- docs/src/tutorials/cmab.ipynb | 11 +++++--- pybandits/model/bnn/network.py | 43 ++++++++++++++++++----------- tests/test_cmab.py | 14 +++++++--- tests/test_model.py | 49 +++++++++++++++++++++++----------- 4 files changed, 79 insertions(+), 38 deletions(-) diff --git a/docs/src/tutorials/cmab.ipynb b/docs/src/tutorials/cmab.ipynb index 0f552600..6e1b74df 100644 --- a/docs/src/tutorials/cmab.ipynb +++ b/docs/src/tutorials/cmab.ipynb @@ -214,7 +214,7 @@ "### 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}$. By default both are `None`, i.e. 0 and 1, so the model trains on raw rewards and the weight priors are in reward units. Fixed constants can be given, alone or together (e.g. only `reward_scale=20` for a centered uplift in dollars), so the default $O(1)$ weight priors fit; they are never re-fitted. With `fit_reward_standardization=True` they are instead fitted on the first batch and then frozen (and must not be given).\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`, neither may be given: both are fitted (mean / std) on the first batch and then frozen.\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", @@ -252,7 +252,7 @@ "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. By default the model is trained on raw rewards. For rewards far from $O(1)$ (e.g. revenue in dollars), pass a fixed `reward_scale` (and optionally `reward_loc`) so the default $O(1)$ weight priors fit, or `fit_reward_standardization=True` to fit both on the first batch. `noise_sigma` (in reward units) sets an initial guess of the noise std." + "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." ] }, { @@ -265,6 +265,8 @@ " 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", @@ -362,8 +364,9 @@ " 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", - " # optionally, fixed values from historical data (or fit_reward_standardization=True to fit on the first batch):\n", - " # reward_loc=..., reward_scale=..., noise_sigma=...,\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", diff --git a/pybandits/model/bnn/network.py b/pybandits/model/bnn/network.py index c5bd1101..39bba8d6 100644 --- a/pybandits/model/bnn/network.py +++ b/pybandits/model/bnn/network.py @@ -146,7 +146,7 @@ class BaseBayesianNeuralNetwork(Model, DNNMixin, ABC): 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. - _output_dim: ClassVar[PositiveInt] = 1 + default_output_dim: ClassVar[PositiveInt] = 1 _logit_var_name: ClassVar[str] = "logit" _prob_var_name: ClassVar[str] = "prob" weight_var_name: ClassVar[str] = "weight" @@ -398,7 +398,7 @@ def create_model_params( 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 default ``_output_dim``. + 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. @@ -414,7 +414,7 @@ def create_model_params( else: _dim_list = [effective_n_features] + hidden_dim_list - _dim_list.append(output_dim or cls._output_dim) + _dim_list.append(output_dim or cls.default_output_dim) layer_params_init = [] for layer_ind in range(len(_dim_list) - 1): @@ -811,7 +811,7 @@ def output_distribution( Parameters ---------- - linear_out : jax.Array of shape (batch, _output_dim) + 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). @@ -838,7 +838,7 @@ def _postprocess_output( Parameters ---------- - linear_out : np.ndarray of shape (n_samples, _output_dim) + 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). @@ -1756,20 +1756,24 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): 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. By default (both None) the targets are the raw rewards, so the weight - priors are in reward units: for rewards far from ``O(1)``, set a fixed ``reward_scale`` (or fit it from data with - ``fit_reward_standardization``) so the default ``O(1)`` weight priors fit. + 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 neither is given and both are fitted on the first update batch (``fit_reward_standardization=True``). + 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, ``reward_loc`` / ``reward_scale`` must not be given. When False, - the values given at cold start are used as fixed constants. Default is False. + and re-fitted after ``reset()``. When True, neither ``reward_loc`` nor ``reward_scale`` may be given. When + False, both must be given, and are used as fixed constants. Default is False. reward_loc : Optional[float] - Location subtracted from the rewards. None means 0 (or fitted, with ``fit_reward_standardization``). + Location subtracted from the rewards. Required when ``fit_reward_standardization`` is False; otherwise must + not be given, and is None until fitted on the first update. reward_scale : Optional[PositiveFloat] - Scale the rewards are divided by. None means 1 (or fitted, with ``fit_reward_standardization``). + Scale the rewards are divided by. Required when ``fit_reward_standardization`` is False; otherwise must not + be given, and is None until fitted on the first update. reward_loc_init : Optional[float] The cold-start value of ``reward_loc``, restored by ``reset()``. Set automatically from ``reward_loc`` at construction. @@ -1805,6 +1809,8 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): -------- >>> 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}, @@ -1825,7 +1831,7 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): supports_continuous_rewards: ClassVar[bool] = True # [mu]; the noise std is the scalar latent ``noise_log_sigma`` - _output_dim: ClassVar[PositiveInt] = 1 + default_output_dim: ClassVar[PositiveInt] = 1 _noise_var_name: ClassVar[str] = "noise_log_sigma" _transfer_learned_keys: ClassVar[Tuple[str, ...]] = ( @@ -1855,20 +1861,25 @@ def validate_reward_standardization(self) -> "GaussianBayesianNeuralNetwork": raise ValueError("reward_loc / reward_scale cannot be given with fit_reward_standardization=True.") if (self.reward_loc is None) != (self.reward_scale is None): raise ValueError("A fitted reward_loc and reward_scale must be either both set or both None.") + elif None in (self.reward_loc, self.reward_scale, self.reward_loc_init, self.reward_scale_init): + 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._output_dim: - raise ValueError(f"The output layer has {output_dim} units, expected {self._output_dim} (mu).") + 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 for an unset or not yet fitted value).""" + """``(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 diff --git a/tests/test_cmab.py b/tests/test_cmab.py index f601c919..ef2465f8 100644 --- a/tests/test_cmab.py +++ b/tests/test_cmab.py @@ -1665,6 +1665,8 @@ def test_cold_start_rejects_delta_with_backbone(delta: float) -> None: _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") @@ -1677,7 +1679,7 @@ def _factory(**kwargs: Any) -> CmabGaussian: n_features=_GAUSSIAN_N_FEATURES, hidden_dim_list=_GAUSSIAN_HIDDEN_DIMS, random_seed=_GAUSSIAN_RANDOM_SEED, - **{**_GAUSSIAN_FAST_FIT_KWARGS, **kwargs}, + **{**_GAUSSIAN_RAW_STANDARDIZATION, **_GAUSSIAN_FAST_FIT_KWARGS, **kwargs}, ) return _factory @@ -1690,9 +1692,12 @@ 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) + 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}) @@ -1770,10 +1775,13 @@ def test_rejects_adaptive_window( 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) + 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]) diff --git a/tests/test_model.py b/tests/test_model.py index 5a7bd8f4..5fb51071 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -2616,6 +2616,8 @@ def test_advi_zero_lr_posterior_equals_prior( _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")] @@ -2640,14 +2642,18 @@ def _noise_sigma_in_reward_units(bnn: GaussianBayesianNeuralNetwork) -> float: @pytest.fixture(scope="module") def make_gaussian_bnn() -> Callable[..., GaussianBayesianNeuralNetwork]: - """Factory: cold-starts a GaussianBayesianNeuralNetwork with the module's shape and extra kwargs.""" + """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, - **kwargs, + **{**standardization, **kwargs}, ) return _factory @@ -2859,24 +2865,34 @@ def test_update_rejects_non_finite_rewards( @pytest.mark.parametrize( "reward_loc, reward_scale", - [(_GAUSSIAN_USER_LOC, None), (None, _GAUSSIAN_USER_SCALE)], - ids=["loc_only", "scale_only"], + [(_GAUSSIAN_USER_LOC, None), (None, _GAUSSIAN_USER_SCALE), (None, None)], + ids=["loc_only", "scale_only", "neither"], ) - def test_fixed_standardization_can_be_partial( + def test_rejects_missing_fixed_values_without_fitted_standardization( self, make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], - rng: np.random.Generator, 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: - """A fixed reward_loc or reward_scale can be given alone: the other one defaults to 0 / 1, never fitted.""" + """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) - loc = 0.0 if reward_loc is None else reward_loc - scale = 1.0 if reward_scale is None else reward_scale - np.testing.assert_allclose(bnn.prepare_rewards(rewards.tolist()), (rewards - loc) / scale, rtol=tol) + 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( @@ -2894,21 +2910,24 @@ def test_rejects_fixed_values_with_fitted_standardization( with pytest.raises(ValidationError): make_gaussian_bnn(fit_reward_standardization=True, reward_loc=reward_loc, reward_scale=reward_scale) - def test_unstandardized_model_trains_and_predicts_on_raw_rewards( + def test_identity_standardization_trains_on_raw_rewards( self, make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], rng: np.random.Generator, n_samples: int = _GAUSSIAN_N_TRAIN, - n_features: int = _GAUSSIAN_N_FEATURES, 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: - """By default no loc / scale are set or fitted: targets are raw and noise_sigma is in raw units.""" - bnn = make_gaussian_bnn(noise_sigma=noise_sigma, **short_fit_kwargs) + """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 is None and bnn.reward_scale is None + 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( From b72458ed55db636c63372bfc997cb4d7d0498e9b Mon Sep 17 00:00:00 2001 From: ronshiff1 Date: Mon, 12 Oct 2026 15:45:37 +0000 Subject: [PATCH 12/14] removed _init params --- docs/src/tutorials/cmab.ipynb | 2 +- pybandits/actions_manager.py | 8 +++++ pybandits/model/bnn/network.py | 54 ++++++++++++---------------------- tests/test_actions_manager.py | 9 +++++- tests/test_model.py | 27 +++++++++++++---- 5 files changed, 57 insertions(+), 43 deletions(-) diff --git a/docs/src/tutorials/cmab.ipynb b/docs/src/tutorials/cmab.ipynb index 6e1b74df..c7a9b4de 100644 --- a/docs/src/tutorials/cmab.ipynb +++ b/docs/src/tutorials/cmab.ipynb @@ -214,7 +214,7 @@ "### 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`, neither may be given: both are fitted (mean / std) on the first batch and then frozen.\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", diff --git a/pybandits/actions_manager.py b/pybandits/actions_manager.py index 13d61a70..7d26e699 100644 --- a/pybandits/actions_manager.py +++ b/pybandits/actions_manager.py @@ -224,6 +224,14 @@ def _check_soft_rewards_delta(self) -> "ActionsManager": 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].""" diff --git a/pybandits/model/bnn/network.py b/pybandits/model/bnn/network.py index 39bba8d6..36a70144 100644 --- a/pybandits/model/bnn/network.py +++ b/pybandits/model/bnn/network.py @@ -1758,7 +1758,7 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): 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 neither is given and both are fitted on the first update batch (``fit_reward_standardization=True``). + 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. @@ -1766,20 +1766,15 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): ---------- 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, neither ``reward_loc`` nor ``reward_scale`` may be given. When - False, both must be given, and are used as fixed constants. Default is False. + 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 must - not be given, and is None until fitted on the first update. + 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 must not - be given, and is None until fitted on the first update. - reward_loc_init : Optional[float] - The cold-start value of ``reward_loc``, restored by ``reset()``. Set automatically from ``reward_loc`` at - construction. - reward_scale_init : Optional[PositiveFloat] - The cold-start value of ``reward_scale``, restored by ``reset()``. Set automatically from ``reward_scale`` at - construction. + 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 @@ -1802,8 +1797,8 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): 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 ``reward_loc`` / ``reward_scale`` to their cold-start values (re-fitted on the - next update with ``fit_reward_standardization``). + 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 -------- @@ -1821,8 +1816,6 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): fit_reward_standardization: bool = False reward_loc: Optional[float] = None reward_scale: Optional[PositiveFloat] = None - reward_loc_init: Optional[float] = None - reward_scale_init: Optional[PositiveFloat] = None n_observations: NonNegativeInt = 0 reward_sum: float = 0.0 noise_sigma: Optional[PositiveFloat] = None @@ -1837,31 +1830,20 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): _transfer_learned_keys: ClassVar[Tuple[str, ...]] = ( "reward_loc", "reward_scale", - "reward_loc_init", - "reward_scale_init", "n_observations", "reward_sum", "noise_log_sigma", ) - @model_validator(mode="before") - @classmethod - def set_standardization_init(cls, data: Any) -> Any: - """Record user-supplied ``reward_loc`` / ``reward_scale`` as their cold-start values (restored by reset).""" - if isinstance(data, dict): - for name in ("reward_loc", "reward_scale"): - if f"{name}_init" not in data and data.get(name) is not None: - data = {**data, f"{name}_init": data[name]} - return data - @model_validator(mode="after") def validate_reward_standardization(self) -> "GaussianBayesianNeuralNetwork": if self.fit_reward_standardization: - if self.reward_loc_init is not None or self.reward_scale_init is not None: - raise ValueError("reward_loc / reward_scale cannot be given with fit_reward_standardization=True.") if (self.reward_loc is None) != (self.reward_scale is None): - raise ValueError("A fitted reward_loc and reward_scale must be either both set or both None.") - elif None in (self.reward_loc, self.reward_scale, self.reward_loc_init, self.reward_scale_init): + 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)." @@ -1910,7 +1892,7 @@ def update_extra_params_from_vi(self, site_mu: dict, site_sigma: dict) -> None: 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=(np.asarray(self.noise_log_sigma.params["sigma"]) * inflation).tolist() + sigma=(self.noise_log_sigma.params["sigma"] * inflation).tolist() ) @staticmethod @@ -2098,8 +2080,8 @@ def _reset(self): self.n_observations = 0 self.reward_sum = 0.0 self.noise_log_sigma = None - self.reward_loc = self.reward_loc_init - self.reward_scale = self.reward_scale_init + if self.fit_reward_standardization: + self.reward_loc = self.reward_scale = None class BaseBayesianNeuralNetworkMO(ModelMO, ABC): 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_model.py b/tests/test_model.py index 5fb51071..cc171cae 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -2897,19 +2897,36 @@ def test_fixed_standardization_maps_rewards( @pytest.mark.parametrize( "reward_loc, reward_scale", - [(_GAUSSIAN_USER_LOC, None), (None, _GAUSSIAN_USER_SCALE), (_GAUSSIAN_USER_LOC, _GAUSSIAN_USER_SCALE)], - ids=["loc_only", "scale_only", "both"], + [(_GAUSSIAN_USER_LOC, None), (None, _GAUSSIAN_USER_SCALE)], + ids=["loc_only", "scale_only"], ) - def test_rejects_fixed_values_with_fitted_standardization( + def test_rejects_partial_fitted_standardization( self, make_gaussian_bnn: Callable[..., GaussianBayesianNeuralNetwork], reward_loc: Optional[float], reward_scale: Optional[float], ) -> None: - """fit_reward_standardization=True fits loc / scale from data, so given values would conflict.""" + """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], @@ -2999,7 +3016,7 @@ def test_reset_after_round_trip_clears_fitted_standardization( 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 and restored.reward_loc_init is None + assert restored.reward_loc == bnn.reward_loc restored.reset() assert restored.reward_loc is None and restored.reward_scale is None From e9b0e42f119de457f5c2531c432669a5123dc657 Mon Sep 17 00:00:00 2001 From: ronshiff1 Date: Mon, 12 Oct 2026 15:52:14 +0000 Subject: [PATCH 13/14] added transfer learning --- pybandits/model/bnn/network.py | 2 + pybandits/transfer.py | 13 +++-- tests/test_transfer.py | 95 +++++++++++++++++++++++++++++++++- 3 files changed, 105 insertions(+), 5 deletions(-) diff --git a/pybandits/model/bnn/network.py b/pybandits/model/bnn/network.py index 36a70144..f69d0d90 100644 --- a/pybandits/model/bnn/network.py +++ b/pybandits/model/bnn/network.py @@ -1834,6 +1834,8 @@ class GaussianBayesianNeuralNetwork(BaseBayesianNeuralNetwork): "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": 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/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) From 052f3c3cc592283b16002887f58e027f839079f2 Mon Sep 17 00:00:00 2001 From: ronshiff1 Date: Mon, 12 Oct 2026 17:13:20 +0000 Subject: [PATCH 14/14] rename arm_rewards --- pybandits/meta_model/cmab_meta_model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pybandits/meta_model/cmab_meta_model.py b/pybandits/meta_model/cmab_meta_model.py index 4b8f8dc1..9ce221c5 100644 --- a/pybandits/meta_model/cmab_meta_model.py +++ b/pybandits/meta_model/cmab_meta_model.py @@ -570,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):