Contextual Multi-Armed Bandit

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

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 BayesianNeuralNetwork, 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

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.91620641 -0.77279409  0.33048932  0.2640768   0.34484474]
 [ 0.5125053  -0.75169093 -0.87127399 -0.21405076 -0.53944607]
 [-0.0537074   0.10669794 -0.50372155 -0.07859156  0.52321436]
 [-0.19836865  0.98750004 -0.8973727  -0.86936949 -0.90326119]
 [-0.1593579  -0.30111161 -0.59061157 -0.08986537  0.21175906]
 [ 0.22240962 -0.1185493  -0.64476461  0.57617029  0.22601114]
 [ 0.86726502 -0.53260971 -0.19891196 -0.02842786 -0.48005051]
 [ 0.53064069  0.0428547   0.53144329  0.48852445 -0.9971394 ]
 [ 0.26625276 -0.17786302  0.79713303  0.77805245 -0.45259756]
 [ 0.7735104  -0.34739189  0.46862325  0.03816798 -0.38249249]]

# 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])
feature_config = FeaturesConfig(n_features=n_features)

update_method = "VI"
update_kwargs = {"num_steps": 100, "batch_size": 128, "optimizer_type": "adam"}

actions = {
    "a1": BayesianNeuralNetwork(
        model_params=model_params,
        feature_config=feature_config,
        update_method=update_method,
        update_kwargs=update_kwargs,
    ),
    "a2": BayesianNeuralNetwork(
        model_params=model_params,
        feature_config=feature_config,
        update_method=update_method,
        update_kwargs=update_kwargs,
    ),
}

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]))

Recommended action: ['a1', 'a2', 'a1', 'a2', 'a2', 'a2', 'a2', 'a2', 'a1', '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, 1, 0, 0, 1, 1, 1, 1, 0]

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)