Contextual Multi-Armed Bandit

For the contextual multi-armed bandit (sMAB) when user information is available (context), we implemented a generalisation of Thompson sampling algorithm (Agrawal and Goyal, 2014) based on PyMC3.

The following notebook contains an example of usage of the class Cmab, which implements the algorithm above.

import numpy as np

from pybandits.cmab import CmabBernoulli
from pybandits.model import BayesianLogisticRegression, BnnLayerParams, BnnParams, StudentTArray

/home/runner/.cache/pypoetry/virtualenvs/pybandits-vYJB-miV-py3.10/lib/python3.10/site-packages/pydantic/_migration.py:283: UserWarning: `pydantic.generics:GenericModel` has been moved to `pydantic.BaseModel`.
  warnings.warn(f'`{import_path}` has been moved to `{new_location}`.')

n_samples = 1000
n_features = 5

First, we need to define the input context matrix \(X\) of size (\(n\_samples, n\_features\)) and the mapping of possible actions \(a_i \in A\) to their associated model.

# context
X = 2 * np.random.random_sample((n_samples, n_features)) - 1  # random float in the interval (-1, 1)
print("X: context matrix of shape (n_samples, n_features)")
print(X[:10])

X: context matrix of shape (n_samples, n_features)
[[ 0.64442861 -0.66783854  0.57177325  0.59325672 -0.2396677 ]
 [ 0.73459484 -0.72518906  0.64395804  0.14733409  0.72497825]
 [ 0.2032226   0.17261862 -0.39321721 -0.45604134 -0.46552433]
 [ 0.3539449  -0.78145525 -0.43865795 -0.5099745   0.66316842]
 [ 0.21236028  0.01069622  0.8277691  -0.48279275  0.20947206]
 [ 0.23633855 -0.54319104 -0.79053374 -0.06539662  0.12987522]
 [-0.65080713  0.95090505 -0.30150012  0.56471149  0.15935843]
 [-0.62085161 -0.58158232  0.05113742  0.70456057 -0.592402  ]
 [ 0.49030487 -0.11631171  0.39211183  0.4678212  -0.15845572]
 [-0.35140496  0.30247346  0.0421965   0.28874923  0.59270882]]

# define action model
bias = StudentTArray.cold_start(mu=1, sigma=2, 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])

actions = {
    "a1": BayesianLogisticRegression(model_params=model_params),
    "a2": BayesianLogisticRegression(model_params=model_params),
}

We can now init the bandit given the mapping of actions \(a_i\) to their model.

# init contextual Multi-Armed Bandit model
cmab = CmabBernoulli(actions=actions)

The predict function below returns the action selected by the bandit at time \(t\): \(a_t = argmax_k P(r=1|\beta_k, x_t)\). The bandit selects one action per each sample of the contect matrix \(X\).

# predict action
pred_actions, _, _ = cmab.predict(X)
print("Recommended action: {}".format(pred_actions[:10]))

/home/runner/.cache/pypoetry/virtualenvs/pybandits-vYJB-miV-py3.10/lib/python3.10/site-packages/pymc/data.py:384: FutureWarning: Data is now always mutable. Specifying the `mutable` kwarg will raise an error in a future release
  warnings.warn(
/home/runner/.cache/pypoetry/virtualenvs/pybandits-vYJB-miV-py3.10/lib/python3.10/site-packages/pytensor/link/c/cmodule.py:2968: UserWarning: PyTensor could not link to a BLAS installation. Operations that might benefit from BLAS will be severely degraded.
This usually happens when PyTensor is installed via pip. We recommend it be installed via conda/mamba/pixi instead.
Alternatively, you can use an experimental backend such as Numba or JAX that perform their own BLAS optimizations, by setting `pytensor.config.mode == 'NUMBA'` or passing `mode='NUMBA'` when compiling a PyTensor function.
For more options and details see https://pytensor.readthedocs.io/en/latest/troubleshooting.html#how-do-i-configure-test-my-blas-library
  warnings.warn(
Sampling: [bias_0, out, weight_0]
/home/runner/.cache/pypoetry/virtualenvs/pybandits-vYJB-miV-py3.10/lib/python3.10/site-packages/pymc/data.py:384: FutureWarning: Data is now always mutable. Specifying the `mutable` kwarg will raise an error in a future release
  warnings.warn(
Sampling: [bias_0, out, weight_0]

Recommended action: ['a1', 'a1', 'a2', 'a1', 'a1', 'a1', 'a2', 'a2', 'a2', 'a2']

Now, we observe the rewards and the context from the environment. In this example rewards and the context are randomly simulated.

# simulate reward from environment
simulated_rewards = np.random.randint(2, size=n_samples).tolist()
print("Simulated rewards: {}".format(simulated_rewards[:10]))

Simulated rewards: [0, 0, 0, 1, 1, 0, 1, 1, 0, 1]

Finally, we update the model providing per each action sample: (i) its context \(x_t\) (ii) the action \(a_t\) selected by the bandit, (iii) the corresponding reward \(r_t\).

# update model
cmab.update(context=X, actions=pred_actions, rewards=simulated_rewards)

/home/runner/.cache/pypoetry/virtualenvs/pybandits-vYJB-miV-py3.10/lib/python3.10/site-packages/pymc/data.py:384: FutureWarning: Data is now always mutable. Specifying the `mutable` kwarg will raise an error in a future release
  warnings.warn(
Initializing NUTS using adapt_diag...
Sequential sampling (2 chains in 1 job)
NUTS: [weight_0, bias_0]
Sampling 2 chains for 500 tune and 1_000 draw iterations (1_000 + 2_000 draws total) took 3 seconds.
We recommend running at least 4 chains for robust computation of convergence diagnostics
/home/runner/.cache/pypoetry/virtualenvs/pybandits-vYJB-miV-py3.10/lib/python3.10/site-packages/pymc/data.py:384: FutureWarning: Data is now always mutable. Specifying the `mutable` kwarg will raise an error in a future release
  warnings.warn(
Initializing NUTS using adapt_diag...
Sequential sampling (2 chains in 1 job)
NUTS: [weight_0, bias_0]
Sampling 2 chains for 500 tune and 1_000 draw iterations (1_000 + 2_000 draws total) took 3 seconds.
We recommend running at least 4 chains for robust computation of convergence diagnostics