<commit_before><commit_msg>Internal Change<commit_after>{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "klGNgWREsvQv"
      },
      "source": [
        "##### Copyright 2018 The TF-Agents Authors."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "lsaQlK8fFQqH"
      },
      "source": [
        "### Get Started\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/agents/blob/master/tf_agents/colabs/dqn_colab.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun 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/tf_agents/colabs/dqn_colab.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "TnE2CgilrngG"
      },
      "outputs": [],
      "source": [
        "# Note: If you haven't installed tf-agents or gym yet, run:\n",
        "!pip install tf-agents\n",
        "!pip install gym"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ZOUOQOrFs3zn"
      },
      "source": [
        "## Introduction"
      ]
    },
    {
      "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. We will walk you through all the components in a Reinforcement Learning (RL) pipeline for training, evaluation and data collection."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "1u9QVVsShC9X"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "sMitx5qSgJk1"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "from tf_agents.agents.dqn import dqn_agent\n",
        "from tf_agents.agents.dqn import q_network\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.environments import trajectory\n",
        "from tf_agents.metrics import metric_utils\n",
        "from tf_agents.metrics import tf_metrics\n",
        "from tf_agents.policies import random_tf_policy\n",
        "from tf_agents.replay_buffers import tf_uniform_replay_buffer\n",
        "from tf_agents.utils import common"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "a4nXjLSTMNmJ"
      },
      "outputs": [],
      "source": [
        "tf.compat.v1.enable_v2_behavior()"
      ]
    },
    {
      "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": [
        "env_name = 'CartPole-v0'  # @param\n",
        "num_iterations = 20000  # @param\n",
        "\n",
        "initial_collect_steps = 1000  # @param\n",
        "collect_steps_per_iteration = 1  # @param\n",
        "replay_buffer_capacity = 100000  # @param\n",
        "\n",
        "fc_layer_params = (100,)\n",
        "\n",
        "batch_size = 64  # @param\n",
        "learning_rate = 1e-3  # @param\n",
        "log_interval = 200  # @param\n",
        "\n",
        "num_eval_episodes = 10  # @param\n",
        "eval_interval = 1000  # @param"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VMsJC3DEgI0x"
      },
      "source": [
        "## Environment\n",
        "\n",
        "Environments in RL represent the task or problem that we are trying to solve. Standard environments can be easily created in TF-Agents using `suites`. We have different `suites` for loading environments from sources such as the OpenAI Gym, Atari, DM Control etc given a string environment name.\n",
        "\n",
        "Now let us load the CartPole environment from the OpenAI Gym suite. A pole is attached by an un-actuated joint to a cart, which moves along a frictionless track. The pole starts upright, and the goal is to prevent it from falling over by applying a horizontal force to the cart."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "pYEz-S9gEv2-"
      },
      "outputs": [],
      "source": [
        "env = suite_gym.load(env_name)\n",
        "\n",
        "# Render environment here."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "B9_lskPOey18"
      },
      "source": [
        "The most important method in environments is the `time_step = environment.step(action)` method which performs one simulation step of the environment given an `action` and returns a `TimeStep` named tuple containing the next observation, reward, etc. Other important methods are `time_step_spec()` and `action_spec()` which return the specifications (types, shapes, bounds) of the `time_step` and `action` respectively."
      ]
    },
    {
      "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\n",
        "print 'Action Spec:'\n",
        "print env.action_spec()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "eJCgJnx3g0yY"
      },
      "source": [
        "So, we see that observation is an array of 4 floats: the position and velocity of the cart, and the angular position and velocity of the pole. Only two discrete actions are available: `0` applies a force on the cart to the left and `1` applies a force to the right (Is it the other way?).\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 = 1\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": "zuUqXAVmecTU"
      },
      "source": [
        "Usually we create two environments: one for training and one for evaluation. Most environments are written in pure python, but they can be easily converted to TensorFlow using the `TFPyEnvironment` wrapper.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Xp-Y4mD6eDhF"
      },
      "outputs": [],
      "source": [
        "train_env = tf_py_environment.TFPyEnvironment(suite_gym.load(env_name))\n",
        "eval_env = tf_py_environment.TFPyEnvironment(suite_gym.load(env_name))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "E9lW_OZYFR8A"
      },
      "source": [
        "## Agent\n",
        "\n",
        "The algorithm that we use to solve an RL problem is represented as an `Agent`. In addition to the DQN agent, TF-Agents provides standard implementations of a variety of `Agents` such as [REINFORCE](http://www-anw.cs.umass.edu/~barto/courses/cs687/williams92simple.pdf), [DDPG](https://arxiv.org/pdf/1509.02971.pdf), [TD3](https://arxiv.org/pdf/1802.09477.pdf), [PPO](https://arxiv.org/abs/1707.06347) and [SAC](https://arxiv.org/abs/1801.01290).\n",
        "\n",
        "The DQN agent can be used in any environment which has a discrete action space. To create a DQN Agent, we first need a `Q-Network` that can learn to predict `Q-Values` (expected return) for all actions given an observation from the environment. \n",
        "\n",
        "We can easily create a `Q-Network` using the specs of the observations and actions. We can specify the layers in the network which, in this example, is the `fc_layer_params` argument set to a tuple of `ints` representing the sizes of each hidden layer.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "TgkdEPg_muzV"
      },
      "outputs": [],
      "source": [
        "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": [
        "We also need an `optimizer` to train the agent and a `train_step_counter` variable to keep track of how many times the agent's networks were updated.\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "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.compat.v2.Variable(0)\n",
        "\n",
        "tf_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=dqn_agent.element_wise_squared_loss,\n",
        "    train_step_counter=train_step_counter)\n",
        "tf_agent.initialize()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "I0KLrEPwkn5x"
      },
      "source": [
        "## Policies\n",
        "\n",
        "In TF-Agents, policies represent the standard notion of policies in RL: given a `time_step` produce an action or a distribution over actions. The main method is `policy_step = policy.step(time_step)` where `policy_step` is a named tuple `PolicyStep(action, state, info)`.  `policy_step.action` is the `action` to be applied to the environment, `state` represents the state for stateful (RNN) policies and `info` may contain auxiliary information such as log probabilities of the actions. \n",
        "\n",
        "Agents contain two policies: the main policy that is used for evaluation/deployment (agent.policy) and another policy that is used for data collection (agent.collect_policy)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "BwY7StuMkuV4"
      },
      "outputs": [],
      "source": [
        "eval_policy = tf_agent.policy\n",
        "collect_policy = tf_agent.collect_policy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "d5a04gY_rD_a"
      },
      "source": [
        "We can also independently create policies that are not part of an agent. For example, a random policy:"
      ]
    },
    {
      "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": "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": 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",
        "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": {
        "colab_type": "text",
        "id": "NLva6g2jdWgr"
      },
      "source": [
        "## Replay Buffer\n",
        "\n",
        "In order to keep track of the data collected from the environment, we will use the TFUniformReplayBuffer. This replay buffer is constructed using specs describing the tensors that are to be stored, which can be obtained from the agent using `tf_agent.collect_data_spec`."
      ]
    },
    {
      "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=tf_agent.collect_data_spec,\n",
        "    batch_size=train_env.batch_size,\n",
        "    max_length=replay_buffer_capacity)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ZGNTDJpZs4NN"
      },
      "source": [
        "For most agents, the `collect_data_spec` is a `Trajectory` named tuple containing the observation, action, reward etc."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "rVD5nQ9ZGo8_"
      },
      "source": [
        "## Data Collection\n",
        "\n",
        "Now let us execute the random policy in the environment for a few steps and record the data (observations, actions, rewards etc) 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):\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",
        "\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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "TujU-PMUsKjS"
      },
      "source": [
        "In order to sample data from the replay buffer, we will create a `tf.data` pipeline which we can feed to the agent for training later. We can specify the `sample_batch_size` to configure the number of items sampled from the replay buffer. We can also optimize the data pipline using parallel calls and prefetching.\n",
        "\n",
        "In order to save space, we only store the current observation in each row of the replay buffer. But since the DQN Agent needs both the current and next observation to compute the loss, we always sample two adjacent rows for each item in the batch by setting `num_steps=2`."
      ]
    },
    {
      "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, sample_batch_size=batch_size, num_steps=2).prefetch(3)\n",
        "\n",
        "iterator = iter(dataset)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "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 ~5 minutes to run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "0pTbJ3PeyF-u"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "%%time\n",
        "\n",
        "# (Optional) Optimize by wrapping some of the code in a graph using TF function.\n",
        "tf_agent.train = common.function(tf_agent.train)\n",
        "\n",
        "# Reset the train step\n",
        "tf_agent.train_step_counter.assign(0)\n",
        "\n",
        "# Evaluate the agent's policy once before training.\n",
        "avg_return = compute_avg_return(eval_env, tf_agent.policy, num_eval_episodes)\n",
        "returns = [avg_return]\n",
        "\n",
        "for _ in range(num_iterations):\n",
        "\n",
        "  # Collect one step using collect_policy and save to the replay buffer.\n",
        "  collect_step(train_env, tf_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 = tf_agent.train(experience)\n",
        "\n",
        "  step = tf_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, tf_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",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "aO-LWCdbbOIC"
      },
      "source": [
        "### Plots\n",
        "\n",
        "We can plot return vs global steps to see the performance of our agent. In `Cartpole-v0`, the environment gives a reward of +1 for every time step the pole stays up, and since the maximum number of steps is 200, the maximum possible return is also 200."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "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=250)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "M7-XpPP99Cy7"
      },
      "source": [
        "### Videos"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "",
        "kind": "local"
      },
      "name": "DQN Colab.ipynb",
      "provenance": [
        {
          "file_id": "1tL_oyGKGsHmtpK056cVQteeadnpaHAtq",
          "timestamp": 1551337271702
        }
      ],
      "toc_visible": true,
      "version": "0.3.2"
    },
    "kernelspec": {
      "display_name": "Python 2",
      "name": "python2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
{'fraction_non_alphanumeric': 0.14994477578970622, 'fraction_numerical': 0.01365142478462558, 'mean_word_length': 2.061401136056262, 'pattern_counts': {'":': 281, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 9, 'lorem ipsum': 0, 'www.': 3, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2120015', 'n_tokens_mistral': 7758, 'n_tokens_neox': 7416, 'n_words': 1781}