06_CMAB_movielens_neural_egreedy_tf.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "06_CMAB_movielens_neural_egreedy_tf.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyMeE3TXgY4AMvKflFl63AeO",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/pstanisl/mlprague-2021/blob/main/06_CMAB_movielens_neural_egreedy_tf.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KfIFc7TXBvQ9"
      },
      "source": [
        "# Contextual Multi-Armed Bandits with neural network for reward prediction\n",
        "\n",
        "#### Installation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qbVPxLuUR25T"
      },
      "source": [
        "!pip install tf-agents -q"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P2MMqktAoX5e"
      },
      "source": [
        "!rm -f ./utils.py\n",
        "!wget --no-check-certificate --no-cache --no-cookies \\\n",
        "    https://raw.githubusercontent.com/pstanisl/mlprague-2021/main/utils.py \\\n",
        "    -O ./utils.py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iohvBdsZB3FL"
      },
      "source": [
        "#### Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rUEquHXHRwwB"
      },
      "source": [
        "import functools\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import random\n",
        "import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import\n",
        "import tensorflow_probability as tfp\n",
        "import zipfile\n",
        "\n",
        "from tqdm.notebook import trange\n",
        "from typing import Iterable, Optional, Sequence, Text, Tuple\n",
        "\n",
        "from tensorflow.python.util import nest\n",
        "from tf_agents.agents import data_converter\n",
        "from tf_agents.agents import tf_agent\n",
        "from tf_agents.bandits.agents import neural_epsilon_greedy_agent as eps_greedy_agent\n",
        "from tf_agents.bandits.agents import greedy_reward_prediction_agent\n",
        "from tf_agents.bandits.agents import utils as bandit_utils\n",
        "from tf_agents.bandits.environments import environment_utilities\n",
        "from tf_agents.bandits.environments import bandit_py_environment\n",
        "from tf_agents.bandits.metrics import tf_metrics as tf_bandit_metrics\n",
        "from tf_agents.bandits.policies import constraints as constr\n",
        "from tf_agents.bandits.policies import greedy_reward_prediction_policy as greedy_reward_policy\n",
        "from tf_agents.bandits.policies import linalg\n",
        "from tf_agents.bandits.policies import linear_bandit_policy\n",
        "from tf_agents.bandits.policies import policy_utilities\n",
        "from tf_agents.drivers import dynamic_step_driver\n",
        "from tf_agents.environments import tf_py_environment\n",
        "from tf_agents.networks import q_network\n",
        "from tf_agents.networks import network\n",
        "from tf_agents.networks import utils\n",
        "from tf_agents.policies import epsilon_greedy_policy\n",
        "from tf_agents.replay_buffers import tf_uniform_replay_buffer\n",
        "from tf_agents.specs import array_spec\n",
        "from tf_agents.trajectories import policy_step\n",
        "from tf_agents.trajectories import time_step as ts\n",
        "from tf_agents.typing import types\n",
        "from tf_agents.utils import common\n",
        "from tf_agents.utils import nest_utils\n",
        "\n",
        "from utils import load_movielens_data, plot_regret"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S7-tyjFvB8Xv"
      },
      "source": [
        "#### Downloading the [MovieLens](https://grouplens.org/datasets/movielens/) (100K) dataset.\n",
        "\n",
        "**Dataset info**\n",
        "\n",
        "MovieLens data sets were collected by the GroupLens Research Project\n",
        "at the University of Minnesota.\n",
        "\n",
        "This data set consists of:\n",
        "* 100,000 ratings (1-5) from 943 users on 1682 movies.\n",
        "* Each user has rated at least 20 movies.\n",
        "* Simple demographic info for the users (age, gender, occupation, zip)\n",
        "\n",
        "The data was collected through the MovieLens web site\n",
        "(movielens.umn.edu) during the seven-month period from September 19th,\n",
        "1997 through April 22nd, 1998. This data has been cleaned up - users\n",
        "who had less than 20 ratings or did not have complete demographic\n",
        "information were removed from this data set. Detailed descriptions of\n",
        "the data file can be found at the end of this file.\n",
        "\n",
        "Neither the University of Minnesota nor any of the researchers\n",
        "involved can guarantee the correctness of the data, its suitability\n",
        "for any particular purpose, or the validity of results based on the\n",
        "use of the data set. The data set may be used for any research\n",
        "purposes under the following conditions:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r8_xdGprnyqb"
      },
      "source": [
        "print(\"Downloading movielens data...\")\n",
        "\n",
        "!wget --no-check-certificate \\\n",
        "    http://files.grouplens.org/datasets/movielens/ml-100k.zip \\\n",
        "    -O ./movielens.zip\n",
        "\n",
        "zip_ref = zipfile.ZipFile('movielens.zip', \"r\")\n",
        "zip_ref.extractall()\n",
        "\n",
        "print(\"Done. Dataset contains:\")\n",
        "print(zip_ref.read('ml-100k/u.info').decode())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mv-RqLRcDKUY"
      },
      "source": [
        "#### Parameters -- Feel Free to Play Around"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D2Tf9d1NcbF9"
      },
      "source": [
        "EPSILON = 0.001 # @param {type: \"number\"}\n",
        "LR = 0.01 # @param {type: \"number\"}\n",
        "NUM_ACTIONS = 20 # @param {type:\"integer\"}\n",
        "RANK_K = 20 # @param {type:\"integer\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HvFAsPqLlmKQ"
      },
      "source": [
        "## Environment\n",
        "\n",
        "Implementation of the environment uses **MovieLens 100K dataset**. As described above, the dataset contains 100000 ratings from 943 users and 1682 movies. The environment can consider only the first $n$ of the dataset's movies. It can be set-up by `num_actions`. The number of \"known\" movies for the environment is equal to actions/arms.\n",
        "\n",
        "> Users without a rating (after selecting first $n$ movies) are removed from the environment."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jW6QuzcKX8gq"
      },
      "source": [
        "class MovieLensPyEnvironment(bandit_py_environment.BanditPyEnvironment):\n",
        "  \"\"\"Implements the MovieLens Bandit environment.\n",
        "  \n",
        "  This environment implements the MovieLens 100K dataset, available at:\n",
        "  https://www.kaggle.com/prajitdatta/movielens-100k-dataset\n",
        "  \n",
        "  This dataset contains 100K ratings from 943 users on 1682 items.\n",
        "  This csv list of:\n",
        "  user id | item id | rating | timestamp.\n",
        "  This environment computes a low-rank matrix factorization (using SVD) of the\n",
        "  data matrix A, such that: A ~= U * V.\n",
        "  \n",
        "  The reward of recommending item `j` to user `i` is provided as A_{ij}.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               data_dir: Text,\n",
        "               rank_k: int,\n",
        "               batch_size: int = 1,\n",
        "               num_movies: int = 20,\n",
        "               name: Optional[Text] = 'movielens'):\n",
        "    \"\"\"Initializes the MovieLens Bandit environment.\n",
        "    Args:\n",
        "      data_dir: (string) Directory where the data lies (in text form).\n",
        "      rank_k : (int) Which rank to use in the matrix factorization.\n",
        "      batch_size: (int) Number of observations generated per call.\n",
        "      num_movies: (int) Only the first `num_movies` movies will be used by the\n",
        "        environment. The rest is cut out from the data.\n",
        "      name: The name of this environment instance.\n",
        "    \"\"\"\n",
        "    self._num_actions = num_movies\n",
        "    self._batch_size = batch_size\n",
        "    self._context_dim = rank_k\n",
        "\n",
        "    # Compute the matrix factorization.\n",
        "    self._data_matrix = load_movielens_data(data_dir)\n",
        "    # Keep only the first items.\n",
        "    self._data_matrix = self._data_matrix[:, :num_movies]\n",
        "    # Filter the users with no iterm rated.\n",
        "    nonzero_users = list(np.nonzero(np.sum(self._data_matrix, axis=1) > 0.0)[0])\n",
        "    self._data_matrix = self._data_matrix[nonzero_users, :]\n",
        "    self._effective_num_users = len(nonzero_users)\n",
        "\n",
        "    # Compute the SVD.\n",
        "    u, s, vh = np.linalg.svd(self._data_matrix, full_matrices=False)\n",
        "\n",
        "    # Keep only the largest singular values.\n",
        "    self._u_hat = u[:, :rank_k] * np.sqrt(s[:rank_k])\n",
        "    self._v_hat = np.transpose(\n",
        "        np.transpose(vh[:rank_k, :]) * np.sqrt(s[:rank_k]))\n",
        "    self._approx_ratings_matrix = np.matmul(self._u_hat, self._v_hat)\n",
        "\n",
        "    self._current_users = np.zeros(batch_size, dtype=np.int32)\n",
        "    self._previous_users = np.zeros(batch_size, dtype=np.int32)\n",
        "\n",
        "    self._action_spec = array_spec.BoundedArraySpec(\n",
        "        shape=(),\n",
        "        dtype=np.int32,\n",
        "        minimum=0,\n",
        "        maximum=self._num_actions - 1,\n",
        "        name='action')\n",
        "    observation_spec = array_spec.ArraySpec(\n",
        "        shape=(self._context_dim,), dtype=np.float64, name='observation')\n",
        "    self._time_step_spec = ts.time_step_spec(observation_spec)\n",
        "    self._observation = np.zeros((self._batch_size, self._context_dim))\n",
        "\n",
        "    self._optimal_action_table = np.argmax(\n",
        "        self._approx_ratings_matrix, axis=1)\n",
        "    self._optimal_reward_table = np.max(\n",
        "        self._approx_ratings_matrix, axis=1)\n",
        "\n",
        "    super(MovieLensPyEnvironment, self).__init__(\n",
        "        observation_spec, self._action_spec)\n",
        "\n",
        "  @property\n",
        "  def batch_size(self):\n",
        "    return self._batch_size\n",
        "\n",
        "  @property\n",
        "  def batched(self):\n",
        "    return True\n",
        "\n",
        "  def _observe(self):\n",
        "    \"\"\"Returns the u vectors of a random sample of users.\"\"\"\n",
        "    sampled_users = random.sample(\n",
        "        range(self._effective_num_users), self._batch_size)\n",
        "    self._previous_users = self._current_users\n",
        "    self._current_users = sampled_users\n",
        "    batched_observations = self._u_hat[sampled_users]\n",
        "    return batched_observations\n",
        "\n",
        "  def _apply_action(self, action):\n",
        "    \"\"\"Computes the reward for the input actions.\"\"\"\n",
        "    rewards = []\n",
        "    for i, j in zip(self._current_users, action):\n",
        "      rewards.append(self._approx_ratings_matrix[i, j])\n",
        "    return np.array(rewards)\n",
        "\n",
        "  def compute_optimal_action(self):\n",
        "    return self._optimal_action_table[self._previous_users]\n",
        "\n",
        "  def compute_optimal_reward(self):\n",
        "    return self._optimal_reward_table[self._previous_users]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SfdzNn4GWey2"
      },
      "source": [
        "Now we are equipped to initialize our environment."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JWLNxLEiT05z"
      },
      "source": [
        "env = MovieLensPyEnvironment('./ml-100k/u.data', RANK_K, 1, num_movies=NUM_ACTIONS)\n",
        "tf_env = tf_py_environment.TFPyEnvironment(env)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VH-uCHRrWbmh"
      },
      "source": [
        "Below we can check what this environment produces."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vcRZbsaSUCEh"
      },
      "source": [
        "print('Observation spec:', tf_env.observation_spec())\n",
        "print('An observation: ', tf_env.reset().observation.numpy())\n",
        "\n",
        "action = tf.zeros(1, dtype=tf.int32)\n",
        "time_step = tf_env.step(action)\n",
        "\n",
        "print(f'For users={env._previous_users}, we selected action={action.numpy()} (optimal={tf_env.compute_optimal_action()})')\n",
        "print(f'For users={env._previous_users}, we received reward={time_step.reward.numpy()} (optimal={tf_env.compute_optimal_reward()})')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qTjTE8ThV85o"
      },
      "source": [
        "## Policy\n",
        "\n",
        "In Reinforcement Learning terminology, policies map an observation from the environment to an action or a distribution over actions. In TF-Agents, observations from the environment are contained in a named tuple `TimeStep('step_type', 'discount', 'reward', 'observation')`, and policies map timesteps to actions or distributions over actions. The policies use `timestep.observation`, `timestep.reward` are ignored during selecting an action.\n",
        "\n",
        "The policy is related to other components in TF-Agents in the following way. In this case, the policy has a neural network to compute actions and from TimeSteps. Agents can contain one or more policies for different purposes, i.e., our `NeuralEpsilonGreedyAgent` (defined below) will use `GreedyRewardPredictionPolicy` for greedy selections and [RandomTFPolicy](https://github.com/tensorflow/agents/blob/master/tf_agents/policies/random_tf_policy.py) otherwise.\n",
        "\n",
        "<br/>\n",
        "\n",
        "TASK: Fill in the missing pieces of code and create:\n",
        "\n",
        "1. prediction of the reward\n",
        "1. choosing the next action\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "10eDwxGDV-n6"
      },
      "source": [
        "class GreedyRewardPredictionPolicy(greedy_reward_policy.GreedyRewardPredictionPolicy):\n",
        "  \"\"\"Simplified implementation of the Greedy policy using tf_agents.Network for\n",
        "  predicting a reward.\n",
        "  \n",
        "  Original implementation can be found here: http://bit.ly/3dIw8Wf\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      time_step_spec: types.TimeStep,\n",
        "      action_spec: types.NestedTensorSpec,\n",
        "      reward_network: types.Network,\n",
        "      name: Optional[Text] = None):\n",
        "    \"\"\"\"Builds a GreedyRewardPredictionPolicy given a reward tf_agents.Network.\n",
        "\n",
        "    This policy takes a tf_agents.Network predicting rewards and generates the\n",
        "    action corresponding to the largest predicted reward.\n",
        "\n",
        "    Args:\n",
        "      time_step_spec: A `TimeStep` spec of the expected time_steps.\n",
        "      action_spec: A nest of BoundedTensorSpec representing the actions.\n",
        "      reward_network: An instance of a `tf_agents.network.Network`,\n",
        "        callable via `network(observation, step_type) -> (output, final_state)`.\n",
        "      name: The name of this policy. All variables in this module will fall\n",
        "        under that name. Defaults to the class name.\n",
        "\n",
        "    Raises:\n",
        "      NotImplementedError: If `action_spec` contains more than one\n",
        "        `BoundedTensorSpec` or the `BoundedTensorSpec` is not valid.\n",
        "    \"\"\"\n",
        "    super(GreedyRewardPredictionPolicy, self).__init__(\n",
        "        time_step_spec,\n",
        "        action_spec,\n",
        "        reward_network,\n",
        "        name=name\n",
        "    )\n",
        "\n",
        "  def _variables(self):\n",
        "    return self._reward_network.variables\n",
        "  \n",
        "  def _distribution(self, time_step, policy_state):\n",
        "    observation = time_step.observation\n",
        "    # Get the predicted rewards from the network\n",
        "    predicted_reward_values, policy_state = # YOUR CODE HERE\n",
        "    batch_size = tf.shape(predicted_reward_values)[0]\n",
        "    \n",
        "    # Check predicted values shape properties\n",
        "    predicted_reward_values.shape.with_rank_at_least(2)\n",
        "    predicted_reward_values.shape.with_rank_at_most(3)\n",
        "    if predicted_reward_values.shape[\n",
        "        -1] is not None and predicted_reward_values.shape[\n",
        "            -1] != self._expected_num_actions:\n",
        "      raise ValueError(\n",
        "          'The number of actions ({}) does not match the reward_network output'\n",
        "          ' size ({}).'.format(self._expected_num_actions,\n",
        "                               predicted_reward_values.shape[1]))\n",
        "    \n",
        "    # Get best actions with argmax\n",
        "    actions = tf.argmax(\n",
        "        # YOUR CODE HERE, \n",
        "        axis=-1, output_type=self.action_spec.dtype) \n",
        "    \n",
        "    # Offset actions with a minimum value from the action_spec\n",
        "    actions += self._action_offset\n",
        "\n",
        "    policy_info = policy_utilities.PolicyInfo()\n",
        "\n",
        "    return policy_step.PolicyStep(\n",
        "        tfp.distributions.Deterministic(loc=actions), policy_state, policy_info)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uo3YPdpkxYsr"
      },
      "source": [
        "## Agent\n",
        "\n",
        "##### Network\n",
        "\n",
        "For our recommender, we will implement QNetwork. It is used in Qlearning for environments with discrete actions, this network maps an observation to value estimates for each possible action. It is feed forward network.\n",
        "\n",
        "It consists from encoding part and Q part. Encoding part creates encoding of the input context vector and the output is this part is input of the last Q layer.\n",
        "\n",
        "Encoding part can cosists of:\n",
        "\n",
        "  * Preprocessing layers\n",
        "  * Preprocessing combiner\n",
        "  * Conv2D\n",
        "  * Flatten\n",
        "  * Dense\n",
        "\n",
        "More information can be found [here](https://www.tensorflow.org/agents/tutorials/8_networks_tutorial#encodingnetwork).\n",
        "\n",
        "<br/>\n",
        "\n",
        "**TASK**: Fill in the missing pieces of code and create:\n",
        "\n",
        "1. An encoding part of the network which contains 1 Flatten layer and 3 Dense layers (i.e., with 50 neurons in each layer).\n",
        "1. In `call` method do the prediction with network."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V-udQcXMjp7s"
      },
      "source": [
        "class QNetwork(network.Network):\n",
        "  \"\"\"Simplified implementation of q_network.QNetwork from tf-agents\n",
        "  \n",
        "  Original implementation can be found here: http://bit.ly/3qVf2bB\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               input_tensor_spec,\n",
        "               action_spec,\n",
        "               activation_fn=tf.keras.activations.relu,\n",
        "               batch_squash=True,\n",
        "               dtype=tf.float32,\n",
        "               name='QNetwork'):\n",
        "    \"\"\"Creates an instance of `QNetwork`.\n",
        "\n",
        "    Args:\n",
        "      input_tensor_spec: A nest of `tensor_spec.TensorSpec` representing the\n",
        "        input observations.\n",
        "      action_spec: A nest of `tensor_spec.BoundedTensorSpec` representing the\n",
        "        actions.\n",
        "      activation_fn: Activation function, e.g. tf.keras.activations.relu.\n",
        "      batch_squash: If True the outer_ranks of the observation are squashed into\n",
        "        the batch dimension. This allow encoding networks to be used with\n",
        "        observations with shape [BxTx...].\n",
        "      dtype: The dtype to use by the convolution and fully connected layers.\n",
        "      name: A string representing the name of the network.\n",
        "\n",
        "    Raises:\n",
        "      ValueError: If `input_tensor_spec` contains more than one observation. Or\n",
        "        if `action_spec` contains more than one action.\n",
        "    \"\"\"\n",
        "    q_network.validate_specs(action_spec, input_tensor_spec)\n",
        "    action_spec = tf.nest.flatten(action_spec)[0]\n",
        "    num_actions = action_spec.maximum - action_spec.minimum + 1\n",
        "\n",
        "    kernel_initializer = tf.compat.v1.variance_scaling_initializer(\n",
        "        scale=2.0, mode='fan_in', distribution='truncated_normal')\n",
        "    \n",
        "    # Create encoder part\n",
        "    layers = []\n",
        "\n",
        "    # YOUR CODE HERE\n",
        "    \n",
        "    \n",
        "    # END OF YOUR CODE\n",
        "\n",
        "    # Create Q part\n",
        "    q_value_layer = tf.keras.layers.Dense(\n",
        "        num_actions,\n",
        "        activation=None,\n",
        "        kernel_initializer=tf.random_uniform_initializer(\n",
        "            minval=-0.03, maxval=0.03),\n",
        "        bias_initializer=tf.constant_initializer(-0.2),\n",
        "        dtype=dtype)\n",
        "\n",
        "    super(QNetwork, self).__init__(\n",
        "        input_tensor_spec=input_tensor_spec,\n",
        "        state_spec=(),\n",
        "        name=name)\n",
        "    \n",
        "    self._batch_squash = batch_squash\n",
        "    self.built = True  # Allow access to self.variables\n",
        "\n",
        "    self._encoder_layers = layers\n",
        "    self._q_value_layer = q_value_layer\n",
        "  \n",
        "  def call(self, observation, step_type=None, network_state=(), training=False):\n",
        "    \"\"\"Runs the given observation through the network.\n",
        "\n",
        "    Args:\n",
        "      observation: The observation to provide to the network.\n",
        "      step_type: The step type for the given observation. See `StepType` in\n",
        "        time_step.py.\n",
        "      network_state: A state tuple to pass to the network, mainly used by RNNs.\n",
        "      training: Whether the output is being used for training.\n",
        "\n",
        "    Returns:\n",
        "      A tuple `(logits, network_state)`.\n",
        "    \"\"\"\n",
        "    del step_type  # unused.\n",
        "\n",
        "    if self._batch_squash:\n",
        "      outer_rank = nest_utils.get_outer_rank(\n",
        "          observation, self.input_tensor_spec)\n",
        "      batch_squash = utils.BatchSquash(outer_rank)\n",
        "      observation = tf.nest.map_structure(batch_squash.flatten, observation)\n",
        "    \n",
        "    state = observation\n",
        "\n",
        "    # HINT: set training=training, see https://keras.io/api/models/model/\n",
        "\n",
        "    # YOUR CODE HERE\n",
        "\n",
        "\n",
        "    # END OF YOUR CODE\n",
        "\n",
        "    if self._batch_squash:\n",
        "      state = tf.nest.map_structure(batch_squash.unflatten, state)\n",
        "\n",
        "    q_value = # YOUR CODE HERE\n",
        "\n",
        "    return q_value, network_state"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EGWJBkCoXKCA"
      },
      "source": [
        "`NeuralEpsilonAgent` use `EpsilonGreedyPolicy` with our `GreedyRewardPredictionPolicy` in greedy situations.\r\n",
        "\r\n",
        "<br/>\r\n",
        "\r\n",
        "**TASK**: Fill in the missing pieces of code (in `_train` and `reward_loss`)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FsgZwX_0MzAf"
      },
      "source": [
        "class NeuralEpsilonGreedyAgent(\n",
        "    greedy_reward_prediction_agent.GreedyRewardPredictionAgent):\n",
        "  \"\"\"Simplified implementation of neural network based epsilon greedy agent.\n",
        "  This agent receives a neural network that it trains to predict rewards. The\n",
        "  action is chosen greedily with respect to the prediction with probability\n",
        "  `1 - epsilon`, and uniformly randomly with probability `epsilon`.\n",
        "\n",
        "  Original implementation can be found here: http://bit.ly/3dIILki\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      time_step_spec: types.TimeStep,\n",
        "      action_spec: types.BoundedTensorSpec,\n",
        "      reward_network: types.Network,\n",
        "      optimizer: types.Optimizer,\n",
        "      epsilon: float,\n",
        "      # Params for training.\n",
        "      error_loss_fn: types.LossFn = tf.compat.v1.losses.mean_squared_error,\n",
        "      # Params for debugging.\n",
        "      debug_summaries: bool = False,\n",
        "      name: Optional[Text] = None):\n",
        "    \"\"\"Creates a Neural Epsilon Greedy Agent.\n",
        "    For more details about the Laplacian smoothing regularization, please see\n",
        "    the documentation of the `GreedyRewardPredictionAgent`.\n",
        "    Args:\n",
        "      time_step_spec: A `TimeStep` spec of the expected time_steps.\n",
        "      action_spec: A nest of `BoundedTensorSpec` representing the actions.\n",
        "      reward_network: A `tf_agents.network.Network` to be used by the agent. The\n",
        "        network will be called with call(observation, step_type) and it is\n",
        "        expected to provide a reward prediction for all actions.\n",
        "        *Note*: when using `observation_and_action_constraint_splitter`, make\n",
        "        sure the `reward_network` is compatible with the network-specific half\n",
        "        of the output of the `observation_and_action_constraint_splitter`. In\n",
        "        particular, `observation_and_action_constraint_splitter` will be called\n",
        "        on the observation before passing to the network.\n",
        "      optimizer: The optimizer to use for training.\n",
        "      epsilon: A float representing the probability of choosing a random action\n",
        "        instead of the greedy action.\n",
        "      error_loss_fn: A function for computing the error loss, taking parameters\n",
        "        labels, predictions, and weights (any function from tf.losses would\n",
        "        work). The default is `tf.losses.mean_squared_error`.)\n",
        "      debug_summaries: A Python bool, default False. When True, debug summaries\n",
        "        are gathered.\n",
        "      name: Python str name of this agent. All variables in this module will\n",
        "        fall under that name. Defaults to the class name.\n",
        "    Raises:\n",
        "      ValueError: If the action spec contains more than one action or or it is\n",
        "      not a bounded scalar int32 spec with minimum 0.\n",
        "    \"\"\"\n",
        "    super(NeuralEpsilonGreedyAgent, self).__init__(\n",
        "        time_step_spec=time_step_spec,\n",
        "        action_spec=action_spec,\n",
        "        reward_network=reward_network,\n",
        "        optimizer=optimizer,\n",
        "        error_loss_fn=error_loss_fn,\n",
        "        debug_summaries=debug_summaries,\n",
        "        name=name)\n",
        "    # Set our custom policy as greedy policy in Epsilon-Greedy scenario\n",
        "    policy = GreedyRewardPredictionPolicy(\n",
        "        time_step_spec,\n",
        "        action_spec,\n",
        "        reward_network)\n",
        "    self._policy = epsilon_greedy_policy.EpsilonGreedyPolicy(\n",
        "        policy, epsilon=epsilon)\n",
        "    self._collect_policy = self._policy\n",
        "    # Registry a method for converting received experience to a trajectory\n",
        "    self._as_trajectory = data_converter.AsTrajectory(\n",
        "        self.data_context, sequence_length=None)\n",
        "\n",
        "  def _loss(self,\n",
        "            experience: types.NestedTensor,\n",
        "            weights: Optional[types.Float] = None,\n",
        "            training: bool = False) -> tf_agent.LossInfo:\n",
        "    \"\"\"Computes loss for training the reward and constraint networks.\n",
        "    Args:\n",
        "      experience: A batch of experience data in the form of a `Trajectory` or\n",
        "        `Transition`.\n",
        "      weights: Optional scalar or elementwise (per-batch-entry) importance\n",
        "        weights.  The output batch loss will be scaled by these weights, and\n",
        "        the final scalar loss is the mean of these values.\n",
        "      training: Whether the loss is being used for training.\n",
        "    Returns:\n",
        "      loss: A `LossInfo` containing the loss for the training step.\n",
        "    Raises:\n",
        "      ValueError:\n",
        "        if the number of actions is greater than 1.\n",
        "    \"\"\"\n",
        "    (observations, actions, rewards) = bandit_utils.process_experience_for_neural_agents(\n",
        "         experience, False, self.training_data_spec)\n",
        "\n",
        "    reward_loss = self.reward_loss(\n",
        "        observations, actions, rewards, weights, training)\n",
        "\n",
        "    self.compute_summaries(reward_loss)\n",
        "\n",
        "    return tf_agent.LossInfo(reward_loss, extra=())\n",
        "\n",
        "  def _variables_to_train(self):\n",
        "    return self._reward_network.trainable_variables\n",
        "  \n",
        "  def _train(self, experience, weights):\n",
        "    \"\"\"Simplified implementation of training GreedyPolicy with NN\n",
        "\n",
        "    Original implementation can be found here: http://bit.ly/3kuyFoJ\n",
        "    \"\"\"\n",
        "    experience = self._as_trajectory(experience)\n",
        "\n",
        "    with tf.GradientTape() as tape:\n",
        "      loss_info = self._loss(# YOUR CODE HERE, weights=weights, training=True) \n",
        "     \n",
        "    variables_to_train = self._variables_to_train()\n",
        "    if not variables_to_train:\n",
        "      logging.info('No variable to train in the agent.')\n",
        "      return loss_info\n",
        "    \n",
        "    grads = tape.gradient(loss_info.loss, variables_to_train)\n",
        "    # Tuple is used for py3, where zip is a generator producing values once.\n",
        "    grads_and_vars = tuple(zip(grads, variables_to_train))\n",
        "\n",
        "    self._optimizer.apply_gradients(# YOUR CODE HERE) \n",
        "    self.train_step_counter.assign_add(1)\n",
        "\n",
        "    return loss_info\n",
        "  \n",
        "  def reward_loss(self,\n",
        "                  observations: types.NestedTensor,\n",
        "                  actions: types.Tensor,\n",
        "                  rewards: types.Tensor,\n",
        "                  weights: Optional[types.Float] = None,\n",
        "                  training: bool = False) -> types.Tensor:\n",
        "    \"\"\"Computes loss for reward prediction training.\n",
        "    Args:\n",
        "      observations: A batch of observations.\n",
        "      actions: A batch of actions.\n",
        "      rewards: A batch of rewards.\n",
        "      weights: Optional scalar or elementwise (per-batch-entry) importance\n",
        "        weights.  The output batch loss will be scaled by these weights, and\n",
        "        the final scalar loss is the mean of these values.\n",
        "      training: Whether the loss is being used for training.\n",
        "    Returns:\n",
        "      loss: A `Tensor` containing the loss for the training step.\n",
        "    Raises:\n",
        "      ValueError:\n",
        "        if the number of actions is greater than 1.\n",
        "    \"\"\"\n",
        "    with tf.name_scope('loss'):\n",
        "      sample_weights = weights if weights is not None else 1\n",
        "\n",
        "      predicted_values, _ = self._reward_network(\n",
        "          observations, training=training)\n",
        "      loss = tf.constant(0.0)\n",
        "\n",
        "      action_predicted_values = common.index_with_actions(\n",
        "          predicted_values,\n",
        "          tf.cast(actions, dtype=tf.int32))\n",
        "      \n",
        "      # Reduction is done outside of the loss function because non-scalar\n",
        "      # weights with unknown shapes may trigger shape validation that fails\n",
        "      # XLA compilation.\n",
        "      loss += tf.reduce_mean(\n",
        "          tf.multiply(\n",
        "              self._error_loss_fn(\n",
        "                  # YOUR CODE HERE,\n",
        "                  # YOUR CODE HERE, \n",
        "                  reduction=tf.compat.v1.losses.Reduction.NONE),\n",
        "              sample_weights))\n",
        "    return loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TGHbwYWTYMvl"
      },
      "source": [
        "Helper function for creating an instance of the `NeuralEpsilonGreedyAgent` with our `QNetwork`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uMo6hdUpD6hO"
      },
      "source": [
        "def get_agent(environment, lr: float = LR, epsilon: float = EPSILON):\n",
        "\n",
        "  network = QNetwork(\n",
        "      input_tensor_spec=environment.time_step_spec().observation,\n",
        "      action_spec=environment.action_spec())\n",
        "  \n",
        "  return NeuralEpsilonGreedyAgent(\n",
        "      time_step_spec=environment.time_step_spec(),\n",
        "      action_spec=environment.action_spec(),\n",
        "      reward_network=network,\n",
        "      optimizer=tf.compat.v1.train.AdamOptimizer(learning_rate=lr),\n",
        "      epsilon=epsilon)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bG8LVv2rW1sK"
      },
      "source": [
        "agent = get_agent(tf_env, lr=LR, epsilon=EPSILON)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "59HNmMTLRppr"
      },
      "source": [
        "Let 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. The `training_data_spec.observation` specificate the structure of the context vector."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HPg7h-hIlTnx"
      },
      "source": [
        "print('training data spec: ', agent.training_data_spec)\n",
        "print('observation spec in training: ', agent.training_data_spec.observation)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nDjYCtBUl2j4"
      },
      "source": [
        "## Training\n",
        "\n",
        "Now we put together all the components that we introduced above: the environment, the policy, and the agent. We run the policy on the environment and output training data with the help of a driver, and train the agent on the data.\n",
        "\n",
        "#### Metrics\n",
        "\n",
        "Important of the training are metrics. If you read some materials you can find, that bandits' most important metric is regret, calculated as the difference between the reward collected by the agent and the expected reward of an oracle policy that has access to the reward functions of the environment. The [RegretMetric](https://github.com/tensorflow/agents/blob/master/tf_agents/bandits/metrics/tf_metrics.py) thus needs a `baseline_reward_fn` function that calculates the best achievable expected reward given an observation. In our example, the optimal reward is computed in the `MovieLensPyEnvironment.compute_optimal_reward` from the approximation of the rating.\n",
        "\n",
        "> In reality, we usually do not have access to an oracle policy, so the regret is hard to get. Thus, the cumulative reward or other metric is often used."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HTmTQNtQpk_a"
      },
      "source": [
        "def get_metrics(environment):\n",
        "  optimal_reward_fn = functools.partial(\n",
        "        environment_utilities.compute_optimal_reward_with_movielens_environment,\n",
        "        environment=tf_env)\n",
        "  optimal_action_fn = functools.partial(\n",
        "        environment_utilities.compute_optimal_action_with_movielens_environment,\n",
        "        environment=tf_env)\n",
        "  \n",
        "  regret_metric = tf_bandit_metrics.RegretMetric(\n",
        "      optimal_reward_fn, \n",
        "      name='regret'\n",
        "  )\n",
        "  suboptimal_arms_metric = tf_bandit_metrics.SuboptimalArmsMetric(\n",
        "      optimal_action_fn,\n",
        "      name='suboptimal_arms'\n",
        "  )\n",
        "  \n",
        "  return [regret_metric, suboptimal_arms_metric]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0UzHUA4SM3co"
      },
      "source": [
        "We will put it all together in `run` function to run the training loop of our implementation of bandits' movie recommendations. The driver below is a helper object and 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 agent's training step. You can find more info about the driver[here](https://www.tensorflow.org/agents/tutorials/4_drivers_tutorial)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P5foepaVcv9Y"
      },
      "source": [
        "def run(\n",
        "    environment, \n",
        "    agent, \n",
        "    iterations, \n",
        "    steps_per_loop,\n",
        "    additional_metrics=()\n",
        "):\n",
        "  replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(\n",
        "      data_spec=agent.policy.trajectory_spec,\n",
        "      batch_size=environment.batch_size,\n",
        "      max_length=steps_per_loop)\n",
        "  \n",
        "  metrics = [] + list(additional_metrics)\n",
        "  ret_metrics = dict([(m.name, []) for m in metrics])\n",
        "\n",
        "  observers = [replay_buffer.add_batch] + metrics\n",
        "\n",
        "  driver = dynamic_step_driver.DynamicStepDriver(\n",
        "      env=environment,\n",
        "      policy=agent.collect_policy,\n",
        "      num_steps=steps_per_loop * environment.batch_size,\n",
        "      observers=observers)\n",
        "\n",
        "  regret_values = []\n",
        "\n",
        "  for _ in trange(num_iterations):\n",
        "    driver.run()\n",
        "    loss_info = agent.train(replay_buffer.gather_all())\n",
        "    replay_buffer.clear()\n",
        "    # Log metrics value\n",
        "    for metric in metrics:\n",
        "      ret_metrics[metric.name].append(metric.result())\n",
        "\n",
        "  return ret_metrics"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ViNag8xHP0wW"
      },
      "source": [
        "Down below is the code for creating all necessary instances. Note that two parameters together specify the number of steps taken. `num_iterations` specifies how many times we run the trainer loop, while the driver will take `steps_per_loop` steps per iteration. The main reason behind keeping both of these parameters is that some operations are done per iteration, while the driver does some in every step. For example, the agent's train function is only called once per iteration. The trade-off here is that if we train more often, our policy is \"fresher\"; on the other hand, training in bigger batches might be more time-efficient. `batch_size` defines how many actions are generated through one step."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9yxB1q29LqWV"
      },
      "source": [
        "batch_size =   128# @param {type:\"integer\"}\n",
        "num_iterations =   1500# @param {type:\"integer\"}\n",
        "steps_per_loop =   8# @param {type:\"integer\"}\n",
        "\n",
        "env = MovieLensPyEnvironment(\n",
        "    './ml-100k/u.data', \n",
        "    rank_k=RANK_K,\n",
        "    batch_size=batch_size, \n",
        "    num_movies=NUM_ACTIONS\n",
        ")\n",
        "tf_env = tf_py_environment.TFPyEnvironment(env)\n",
        "tf_env.reset()\n",
        "\n",
        "agent = get_agent(tf_env, lr=LR, epsilon=EPSILON)\n",
        "\n",
        "additional_metrics = get_metrics(tf_env)\n",
        "\n",
        "metrics = run(\n",
        "    tf_env, \n",
        "    agent, \n",
        "    iterations=num_iterations,\n",
        "    steps_per_loop=steps_per_loop,\n",
        "    additional_metrics=additional_metrics\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_gvfBsUCDi2U"
      },
      "source": [
        "Now let's see the result. After running the last code snippet, the resulting plot (hopefully) shows that the average regret is going down as the agent is trained and the policy gets better in figuring out what the right action is, given the observation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jqo41BBZXLZh"
      },
      "source": [
        "plot_regret(metrics['regret'], {'algorithm': 'NeuralEpsGreedy'})"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XYdibk130cGM"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}