tf_agents/colabs/6_reinforce_tutorial.ipynb

{
 "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",
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/agents/blob/master/tf_agents/colabs/6_reinforce_tutorial.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://github.com/tensorflow/agents/blob/master/tf_agents/colabs/6_reinforce_tutorial.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
    "  </td>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KEHR2Ui-lo8O"
   },
   "outputs": [],
   "source": [
    "# Note: If you haven't installed the following dependencies, run:\n",
    "!apt-get install xvfb\n",
    "!pip install 'gym==0.10.11'\n",
    "!pip install 'imageio==2.4.0'\n",
    "!pip install PILLOW\n",
    "!pip install pyglet\n",
    "!pip install pyvirtualdisplay\n",
    "!pip install tf-agents-nightly\n",
    "!pip install tf-nightly"
   ]
  },
  {
   "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 [REINFORCE](http://www-anw.cs.umass.edu/~barto/courses/cs687/williams92simple.pdf) agent on the Cartpole environment using the TF-Agents library, similar to the [DQN tutorial](1_dqn_tutorial.ipynb).\n",
    "\n",
    "![Cartpole environment](images/cartpole.png)\n",
    "\n",
    "We will walk you through all the components in a Reinforcement Learning (RL) pipeline for training, evaluation and data collection.\n"
   ]
  },
  {
   "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": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import base64\n",
    "import imageio\n",
    "import IPython\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import PIL.Image\n",
    "import pyvirtualdisplay\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "from tf_agents.agents.reinforce import reinforce_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 actor_distribution_network\n",
    "from tf_agents.replay_buffers import tf_uniform_replay_buffer\n",
    "from tf_agents.trajectories import trajectory\n",
    "from tf_agents.utils import common\n",
    "\n",
    "tf.compat.v1.enable_v2_behavior()\n",
    "\n",
    "\n",
    "# Set up a virtual display for rendering OpenAI gym environments.\n",
    "display = pyvirtualdisplay.Display(visible=0, size=(1400, 900)).start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "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 = 250  # @param\n",
    "collect_episodes_per_iteration = 2  # @param\n",
    "replay_buffer_capacity = 2000  # @param\n",
    "\n",
    "fc_layer_params = (100,)\n",
    "\n",
    "learning_rate = 1e-3  # @param\n",
    "log_interval = 25  # @param\n",
    "num_eval_episodes = 10  # @param\n",
    "eval_interval = 50  # @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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "pYEz-S9gEv2-"
   },
   "outputs": [],
   "source": [
    "env = suite_gym.load(env_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IIHYVBkuvPNw"
   },
   "source": [
    "We 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 `time_step = environment.step(action)` statement takes `action` in the environment.  The `TimeStep` tuple returned contains the environment's next observation and reward for that action. The `time_step_spec()` and `action_spec()` methods in the environment 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. Since only two actions are possible (move left or move right), the `action_spec` is a scalar where 0 means \"move left\" and 1 means \"move right.\""
   ]
  },
  {
   "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. The original environment's API uses numpy arrays, the `TFPyEnvironment` converts these to/from `Tensors` for you to more easily interact with TensorFlow policies and agents.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Xp-Y4mD6eDhF"
   },
   "outputs": [],
   "source": [
    "train_py_env = suite_gym.load(env_name)\n",
    "eval_py_env = suite_gym.load(env_name)\n",
    "\n",
    "train_env = tf_py_environment.TFPyEnvironment(train_py_env)\n",
    "eval_env = tf_py_environment.TFPyEnvironment(eval_py_env)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "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 REINFORCE agent, TF-Agents provides standard implementations of a variety of `Agents` such as [DQN](https://storage.googleapis.com/deepmind-media/dqn/DQNNaturePaper.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",
    "To create a REINFORCE Agent, we first need an `Actor Network` that can learn to predict the action given an observation from the environment.\n",
    "\n",
    "We can easily create an `Actor 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 (see the Hyperparameters section above).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TgkdEPg_muzV"
   },
   "outputs": [],
   "source": [
    "actor_net = actor_distribution_network.ActorDistributionNetwork(\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 network we just created, and a `train_step_counter` variable to keep track of how many times the network was 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 = reinforce_agent.ReinforceAgent(\n",
    "    train_env.time_step_spec(),\n",
    "    train_env.action_spec(),\n",
    "    actor_network=actor_net,\n",
    "    optimizer=optimizer,\n",
    "    normalize_returns=True,\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)`.  The `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": "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",
    "# 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",
    "As REINFORCE learns from whole episodes, we define a function to collect an episode using the given data collection policy and save the data (observations, actions, rewards etc.) as trajectories in the replay buffer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wr1KSAEGG4h9"
   },
   "outputs": [],
   "source": [
    "#@test {\"skip\": true}\n",
    "\n",
    "def collect_episode(environment, policy, num_episodes):\n",
    "\n",
    "  episode_counter = 0\n",
    "  environment.reset()\n",
    "\n",
    "  while episode_counter < num_episodes:\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",
    "    if traj.is_boundary():\n",
    "      episode_counter += 1\n",
    "\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": "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 ~3 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 a few episodes using collect_policy and save to the replay buffer.\n",
    "  collect_episode(\n",
    "      train_env, tf_agent.collect_policy, collect_episodes_per_iteration)\n",
    "\n",
    "  # Use data from the buffer and update the agent's network.\n",
    "  experience = replay_buffer.gather_all()\n",
    "  train_loss = tf_agent.train(experience)\n",
    "  replay_buffer.clear()\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9pGfGxSH32gn"
   },
   "source": [
    "It is helpful to visualize the performance of an agent by rendering the environment at each step. Before we do that, let us first create a function to embed videos in this colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 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",
    "  <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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9c_PH-pX4Pr5"
   },
   "source": [
    "The following code visualizes the agent's policy for a few episodes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "owOVWB158NlF"
   },
   "outputs": [],
   "source": [
    "num_episodes = 3\n",
    "video_filename = 'imageio.mp4'\n",
    "with imageio.get_writer(video_filename, fps=60) as video:\n",
    "  for _ in range(num_episodes):\n",
    "    time_step = eval_env.reset()\n",
    "    video.append_data(eval_py_env.render())\n",
    "    while not time_step.is_last():\n",
    "      action_step = tf_agent.policy.action(time_step)\n",
    "      time_step = eval_env.step(action_step.action)\n",
    "      video.append_data(eval_py_env.render())\n",
    "\n",
    "embed_mp4(video_filename)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "REINFORCE Tutorial.ipynb",
   "toc_visible": true,
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 2",
   "name": "python2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}