t81_558_class_12_04_atari.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "anaconda-cloud": {},
    "colab": {
      "name": "Copy of t81_558_class_12_04_atari.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3.6 (tensorflow)",
      "language": "python",
      "name": "tensorflow"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.7"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "klGNgWREsvQv"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jeffheaton/t81_558_deep_learning/blob/master/t81_558_class_12_04_atari.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yzzbqc-JS2z9"
      },
      "source": [
        "# T81-558: Applications of Deep Neural Networks\n",
        "**Module 12: Reinforcement Learning**\n",
        "* Instructor: [Jeff Heaton](https://sites.wustl.edu/jeffheaton/), McKelvey School of Engineering, [Washington University in St. Louis](https://engineering.wustl.edu/Programs/Pages/default.aspx)\n",
        "* For more information visit the [class website](https://sites.wustl.edu/jeffheaton/t81-558/)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pCrXGd_CS6eB"
      },
      "source": [
        "# Module 12 Video Material\n",
        "\n",
        "* Part 12.1: Introduction to the OpenAI Gym [[Video]](https://www.youtube.com/watch?v=_KbUxgyisjM&list=PLjy4p-07OYzulelvJ5KVaT2pDlxivl_BN) [[Notebook]](t81_558_class_12_01_ai_gym.ipynb)\n",
        "* Part 12.2: Introduction to Q-Learning [[Video]](https://www.youtube.com/watch?v=A3sYFcJY3lA&list=PLjy4p-07OYzulelvJ5KVaT2pDlxivl_BN) [[Notebook]](t81_558_class_12_02_qlearningreinforcement.ipynb)\n",
        "* Part 12.3: Keras Q-Learning in the OpenAI Gym [[Video]](https://www.youtube.com/watch?v=qy1SJmsRhvM&list=PLjy4p-07OYzulelvJ5KVaT2pDlxivl_BN) [[Notebook]](t81_558_class_12_03_keras_reinforce.ipynb)\n",
        "* **Part 12.4: Atari Games with Keras Neural Networks** [[Video]](https://www.youtube.com/watch?v=co0SwPWoZh0&list=PLjy4p-07OYzulelvJ5KVaT2pDlxivl_BN) [[Notebook]](t81_558_class_12_04_atari.ipynb)\n",
        "* Part 12.5: Application of Reinforcement Learning [[Video]](https://www.youtube.com/watch?v=1jQPP3RfwMI&list=PLjy4p-07OYzulelvJ5KVaT2pDlxivl_BN) [[Notebook]](t81_558_class_12_05_apply_rl.ipynb)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pmDI-h7cI0tI"
      },
      "source": [
        "# Google CoLab Instructions\n",
        "\n",
        "The following code ensures that Google CoLab is running the correct version of TensorFlow, and has the necessary Python libraries installed."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9KQhYThvTCQC"
      },
      "source": [
        "try:\n",
        "    from google.colab import drive\n",
        "    %tensorflow_version 2.x\n",
        "    COLAB = True\n",
        "    print(\"Note: using Google CoLab\")\n",
        "except:\n",
        "    print(\"Note: not using Google CoLab\")\n",
        "    COLAB = False\n",
        "\n",
        "if COLAB:\n",
        "  !sudo apt-get install -y xvfb ffmpeg\n",
        "  !pip install -q 'gym==0.10.11'\n",
        "  !pip install -q 'imageio==2.4.0'\n",
        "  !pip install -q PILLOW\n",
        "  !pip install -q 'pyglet==1.3.2'\n",
        "  !pip install -q pyvirtualdisplay\n",
        "  !pip install -q --upgrade tensorflow-probability\n",
        "  !pip install -q tf-agents"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lsaQlK8fFQqH"
      },
      "source": [
        "# Part 12.4: Atari Games with Keras Neural Networks\n",
        "\n",
        "\n",
        "The Atari 2600 is a home video game console from Atari, Inc. Released on September 11, 1977. It is credited with popularizing the use of microprocessor-based hardware and games stored on ROM cartridges instead of dedicated hardware with games physically built into the unit. The 2600 was bundled with two joystick controllers, a conjoined pair of paddle controllers, and a game cartridge: initially [Combat](https://en.wikipedia.org/wiki/Combat_(Atari_2600)), and later [Pac-Man](https://en.wikipedia.org/wiki/Pac-Man_(Atari_2600)).\n",
        "\n",
        "Atari emulators are popular and allow many of the old Atari video games to be played on modern computers.  They are even available as JavaScript.\n",
        "\n",
        "* [Virtual Atari](http://www.virtualatari.org/listP.html)\n",
        "\n",
        "Atari games have become popular benchmarks for AI systems, particularly reinforcement learning.  OpenAI Gym internally uses the [Stella Atari Emulator](https://stella-emu.github.io/). The Atari 2600 is shown in Figure 12.ATARI.\n",
        "\n",
        "**Figure 12.ATARI: The Atari 2600**\n",
        "![Atari 2600 Console](https://raw.githubusercontent.com/jeffheaton/t81_558_deep_learning/master/images/atari-1.png \"Atari 2600 Console\")\n",
        "\n",
        "### Actual Atari 2600 Specs\n",
        "\n",
        "* CPU: 1.19 MHz MOS Technology 6507\n",
        "* Audio + Video processor: Television Interface Adapter (TIA)\n",
        "* Playfield resolution: 40 x 192 pixels (NTSC). Uses a 20-pixel register that is mirrored or copied, left side to right side, to achieve the width of 40 pixels.\n",
        "* Player sprites: 8 x 192 pixels (NTSC). Player, ball, and missile sprites use pixels that are 1/4 the width of playfield pixels (unless stretched).\n",
        "* Ball and missile sprites: 1 x 192 pixels (NTSC).\n",
        "* Maximum resolution: 160 x 192 pixels (NTSC). Max resolution is only somewhat achievable with programming tricks that combine sprite pixels with playfield pixels.\n",
        "* 128 colors (NTSC). 128 possible on screen. Max of 4 per line: background, playfield, player0 sprite, and player1 sprite. Palette switching between lines is common. Palette switching mid line is possible but not common due to resource limitations.\n",
        "* 2 channels of 1-bit monaural sound with 4-bit volume control.\n",
        "\n",
        "### OpenAI Lab Atari Pong\n",
        "\n",
        "OpenAI Gym can be used with Windows; however, it requires a special [installation procedure](https://towardsdatascience.com/how-to-install-openai-gym-in-a-windows-environment-338969e24d30). \n",
        "\n",
        "This chapter demonstrates playing [Atari Pong](https://github.com/wau/keras-rl2/blob/master/examples/dqn_atari.py). Pong is a two-dimensional sports game that simulates table tennis. The player controls an in-game paddle by moving it vertically across the left or right side of the screen. They can compete against another player controlling a second paddle on the opposing side. Players use the paddles to hit a ball back and forth. The goal is for each player to reach eleven points before the opponent; you earn points when one fails to return it to the other.  For the Atari 2600 version of Pong, a computer player (controlled by the 2600) is the opposing player.\n",
        "\n",
        "This section shows how to adapt TF-Agents to an Atari game.  Some changes are necessary when compared to the pole-cart game presented earlier in this chapter. You can quickly adapt this example to any Atari game by simply changing the environment name.  However, I tuned the code presented here for Pong, and it may not perform as well for other games.  Some tuning will likely be necessary to produce a good agent for other games.\n",
        "\n",
        "We begin by importing the needed Python packages."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sMitx5qSgJk1"
      },
      "source": [
        "import base64\n",
        "import imageio\n",
        "import IPython\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import PIL.Image\n",
        "import pyvirtualdisplay\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from tf_agents.agents.dqn import dqn_agent\n",
        "from tf_agents.drivers import dynamic_step_driver\n",
        "from tf_agents.environments import suite_gym, suite_atari\n",
        "from tf_agents.environments import tf_py_environment, batched_py_environment\n",
        "from tf_agents.eval import metric_utils\n",
        "from tf_agents.metrics import tf_metrics\n",
        "from tf_agents.networks import q_network, 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",
        "from tf_agents.agents.categorical_dqn import categorical_dqn_agent\n",
        "from tf_agents.networks import categorical_q_network\n",
        "\n",
        "from tf_agents.specs import tensor_spec\n",
        "from tf_agents.trajectories import time_step as ts"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J6HsdS5GbSjd"
      },
      "source": [
        "# Set up a virtual display for rendering OpenAI gym environments.\n",
        "display = pyvirtualdisplay.Display(visible=0, size=(1400, 900)).start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LmC0NDhdLIKY"
      },
      "source": [
        "## Hyperparameters\n",
        "\n",
        "The hyperparameter names are the same as the previous DQN example; however, I tuned the numeric values for the more complex Atari game."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HC1kNrOsLSIZ"
      },
      "source": [
        "num_iterations = 250000 \n",
        "\n",
        "initial_collect_steps = 200  \n",
        "collect_steps_per_iteration = 10 \n",
        "replay_buffer_max_length = 100000\n",
        "\n",
        "batch_size =   32\n",
        "learning_rate = 2.5e-3\n",
        "log_interval =   5000\n",
        "\n",
        "num_eval_episodes = 5  \n",
        "eval_interval = 25000  "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XZUoWspvOJqB"
      },
      "source": [
        "The algorithm needs more iterations for an Atari game.  I also found that increasing the number of collection steps helped the algorithm to train."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VMsJC3DEgI0x"
      },
      "source": [
        "## Atari Environment's\n",
        "\n",
        "You must handle Atari environments differently than games like cart-poll.  Atari games typically use their 2D displays as the environment state.  AI Gym represents Atari games as either a 3D (height by width by color) state spaced based on their screens, or a vector representing the state of the game's computer RAM. To preprocess Atari games for greater computational efficiency, we generally skip several frames, decrease the resolution, and discard color information.  The following code shows how we can set up an Atari environment."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HWCaUF9mU53o"
      },
      "source": [
        "! wget http://www.atarimania.com/roms/Roms.rar\n",
        "! mkdir /content/ROM/\n",
        "! unrar e /content/Roms.rar /content/ROM/\n",
        "! python -m atari_py.import_roms /content/ROM/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pYEz-S9gEv2-"
      },
      "source": [
        "#env_name = 'Breakout-v4'\n",
        "env_name = 'Pong-v0'\n",
        "#env_name = 'BreakoutDeterministic-v4'\n",
        "#env = suite_gym.load(env_name)\n",
        "\n",
        "# AtariPreprocessing runs 4 frames at a time, max-pooling over the last 2\n",
        "# frames. We need to account for this when computing things like update\n",
        "# intervals.\n",
        "ATARI_FRAME_SKIP = 4\n",
        "\n",
        "max_episode_frames=108000  # ALE frames\n",
        "\n",
        "env = suite_atari.load(\n",
        "    env_name,\n",
        "    max_episode_steps=max_episode_frames / ATARI_FRAME_SKIP,\n",
        "    gym_env_wrappers=suite_atari.DEFAULT_ATARI_GYM_WRAPPERS_WITH_STACKING)\n",
        "#env = batched_py_environment.BatchedPyEnvironment([env])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IIHYVBkuvPNw"
      },
      "source": [
        "We can now reset the environment and display one step.  The following image shows how the Pong game environment appears to a user."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RlO7WIQHu_7D"
      },
      "source": [
        "env.reset()\n",
        "PIL.Image.fromarray(env.render())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b_lHcIcqUaqB"
      },
      "source": [
        "We are now ready to load and wrap the two environments for TF-Agents. The algorithm uses the first environment for evaluation, and the second to train."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N7brXNIGWXjC"
      },
      "source": [
        "train_py_env = suite_atari.load(\n",
        "    env_name,\n",
        "    max_episode_steps=max_episode_frames / ATARI_FRAME_SKIP,\n",
        "    gym_env_wrappers=suite_atari.DEFAULT_ATARI_GYM_WRAPPERS_WITH_STACKING)\n",
        "\n",
        "eval_py_env = suite_atari.load(\n",
        "    env_name,\n",
        "    max_episode_steps=max_episode_frames / ATARI_FRAME_SKIP,\n",
        "    gym_env_wrappers=suite_atari.DEFAULT_ATARI_GYM_WRAPPERS_WITH_STACKING)\n",
        "\n",
        "train_env = tf_py_environment.TFPyEnvironment(train_py_env)\n",
        "eval_env = tf_py_environment.TFPyEnvironment(eval_py_env)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E9lW_OZYFR8A"
      },
      "source": [
        "## Agent\n",
        "\n",
        "I used the following class, from TF-Agents examples, to wrap the regular Q-network class.  The AtariQNetwork class ensures that the pixel values from the Atari screen are divided by 255.  This division assists the neural network by normalizing the pixel values to between 0 and 1."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EoLEdYvzUHeX"
      },
      "source": [
        "# AtariPreprocessing runs 4 frames at a time, max-pooling over the last 2\n",
        "# frames. We need to account for this when computing things like update\n",
        "# intervals.\n",
        "ATARI_FRAME_SKIP = 4\n",
        "\n",
        "class AtariCategoricalQNetwork(network.Network):\n",
        "  \"\"\"CategoricalQNetwork subclass that divides observations by 255.\"\"\"\n",
        "\n",
        "  def __init__(self, input_tensor_spec, action_spec, **kwargs):\n",
        "    super(AtariCategoricalQNetwork, self).__init__(\n",
        "        input_tensor_spec, state_spec=())\n",
        "    input_tensor_spec = tf.TensorSpec(\n",
        "        dtype=tf.float32, shape=input_tensor_spec.shape)\n",
        "    self._categorical_q_network = categorical_q_network.CategoricalQNetwork(\n",
        "        input_tensor_spec, action_spec, **kwargs)\n",
        "\n",
        "  @property\n",
        "  def num_atoms(self):\n",
        "    return self._categorical_q_network.num_atoms\n",
        "\n",
        "  def call(self, observation, step_type=None, network_state=()):\n",
        "    state = tf.cast(observation, tf.float32)\n",
        "    # We divide the grayscale pixel values by 255 here rather than storing\n",
        "    # normalized values beause uint8s are 4x cheaper to store than float32s.\n",
        "    # TODO(b/129805821): handle the division by 255 for train_eval_atari.py in\n",
        "    # a preprocessing layer instead.\n",
        "    state = state / 255\n",
        "    return self._categorical_q_network(\n",
        "        state, step_type=step_type, network_state=network_state)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l--Jj22eVRZD"
      },
      "source": [
        "Next, we introduce two hyperparameters that are specific to the neural network we are about to define."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TgkdEPg_muzV"
      },
      "source": [
        "fc_layer_params = (512,)\n",
        "conv_layer_params=((32, (8, 8), 4), (64, (4, 4), 2), (64, (3, 3), 1))\n",
        "\n",
        "q_net = AtariCategoricalQNetwork(\n",
        "            train_env.observation_spec(),\n",
        "            train_env.action_spec(),\n",
        "            conv_layer_params=conv_layer_params,\n",
        "            fc_layer_params=fc_layer_params)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z62u55hSmviJ"
      },
      "source": [
        "Convolutional neural networks usually are made up of several alternating pairs of convolution and max-pooling layers, ultimately culminating in one or more dense layers.  These layers are the same types as previously seen in this course.  The QNetwork accepts two parameters that define the convolutional neural network structure. \n",
        "\n",
        "The more simple of the two parameters is **fc_layer_params**.  This parameter specifies the size of each of the dense layers.  A tuple specifies the size of each of the layers in a list. \n",
        "\n",
        "The second parameter, named **conv_layer_params**, is a list of convolution layers parameters, where each item is a length-three tuple indicating (filters, kernel_size, stride).  This implementation of QNetwork supports only convolution layers.  If you desire a more complex convolutional neural network, you must define your variant of the QNetwork.\n",
        "\n",
        "The QNetwork defined here is not the agent, instead, the QNetwork is used by the DQN agent to implement the actual neural network. This allows flexibility as you can set your own class if needed.\n",
        "\n",
        "Next, we define the optimizer.  For this example, I used RMSPropOptimizer.  However, AdamOptimizer is another popular choice.  We also create the DQN agent and reference the Q-network we just created."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "jbY4yrjTEyc9"
      },
      "source": [
        "optimizer = tf.compat.v1.train.RMSPropOptimizer(\n",
        "    learning_rate=learning_rate,\n",
        "    decay=0.95,\n",
        "    momentum=0.0,\n",
        "    epsilon=0.00001,\n",
        "    centered=True)\n",
        "\n",
        "train_step_counter = tf.Variable(0)\n",
        "\n",
        "observation_spec = tensor_spec.from_spec(train_env.observation_spec())\n",
        "time_step_spec = ts.time_step_spec(observation_spec)\n",
        "\n",
        "action_spec = tensor_spec.from_spec(train_env.action_spec())\n",
        "target_update_period=32000  # ALE frames\n",
        "update_period=16  # ALE frames\n",
        "_update_period = update_period / ATARI_FRAME_SKIP\n",
        "\n",
        "\n",
        "agent = categorical_dqn_agent.CategoricalDqnAgent(\n",
        "    time_step_spec,\n",
        "    action_spec,\n",
        "    categorical_q_network=q_net,\n",
        "    optimizer=optimizer,\n",
        "    #epsilon_greedy=epsilon,\n",
        "    n_step_update=1.0,\n",
        "    target_update_tau=1.0,\n",
        "    target_update_period=(\n",
        "        target_update_period / ATARI_FRAME_SKIP / _update_period),\n",
        "    gamma=0.99,\n",
        "    reward_scale_factor=1.0,\n",
        "    gradient_clipping=None,\n",
        "    debug_summaries=False,\n",
        "    summarize_grads_and_vars=False)\n",
        "\n",
        "\"\"\"\n",
        "agent = dqn_agent.DqnAgent(\n",
        "    time_step_spec,\n",
        "    action_spec,\n",
        "    q_network=q_net,\n",
        "    optimizer=optimizer,\n",
        "    epsilon_greedy=0.01,\n",
        "    n_step_update=1.0,\n",
        "    target_update_tau=1.0,\n",
        "    target_update_period=(\n",
        "        target_update_period / ATARI_FRAME_SKIP / _update_period),\n",
        "    td_errors_loss_fn=common.element_wise_huber_loss,\n",
        "    gamma=0.99,\n",
        "    reward_scale_factor=1.0,\n",
        "    gradient_clipping=None,\n",
        "    debug_summaries=False,\n",
        "    summarize_grads_and_vars=False,\n",
        "    train_step_counter=_global_step)\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "agent.initialize()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "EQOs8OT0Y2w6"
      },
      "source": [
        "q_net.input_tensor_spec"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "_073wq7RlmTi"
      },
      "source": [
        "train_env.observation_spec()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "OviZu4Bml9BO"
      },
      "source": [
        "train_py_env.observation_spec()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "hVwGgnS-nwHV"
      },
      "source": [
        "train_py_env"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "94rCXQtbUbXv"
      },
      "source": [
        "## Metrics and Evaluation\n",
        "\n",
        "There are many different ways to measure the effectiveness of a model trained with reinforcement learning.  The loss function of the internal Q-network is not a good measure of the entire DQN algorithm's overall fitness.  The network loss function measures how close the Q-network was fit to the collected data and did not indicate how effective the DQN is in maximizing rewards. The method used for this example tracks the average reward received over several episodes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "bitzHo5_UbXy"
      },
      "source": [
        "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",
        "# See also the metrics module for standard implementations of \n",
        "# different metrics.\n",
        "# https://github.com/tensorflow/agents/tree/master/tf_agents/metrics"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NLva6g2jdWgr"
      },
      "source": [
        "## Replay Buffer\n",
        "\n",
        "DQN works by training a neural network to predict the Q-values for every possible environment-state.  A neural network needs training data, so the algorithm accumulates this training data as it runs episodes. The replay buffer is where this data is stored.  Only the most recent episodes are stored, older episode data rolls off the queue as the queue accumulates new data."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "vX2zGUWJGWAl"
      },
      "source": [
        "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_max_length)\n",
        "\n",
        "# Dataset generates trajectories with shape [Bx2x...]\n",
        "dataset = replay_buffer.as_dataset(\n",
        "    num_parallel_calls=3, \n",
        "    sample_batch_size=batch_size, \n",
        "    num_steps=2).prefetch(3)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rVD5nQ9ZGo8_"
      },
      "source": [
        "## Random Collection\n",
        "\n",
        "The algorithm must prime the pump.  Training cannot begin on an empty replay buffer.  The following code performs a predefined number of steps to generate initial training data."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "wr1KSAEGG4h9"
      },
      "source": [
        "random_policy = random_tf_policy.RandomTFPolicy(train_env.time_step_spec(),\n",
        "                                                train_env.action_spec())\n",
        "\n",
        "def collect_step(environment, policy, buffer):\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",
        "  buffer.add_batch(traj)\n",
        "\n",
        "def collect_data(env, policy, buffer, steps):\n",
        "  for _ in range(steps):\n",
        "    collect_step(env, policy, buffer)\n",
        "\n",
        "collect_data(train_env, random_policy, replay_buffer, \\\n",
        "             steps=initial_collect_steps)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hBc9lj9VWWtZ"
      },
      "source": [
        "## Training the agent\n",
        "\n",
        "We are now ready to train the DQN. This process can take many hours, depending on how many episodes you wish to run through.  As training occurs, this code will update on both the loss and average return.  As training becomes more successful, the average return should increase. The losses reported reflecting the average loss for individual training batches."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "0pTbJ3PeyF-u"
      },
      "source": [
        "iterator = iter(dataset)\n",
        "\n",
        "# (Optional) Optimize by wrapping some of the code in a graph \n",
        "# 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, \\\n",
        "                                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, replay_buffer)\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).loss\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))\n",
        "\n",
        "  if step % eval_interval == 0:\n",
        "    avg_return = compute_avg_return(eval_env, agent.policy, \\\n",
        "                                    num_eval_episodes)\n",
        "    print('step = {0}: Average Return = {1}'.format(step, avg_return))\n",
        "    returns.append(avg_return)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "68jNcA_TiJDq"
      },
      "source": [
        "## Visualization\n",
        "\n",
        "The notebook can plot the average return over training iterations.  The average return should increase as the program performs more training iterations."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "NxtL1mbOYCVO"
      },
      "source": [
        "iterations = range(0, num_iterations + 1, eval_interval)\n",
        "plt.plot(iterations, returns)\n",
        "plt.ylabel('Average Return')\n",
        "plt.xlabel('Iterations')\n",
        "plt.ylim(top=10)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M7-XpPP99Cy7"
      },
      "source": [
        "### Videos\n",
        "\n",
        "We now have a trained model and observed its training progress on a graph.  Perhaps the most compelling way to view an Atari game's results is a video that allows us to see the agent play the game. The following functions are defined so that we can watch the agent play the game in the notebook."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "ULaGr8pvOKbl"
      },
      "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",
        "  <video width=\"640\" height=\"480\" controls>\n",
        "    <source src=\"data:video/mp4;base64,{0}\" type=\"video/mp4\">\n",
        "  Your browser does not support the video tag.\n",
        "  </video>'''.format(b64.decode())\n",
        "\n",
        "  return IPython.display.HTML(tag)\n",
        "\n",
        "def create_policy_eval_video(policy, filename, num_episodes=5, fps=30):\n",
        "  filename = filename + \".mp4\"\n",
        "  with imageio.get_writer(filename, fps=fps) 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 = policy.action(time_step)\n",
        "        time_step = eval_env.step(action_step.action)\n",
        "        video.append_data(eval_py_env.render())\n",
        "  return embed_mp4(filename)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9c_PH-pX4Pr5"
      },
      "source": [
        "First, we will observe the trained agent play the game."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "owOVWB158NlF"
      },
      "source": [
        "create_policy_eval_video(agent.policy, \"trained-agent\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "povaAOcZygLw"
      },
      "source": [
        "For comparison, we observe a random agent play.  While the trained agent is far from perfect, it does outperform the random agent by a considerable amount."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "pJZIdC37yNH4"
      },
      "source": [
        "create_policy_eval_video(random_policy, \"random-agent\")"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}