cMAB Simulation

This notebook shows a simulation framework for the contextual multi-armed bandit (cMAB). It allows to study the behaviour of the bandit algoritm, to evaluate results and to run experiments on simulated data under different context, reward and action settings.

[1]:
from sklearn.datasets import make_classification

from pybandits.cmab import CmabBernoulli
from pybandits.cmab_simulator import CmabSimulator
from pybandits.model import BayesianLogisticRegression, BnnLayerParams, BnnParams, FeaturesConfig, StudentTArray
/home/runner/.cache/pypoetry/virtualenvs/pybandits-vYJB-miV-py3.10/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

First we need to define the simulation parameters. The parameters are split into two parts. The general parameters contain:

  • Number of update rounds

  • Number of samples per batch of update round

  • Seed for reproducibility

  • Verbosity enabler

  • Visualization enabler

The problem definition parameters contain:

  • Number of groups

  • Number of features

Data are processed in batches of size n>=1. Per each batch of simulated samples, the cMAB selects one action and collects the corresponding simulated reward for each sample. Then, prior parameters are updated based on returned rewards from recommended actions.

[2]:
# general simulator parameters
n_updates = 10
batch_size = 100
random_seed = None
verbose = True
visualize = True
[3]:
# problem definition simulation parameters
n_groups = 3
n_features = 5

Next, we initialize the context matrix \(X\) and the groups of samples. Samples that belong to the same group have features that come from the same distribution. Then, the action model and the cMAB are defined. We define three actions, each with a Bayesian Logistic Regression model. The model is defined by a Student-T prior for the intercept and a Student-T prior for each feature coefficient.

[4]:
# init context matrix and groups

context, group = make_classification(
    n_samples=batch_size * n_updates, n_features=n_features, n_informative=n_features, n_redundant=0, n_classes=n_groups
)
group = [str(g) for g in group]
[5]:
# define action model


def create_blr(n_features, bias_mu, bias_sigma, update_method, update_kwargs):
    """Create a BayesianLogisticRegression with given parameters."""
    bias = StudentTArray.cold_start(mu=bias_mu, sigma=bias_sigma, shape=1)
    weight = StudentTArray.cold_start(shape=(n_features, 1))
    layer_params = BnnLayerParams(weight=weight, bias=bias)
    model_params = BnnParams(bnn_layer_params=[layer_params])
    feature_config = FeaturesConfig(n_features=n_features)
    return BayesianLogisticRegression(
        model_params=model_params,
        feature_config=feature_config,
        update_method=update_method,
        update_kwargs=update_kwargs,
    )


update_method = "VI"
update_kwargs = {"num_steps": 100, "batch_size": 128, "optimizer_type": "adam"}
blr_kwargs = dict(
    n_features=n_features, bias_mu=1, bias_sigma=2, update_method=update_method, update_kwargs=update_kwargs
)
actions = {
    "a1": create_blr(**blr_kwargs),
    "a2": create_blr(**blr_kwargs),
    "a3": create_blr(**blr_kwargs),
}
# init contextual Multi-Armed Bandit model
cmab = CmabBernoulli(actions=actions)

Finally, we need to define the probabilities of positive rewards per each action/group, i.e. the ground truth (‘Action A’: 0.8 for group ‘0’ means that if the bandits selects ‘Action A’ for samples that belong to group ‘0’, then the environment will return a positive reward with 80% probability).

[6]:
# init probability of rewards randomly using splines
probs_reward = None

Now, we initialize the cMAB as shown in the previous notebook and the CmabSimulator with the parameters set above.

[7]:
# init simulation
cmab_simulator = CmabSimulator(
    mab=cmab,
    group=group,
    batch_size=batch_size,
    n_updates=n_updates,
    probs_reward=probs_reward,
    context=context,
    verbose=verbose,
)

Now, we can start simulation process by executing run() which performs the following steps:

For i=0 to n_updates:
    Extract batch[i] of samples from X
    Model recommends the best actions as the action with the highest reward probability to each simulated sample in batch[i] and collect corresponding simulated rewards
    Model priors are updated using information from recommended actions and returned rewards

Finally, we can visualize the results of the simulation. As defined in the ground truth: ‘a2’ was the action recommended the most for samples that belong to group ‘0’, ‘a1’ to group ‘1’ and both ‘a1’ and ‘a3’ to group ‘2’.

[8]:
cmab_simulator.run()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:335: FutureWarning: The behavior of DataFrame concatenation with empty or all-NA entries is deprecated. In a future version, this will no longer exclude empty or all-NA columns when determining the result dtypes. To retain the old behavior, exclude the relevant entries before the concat operation.
  self._results = pd.concat((self._results, batch_results), ignore_index=True)
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 36. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 30. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 34. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 19. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 49. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 32. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 37. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 36. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 27. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 31. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 38. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 40. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 39. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 21. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 31. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 36. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 33. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 38. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 24. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 36. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 32. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 58. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 17. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 25. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: overflow encountered in exp
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/simulator.py:229: RuntimeWarning: invalid value encountered in scalar divide
  return np.where(s >= 0, 1 / (1 + np.exp(-s)), np.exp(s) / (1 + np.exp(s))).item()
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 47. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 30. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
/home/runner/work/pybandits/pybandits/pybandits/model.py:1513: UserWarning: subsample_size does not match len(subsample), 128 vs 23. Did you accidentally use different subsample_size in the model and guide?
  with numpyro.plate("data", N, subsample_size=batch_size) as idx:
2026-03-29 18:47:05.932 | INFO     | pybandits.simulator:_print_results:541 - Simulation results (first 10 observations):

2026-03-29 18:47:05.953 | INFO     | pybandits.simulator:_print_results:542 - Count of actions selected by the bandit:

2026-03-29 18:47:05.955 | INFO     | pybandits.simulator:_print_results:543 - Observed proportion of positive rewards for each action:

Furthermore, we can examine the number of times each action was selected and the proportion of positive rewards for each action.

[9]:
cmab_simulator.selected_actions_count
[9]:
action a1 a2 a3 cum_a1 cum_a2 cum_a3
group batch
0 0.0 11 10 4 11 10 4
1 0.0 7 14 11 7 14 11
2 0.0 12 12 19 12 12 19
0 1.0 8 20 3 19 30 7
1 1.0 9 5 13 16 19 24
2 1.0 2 24 16 14 36 35
0 2.0 8 22 2 27 52 9
1 2.0 20 7 13 36 26 37
2 2.0 9 7 12 23 43 47
0 3.0 10 15 7 37 67 16
1 3.0 18 6 14 54 32 51
2 3.0 3 17 10 26 60 57
0 4.0 16 17 4 53 84 20
1 4.0 18 2 9 72 34 60
2 4.0 6 20 8 32 80 65
0 5.0 11 15 11 64 99 31
1 5.0 11 5 12 83 39 72
2 5.0 9 16 10 41 96 75
0 6.0 15 15 9 79 114 40
1 6.0 17 8 6 100 47 78
2 6.0 6 15 9 47 111 84
0 7.0 13 5 11 92 119 51
1 7.0 16 11 12 116 58 90
2 7.0 7 16 9 54 127 93
0 8.0 24 6 11 116 125 62
1 8.0 21 4 5 137 62 95
2 8.0 13 7 9 67 134 102
0 9.0 19 3 10 135 128 72
1 9.0 23 8 9 160 70 104
2 9.0 5 19 4 72 153 106
0 total 135 128 72 135 128 72
1 total 160 70 104 160 70 104
2 total 72 153 106 72 153 106
total total 367 351 282 367 351 282
[10]:
cmab_simulator.positive_reward_proportion
[10]:
proportion
action group
a1 0 0.874074
1 0.43125
2 0.555556
a2 0 0.59375
1 0.414286
2 0.836601
a3 0 0.680556
1 0.144231
2 0.226415