docs/tutorials/per_arm_bandits_tutorial.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nPjtEgqN4SjA"
      },
      "source": [
        "##### Copyright 2023 The TF-Agents Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6AZJOyCA4NpL"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GdBl31Dqwomt"
      },
      "source": [
        "# A Tutorial on Multi-Armed Bandits with Per-Arm Features"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y2vvG61d35bG"
      },
      "source": [
        "### Get Started\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/agents/tutorials/per_arm_bandits_tutorial\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003e\n",
        "    View on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/agents/blob/master/docs/tutorials/per_arm_bandits_tutorial.ipynb\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003e\n",
        "    Run in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/agents/blob/master/docs/tutorials/per_arm_bandits_tutorial.ipynb\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003e\n",
        "    View source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/agents/docs/tutorials/per_arm_bandits_tutorial.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ddRJNIiEwu9O"
      },
      "source": [
        "This tutorial is a step-by-step guide on how to use the TF-Agents library for contextual bandits problems where the actions (arms) have their own features, such as a list of movies represented by features (genre, year of release, ...)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I6mUk-hZa3pB"
      },
      "source": [
        "### Prerequisite\n",
        "\n",
        "It is assumed that the reader is somewhat familiar with the Bandit library of TF-Agents, in particular, has worked through the [tutorial for Bandits in TF-Agents](https://github.com/tensorflow/agents/tree/master/docs/tutorials/bandits_tutorial.ipynb) before reading this tutorial."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kscmAIA5xtJW"
      },
      "source": [
        "## Multi-Armed Bandits with Arm Features\n",
        "In the \"classic\" Contextual Multi-Armed Bandits setting, an agent receives a context vector (aka observation) at every time step and has to choose from a finite set of numbered actions (arms) so as to maximize its cumulative reward."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BDkno4bQ1vPE"
      },
      "source": [
        "\n",
        "Now consider the scenario where an agent recommends to a user the next movie to watch. Every time a decision has to be made, the agent receives as context some information about the user (watch history, genre preference, etc...), as well as the list of movies to choose from.\n",
        "\n",
        "We could try to formulate this problem by having the user information as the context and the arms would be `movie_1, movie_2, ..., movie_K`, but this approach has multiple shortcomings:\n",
        "\n",
        "\n",
        "*   The number of actions would have to be all the movies in the system and it is cumbersome to add a new movie.\n",
        "*   The agent has to learn a model for every single movie.\n",
        "*   Similarity between movies is not taken into account.\n",
        "\n",
        "Instead of numbering the movies, we can do something more intuitive: we can represent movies with a set of features including genre, length, cast, rating, year, etc. The advantages of this approach are manifold:\n",
        "\n",
        "*   Generalisation across movies.\n",
        "*   The agent learns just one reward function that models reward with user and movie features.\n",
        "*   Easy to remove from, or introduce new movies to the system.\n",
        "\n",
        "In this new setting, the number of actions does not even have to be the same in every time step.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hMXxIHXNRP5_"
      },
      "source": [
        "## Per-Arm Bandits in TF-Agents\n",
        "The TF-Agents Bandit suite is developed so that one can use it for the per-arm case as well. There are per-arm environments, and also most of the policies and agents can operate in per-arm mode.\n",
        "\n",
        "Before we dive into coding an example, we need the necessery imports."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cl5_CCIWSFvn"
      },
      "source": [
        "### Installation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QxiNIm5XSIIp"
      },
      "outputs": [],
      "source": [
        "!pip install tf-agents\n",
        "!pip install tf-keras"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WPuD0bMEY9Iz"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "# Keep using keras-2 (tf-keras) rather than keras-3 (keras).\n",
        "os.environ['TF_USE_LEGACY_KERAS'] = '1'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EDKNu5JTSDmf"
      },
      "source": [
        "### Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BbQXsoeKR2ui"
      },
      "outputs": [],
      "source": [
        "import functools\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "from tf_agents.bandits.agents import lin_ucb_agent\n",
        "from tf_agents.bandits.environments import stationary_stochastic_per_arm_py_environment as p_a_env\n",
        "from tf_agents.bandits.metrics import tf_metrics as tf_bandit_metrics\n",
        "from tf_agents.drivers import dynamic_step_driver\n",
        "from tf_agents.environments import tf_py_environment\n",
        "from tf_agents.replay_buffers import tf_uniform_replay_buffer\n",
        "from tf_agents.specs import tensor_spec\n",
        "from tf_agents.trajectories import time_step as ts\n",
        "\n",
        "nest = tf.nest"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q4cVyq3JMM7Z"
      },
      "source": [
        "### Parameters -- Feel Free to Play Around"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bfna8xm0MSCn"
      },
      "outputs": [],
      "source": [
        "# The dimension of the global features.\n",
        "GLOBAL_DIM = 40  #@param {type:\"integer\"}\n",
        "# The elements of the global feature will be integers in [-GLOBAL_BOUND, GLOBAL_BOUND).\n",
        "GLOBAL_BOUND = 10  #@param {type:\"integer\"}\n",
        "# The dimension of the per-arm features.\n",
        "PER_ARM_DIM = 50  #@param {type:\"integer\"}\n",
        "# The elements of the PER-ARM feature will be integers in [-PER_ARM_BOUND, PER_ARM_BOUND).\n",
        "PER_ARM_BOUND = 6  #@param {type:\"integer\"}\n",
        "# The variance of the Gaussian distribution that generates the rewards.\n",
        "VARIANCE = 100.0  #@param {type: \"number\"}\n",
        "# The elements of the linear reward parameter will be integers in [-PARAM_BOUND, PARAM_BOUND).\n",
        "PARAM_BOUND = 10  #@param {type: \"integer\"}\n",
        "\n",
        "NUM_ACTIONS = 70  #@param {type:\"integer\"}\n",
        "BATCH_SIZE = 20  #@param {type:\"integer\"}\n",
        "\n",
        "# Parameter for linear reward function acting on the\n",
        "# concatenation of global and per-arm features.\n",
        "reward_param = list(np.random.randint(\n",
        "      -PARAM_BOUND, PARAM_BOUND, [GLOBAL_DIM + PER_ARM_DIM]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N-mEz1HvRIBC"
      },
      "source": [
        "### A Simple Per-Arm Environment\n",
        "\n",
        "The stationary stochastic environment, explained in the other [tutorial](https://github.com/tensorflow/agents/tree/master/docs/tutorials/bandits_tutorial.ipynb), has a per-arm counterpart.\n",
        "\n",
        "To initialize the per-arm environment, one has to define functions that generate\n",
        "*   *global and per-arm features*: These functions have no input parameters and generate a single (global or per-arm) feature vector when called.\n",
        "*   *rewards*: This function takes as parameter the concatenation of a global and a per-arm feature vector, and generates a reward. Basically this is the function that the agent will have to \"guess\". It is worth noting here that in the per-arm case the reward function is identical for every arm. This is a fundamental difference from the classic bandit case, where the agent has to estimate reward functions for each arm independently.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kfTa5Y4ZYjhO"
      },
      "outputs": [],
      "source": [
        "def global_context_sampling_fn():\n",
        "  \"\"\"This function generates a single global observation vector.\"\"\"\n",
        "  return np.random.randint(\n",
        "      -GLOBAL_BOUND, GLOBAL_BOUND, [GLOBAL_DIM]).astype(np.float32)\n",
        "\n",
        "def per_arm_context_sampling_fn():\n",
        "  \"\"\"\"This function generates a single per-arm observation vector.\"\"\"\n",
        "  return np.random.randint(\n",
        "      -PER_ARM_BOUND, PER_ARM_BOUND, [PER_ARM_DIM]).astype(np.float32)\n",
        "\n",
        "def linear_normal_reward_fn(x):\n",
        "  \"\"\"This function generates a reward from the concatenated global and per-arm observations.\"\"\"\n",
        "  mu = np.dot(x, reward_param)\n",
        "  return np.random.normal(mu, VARIANCE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P2vpRPhheTo5"
      },
      "source": [
        "Now we are equipped to initialize our environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ny-yikfQQi9l"
      },
      "outputs": [],
      "source": [
        "per_arm_py_env = p_a_env.StationaryStochasticPerArmPyEnvironment(\n",
        "    global_context_sampling_fn,\n",
        "    per_arm_context_sampling_fn,\n",
        "    NUM_ACTIONS,\n",
        "    linear_normal_reward_fn,\n",
        "    batch_size=BATCH_SIZE\n",
        ")\n",
        "per_arm_tf_env = tf_py_environment.TFPyEnvironment(per_arm_py_env)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iIzFl8HiAIxg"
      },
      "source": [
        "Below we can check what this environment produces."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J8ZVqMU5AOzs"
      },
      "outputs": [],
      "source": [
        "print('observation spec: ', per_arm_tf_env.observation_spec())\n",
        "print('\\nAn observation: ', per_arm_tf_env.reset().observation)\n",
        "\n",
        "action = tf.zeros(BATCH_SIZE, dtype=tf.int32)\n",
        "time_step = per_arm_tf_env.step(action)\n",
        "print('\\nRewards after taking an action: ', time_step.reward)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eIlCGssRAzIr"
      },
      "source": [
        "We see that the observation spec is a dictionary with two elements:\n",
        "\n",
        "\n",
        "*   One with key `'global'`: this is the global context part, with shape matching the parameter `GLOBAL_DIM`.\n",
        "*   One with key `'per_arm'`: this is the per-arm context, and its shape is `[NUM_ACTIONS, PER_ARM_DIM]`. This part is the placeholder for the arm features for every arm in a time step.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eTpWbNVeS6ci"
      },
      "source": [
        "### The LinUCB Agent"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q83tgZR9LLUx"
      },
      "source": [
        "The LinUCB agent implements the identically named Bandit algorithm, which estimates the parameter of the linear reward function while also maintains a confidence ellipsoid around the estimate. The agent chooses the arm that has the highest estimated expected reward, assuming that the parameter lies within the confidence ellipsoid.\n",
        "\n",
        "Creating an agent requires the knowledge of the observation and the action specification. When defining the agent, we set the boolean parameter `accepts_per_arm_features` set to `True`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MqzA9Zi0Q2No"
      },
      "outputs": [],
      "source": [
        "observation_spec = per_arm_tf_env.observation_spec()\n",
        "time_step_spec = ts.time_step_spec(observation_spec)\n",
        "action_spec = tensor_spec.BoundedTensorSpec(\n",
        "    dtype=tf.int32, shape=(), minimum=0, maximum=NUM_ACTIONS - 1)\n",
        "\n",
        "agent = lin_ucb_agent.LinearUCBAgent(time_step_spec=time_step_spec,\n",
        "                                     action_spec=action_spec,\n",
        "                                     accepts_per_arm_features=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MaQlv1lpC-oc"
      },
      "source": [
        "### The Flow of Training Data\n",
        "\n",
        "This section gives a sneak peek into the mechanics of how per-arm features go from the policy to training. Feel free to jump to the next section (Defining the Regret Metric) and come back here later if interested."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gUSy2IFK5NpU"
      },
      "source": [
        "First, let us have a look at the data specification in the agent. The `training_data_spec` attribute of the agent specifies what elements and structure the training data should have."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YQRZd43o5M0j"
      },
      "outputs": [],
      "source": [
        "print('training data spec: ', agent.training_data_spec)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UyEFAHUg-m9V"
      },
      "source": [
        "If we have a closer look to the `observation` part of the spec, we see that it does not contain per-arm features!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lTBR1vgG-2AM"
      },
      "outputs": [],
      "source": [
        "print('observation spec in training: ', agent.training_data_spec.observation)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eDtmUgeJA_DN"
      },
      "source": [
        "What happened to the per-arm features? To answer this question, first we note that when the LinUCB agent trains, it does not need the per-arm features of **all** arms, it only needs those of the **chosen** arm. Hence, it makes sense to drop the tensor of shape `[BATCH_SIZE, NUM_ACTIONS, PER_ARM_DIM]`, as it is very wasteful, especially if the number of actions is large.\n",
        "\n",
        "But still, the per-arm features of the chosen arm must be somewhere! To this end, we make sure that the LinUCB policy stores the features of the chosen arm within the `policy_info` field of the training data:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C0fHvLm0Cpq9"
      },
      "outputs": [],
      "source": [
        "print('chosen arm features: ', agent.training_data_spec.policy_info.chosen_arm_features)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4psGJKGIDVgN"
      },
      "source": [
        "We see from the shape that the `chosen_arm_features` field has only the feature vector of one arm, and that will be the chosen arm. Note that the `policy_info`, and with it the `chosen_arm_features`, is part of the training data, as we saw from inspecting the training data spec, and thus it is available at training time."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ancYxxEHc-6Q"
      },
      "source": [
        "### Defining the Regret Metric"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4j2RTrBfzVJQ"
      },
      "source": [
        "Before starting the training loop, we define some utility functions that help calculate the regret of our agent. These functions help determining the optimal expected reward given the set of actions (given by their arm features) and the linear parameter that is hidden from the agent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N1_kJAOS6VPo"
      },
      "outputs": [],
      "source": [
        "def _all_rewards(observation, hidden_param):\n",
        "  \"\"\"Outputs rewards for all actions, given an observation.\"\"\"\n",
        "  hidden_param = tf.cast(hidden_param, dtype=tf.float32)\n",
        "  global_obs = observation['global']\n",
        "  per_arm_obs = observation['per_arm']\n",
        "  num_actions = tf.shape(per_arm_obs)[1]\n",
        "  tiled_global = tf.tile(\n",
        "      tf.expand_dims(global_obs, axis=1), [1, num_actions, 1])\n",
        "  concatenated = tf.concat([tiled_global, per_arm_obs], axis=-1)\n",
        "  rewards = tf.linalg.matvec(concatenated, hidden_param)\n",
        "  return rewards\n",
        "\n",
        "def optimal_reward(observation):\n",
        "  \"\"\"Outputs the maximum expected reward for every element in the batch.\"\"\"\n",
        "  return tf.reduce_max(_all_rewards(observation, reward_param), axis=1)\n",
        "\n",
        "regret_metric = tf_bandit_metrics.RegretMetric(optimal_reward)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i06WDbjrUSog"
      },
      "source": [
        "Now we are all set for starting our bandit training loop. The driver below takes care of choosing actions using the policy, storing rewards of chosen actions in the replay buffer, calculating the predefined regret metric, and executing the training step of the agent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j2Iimtmkzs0-"
      },
      "outputs": [],
      "source": [
        "num_iterations = 20 # @param\n",
        "steps_per_loop = 1 # @param\n",
        "\n",
        "replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(\n",
        "    data_spec=agent.policy.trajectory_spec,\n",
        "    batch_size=BATCH_SIZE,\n",
        "    max_length=steps_per_loop)\n",
        "\n",
        "observers = [replay_buffer.add_batch, regret_metric]\n",
        "\n",
        "driver = dynamic_step_driver.DynamicStepDriver(\n",
        "    env=per_arm_tf_env,\n",
        "    policy=agent.collect_policy,\n",
        "    num_steps=steps_per_loop * BATCH_SIZE,\n",
        "    observers=observers)\n",
        "\n",
        "regret_values = []\n",
        "\n",
        "for _ in range(num_iterations):\n",
        "  driver.run()\n",
        "  loss_info = agent.train(replay_buffer.gather_all())\n",
        "  replay_buffer.clear()\n",
        "  regret_values.append(regret_metric.result())\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yG5VMgSlUqYS"
      },
      "source": [
        " Now let's see the result. If we did everything right, the agent is able to estimate the linear reward function well, and thus the policy can pick actions whose expected reward is close to that of the optimal. This is indicated by our above defined regret metric, which goes down and approaches zero."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p4DOYwhMSUVh"
      },
      "outputs": [],
      "source": [
        "plt.plot(regret_values)\n",
        "plt.title('Regret of LinUCB on the Linear per-arm environment')\n",
        "plt.xlabel('Number of Iterations')\n",
        "_ = plt.ylabel('Average Regret')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t7ZgHgOx5Ojq"
      },
      "source": [
        "### What's Next?\n",
        "The above example is [implemented](https://github.com/tensorflow/agents/blob/master/tf_agents/bandits/agents/examples/v2/train_eval_per_arm_stationary_linear.py) in our codebase where you can choose from other agents as well, including the [Neural epsilon-Greedy agent](https://github.com/tensorflow/agents/blob/master/tf_agents/bandits/agents/neural_epsilon_greedy_agent.py)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "t7ZgHgOx5Ojq"
      ],
      "name": "per_arm_tutorial.ipynb",
      "provenance": [
        {
          "file_id": "1FGIvkvTjJ_yQ4sthnCKxMDtDyQjhn9GF",
          "timestamp": 1600257099961
        }
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}