docs/tutorials/4_drivers_tutorial.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "beObUOFyuRjT"
      },
      "source": [
        "##### Copyright 2023 The TF-Agents Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "nQnmcm0oI1Q-"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s6D70EeAZe-Q"
      },
      "source": [
        "# Drivers\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/agents/tutorials/4_drivers_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/4_drivers_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/4_drivers_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/4_drivers_tutorial.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8aPHF9kXFggA"
      },
      "source": [
        "## Introduction\n",
        "\n",
        "A common pattern in reinforcement learning is to execute a policy in an environment for a specified number of steps or episodes. This happens, for example, during data collection, evaluation and generating a video of the agent.\n",
        "\n",
        "While this is relatively straightforward to write in python, it is much more complex to write and debug in TensorFlow because it involves `tf.while` loops, `tf.cond` and `tf.control_dependencies`. Therefore we abstract this notion of a run loop into a class called `driver`, and provide well tested implementations both in Python and TensorFlow.\n",
        "\n",
        "Additionally, the data encountered by the driver at each step is saved in a named tuple called Trajectory and broadcast to a set of observers such as replay buffers and metrics. This data includes the observation from the environment, the action recommended by the policy, the reward obtained, the type of the current and the next step, etc."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t7PM1QfMZqkS"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0w-Ykwl1bn4v"
      },
      "source": [
        "If you haven't installed tf-agents or gym yet, run:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TnE2CgilrngG"
      },
      "outputs": [],
      "source": [
        "!pip install tf-agents\n",
        "!pip install tf-keras"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WPuD0bMEY9Iz"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "# Keep using keras-2 (tf-keras) rather than keras-3 (keras).\n",
        "os.environ['TF_USE_LEGACY_KERAS'] = '1'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "whYNP894FSkA"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import\n",
        "from __future__ import division\n",
        "from __future__ import print_function\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "\n",
        "from tf_agents.environments import suite_gym\n",
        "from tf_agents.environments import tf_py_environment\n",
        "from tf_agents.policies import random_py_policy\n",
        "from tf_agents.policies import random_tf_policy\n",
        "from tf_agents.metrics import py_metrics\n",
        "from tf_agents.metrics import tf_metrics\n",
        "from tf_agents.drivers import py_driver\n",
        "from tf_agents.drivers import dynamic_episode_driver"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9V7DEcB8IeiQ"
      },
      "source": [
        "## Python Drivers\n",
        "\n",
        "The `PyDriver` class takes a python environment, a python policy and a list of observers to update at each step. The main method is `run()`, which steps the environment using actions from the policy until at least one of the following termination criteria is met: The number of steps reaches `max_steps` or the number of episodes reaches `max_episodes`.\n",
        "\n",
        "The implementation is roughly as follows:\n",
        "\n",
        "\n",
        "```python\n",
        "class PyDriver(object):\n",
        "\n",
        "  def __init__(self, env, policy, observers, max_steps=1, max_episodes=1):\n",
        "    self._env = env\n",
        "    self._policy = policy\n",
        "    self._observers = observers or []\n",
        "    self._max_steps = max_steps or np.inf\n",
        "    self._max_episodes = max_episodes or np.inf\n",
        "\n",
        "  def run(self, time_step, policy_state=()):\n",
        "    num_steps = 0\n",
        "    num_episodes = 0\n",
        "    while num_steps \u003c self._max_steps and num_episodes \u003c self._max_episodes:\n",
        "\n",
        "      # Compute an action using the policy for the given time_step\n",
        "      action_step = self._policy.action(time_step, policy_state)\n",
        "\n",
        "      # Apply the action to the environment and get the next step\n",
        "      next_time_step = self._env.step(action_step.action)\n",
        "\n",
        "      # Package information into a trajectory\n",
        "      traj = trajectory.Trajectory(\n",
        "         time_step.step_type,\n",
        "         time_step.observation,\n",
        "         action_step.action,\n",
        "         action_step.info,\n",
        "         next_time_step.step_type,\n",
        "         next_time_step.reward,\n",
        "         next_time_step.discount)\n",
        "\n",
        "      for observer in self._observers:\n",
        "        observer(traj)\n",
        "\n",
        "      # Update statistics to check termination\n",
        "      num_episodes += np.sum(traj.is_last())\n",
        "      num_steps += np.sum(~traj.is_boundary())\n",
        "\n",
        "      time_step = next_time_step\n",
        "      policy_state = action_step.state\n",
        "\n",
        "    return time_step, policy_state\n",
        "\n",
        "```\n",
        "\n",
        "Now, let us run through the example of running a random policy on the CartPole environment, saving the results to a replay buffer and computing some metrics."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dj4_-77_5ExP"
      },
      "outputs": [],
      "source": [
        "env = suite_gym.load('CartPole-v0')\n",
        "policy = random_py_policy.RandomPyPolicy(time_step_spec=env.time_step_spec(), \n",
        "                                         action_spec=env.action_spec())\n",
        "replay_buffer = []\n",
        "metric = py_metrics.AverageReturnMetric()\n",
        "observers = [replay_buffer.append, metric]\n",
        "driver = py_driver.PyDriver(\n",
        "    env, policy, observers, max_steps=20, max_episodes=1)\n",
        "\n",
        "initial_time_step = env.reset()\n",
        "final_time_step, _ = driver.run(initial_time_step)\n",
        "\n",
        "print('Replay Buffer:')\n",
        "for traj in replay_buffer:\n",
        "  print(traj)\n",
        "\n",
        "print('Average Return: ', metric.result())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X3Yrxg36Ik1x"
      },
      "source": [
        "## TensorFlow Drivers\n",
        "\n",
        "We also have drivers in TensorFlow which are functionally similar to Python drivers, but use TF environments, TF policies, TF observers etc. We currently have 2 TensorFlow drivers: `DynamicStepDriver`, which terminates after a given number of (valid) environment steps and `DynamicEpisodeDriver`, which terminates after a given number of episodes. Let us look at an example of the DynamicEpisode in action.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WC4ba3ObSceA"
      },
      "outputs": [],
      "source": [
        "env = suite_gym.load('CartPole-v0')\n",
        "tf_env = tf_py_environment.TFPyEnvironment(env)\n",
        "\n",
        "tf_policy = random_tf_policy.RandomTFPolicy(action_spec=tf_env.action_spec(),\n",
        "                                            time_step_spec=tf_env.time_step_spec())\n",
        "\n",
        "\n",
        "num_episodes = tf_metrics.NumberOfEpisodes()\n",
        "env_steps = tf_metrics.EnvironmentSteps()\n",
        "observers = [num_episodes, env_steps]\n",
        "driver = dynamic_episode_driver.DynamicEpisodeDriver(\n",
        "    tf_env, tf_policy, observers, num_episodes=2)\n",
        "\n",
        "# Initial driver.run will reset the environment and initialize the policy.\n",
        "final_time_step, policy_state = driver.run()\n",
        "\n",
        "print('final_time_step', final_time_step)\n",
        "print('Number of Steps: ', env_steps.result().numpy())\n",
        "print('Number of Episodes: ', num_episodes.result().numpy())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Sz5jhHnU0fX1"
      },
      "outputs": [],
      "source": [
        "# Continue running from previous state\n",
        "final_time_step, _ = driver.run(final_time_step, policy_state)\n",
        "\n",
        "print('final_time_step', final_time_step)\n",
        "print('Number of Steps: ', env_steps.result().numpy())\n",
        "print('Number of Episodes: ', num_episodes.result().numpy())"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "TF-Agents Drivers Tutorial.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}