docs/tutorials/1_dqn_tutorial.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "klGNgWREsvQv"
      },
      "source": [
        "##### Copyright 2018 The TF-Agents Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "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": {
        "colab_type": "text",
        "id": "pmDI-h7cI0tI"
      },
      "source": [
        "# Train a Deep Q Network with TF-Agents\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/agents/tutorials/1_dqn_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/1_dqn_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/1_dqn_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/1_dqn_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": {
        "colab_type": "text",
        "id": "lsaQlK8fFQqH"
      },
      "source": [
        "## Introduction\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "cKOCZlhUgXVK"
      },
      "source": [
        "This example shows how to train a [DQN (Deep Q Networks)](https://storage.googleapis.com/deepmind-media/dqn/DQNNaturePaper.pdf)  agent on the Cartpole environment using the TF-Agents library.\n",
        "\n",
        "![Cartpole environment](https://raw.githubusercontent.com/tensorflow/agents/master/docs/tutorials/images/cartpole.png)\n",
        "\n",
        "It will walk you through all the components in a Reinforcement Learning (RL) pipeline for training, evaluation and data collection.\n",
        "\n",
        "\n",
        "To run this code live, click the 'Run in Google Colab' link above.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "1u9QVVsShC9X"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kNrNXKI7bINP"
      },
      "source": [
        "If you haven't installed the following dependencies, run:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "KEHR2Ui-lo8O"
      },
      "outputs": [],
      "source": [
        "!sudo apt-get install -y xvfb ffmpeg\n",
        "!pip install 'gym==0.10.11'\n",
        "!pip install 'imageio==2.4.0'\n",
        "!pip install PILLOW\n",
        "!pip install 'pyglet==1.3.2'\n",
        "!pip install pyvirtualdisplay\n",
        "!pip install --upgrade tensorflow-probability\n",
        "!pip install tf-agents"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "sMitx5qSgJk1"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import, division, print_function\n",
        "\n",
        "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\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 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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "J6HsdS5GbSjd"
      },
      "outputs": [],
      "source": [
        "tf.compat.v1.enable_v2_behavior()\n",
        "\n",
        "# Set up a virtual display for rendering OpenAI gym environments.\n",
        "display = pyvirtualdisplay.Display(visible=0, size=(1400, 900)).start()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "NspmzG4nP3b9"
      },
      "outputs": [],
      "source": [
        "tf.version.VERSION"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "LmC0NDhdLIKY"
      },
      "source": [
        "## Hyperparameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "HC1kNrOsLSIZ"
      },
      "outputs": [],
      "source": [
        "num_iterations = 20000 # @param {type:\"integer\"}\n",
        "\n",
        "initial_collect_steps = 1000  # @param {type:\"integer\"} \n",
        "collect_steps_per_iteration = 1  # @param {type:\"integer\"}\n",
        "replay_buffer_max_length = 100000  # @param {type:\"integer\"}\n",
        "\n",
        "batch_size = 64  # @param {type:\"integer\"}\n",
        "learning_rate = 1e-3  # @param {type:\"number\"}\n",
        "log_interval = 200  # @param {type:\"integer\"}\n",
        "\n",
        "num_eval_episodes = 10  # @param {type:\"integer\"}\n",
        "eval_interval = 1000  # @param {type:\"integer\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VMsJC3DEgI0x"
      },
      "source": [
        "## Environment\n",
        "\n",
        "In Reinforcement Learning (RL), an environment represents the task or problem to be solved. Standard environments can be created in TF-Agents using `tf_agents.environments` suites. TF-Agents has suites for loading environments from sources such as the OpenAI Gym, Atari, and DM Control.\n",
        "\n",
        "Load the CartPole environment from the OpenAI Gym suite. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "pYEz-S9gEv2-"
      },
      "outputs": [],
      "source": [
        "env_name = 'CartPole-v0'\n",
        "env = suite_gym.load(env_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "IIHYVBkuvPNw"
      },
      "source": [
        "You can render this environment to see how it looks. A free-swinging pole is attached to a cart.  The goal is to move the cart right or left in order to keep the pole pointing up."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "RlO7WIQHu_7D"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "env.reset()\n",
        "PIL.Image.fromarray(env.render())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "B9_lskPOey18"
      },
      "source": [
        "The `environment.step` method takes an `action` in the environment and returns a `TimeStep` tuple containing the next observation of the environment and the reward for the action.\n",
        "\n",
        "The `time_step_spec()` method returns the specification for the `TimeStep` tuple. Its `observation` attribute shows the shape of observations, the data types, and the ranges of allowed values. The `reward` attribute shows the same details for the reward.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "exDv57iHfwQV"
      },
      "outputs": [],
      "source": [
        "print('Observation Spec:')\n",
        "print(env.time_step_spec().observation)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "UxiSyCbBUQPi"
      },
      "outputs": [],
      "source": [
        "print('Reward Spec:')\n",
        "print(env.time_step_spec().reward)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "b_lHcIcqUaqB"
      },
      "source": [
        "The `action_spec()` method returns the shape, data types, and allowed values of valid actions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "bttJ4uxZUQBr"
      },
      "outputs": [],
      "source": [
        "print('Action Spec:')\n",
        "print(env.action_spec())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "eJCgJnx3g0yY"
      },
      "source": [
        "In the Cartpole environment:\n",
        "\n",
        "-   `observation` is an array of 4 floats: \n",
        "    -   the position and velocity of the cart\n",
        "    -   the angular position and velocity of the pole \n",
        "-   `reward` is a scalar float value\n",
        "-   `action` is a scalar integer with only two possible values:\n",
        "    -   `0` — \"move left\"\n",
        "    -   `1` — \"move right\"\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "V2UGR5t_iZX-"
      },
      "outputs": [],
      "source": [
        "time_step = env.reset()\n",
        "print('Time step:')\n",
        "print(time_step)\n",
        "\n",
        "action = np.array(1, dtype=np.int32)\n",
        "\n",
        "next_time_step = env.step(action)\n",
        "print('Next time step:')\n",
        "print(next_time_step)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4JSc9GviWUBK"
      },
      "source": [
        "Usually two environments are instantiated: one for training and one for evaluation. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "N7brXNIGWXjC"
      },
      "outputs": [],
      "source": [
        "train_py_env = suite_gym.load(env_name)\n",
        "eval_py_env = suite_gym.load(env_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "zuUqXAVmecTU"
      },
      "source": [
        "The Cartpole environment, like most environments, is written in pure Python. This is converted to TensorFlow using the `TFPyEnvironment` wrapper.\n",
        "\n",
        "The original environment's API uses Numpy arrays. The `TFPyEnvironment` converts these to `Tensors` to make it compatible with Tensorflow agents and policies.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Xp-Y4mD6eDhF"
      },
      "outputs": [],
      "source": [
        "train_env = tf_py_environment.TFPyEnvironment(train_py_env)\n",
        "eval_env = tf_py_environment.TFPyEnvironment(eval_py_env)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "E9lW_OZYFR8A"
      },
      "source": [
        "## Agent\n",
        "\n",
        "The algorithm used to solve an RL problem is represented by an `Agent`. TF-Agents provides standard implementations of a variety of `Agents`, including:\n",
        "\n",
        "-   [DQN](https://storage.googleapis.com/deepmind-media/dqn/DQNNaturePaper.pdf) (used in this tutorial)\n",
        "-   [REINFORCE](http://www-anw.cs.umass.edu/~barto/courses/cs687/williams92simple.pdf)\n",
        "-   [DDPG](https://arxiv.org/pdf/1509.02971.pdf)\n",
        "-   [TD3](https://arxiv.org/pdf/1802.09477.pdf)\n",
        "-   [PPO](https://arxiv.org/abs/1707.06347)\n",
        "-   [SAC](https://arxiv.org/abs/1801.01290).\n",
        "\n",
        "The DQN agent can be used in any environment which has a discrete action space.\n",
        "\n",
        "At the heart of a DQN Agent is a `QNetwork`, a neural network model that can learn to predict `QValues` (expected returns) for all actions, given an observation from the environment.\n",
        "\n",
        "Use `tf_agents.networks.q_network` to create a `QNetwork`, passing in the `observation_spec`, `action_spec`, and a tuple describing the number and size of the model's hidden layers.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "TgkdEPg_muzV"
      },
      "outputs": [],
      "source": [
        "fc_layer_params = (100,)\n",
        "\n",
        "q_net = q_network.QNetwork(\n",
        "    train_env.observation_spec(),\n",
        "    train_env.action_spec(),\n",
        "    fc_layer_params=fc_layer_params)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "z62u55hSmviJ"
      },
      "source": [
        "Now use `tf_agents.agents.dqn.dqn_agent` to instantiate a `DqnAgent`. In addition to the `time_step_spec`, `action_spec` and the QNetwork, the agent constructor also requires an optimizer (in this case, `AdamOptimizer`), a loss function, and an integer step counter."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "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 = dqn_agent.DqnAgent(\n",
        "    train_env.time_step_spec(),\n",
        "    train_env.action_spec(),\n",
        "    q_network=q_net,\n",
        "    optimizer=optimizer,\n",
        "    td_errors_loss_fn=common.element_wise_squared_loss,\n",
        "    train_step_counter=train_step_counter)\n",
        "\n",
        "agent.initialize()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "I0KLrEPwkn5x"
      },
      "source": [
        "## Policies\n",
        "\n",
        "A policy defines the way an agent acts in an environment. Typically, the goal of reinforcement learning is to train the underlying model until the policy produces the desired outcome.\n",
        "\n",
        "In this tutorial:\n",
        "\n",
        "-   The desired outcome is keeping the pole balanced upright over the cart.\n",
        "-   The policy returns an action (left or right) for each `time_step` observation.\n",
        "\n",
        "Agents contain two policies: \n",
        "\n",
        "-   `agent.policy` — The main policy that is used for evaluation and deployment.\n",
        "-   `agent.collect_policy` — A second policy that is used for data collection.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "BwY7StuMkuV4"
      },
      "outputs": [],
      "source": [
        "eval_policy = agent.policy\n",
        "collect_policy = agent.collect_policy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "2Qs1Fl3dV0ae"
      },
      "source": [
        "Policies can be created independently of agents. For example, use `tf_agents.policies.random_tf_policy` to create a policy which will randomly select an action for each `time_step`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "HE37-UCIrE69"
      },
      "outputs": [],
      "source": [
        "random_policy = random_tf_policy.RandomTFPolicy(train_env.time_step_spec(),\n",
        "                                                train_env.action_spec())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "dOlnlRRsUbxP"
      },
      "source": [
        "To get an action from a policy, call the `policy.action(time_step)` method. The `time_step` contains the observation from the environment. This method returns a `PolicyStep`, which is a named tuple with three components:\n",
        "\n",
        "-   `action` — the action to be taken (in this case, `0` or `1`)\n",
        "-   `state` — used for stateful (that is, RNN-based) policies\n",
        "-   `info` — auxiliary data, such as log probabilities of actions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "5gCcpXswVAxk"
      },
      "outputs": [],
      "source": [
        "example_environment = tf_py_environment.TFPyEnvironment(\n",
        "    suite_gym.load('CartPole-v0'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "D4DHZtq3Ndis"
      },
      "outputs": [],
      "source": [
        "time_step = example_environment.reset()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "PRFqAUzpNaAW"
      },
      "outputs": [],
      "source": [
        "random_policy.action(time_step)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "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. Several episodes are run, creating an average return.\n",
        "\n",
        "The following function computes the average return of a policy, given the policy, environment, and a number of episodes.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "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",
        "# See also the metrics module for standard implementations of different metrics.\n",
        "# https://github.com/tensorflow/agents/tree/master/tf_agents/metrics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "_snCVvq5Z8lJ"
      },
      "source": [
        "Running this computation on the `random_policy` shows a baseline performance in the environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "9bgU6Q6BZ8Bp"
      },
      "outputs": [],
      "source": [
        "compute_avg_return(eval_env, random_policy, num_eval_episodes)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NLva6g2jdWgr"
      },
      "source": [
        "## Replay Buffer\n",
        "\n",
        "The replay buffer keeps track of data collected from the environment. This tutorial uses `tf_agents.replay_buffers.tf_uniform_replay_buffer.TFUniformReplayBuffer`, as it is the most common. \n",
        "\n",
        "The constructor requires the specs for the data it will be collecting. This is available from the agent using the `collect_data_spec` method. The batch size and maximum buffer length are also required.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "vX2zGUWJGWAl"
      },
      "outputs": [],
      "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)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ZGNTDJpZs4NN"
      },
      "source": [
        "For most agents, `collect_data_spec` is a named tuple called `Trajectory`, containing the specs for observations, actions, rewards, and other items."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "_IZ-3HcqgE1z"
      },
      "outputs": [],
      "source": [
        "agent.collect_data_spec"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "sy6g1tGcfRlw"
      },
      "outputs": [],
      "source": [
        "agent.collect_data_spec._fields"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "rVD5nQ9ZGo8_"
      },
      "source": [
        "## Data Collection\n",
        "\n",
        "Now execute the random policy in the environment for a few steps, recording the data in the replay buffer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "wr1KSAEGG4h9"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\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, steps=100)\n",
        "\n",
        "# This loop is so common in RL, that we provide standard implementations. \n",
        "# For more details see the drivers module.\n",
        "# https://github.com/tensorflow/agents/blob/master/tf_agents/docs/python/tf_agents/drivers.md"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "84z5pQJdoKxo"
      },
      "source": [
        "The replay buffer is now a collection of Trajectories."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "4wZnLu2ViO4E"
      },
      "outputs": [],
      "source": [
        "# For the curious:\n",
        "# Uncomment to peel one of these off and inspect it.\n",
        "# iter(replay_buffer.as_dataset()).next()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "TujU-PMUsKjS"
      },
      "source": [
        "The agent needs access to the replay buffer. This is provided by creating an iterable `tf.data.Dataset` pipeline which will feed data to the agent.\n",
        "\n",
        "Each row of the replay buffer only stores a single observation step. But since the DQN Agent needs both the current and next observation to compute the loss, the dataset pipeline will sample two adjacent rows for each item in the batch (`num_steps=2`).\n",
        "\n",
        "This dataset is also optimized by running parallel calls and prefetching data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ba7bilizt_qW"
      },
      "outputs": [],
      "source": [
        "# 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)\n",
        "\n",
        "\n",
        "dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "K13AST-2ppOq"
      },
      "outputs": [],
      "source": [
        "iterator = iter(dataset)\n",
        "\n",
        "print(iterator)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Th5w5Sff0b16"
      },
      "outputs": [],
      "source": [
        "# For the curious:\n",
        "# Uncomment to see what the dataset iterator is feeding to the agent.\n",
        "# Compare this representation of replay data \n",
        "# to the collection of individual trajectories shown earlier.\n",
        "\n",
        "# iterator.next()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hBc9lj9VWWtZ"
      },
      "source": [
        "## Training the agent\n",
        "\n",
        "Two things must happen during the training loop:\n",
        "\n",
        "-   collect data from the environment\n",
        "-   use that data to train the agent's neural network(s)\n",
        "\n",
        "This example also periodicially evaluates the policy and prints the current score.\n",
        "\n",
        "The following will take ~5 minutes to run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "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, 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, num_eval_episodes)\n",
        "    print('step = {0}: Average Return = {1}'.format(step, avg_return))\n",
        "    returns.append(avg_return)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "68jNcA_TiJDq"
      },
      "source": [
        "## Visualization\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "aO-LWCdbbOIC"
      },
      "source": [
        "### Plots\n",
        "\n",
        "Use `matplotlib.pyplot` to chart how the policy improved during training.\n",
        "\n",
        "One iteration of `Cartpole-v0` consists of 200 time steps. The environment gives a reward of `+1` for each step the pole stays up, so the maximum return for one episode is 200. The charts shows the return increasing towards that maximum each time it is evaluated during training. (It may be a little unstable and not increase monotonically each time.)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "NxtL1mbOYCVO"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "\n",
        "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=250)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "M7-XpPP99Cy7"
      },
      "source": [
        "### Videos"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "9pGfGxSH32gn"
      },
      "source": [
        "Charts are nice. But more exciting is seeing an agent actually performing a task in an environment. \n",
        "\n",
        "First, create a function to embed videos in the notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "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": {
        "colab_type": "text",
        "id": "9c_PH-pX4Pr5"
      },
      "source": [
        "Now iterate through a few episodes of the Cartpole game with the agent. The underlying Python environment (the one \"inside\" the TensorFlow environment wrapper) provides a `render()` method, which outputs an image of the environment state. These can be collected into a video."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "owOVWB158NlF"
      },
      "outputs": [],
      "source": [
        "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)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "create_policy_eval_video(agent.policy, \"trained-agent\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "povaAOcZygLw"
      },
      "source": [
        "For fun, compare the trained agent (above) to an agent moving randomly. (It does not do as well.)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "pJZIdC37yNH4"
      },
      "outputs": [],
      "source": [
        "create_policy_eval_video(random_policy, \"random-agent\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "DQN 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
}