docs/tutorials/9_c51_tutorial.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "klGNgWREsvQv"
      },
      "source": [
        "##### Copyright 2023 The TF-Agents Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "nQnmcm0oI1Q-"
      },
      "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": "oMaGpi7TciQs"
      },
      "source": [
        "# DQN C51/Rainbow\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/agents/tutorials/9_c51_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/9_c51_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/9_c51_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/9_c51_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": "ZOUOQOrFs3zn"
      },
      "source": [
        "## Introduction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cKOCZlhUgXVK"
      },
      "source": [
        "This example shows how to train a [Categorical DQN (C51)](https://arxiv.org/pdf/1707.06887.pdf) agent on the Cartpole environment using the TF-Agents library.\n",
        "\n",
        "![Cartpole environment](https://github.com/tensorflow/agents/blob/master/docs/tutorials/images/cartpole.png?raw=1)\n",
        "\n",
        "Make sure you take a look through the [DQN tutorial](https://github.com/tensorflow/agents/blob/master/docs/tutorials/1_dqn_tutorial.ipynb) as a prerequisite. This tutorial will assume familiarity with the DQN tutorial; it will mainly focus on the differences between DQN and C51.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lsaQlK8fFQqH"
      },
      "source": [
        "## Setup\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-NzBsZzPcyBm"
      },
      "source": [
        "If you haven't installed tf-agents yet, run:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KEHR2Ui-lo8O"
      },
      "outputs": [],
      "source": [
        "!sudo apt-get update\n",
        "!sudo apt-get install -y xvfb ffmpeg freeglut3-dev\n",
        "!pip install 'imageio==2.4.0'\n",
        "!pip install pyvirtualdisplay\n",
        "!pip install tf-agents\n",
        "!pip install pyglet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sMitx5qSgJk1"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import\n",
        "from __future__ import division\n",
        "from __future__ import print_function\n",
        "\n",
        "import base64\n",
        "import imageio\n",
        "import IPython\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "import PIL.Image\n",
        "import pyvirtualdisplay\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from tf_agents.agents.categorical_dqn import categorical_dqn_agent\n",
        "from tf_agents.drivers import dynamic_step_driver\n",
        "from tf_agents.environments import suite_gym\n",
        "from tf_agents.environments import tf_py_environment\n",
        "from tf_agents.eval import metric_utils\n",
        "from tf_agents.metrics import tf_metrics\n",
        "from tf_agents.networks import categorical_q_network\n",
        "from tf_agents.policies import random_tf_policy\n",
        "from tf_agents.replay_buffers import tf_uniform_replay_buffer\n",
        "from tf_agents.trajectories import trajectory\n",
        "from tf_agents.utils import common\n",
        "\n",
        "# Set up a virtual display for rendering OpenAI gym environments.\n",
        "display = pyvirtualdisplay.Display(visible=0, size=(1400, 900)).start()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LmC0NDhdLIKY"
      },
      "source": [
        "## Hyperparameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HC1kNrOsLSIZ"
      },
      "outputs": [],
      "source": [
        "env_name = \"CartPole-v1\" # @param {type:\"string\"}\n",
        "num_iterations = 15000 # @param {type:\"integer\"}\n",
        "\n",
        "initial_collect_steps = 1000  # @param {type:\"integer\"} \n",
        "collect_steps_per_iteration = 1  # @param {type:\"integer\"}\n",
        "replay_buffer_capacity = 100000  # @param {type:\"integer\"}\n",
        "\n",
        "fc_layer_params = (100,)\n",
        "\n",
        "batch_size = 64  # @param {type:\"integer\"}\n",
        "learning_rate = 1e-3  # @param {type:\"number\"}\n",
        "gamma = 0.99\n",
        "log_interval = 200  # @param {type:\"integer\"}\n",
        "\n",
        "num_atoms = 51  # @param {type:\"integer\"}\n",
        "min_q_value = -20  # @param {type:\"integer\"}\n",
        "max_q_value = 20  # @param {type:\"integer\"}\n",
        "n_step_update = 2  # @param {type:\"integer\"}\n",
        "\n",
        "num_eval_episodes = 10  # @param {type:\"integer\"}\n",
        "eval_interval = 1000  # @param {type:\"integer\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VMsJC3DEgI0x"
      },
      "source": [
        "## Environment\n",
        "\n",
        "Load the environment as before, with one for training and one for evaluation. Here we use CartPole-v1 (vs. CartPole-v0 in the DQN tutorial), which has a larger max reward of 500 rather than 200."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xp-Y4mD6eDhF"
      },
      "outputs": [],
      "source": [
        "train_py_env = suite_gym.load(env_name)\n",
        "eval_py_env = suite_gym.load(env_name)\n",
        "\n",
        "train_env = tf_py_environment.TFPyEnvironment(train_py_env)\n",
        "eval_env = tf_py_environment.TFPyEnvironment(eval_py_env)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E9lW_OZYFR8A"
      },
      "source": [
        "## Agent\n",
        "\n",
        "C51 is a Q-learning algorithm based on DQN. Like DQN, it can be used on any environment with a discrete action space.\n",
        "\n",
        "The main difference between C51 and DQN is that rather than simply predicting the Q-value for each state-action pair, C51 predicts a histogram model for the probability distribution of the Q-value:\n",
        "\n",
        "![Example C51 Distribution](images/c51_distribution.png)\n",
        "\n",
        "By learning the distribution rather than simply the expected value, the algorithm is able to stay more stable during training, leading to improved final performance. This is particularly true in situations with bimodal or even multimodal value distributions, where a single average does not provide an accurate picture.\n",
        "\n",
        "In order to train on probability distributions rather than on values, C51 must perform some complex distributional computations in order to calculate its loss function. But don't worry, all of this is taken care of for you in TF-Agents!\n",
        "\n",
        "To create a C51 Agent, we first need to create a `CategoricalQNetwork`. The API of the `CategoricalQNetwork` is the same as that of the `QNetwork`, except that there is an additional argument `num_atoms`. This represents the number of support points in our probability distribution estimates. (The above image includes 10 support points, each represented by a vertical blue bar.) As you can tell from the name, the default number of atoms is 51.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TgkdEPg_muzV"
      },
      "outputs": [],
      "source": [
        "categorical_q_net = categorical_q_network.CategoricalQNetwork(\n",
        "    train_env.observation_spec(),\n",
        "    train_env.action_spec(),\n",
        "    num_atoms=num_atoms,\n",
        "    fc_layer_params=fc_layer_params)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z62u55hSmviJ"
      },
      "source": [
        "We also need an `optimizer` to train the network we just created, and a `train_step_counter` variable to keep track of how many times the network was updated.\n",
        "\n",
        "Note that one other significant difference from vanilla `DqnAgent` is that we now need to specify `min_q_value` and `max_q_value` as arguments. These specify the most extreme values of the support (in other words, the most extreme of the 51 atoms on either side). Make sure to choose these appropriately for your particular environment. Here we use -20 and 20."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jbY4yrjTEyc9"
      },
      "outputs": [],
      "source": [
        "optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=learning_rate)\n",
        "\n",
        "train_step_counter = tf.Variable(0)\n",
        "\n",
        "agent = categorical_dqn_agent.CategoricalDqnAgent(\n",
        "    train_env.time_step_spec(),\n",
        "    train_env.action_spec(),\n",
        "    categorical_q_network=categorical_q_net,\n",
        "    optimizer=optimizer,\n",
        "    min_q_value=min_q_value,\n",
        "    max_q_value=max_q_value,\n",
        "    n_step_update=n_step_update,\n",
        "    td_errors_loss_fn=common.element_wise_squared_loss,\n",
        "    gamma=gamma,\n",
        "    train_step_counter=train_step_counter)\n",
        "agent.initialize()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L7O7F_HqiQ1G"
      },
      "source": [
        "One last thing to note is that we also added an argument to use n-step updates with $n$ = 2. In single-step Q-learning ($n$ = 1), we only compute the error between the Q-values at the current time step and the next time step using the single-step return (based on the Bellman optimality equation). The single-step return is defined as:\n",
        "\n",
        "$G_t = R_{t + 1} + \\gamma V(s_{t + 1})$\n",
        "\n",
        "where we define $V(s) = \\max_a{Q(s, a)}$.\n",
        "\n",
        "N-step updates involve expanding the standard single-step return function $n$ times:\n",
        "\n",
        "$G_t^n = R_{t + 1} + \\gamma R_{t + 2} + \\gamma^2 R_{t + 3} + \\dots + \\gamma^n V(s_{t + n})$\n",
        "\n",
        "N-step updates enable the agent to bootstrap from further in the future, and with the right value of $n$, this often leads to faster learning.\n",
        "\n",
        "Although C51 and n-step updates are often combined with prioritized replay to form the core of the [Rainbow agent](https://arxiv.org/pdf/1710.02298.pdf), we saw no measurable improvement from implementing prioritized replay. Moreover, we find that when combining our C51 agent with n-step updates alone, our agent performs as well as other Rainbow agents on the sample of Atari environments we've tested."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "94rCXQtbUbXv"
      },
      "source": [
        "## Metrics and Evaluation\n",
        "\n",
        "The most common metric used to evaluate a policy is the average return. The return is the sum of rewards obtained while running a policy in an environment for an episode, and we usually average this over a few episodes. We can compute the average return metric as follows.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bitzHo5_UbXy"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "def compute_avg_return(environment, policy, num_episodes=10):\n",
        "\n",
        "  total_return = 0.0\n",
        "  for _ in range(num_episodes):\n",
        "\n",
        "    time_step = environment.reset()\n",
        "    episode_return = 0.0\n",
        "\n",
        "    while not time_step.is_last():\n",
        "      action_step = policy.action(time_step)\n",
        "      time_step = environment.step(action_step.action)\n",
        "      episode_return += time_step.reward\n",
        "    total_return += episode_return\n",
        "\n",
        "  avg_return = total_return / num_episodes\n",
        "  return avg_return.numpy()[0]\n",
        "\n",
        "\n",
        "random_policy = random_tf_policy.RandomTFPolicy(train_env.time_step_spec(),\n",
        "                                                train_env.action_spec())\n",
        "\n",
        "compute_avg_return(eval_env, random_policy, num_eval_episodes)\n",
        "\n",
        "# Please also see the metrics module for standard implementations of different\n",
        "# metrics."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NLva6g2jdWgr"
      },
      "source": [
        "## Data Collection\n",
        "\n",
        "As in the DQN tutorial, set up the replay buffer and the initial data collection with the random policy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wr1KSAEGG4h9"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(\n",
        "    data_spec=agent.collect_data_spec,\n",
        "    batch_size=train_env.batch_size,\n",
        "    max_length=replay_buffer_capacity)\n",
        "\n",
        "def collect_step(environment, policy):\n",
        "  time_step = environment.current_time_step()\n",
        "  action_step = policy.action(time_step)\n",
        "  next_time_step = environment.step(action_step.action)\n",
        "  traj = trajectory.from_transition(time_step, action_step, next_time_step)\n",
        "\n",
        "  # Add trajectory to the replay buffer\n",
        "  replay_buffer.add_batch(traj)\n",
        "\n",
        "for _ in range(initial_collect_steps):\n",
        "  collect_step(train_env, random_policy)\n",
        "\n",
        "# This loop is so common in RL, that we provide standard implementations of\n",
        "# these. For more details see the drivers module.\n",
        "\n",
        "# Dataset generates trajectories with shape [BxTx...] where\n",
        "# T = n_step_update + 1.\n",
        "dataset = replay_buffer.as_dataset(\n",
        "    num_parallel_calls=3, sample_batch_size=batch_size,\n",
        "    num_steps=n_step_update + 1).prefetch(3)\n",
        "\n",
        "iterator = iter(dataset)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hBc9lj9VWWtZ"
      },
      "source": [
        "## Training the agent\n",
        "\n",
        "The training loop involves both collecting data from the environment and optimizing the agent's networks. Along the way, we will occasionally evaluate the agent's policy to see how we are doing.\n",
        "\n",
        "The following will take ~7 minutes to run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0pTbJ3PeyF-u"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "try:\n",
        "  %%time\n",
        "except:\n",
        "  pass\n",
        "\n",
        "# (Optional) Optimize by wrapping some of the code in a graph using TF function.\n",
        "agent.train = common.function(agent.train)\n",
        "\n",
        "# Reset the train step\n",
        "agent.train_step_counter.assign(0)\n",
        "\n",
        "# Evaluate the agent's policy once before training.\n",
        "avg_return = compute_avg_return(eval_env, agent.policy, num_eval_episodes)\n",
        "returns = [avg_return]\n",
        "\n",
        "for _ in range(num_iterations):\n",
        "\n",
        "  # Collect a few steps using collect_policy and save to the replay buffer.\n",
        "  for _ in range(collect_steps_per_iteration):\n",
        "    collect_step(train_env, agent.collect_policy)\n",
        "\n",
        "  # Sample a batch of data from the buffer and update the agent's network.\n",
        "  experience, unused_info = next(iterator)\n",
        "  train_loss = agent.train(experience)\n",
        "\n",
        "  step = agent.train_step_counter.numpy()\n",
        "\n",
        "  if step % log_interval == 0:\n",
        "    print('step = {0}: loss = {1}'.format(step, train_loss.loss))\n",
        "\n",
        "  if step % eval_interval == 0:\n",
        "    avg_return = compute_avg_return(eval_env, agent.policy, num_eval_episodes)\n",
        "    print('step = {0}: Average Return = {1:.2f}'.format(step, avg_return))\n",
        "    returns.append(avg_return)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "68jNcA_TiJDq"
      },
      "source": [
        "## Visualization\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aO-LWCdbbOIC"
      },
      "source": [
        "### Plots\n",
        "\n",
        "We can plot return vs global steps to see the performance of our agent. In `Cartpole-v1`, the environment gives a reward of +1 for every time step the pole stays up, and since the maximum number of steps is 500, the maximum possible return is also 500."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NxtL1mbOYCVO"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "\n",
        "steps = range(0, num_iterations + 1, eval_interval)\n",
        "plt.plot(steps, returns)\n",
        "plt.ylabel('Average Return')\n",
        "plt.xlabel('Step')\n",
        "plt.ylim(top=550)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M7-XpPP99Cy7"
      },
      "source": [
        "### Videos"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9pGfGxSH32gn"
      },
      "source": [
        "It is helpful to visualize the performance of an agent by rendering the environment at each step. Before we do that, let us first create a function to embed videos in this colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ULaGr8pvOKbl"
      },
      "outputs": [],
      "source": [
        "def embed_mp4(filename):\n",
        "  \"\"\"Embeds an mp4 file in the notebook.\"\"\"\n",
        "  video = open(filename,'rb').read()\n",
        "  b64 = base64.b64encode(video)\n",
        "  tag = '''\n",
        "  \u003cvideo width=\"640\" height=\"480\" controls\u003e\n",
        "    \u003csource src=\"data:video/mp4;base64,{0}\" type=\"video/mp4\"\u003e\n",
        "  Your browser does not support the video tag.\n",
        "  \u003c/video\u003e'''.format(b64.decode())\n",
        "\n",
        "  return IPython.display.HTML(tag)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9c_PH-pX4Pr5"
      },
      "source": [
        "The following code visualizes the agent's policy for a few episodes:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "owOVWB158NlF"
      },
      "outputs": [],
      "source": [
        "num_episodes = 3\n",
        "video_filename = 'imageio.mp4'\n",
        "with imageio.get_writer(video_filename, fps=60) as video:\n",
        "  for _ in range(num_episodes):\n",
        "    time_step = eval_env.reset()\n",
        "    video.append_data(eval_py_env.render())\n",
        "    while not time_step.is_last():\n",
        "      action_step = agent.policy.action(time_step)\n",
        "      time_step = eval_env.step(action_step.action)\n",
        "      video.append_data(eval_py_env.render())\n",
        "\n",
        "embed_mp4(video_filename)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "exziB27hY8ia"
      },
      "source": [
        "C51 tends to do slightly better than DQN on CartPole-v1, but the difference between the two agents becomes more and more significant in increasingly complex environments. For example, on the full Atari 2600 benchmark, C51 demonstrates a mean score improvement of 126% over DQN after normalizing with respect to a random agent. Additional improvements can be gained by including n-step updates.\n",
        "\n",
        "For a deeper dive into the C51 algorithm, see [A Distributional Perspective on Reinforcement Learning (2017)](https://arxiv.org/pdf/1707.06887.pdf)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "DQN C51/Rainbow Tutorial.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "pycharm": {
      "stem_cell": {
        "cell_type": "raw",
        "metadata": {
          "collapsed": false
        },
        "source": []
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}