Skip to content

gaussian cmab #165

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ronshiff1 merged 14 commits into from
Oct 12, 2026
Merged

gaussian cmab #165

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
1042258
added continous cmab
ronshiff1 Oct 9, 2026
8201a70
updates continous cmab
ronshiff1 Oct 9, 2026
cc2e3f9
changed normlizing to false in default + updated tutorial
ronshiff1 Oct 9, 2026
1c4dd2f
fix docstring
ronshiff1 Oct 9, 2026
b3da4be
change in loc/scalre mechansim
ronshiff1 Oct 9, 2026
0e4fe94
skill updated
ronshiff1 Oct 12, 2026
7a415f6
remove skill
ronshiff1 Oct 12, 2026
6e38025
tutorial moved
ronshiff1 Oct 12, 2026
6cf277f
tests mobed to old files
ronshiff1 Oct 12, 2026
452c7ce
some PR fixes
ronshiff1 Oct 12, 2026
d5b843a
fix scaling + deafult output size
ronshiff1 Oct 12, 2026
b72458e
removed _init params
ronshiff1 Oct 12, 2026
e9b0e42
added transfer learning
ronshiff1 Oct 12, 2026
052f3c3
rename arm_rewards
ronshiff1 Oct 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
232 changes: 220 additions & 12 deletions docs/src/tutorials/cmab.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,7 @@
"cells": [
{
"cell_type": "markdown",
"metadata": {
"pycharm": {
"name": "#%% md\n"
}
},
"metadata": {},
"source": [
"# Contextual Multi-Armed Bandit"
]
Expand All @@ -29,22 +25,42 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"The following notebook contains an example of usage of the class Cmab, which implements the algorithm above."
"The following notebook contains an example of usage of the classes `CmabBernoulli` (binary rewards) and `CmabGaussian` (continuous, real-valued rewards), which implement the algorithm above."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"pybandits supports three kinds of reward for contextual bandits:\n",
"\n",
"| reward | type | model | bandit |\n",
"|---|---|---|---|\n",
"| binary 0 / 1 | `BinaryReward` | `BayesianNeuralNetwork` (Bernoulli) | `CmabBernoulli` (and the CC / MO / DP variants) |\n",
"| soft label in $[0, 1]$ | `SoftReward` (`use_soft_rewards=True`) | `BayesianNeuralNetwork` (Bernoulli on fractional targets) | the same |\n",
"| real-valued | `ContinuousReward` | `GaussianBayesianNeuralNetwork` | `CmabGaussian` |\n",
"\n",
"The bandit entry points accept any of these reward types, and the actions manager checks each reward against what its models support."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Binary rewards: `CmabBernoulli`\n",
"\n",
"Use `CmabBernoulli` when the reward is binary (or a soft reward in $[0, 1]$), e.g. a click or a purchase."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"is_executing": false
}
},
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"\n",
"from pybandits.cmab import CmabBernoulli\n",
"from pybandits.cmab import CmabBernoulli, CmabGaussian\n",
"from pybandits.model import BayesianNeuralNetwork, BnnLayerParams, BnnParams, FeaturesConfig, StudentTArray"
]
},
Expand Down Expand Up @@ -179,6 +195,198 @@
"# update model\n",
"cmab.update(context=X, actions=pred_actions, rewards=simulated_rewards)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Continuous rewards: `CmabGaussian`\n",
"\n",
"Use `CmabGaussian` when the reward is real-valued, e.g. revenue or an uplift $R - E_0(x)$ against a control baseline computed outside the bandit. Each action is a `GaussianBayesianNeuralNetwork`: a BNN with a single mean head $\\mu_k(x)$ and a learned scalar noise std $\\sigma_k$, trained with a Gaussian likelihood $r \\sim \\mathcal{N}(\\mu_k(x), \\sigma_k)$. At prediction time the weights are sampled from the posterior and the bandit selects $a_t = argmax_k \\mu_k(x_t)$.\n",
"\n",
"Unlike `CmabBernoulli`, `predict` returns the sampled reward means in place of the probabilities, and the noise std $\\sigma_k$ (for monitoring only) in place of the weighted sums. The adaptive window (`delta`) is not supported, since its change-point test assumes binary rewards."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### The Gaussian model\n",
"\n",
"- **$\\sigma_k$ is a single noise std per action,** not a function of the context. It is a latent ($\\log \\sigma_k \\sim \\mathcal{N}$) whose posterior is stored in the model state (`noise_log_sigma`) and carried across updates, like the weights. Its initial prior is centered on `noise_sigma` (in reward units) if given, otherwise on the std of the first update batch.\n",
"- **Training is on the targets** $(r - \\text{reward\\_loc}) / \\text{reward\\_scale}$. They are set in exactly one of two ways. With `fit_reward_standardization=False` (the default), both must be given as fixed constants and are never re-fitted: choose them so the default $O(1)$ weight priors fit (e.g. `reward_loc=0, reward_scale=20` for a centered uplift in dollars), or use `reward_loc=0, reward_scale=1` to train on raw rewards. With `fit_reward_standardization=True`, leave both unset: they are fitted (mean / std) on the first batch, then frozen until `reset()`.\n",
"- **Thompson sampling** draws the weights from the posterior and ranks actions on $\\mu_k(x)$.\n",
"- **`reset()`** returns the model to its cold-start state, including the standardization: fixed values are restored, fitted ones are re-fitted on the next update.\n",
"\n",
"#### Why one $\\sigma$ and not $\\sigma(x)$\n",
"\n",
"With a single noise level, the fit of $\\mu$ is a (Bayesian) least-squares fit, whose target is $E[r \\mid x]$ whatever the shape of the noise. A context-dependent $\\sigma(x)$ weights each residual by $1/\\sigma(x)^2$. On zero-inflated or skewed rewards (e.g. revenue), the model can then explain large rewards as noise instead of raising $\\mu$, which biases $\\mu$ low exactly where the rewards are most variable. That bias is what Thompson sampling would rank on."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We simulate a two-action environment where the expected reward of each action depends linearly on the first context feature, plus Gaussian noise. Action `a1` is best when $x_0 > 0.25$ and `a2` otherwise."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"rng = np.random.default_rng(0)\n",
"noise_std = 0.5\n",
"\n",
"# expected reward of each action as a function of the context\n",
"true_mean = {\n",
" \"a1\": lambda X: 1.0 + 2.0 * X[:, 0],\n",
" \"a2\": lambda X: 2.0 - 2.0 * X[:, 0],\n",
"}\n",
"\n",
"X_gauss = 2 * rng.random((n_samples, n_features)) - 1 # random float in the interval (-1, 1)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The bandit is built with `cold_start`, which creates one `GaussianBayesianNeuralNetwork` per action id. We use a narrow prior on the weights and a larger optimizer step size, which make the unbounded Gaussian likelihood fit much faster. The rewards here are already $O(1)$, so we train on the raw rewards with the fixed standardization `reward_loc=0.0, reward_scale=1.0`. For rewards far from $O(1)$ (e.g. revenue in dollars), pass fixed values that bring them to $O(1)$, or set `fit_reward_standardization=True` (and no `reward_loc` / `reward_scale`) to fit both on the first batch. `noise_sigma` (in reward units) sets an initial guess of the noise std."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"cmab_gauss = CmabGaussian.cold_start(\n",
" action_ids={\"a1\", \"a2\"},\n",
" n_features=n_features,\n",
" hidden_dim_list=[16],\n",
" reward_loc=0.0,\n",
" reward_scale=1.0,\n",
" dist_type=\"normal\",\n",
" dist_params_init={\"mu\": 0, \"sigma\": 0.1},\n",
" update_kwargs={\"num_steps\": 1000, \"optimizer_kwargs\": {\"step_size\": 3e-3}},\n",
" random_seed=0,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Before any update, the actions are selected from the prior. `predict` returns, for each sample, the selected action, the sampled mean reward of every action and its noise std."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"pred_actions_g, mus, sigmas = cmab_gauss.predict(X_gauss)\n",
"print(\"Recommended action: {}\".format(pred_actions_g[:10]))\n",
"print(\"Sampled mean rewards: {}\".format(mus[0]))\n",
"print(\"Noise std: {}\".format(sigmas[0]))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We now observe the (continuous) rewards of the selected actions. Here they are simulated from the true model above."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"simulated_rewards_g = [\n",
" float(true_mean[a](X_gauss[i : i + 1])[0] + rng.normal(scale=noise_std)) for i, a in enumerate(pred_actions_g)\n",
"]\n",
"print(\"Simulated rewards: {}\".format(np.round(simulated_rewards_g[:10], 2)))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finally, we update the model providing per each sample the context $x_t$, the selected action $a_t$ and the corresponding reward $r_t$."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"cmab_gauss.update(context=X_gauss, actions=pred_actions_g, rewards=simulated_rewards_g)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"After the update, the bandit selects the best action for each context far more often. We compare its recommendations on fresh contexts with the truly best action."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"X_test = 2 * rng.random((500, n_features)) - 1\n",
"pred_actions_test, mus_test, sigmas_test = cmab_gauss.predict(X_test)\n",
"best_actions = np.where(true_mean[\"a1\"](X_test) > true_mean[\"a2\"](X_test), \"a1\", \"a2\")\n",
"print(\"Share of best-action recommendations: {:.2f}\".format(np.mean(np.array(pred_actions_test) == best_actions)))\n",
"print(\"Sampled mean rewards (first context): {}\".format(mus_test[0]))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Recommended use: uplift against a control baseline\n",
"\n",
"For rewards with a lot of player-to-player variance (e.g. revenue), train on the **uplift** $r = R - E_0(x)$, where $E_0(x)$ is the expected reward under a control policy, from a separate baseline model. $E_0$ doesn't depend on the action, so the ranking of actions is unchanged, and the shared variance is removed (a control variate). The posterior then contracts much faster with per-action data volumes. Note that $r$ can be negative.\n",
"\n",
"```python\n",
"mab = CmabGaussian.cold_start(\n",
" action_ids={\"a1\", \"a2\", \"a3\"},\n",
" n_features=n_features,\n",
" hidden_dim_list=[16],\n",
" dist_type=\"normal\",\n",
" dist_params_init={\"mu\": 0, \"sigma\": 0.1}, # a narrow prior fits much faster than the default sigma=1\n",
" update_kwargs={\"num_steps\": 400, \"optimizer_kwargs\": {\"step_size\": 3e-3}},\n",
" reward_loc=0.0, # fixed values, e.g. from historical data;\n",
" reward_scale=..., # or fit_reward_standardization=True instead of both, to fit them on the first batch\n",
" # noise_sigma=..., # optional initial guess of the noise std, in reward units\n",
" # decay_factor=..., # if the reward process drifts\n",
")\n",
"actions, mu, sigma = mab.predict(context=X)\n",
"mab.update(actions=actions, rewards=(R - E0).tolist(), context=X)\n",
"```\n",
"\n",
"A shared backbone (`backbone_hidden_dims=...`) works as for `CmabBernoulli`; each action keeps its own $\\sigma_k$ and standardization.\n",
"\n",
"### Alternatives that were evaluated and rejected\n",
"\n",
"- **$\\sigma(x)$ noise head:** biases $\\mu$ on skewed rewards (see above).\n",
"- **Bernoulli on a $[0, 1]$ rescaling of a continuous reward:** the mean is consistent, but the Bernoulli variance $p(1 - p)$ can overstate the real noise by orders of magnitude. Each observation then carries far too little weight, and the model learns little beyond an intercept.\n",
"- **Hurdle $\\times$ log-normal amount:** a good fit of the distribution's shape, but the back-transform $\\exp(m + \\sigma^2/2)$ can bias the level, and Thompson draws of it have extreme tails. If a split of the effect into conversion vs. amount is needed, a Bernoulli $\\times$ Gaussian amount (on the raw scale) is the better hurdle.\n",
"\n",
"### Not supported (yet)\n",
"\n",
"- Gaussian variants of the cost-control, multi-objective, dynamic-pricing and quantitative bandits.\n",
"- The adaptive window (`delta`) with continuous rewards.\n",
"- Offline policy evaluation and the simulators with continuous rewards."
]
}
],
"metadata": {
Expand Down
55 changes: 48 additions & 7 deletions pybandits/actions_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@

from pybandits.base import (
ActionId,
AnyReward,
PositiveProbability,
PyBanditsBaseModel,
Reward,
Expand All @@ -78,6 +79,7 @@
BetaDP,
BetaMO,
BetaMOCC,
GaussianBayesianNeuralNetwork,
Model,
ModelMO,
)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -217,8 +220,46 @@ def _check_soft_rewards_delta(self) -> "ActionsManager":
"""The adaptive window's change-point test and memory checks assume binary rewards."""
if self.delta is not None and self.use_soft_rewards:
raise ValueError("Soft rewards are not supported with the adaptive window (delta).")
if self.delta is not None and self.continuous_rewards:
raise ValueError("Continuous rewards are not supported with the adaptive window (delta).")
return self

@model_validator(mode="after")
def _check_uniform_reward_support(self) -> "ActionsManager":
"""All action models must take the same kind of reward: either all continuous or none."""
supports = {model.supports_continuous_rewards for model in self.actions.values()}
if len(supports) > 1:
raise ValueError("Action models must either all take continuous rewards or none of them.")
return self

@property
def continuous_rewards(self) -> bool:
"""Whether the action models take real-valued rewards rather than binary / soft ones in [0, 1]."""
return all(model.supports_continuous_rewards for model in self.actions.values())
Comment thread
ronshiff1 marked this conversation as resolved.

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]

@shaharbar1 shaharbar1 Oct 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

dtype float...

@ronshiff1 ronshiff1 Oct 12, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's a check for the input. it should be ok (unless the reward is above - 3.4e38 - the test will pass, but since later we use 32 bit, it will appear as inf. but this kind of input should not happen

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],
Expand Down Expand Up @@ -293,10 +334,10 @@ def sample_proba(
def update(
self,
actions: List[ActionId],
rewards: Union[List[Reward], List[List[Reward]]],
rewards: Union[List[AnyReward], List[List[AnyReward]]],
quantities: Optional[List[Union[float, List[float], None]]] = None,
actions_memory: Optional[List[ActionId]] = None,
rewards_memory: Optional[Union[List[Reward], List[List[Reward]]]] = None,
rewards_memory: Optional[Union[List[AnyReward], List[List[AnyReward]]]] = None,
**kwargs,
):
"""
Expand All @@ -317,8 +358,7 @@ def update(
List of previously collected rewards.
"""
self.actions_with_change.clear()
if not self.use_soft_rewards and not all(np.isin(r, (0, 1)).all() for r in (rewards, rewards_memory or [])):
raise ValueError("Rewards must be binary (0/1); set use_soft_rewards=True to pass soft rewards in [0, 1].")
self._check_rewards(rewards, rewards_memory or [])
if self.delta is None and (actions_memory or rewards_memory):
raise AttributeError("Adaptive window size is not set, so memory should not be provided.")
if self.delta is not None and (actions_memory is None or rewards_memory is None):
Expand Down Expand Up @@ -851,11 +891,11 @@ def _check_backbone_delta(self) -> "CmabActionsManager":
def update(
self,
actions: List[ActionId],
rewards: Union[List[Reward], List[List[Reward]]],
rewards: Union[List[AnyReward], List[List[AnyReward]]],
quantities: Optional[List[Union[float, List[float], None]]],
context: np.ndarray,
actions_memory: Optional[List[ActionId]] = None,
rewards_memory: Optional[Union[List[Reward], List[List[Reward]]]] = None,
rewards_memory: Optional[Union[List[AnyReward], List[List[AnyReward]]]] = None,
context_memory: Optional[np.ndarray] = None,
):
"""
Expand Down Expand Up @@ -967,3 +1007,4 @@ def _update_actions(
CmabActionsManagerDP = CmabActionsManager[Union[BayesianNeuralNetworkDP, QuantitativeBayesianNeuralNetworkDP]]
CmabActionsManagerMO = CmabActionsManager[BayesianNeuralNetworkMO]
CmabActionsManagerMOCC = CmabActionsManager[BayesianNeuralNetworkMOCC]
CmabActionsManagerGaussian = CmabActionsManager[GaussianBayesianNeuralNetwork]
Loading
Loading