{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Contextual Multi-Armed Bandit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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](https://arxiv.org/pdf/1209.3352.pdf)) based on NumPyro."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![title](img/cmab.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from pybandits.cmab import CmabBernoulli, CmabGaussian\n",
    "from pybandits.model import BayesianNeuralNetwork, BnnLayerParams, BnnParams, FeaturesConfig, StudentTArray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_samples = 1000\n",
    "n_features = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# context\n",
    "X = 2 * np.random.random_sample((n_samples, n_features)) - 1  # random float in the interval (-1, 1)\n",
    "print(\"X: context matrix of shape (n_samples, n_features)\")\n",
    "print(X[:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define action model\n",
    "bias = StudentTArray.cold_start(mu=1, sigma=2, shape=1)\n",
    "weight = StudentTArray.cold_start(shape=(n_features, 1))\n",
    "layer_params = BnnLayerParams(weight=weight, bias=bias)\n",
    "model_params = BnnParams(bnn_layer_params=[layer_params])\n",
    "feature_config = FeaturesConfig(n_features=n_features)\n",
    "\n",
    "update_kwargs = {\"num_steps\": 100, \"batch_size\": 128, \"optimizer_type\": \"adam\"}\n",
    "\n",
    "actions = {\n",
    "    \"a1\": BayesianNeuralNetwork(\n",
    "        model_params=model_params,\n",
    "        feature_config=feature_config,\n",
    "        update_kwargs=update_kwargs,\n",
    "    ),\n",
    "    \"a2\": BayesianNeuralNetwork(\n",
    "        model_params=model_params,\n",
    "        feature_config=feature_config,\n",
    "        update_kwargs=update_kwargs,\n",
    "    ),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now init the bandit given the mapping of actions $a_i$ to their model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# init contextual Multi-Armed Bandit model\n",
    "cmab = CmabBernoulli(actions=actions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# predict action\n",
    "pred_actions, _, _ = cmab.predict(X)\n",
    "print(\"Recommended action: {}\".format(pred_actions[:10]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we observe the rewards and the context from the environment. In this example rewards and the context are randomly simulated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simulate reward from environment\n",
    "simulated_rewards = np.random.randint(2, size=n_samples).tolist()\n",
    "print(\"Simulated rewards: {}\".format(simulated_rewards[:10]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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": {
  "hide_input": false,
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}