docs/tutorials/2_environments_tutorial.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ma19Ks2CTDbZ"
      },
      "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": "XljsiF6lYkuV"
      },
      "source": [
        "# Environments\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/agents/tutorials/2_environments_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/2_environments_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/2_environments_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/2_environments_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": "9h3B-YBHopJI"
      },
      "source": [
        "## Introduction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n9c6vCPGovOM"
      },
      "source": [
        "The goal of Reinforcement Learning (RL) is to design agents that learn by interacting with an environment. In the standard RL setting, the agent receives an observation at every time step and chooses an action. The action is applied to the environment and the environment returns a reward and a new observation. The agent trains a policy to choose actions to maximize the sum of rewards, also known as return.\n",
        "\n",
        "In TF-Agents, environments can be implemented either in Python or TensorFlow. Python environments are usually easier to implement, understand, and debug, but TensorFlow environments are more efficient and allow natural parallelization. The most common workflow is to implement an environment in Python and use one of our wrappers to automatically convert it into TensorFlow.\n",
        "\n",
        "Let us look at Python environments first. TensorFlow environments follow a very similar API."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4_16bQF0anmE"
      },
      "source": [
        "## Setup\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qax00bg2a4Jj"
      },
      "source": [
        "If you haven't installed tf-agents or gym yet, run:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KKU2iY_7at8Y"
      },
      "outputs": [],
      "source": [
        "!pip install tf-agents[reverb]\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": "1ZAoFNwnRbKK"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import\n",
        "from __future__ import division\n",
        "from __future__ import print_function\n",
        "\n",
        "import abc\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "\n",
        "from tf_agents.environments import py_environment\n",
        "from tf_agents.environments import tf_environment\n",
        "from tf_agents.environments import tf_py_environment\n",
        "from tf_agents.environments import utils\n",
        "from tf_agents.specs import array_spec\n",
        "from tf_agents.environments import wrappers\n",
        "from tf_agents.environments import suite_gym\n",
        "from tf_agents.trajectories import time_step as ts"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x-y4p9i9UURn"
      },
      "source": [
        "## Python Environments"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JPSwHONKMNv9"
      },
      "source": [
        "Python environments have a `step(action) -\u003e next_time_step` method that applies an action to the environment, and returns the following information about the next step:\n",
        "1. `observation`: This is the part of the environment state that the agent can observe to choose its actions at the next step.\n",
        "2. `reward`: The agent is learning to maximize the sum of these rewards across multiple steps.\n",
        "3. `step_type`: Interactions with the environment are usually part of a sequence/episode. e.g. multiple moves in a game of chess. step_type can be either `FIRST`, `MID` or `LAST` to indicate whether this time step is the first, intermediate or last step in a sequence.\n",
        "4. `discount`: This is a float representing how much to weight the reward at the next time step relative to the reward at the current time step.\n",
        "\n",
        "These are grouped into a named tuple `TimeStep(step_type, reward, discount, observation)`.\n",
        "\n",
        "The interface that all Python environments must implement is in `environments/py_environment.PyEnvironment`. The main methods are:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GlD2Dd2vUTtg"
      },
      "outputs": [],
      "source": [
        "class PyEnvironment(object):\n",
        "\n",
        "  def reset(self):\n",
        "    \"\"\"Return initial_time_step.\"\"\"\n",
        "    self._current_time_step = self._reset()\n",
        "    return self._current_time_step\n",
        "\n",
        "  def step(self, action):\n",
        "    \"\"\"Apply action and return new time_step.\"\"\"\n",
        "    if self._current_time_step is None:\n",
        "        return self.reset()\n",
        "    self._current_time_step = self._step(action)\n",
        "    return self._current_time_step\n",
        "\n",
        "  def current_time_step(self):\n",
        "    return self._current_time_step\n",
        "\n",
        "  def time_step_spec(self):\n",
        "    \"\"\"Return time_step_spec.\"\"\"\n",
        "\n",
        "  @abc.abstractmethod\n",
        "  def observation_spec(self):\n",
        "    \"\"\"Return observation_spec.\"\"\"\n",
        "\n",
        "  @abc.abstractmethod\n",
        "  def action_spec(self):\n",
        "    \"\"\"Return action_spec.\"\"\"\n",
        "\n",
        "  @abc.abstractmethod\n",
        "  def _reset(self):\n",
        "    \"\"\"Return initial_time_step.\"\"\"\n",
        "\n",
        "  @abc.abstractmethod\n",
        "  def _step(self, action):\n",
        "    \"\"\"Apply action and return new time_step.\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zfF8koryiGPR"
      },
      "source": [
        "In addition to the `step()` method, environments also provide a `reset()` method that starts a new sequence and provides an initial `TimeStep`. It is not necessary to call the `reset` method explicitly. We assume that environments reset automatically, either when they get to the end of an episode or when step() is called the first time.\n",
        "\n",
        "Note that subclasses do not implement `step()` or `reset()` directly. They instead override the `_step()` and `_reset()` methods. The time steps returned from  these methods will be cached and exposed through `current_time_step()`.\n",
        "\n",
        "The `observation_spec` and the `action_spec` methods return a nest of `(Bounded)ArraySpecs` that describe the name, shape, datatype and ranges of the observations and actions respectively.\n",
        "\n",
        "In TF-Agents we repeatedly refer to nests which are defined as any tree like structure composed of lists, tuples, named-tuples, or dictionaries. These can be arbitrarily composed to maintain structure of observations and actions. We have found this to be very useful for more complex environments where you have many observations and actions."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r63R-RbjcIRw"
      },
      "source": [
        "### Using Standard Environments\n",
        "\n",
        "TF Agents has built-in wrappers for many standard environments like the OpenAI Gym, DeepMind-control and Atari, so that they follow our `py_environment.PyEnvironment` interface. These wrapped evironments can be easily loaded using our environment suites. Let's load the CartPole environment from the OpenAI gym and look at the action and time_step_spec."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1kBPE5T-nb2-"
      },
      "outputs": [],
      "source": [
        "environment = suite_gym.load('CartPole-v0')\n",
        "print('action_spec:', environment.action_spec())\n",
        "print('time_step_spec.observation:', environment.time_step_spec().observation)\n",
        "print('time_step_spec.step_type:', environment.time_step_spec().step_type)\n",
        "print('time_step_spec.discount:', environment.time_step_spec().discount)\n",
        "print('time_step_spec.reward:', environment.time_step_spec().reward)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vWXOC863Apo_"
      },
      "source": [
        "So we see that the environment expects actions of type `int64` in [0, 1] and returns `TimeSteps` where the observations are a `float32` vector of length 4 and discount factor is a `float32` in [0.0, 1.0]. Now, let's try to take a fixed action `(1,)` for a whole episode."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AzIbOJ0-0y12"
      },
      "outputs": [],
      "source": [
        "action = np.array(1, dtype=np.int32)\n",
        "time_step = environment.reset()\n",
        "print(time_step)\n",
        "while not time_step.is_last():\n",
        "  time_step = environment.step(action)\n",
        "  print(time_step)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4xAbBl4_PMtA"
      },
      "source": [
        "### Creating your own Python Environment\n",
        "\n",
        "For many clients, a common use case is to apply one of the standard agents (see agents/) in TF-Agents to their problem. To do this, they have to frame their problem as an environment. So let us look at how to implement an environment in Python.\n",
        "\n",
        "Let's say we want to train an agent to play the following (Black Jack inspired) card game:\n",
        "\n",
        "1. The game is played using an infinite deck of cards numbered 1...10.\n",
        "2. At every turn the agent can do 2 things: get a new random card, or stop the current round.\n",
        "3. The goal is to get the sum of your cards as close to 21 as possible at the end of the round, without going over.\n",
        "\n",
        "An environment that represents the game could look like this:\n",
        "\n",
        "1. Actions: We have 2 actions. Action 0: get a new card, and Action 1: terminate the current round.\n",
        "2. Observations: Sum of the cards in the current round.\n",
        "3. Reward: The objective is to get as close to 21 as possible without going over, so we can achieve this using the following reward at the end of the round:\n",
        "   sum_of_cards - 21 if sum_of_cards \u003c= 21, else -21\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9HD0cDykPL6I"
      },
      "outputs": [],
      "source": [
        "class CardGameEnv(py_environment.PyEnvironment):\n",
        "\n",
        "  def __init__(self):\n",
        "    self._action_spec = array_spec.BoundedArraySpec(\n",
        "        shape=(), dtype=np.int32, minimum=0, maximum=1, name='action')\n",
        "    self._observation_spec = array_spec.BoundedArraySpec(\n",
        "        shape=(1,), dtype=np.int32, minimum=0, name='observation')\n",
        "    self._state = 0\n",
        "    self._episode_ended = False\n",
        "\n",
        "  def action_spec(self):\n",
        "    return self._action_spec\n",
        "\n",
        "  def observation_spec(self):\n",
        "    return self._observation_spec\n",
        "\n",
        "  def _reset(self):\n",
        "    self._state = 0\n",
        "    self._episode_ended = False\n",
        "    return ts.restart(np.array([self._state], dtype=np.int32))\n",
        "\n",
        "  def _step(self, action):\n",
        "\n",
        "    if self._episode_ended:\n",
        "      # The last action ended the episode. Ignore the current action and start\n",
        "      # a new episode.\n",
        "      return self.reset()\n",
        "\n",
        "    # Make sure episodes don't go on forever.\n",
        "    if action == 1:\n",
        "      self._episode_ended = True\n",
        "    elif action == 0:\n",
        "      new_card = np.random.randint(1, 11)\n",
        "      self._state += new_card\n",
        "    else:\n",
        "      raise ValueError('`action` should be 0 or 1.')\n",
        "\n",
        "    if self._episode_ended or self._state \u003e= 21:\n",
        "      reward = self._state - 21 if self._state \u003c= 21 else -21\n",
        "      return ts.termination(np.array([self._state], dtype=np.int32), reward)\n",
        "    else:\n",
        "      return ts.transition(\n",
        "          np.array([self._state], dtype=np.int32), reward=0.0, discount=1.0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LYEwyX7QsqeX"
      },
      "source": [
        "Let's make sure we did everything correctly defining the above environment. When creating your own environment you must make sure the observations and time_steps generated follow the correct shapes and types as defined in your specs. These are used to generate the TensorFlow graph and as such can create hard to debug problems if we get them wrong.\n",
        "\n",
        "To validate our environment we will use a random policy to generate actions and we will iterate over 5 episodes to make sure things are working as intended. An error is raised if we receive a time_step that does not follow the environment specs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Hhm-5R7spVx"
      },
      "outputs": [],
      "source": [
        "environment = CardGameEnv()\n",
        "utils.validate_py_environment(environment, episodes=5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_36eM7MvkNOg"
      },
      "source": [
        "Now that we know the environment is working as intended, let's run this environment using a fixed policy: ask for 3 cards and then end the round."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FILylafAkMEx"
      },
      "outputs": [],
      "source": [
        "get_new_card_action = np.array(0, dtype=np.int32)\n",
        "end_round_action = np.array(1, dtype=np.int32)\n",
        "\n",
        "environment = CardGameEnv()\n",
        "time_step = environment.reset()\n",
        "print(time_step)\n",
        "cumulative_reward = time_step.reward\n",
        "\n",
        "for _ in range(3):\n",
        "  time_step = environment.step(get_new_card_action)\n",
        "  print(time_step)\n",
        "  cumulative_reward += time_step.reward\n",
        "\n",
        "time_step = environment.step(end_round_action)\n",
        "print(time_step)\n",
        "cumulative_reward += time_step.reward\n",
        "print('Final Reward = ', cumulative_reward)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_vBLPN3ioyGx"
      },
      "source": [
        "### Environment Wrappers\n",
        "\n",
        "An environment wrapper takes a Python environment and returns a modified version of the environment. Both the original environment and the modified environment are instances of `py_environment.PyEnvironment`, and multiple wrappers can be chained together.\n",
        "\n",
        "Some common wrappers can be found in `environments/wrappers.py`. For example:\n",
        "\n",
        "1. `ActionDiscretizeWrapper`: Converts a continuous action space to a discrete action space.\n",
        "2. `RunStats`: Captures run statistics of the environment such as number of steps taken, number of episodes completed etc.\n",
        "3. `TimeLimit`: Terminates the episode after a fixed number of steps.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_8aIybRdnFfb"
      },
      "source": [
        "#### Example 1: Action Discretize Wrapper"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YIaxJRUpvfyc"
      },
      "source": [
        "InvertedPendulum is a PyBullet environment that accepts continuous actions in the range `[-2, 2]`. If we want to train a discrete action agent such as DQN on this environment, we have to discretize (quantize) the action space. This is exactly what the `ActionDiscretizeWrapper` does. Compare the `action_spec` before and after wrapping:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AJxEoZ4HoyjR"
      },
      "outputs": [],
      "source": [
        "env = suite_gym.load('Pendulum-v1')\n",
        "print('Action Spec:', env.action_spec())\n",
        "\n",
        "discrete_action_env = wrappers.ActionDiscretizeWrapper(env, num_actions=5)\n",
        "print('Discretized Action Spec:', discrete_action_env.action_spec())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "njFjW8bmwTWJ"
      },
      "source": [
        "The wrapped `discrete_action_env` is an instance of `py_environment.PyEnvironment` and can be treated like a regular Python environment.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8l5dwAhsP_F_"
      },
      "source": [
        "## TensorFlow Environments"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iZG39AjBkTjr"
      },
      "source": [
        "The interface for TF environments is defined in `environments/tf_environment.TFEnvironment` and looks very similar to the Python environments. TF Environments differ from Python envs in a couple of ways:\n",
        "\n",
        "* They generate tensor objects instead of arrays\n",
        "* TF environments add a batch dimension to the tensors generated when compared to the specs.\n",
        "\n",
        "Converting the Python environments into TFEnvs allows tensorflow to parallelize operations. For example, one could define a `collect_experience_op` that collects data from the environment and adds to a `replay_buffer`, and a `train_op` that reads from the `replay_buffer` and trains the agent, and run them in parallel naturally in TensorFlow."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WKBDDZqKTxsL"
      },
      "outputs": [],
      "source": [
        "class TFEnvironment(object):\n",
        "\n",
        "  def time_step_spec(self):\n",
        "    \"\"\"Describes the `TimeStep` tensors returned by `step()`.\"\"\"\n",
        "\n",
        "  def observation_spec(self):\n",
        "    \"\"\"Defines the `TensorSpec` of observations provided by the environment.\"\"\"\n",
        "\n",
        "  def action_spec(self):\n",
        "    \"\"\"Describes the TensorSpecs of the action expected by `step(action)`.\"\"\"\n",
        "\n",
        "  def reset(self):\n",
        "    \"\"\"Returns the current `TimeStep` after resetting the Environment.\"\"\"\n",
        "    return self._reset()\n",
        "\n",
        "  def current_time_step(self):\n",
        "    \"\"\"Returns the current `TimeStep`.\"\"\"\n",
        "    return self._current_time_step()\n",
        "\n",
        "  def step(self, action):\n",
        "    \"\"\"Applies the action and returns the new `TimeStep`.\"\"\"\n",
        "    return self._step(action)\n",
        "\n",
        "  @abc.abstractmethod\n",
        "  def _reset(self):\n",
        "    \"\"\"Returns the current `TimeStep` after resetting the Environment.\"\"\"\n",
        "\n",
        "  @abc.abstractmethod\n",
        "  def _current_time_step(self):\n",
        "    \"\"\"Returns the current `TimeStep`.\"\"\"\n",
        "\n",
        "  @abc.abstractmethod\n",
        "  def _step(self, action):\n",
        "    \"\"\"Applies the action and returns the new `TimeStep`.\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tFkBIA92ThWf"
      },
      "source": [
        "The `current_time_step()` method returns the current time_step and initializes the environment if needed.\n",
        "\n",
        "The `reset()` method forces a reset in the environment and returns the current_step.\n",
        "\n",
        "If the `action` doesn't depend on the previous `time_step` a `tf.control_dependency` is needed in `Graph` mode.\n",
        "\n",
        "For now, let us look at how `TFEnvironments` are created."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S6wS3AaLdVLT"
      },
      "source": [
        "### Creating your own TensorFlow Environment\n",
        "\n",
        "This is more complicated than creating environments in Python, so we will not cover it in this colab. An example is available [here](https://github.com/tensorflow/agents/blob/master/tf_agents/environments/tf_environment_test.py). The more common use case is to implement your environment in Python and wrap it in TensorFlow using our `TFPyEnvironment` wrapper (see below)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V_Ny2lb-dU5R"
      },
      "source": [
        "### Wrapping a Python Environment in TensorFlow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lv4-UcurZ8nb"
      },
      "source": [
        "We can easily wrap any Python environment into a TensorFlow environment using the `TFPyEnvironment` wrapper."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UYerqyNfnVRL"
      },
      "outputs": [],
      "source": [
        "env = suite_gym.load('CartPole-v0')\n",
        "tf_env = tf_py_environment.TFPyEnvironment(env)\n",
        "\n",
        "print(isinstance(tf_env, tf_environment.TFEnvironment))\n",
        "print(\"TimeStep Specs:\", tf_env.time_step_spec())\n",
        "print(\"Action Specs:\", tf_env.action_spec())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z3WFrnX9CNpC"
      },
      "source": [
        "Note the specs are now of type: `(Bounded)TensorSpec`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vQPvC1ARYALj"
      },
      "source": [
        "### Usage Examples"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ov7EIrk8dKUU"
      },
      "source": [
        "#### Simple Example"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gdvFqUqbdB7u"
      },
      "outputs": [],
      "source": [
        "env = suite_gym.load('CartPole-v0')\n",
        "\n",
        "tf_env = tf_py_environment.TFPyEnvironment(env)\n",
        "# reset() creates the initial time_step after resetting the environment.\n",
        "time_step = tf_env.reset()\n",
        "num_steps = 3\n",
        "transitions = []\n",
        "reward = 0\n",
        "for i in range(num_steps):\n",
        "  action = tf.constant([i % 2])\n",
        "  # applies the action and returns the new TimeStep.\n",
        "  next_time_step = tf_env.step(action)\n",
        "  transitions.append([time_step, action, next_time_step])\n",
        "  reward += next_time_step.reward\n",
        "  time_step = next_time_step\n",
        "\n",
        "np_transitions = tf.nest.map_structure(lambda x: x.numpy(), transitions)\n",
        "print('\\n'.join(map(str, np_transitions)))\n",
        "print('Total reward:', reward.numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kWs48LNsdLnc"
      },
      "source": [
        "#### Whole Episodes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t561kUXMk-KM"
      },
      "outputs": [],
      "source": [
        "env = suite_gym.load('CartPole-v0')\n",
        "tf_env = tf_py_environment.TFPyEnvironment(env)\n",
        "\n",
        "time_step = tf_env.reset()\n",
        "rewards = []\n",
        "steps = []\n",
        "num_episodes = 5\n",
        "\n",
        "for _ in range(num_episodes):\n",
        "  episode_reward = 0\n",
        "  episode_steps = 0\n",
        "  while not time_step.is_last():\n",
        "    action = tf.random.uniform([1], 0, 2, dtype=tf.int32)\n",
        "    time_step = tf_env.step(action)\n",
        "    episode_steps += 1\n",
        "    episode_reward += time_step.reward.numpy()\n",
        "  rewards.append(episode_reward)\n",
        "  steps.append(episode_steps)\n",
        "  time_step = tf_env.reset()\n",
        "\n",
        "num_steps = np.sum(steps)\n",
        "avg_length = np.mean(steps)\n",
        "avg_reward = np.mean(rewards)\n",
        "\n",
        "print('num_episodes:', num_episodes, 'num_steps:', num_steps)\n",
        "print('avg_length', avg_length, 'avg_reward:', avg_reward)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "TF-Agents Environments Tutorial.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}