docs/tutorials/7_SAC_minitaur_tutorial.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "klGNgWREsvQv"
      },
      "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": "lsaQlK8fFQqH"
      },
      "source": [
        "# SAC minitaur with the Actor-Learner API\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/agents/tutorials/7_SAC_minitaur_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/7_SAC_minitaur_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/7_SAC_minitaur_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/7_SAC_minitaur_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\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZOUOQOrFs3zn"
      },
      "source": [
        "## Introduction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cKOCZlhUgXVK"
      },
      "source": [
        "This example shows how to train a [Soft Actor Critic](https://arxiv.org/abs/1812.05905) agent on the [Minitaur](https://github.com/bulletphysics/bullet3/blob/master/examples/pybullet/gym/pybullet_envs/bullet/minitaur.py) environment.\n",
        "\n",
        "If you've worked through the [DQN Colab](https://github.com/tensorflow/agents/blob/master/docs/tutorials/1_dqn_tutorial.ipynb) this should feel very familiar. Notable changes include:\n",
        "\n",
        "  * Changing the agent from DQN to SAC.\n",
        "  * Training on Minitaur which is a much more complex environment than CartPole. The Minitaur environment aims to train a quadruped robot to move forward.\n",
        "  * Using the TF-Agents Actor-Learner API for distributed Reinforcement Learning.\n",
        "\n",
        "The API supports both distributed data collection using an experience replay buffer and variable container (parameter server) and distributed training across multiple devices. The API is designed to be very simple and modular. We utilize [Reverb](https://deepmind.com/research/open-source/Reverb) for both replay buffer and variable container and [TF DistributionStrategy API](https://www.tensorflow.org/guide/distributed_training) for distributed training on GPUs and TPUs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9vUQms4DAY5A"
      },
      "source": [
        "If you haven't installed the following dependencies, run:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fskoLlB-AZ9j"
      },
      "outputs": [],
      "source": [
        "!sudo apt-get update\n",
        "!sudo apt-get install -y xvfb ffmpeg\n",
        "!pip install 'imageio==2.4.0'\n",
        "!pip install matplotlib\n",
        "!pip install tf-agents[reverb]\n",
        "!pip install pybullet\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": "markdown",
      "metadata": {
        "id": "1u9QVVsShC9X"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nNV5wyH3dyMl"
      },
      "source": [
        "First we will import the different tools that we need."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sMitx5qSgJk1"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "import imageio\n",
        "import IPython\n",
        "import matplotlib.pyplot as plt\n",
        "import os\n",
        "import reverb\n",
        "import tempfile\n",
        "import PIL.Image\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from tf_agents.agents.ddpg import critic_network\n",
        "from tf_agents.agents.sac import sac_agent\n",
        "from tf_agents.agents.sac import tanh_normal_projection_network\n",
        "from tf_agents.environments import suite_pybullet\n",
        "from tf_agents.metrics import py_metrics\n",
        "from tf_agents.networks import actor_distribution_network\n",
        "from tf_agents.policies import greedy_policy\n",
        "from tf_agents.policies import py_tf_eager_policy\n",
        "from tf_agents.policies import random_py_policy\n",
        "from tf_agents.replay_buffers import reverb_replay_buffer\n",
        "from tf_agents.replay_buffers import reverb_utils\n",
        "from tf_agents.train import actor\n",
        "from tf_agents.train import learner\n",
        "from tf_agents.train import triggers\n",
        "from tf_agents.train.utils import spec_utils\n",
        "from tf_agents.train.utils import strategy_utils\n",
        "from tf_agents.train.utils import train_utils\n",
        "\n",
        "tempdir = tempfile.gettempdir()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LmC0NDhdLIKY"
      },
      "source": [
        "## Hyperparameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HC1kNrOsLSIZ"
      },
      "outputs": [],
      "source": [
        "env_name = \"MinitaurBulletEnv-v0\" # @param {type:\"string\"}\n",
        "\n",
        "# Use \"num_iterations = 1e6\" for better results (2 hrs)\n",
        "# 1e5 is just so this doesn't take too long (1 hr)\n",
        "num_iterations = 100000 # @param {type:\"integer\"}\n",
        "\n",
        "initial_collect_steps = 10000 # @param {type:\"integer\"}\n",
        "collect_steps_per_iteration = 1 # @param {type:\"integer\"}\n",
        "replay_buffer_capacity = 10000 # @param {type:\"integer\"}\n",
        "\n",
        "batch_size = 256 # @param {type:\"integer\"}\n",
        "\n",
        "critic_learning_rate = 3e-4 # @param {type:\"number\"}\n",
        "actor_learning_rate = 3e-4 # @param {type:\"number\"}\n",
        "alpha_learning_rate = 3e-4 # @param {type:\"number\"}\n",
        "target_update_tau = 0.005 # @param {type:\"number\"}\n",
        "target_update_period = 1 # @param {type:\"number\"}\n",
        "gamma = 0.99 # @param {type:\"number\"}\n",
        "reward_scale_factor = 1.0 # @param {type:\"number\"}\n",
        "\n",
        "actor_fc_layer_params = (256, 256)\n",
        "critic_joint_fc_layer_params = (256, 256)\n",
        "\n",
        "log_interval = 5000 # @param {type:\"integer\"}\n",
        "\n",
        "num_eval_episodes = 20 # @param {type:\"integer\"}\n",
        "eval_interval = 10000 # @param {type:\"integer\"}\n",
        "\n",
        "policy_save_interval = 5000 # @param {type:\"integer\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "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's load the Minitaur environment from the Pybullet suite."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RlO7WIQHu_7D"
      },
      "outputs": [],
      "source": [
        "env = suite_pybullet.load(env_name)\n",
        "env.reset()\n",
        "PIL.Image.fromarray(env.render())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gY179d1xlmoM"
      },
      "source": [
        "In this environment the goal is for the agent to train a policy that will control the Minitaur robot and have it move forward as fast as possible. Episodes last 1000 steps and the return will be the sum of rewards throughout the episode.\n",
        "\n",
        "Let's look at the information the environment provides as an `observation` which the policy will use to generate `actions`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "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": {
        "id": "Wg5ysVTnctIm"
      },
      "source": [
        "The observation is fairly complex. We receive 28 values representing the angles, velocities, and torques for all the motors. In return the environment expects 8 values for the actions between `[-1, 1]`. These are the desired motor angles.\n",
        "\n",
        "Usually we create two environments: one for collecting data during training and one for evaluation. The environments are written in pure python and use numpy arrays, which the Actor Learner API directly consumes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xp-Y4mD6eDhF"
      },
      "outputs": [],
      "source": [
        "collect_env = suite_pybullet.load(env_name)\n",
        "eval_env = suite_pybullet.load(env_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Da-z2yF66FR9"
      },
      "source": [
        "## Distribution Strategy\n",
        "We use the DistributionStrategy API to enable running the train step computation across multiple devices such as multiple GPUs or TPUs using data parallelism. The train step:\n",
        "* Receives a batch of training data\n",
        "* Splits it across the devices\n",
        "* Computes the forward step\n",
        "* Aggregates and computes the MEAN of the loss\n",
        "* Computes the backward step and performs a gradient variable update\n",
        "\n",
        "With TF-Agents Learner API and DistributionStrategy API it is quite easy to switch between running the train step on GPUs (using MirroredStrategy) to TPUs (using TPUStrategy) without changing any of the training logic below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wGREYZCaDB1h"
      },
      "source": [
        "### Enabling the GPU\n",
        "If you want to try running on a GPU, you'll first need to enable GPUs for the notebook:\n",
        "\n",
        "* Navigate to Edit→Notebook Settings\n",
        "* Select GPU from the Hardware Accelerator drop-down"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5ZuvwDV66Mn1"
      },
      "source": [
        "### Picking a strategy\n",
        "Use `strategy_utils` to generate a strategy. Under the hood, passing the parameter:\n",
        "* `use_gpu = False` returns `tf.distribute.get_strategy()`, which uses CPU\n",
        "* `use_gpu = True` returns `tf.distribute.MirroredStrategy()`, which uses all GPUs that are visible to TensorFlow on one machine"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ff5ZZRZI15ds"
      },
      "outputs": [],
      "source": [
        "use_gpu = True #@param {type:\"boolean\"}\n",
        "\n",
        "strategy = strategy_utils.get_strategy(tpu=False, use_gpu=use_gpu)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fMn5FTs5kHvt"
      },
      "source": [
        "All variables and Agents need to be created under `strategy.scope()`, as you'll see below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E9lW_OZYFR8A"
      },
      "source": [
        "## Agent\n",
        "\n",
        "To create an SAC Agent, we first need to create the networks that it will train. SAC is an actor-critic agent, so we will need two networks.\n",
        "\n",
        "The critic will give us value estimates for `Q(s,a)`. That is, it will recieve as input an observation and an action, and it will give us an estimate of how good that action was for the given state.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TgkdEPg_muzV"
      },
      "outputs": [],
      "source": [
        "observation_spec, action_spec, time_step_spec = (\n",
        "      spec_utils.get_tensor_specs(collect_env))\n",
        "\n",
        "with strategy.scope():\n",
        "  critic_net = critic_network.CriticNetwork(\n",
        "        (observation_spec, action_spec),\n",
        "        observation_fc_layer_params=None,\n",
        "        action_fc_layer_params=None,\n",
        "        joint_fc_layer_params=critic_joint_fc_layer_params,\n",
        "        kernel_initializer='glorot_uniform',\n",
        "        last_kernel_initializer='glorot_uniform')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pYy4AH4V7Ph4"
      },
      "source": [
        "We will use this critic to train an `actor` network which will allow us to generate actions given an observation.\n",
        "\n",
        "The `ActorNetwork` will predict parameters for a tanh-squashed [MultivariateNormalDiag](https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/MultivariateNormalDiag) distribution. This distribution will then be sampled, conditioned on the current observation, whenever we need to generate actions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TB5Y3Oub4u7f"
      },
      "outputs": [],
      "source": [
        "with strategy.scope():\n",
        "  actor_net = actor_distribution_network.ActorDistributionNetwork(\n",
        "      observation_spec,\n",
        "      action_spec,\n",
        "      fc_layer_params=actor_fc_layer_params,\n",
        "      continuous_projection_net=(\n",
        "          tanh_normal_projection_network.TanhNormalProjectionNetwork))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z62u55hSmviJ"
      },
      "source": [
        "With these networks at hand we can now instantiate the agent.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jbY4yrjTEyc9"
      },
      "outputs": [],
      "source": [
        "with strategy.scope():\n",
        "  train_step = train_utils.create_train_step()\n",
        "\n",
        "  tf_agent = sac_agent.SacAgent(\n",
        "        time_step_spec,\n",
        "        action_spec,\n",
        "        actor_network=actor_net,\n",
        "        critic_network=critic_net,\n",
        "        actor_optimizer=tf.keras.optimizers.Adam(\n",
        "            learning_rate=actor_learning_rate),\n",
        "        critic_optimizer=tf.keras.optimizers.Adam(\n",
        "            learning_rate=critic_learning_rate),\n",
        "        alpha_optimizer=tf.keras.optimizers.Adam(\n",
        "            learning_rate=alpha_learning_rate),\n",
        "        target_update_tau=target_update_tau,\n",
        "        target_update_period=target_update_period,\n",
        "        td_errors_loss_fn=tf.math.squared_difference,\n",
        "        gamma=gamma,\n",
        "        reward_scale_factor=reward_scale_factor,\n",
        "        train_step_counter=train_step)\n",
        "\n",
        "  tf_agent.initialize()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NLva6g2jdWgr"
      },
      "source": [
        "## Replay Buffer\n",
        "\n",
        "In order to keep track of the data collected from the environment, we will use [Reverb](https://deepmind.com/research/open-source/Reverb), an efficient, extensible, and easy-to-use replay system by Deepmind. It stores experience data collected by the Actors and consumed by the Learner during training.\n",
        "\n",
        "In this tutorial, this is less important than `max_size` -- but in a distributed setting with async collection and training, you will probably want to experiment with `rate_limiters.SampleToInsertRatio`, using a samples_per_insert somewhere between 2 and 1000. For example:\n",
        "```\n",
        "rate_limiter=reverb.rate_limiters.SampleToInsertRatio(samples_per_insert=3.0, min_size_to_sample=3, error_buffer=3.0)\n",
        "```\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vX2zGUWJGWAl"
      },
      "outputs": [],
      "source": [
        "table_name = 'uniform_table'\n",
        "table = reverb.Table(\n",
        "    table_name,\n",
        "    max_size=replay_buffer_capacity,\n",
        "    sampler=reverb.selectors.Uniform(),\n",
        "    remover=reverb.selectors.Fifo(),\n",
        "    rate_limiter=reverb.rate_limiters.MinSize(1))\n",
        "\n",
        "reverb_server = reverb.Server([table])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LRNvAnkO7JK2"
      },
      "source": [
        "The 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`.\n",
        "\n",
        "Since the SAC Agent needs both the current and next observation to compute the loss, we set `sequence_length=2`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xVLUxyUo7HQR"
      },
      "outputs": [],
      "source": [
        "reverb_replay = reverb_replay_buffer.ReverbReplayBuffer(\n",
        "    tf_agent.collect_data_spec,\n",
        "    sequence_length=2,\n",
        "    table_name=table_name,\n",
        "    local_server=reverb_server)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rVD5nQ9ZGo8_"
      },
      "source": [
        "Now we generate a TensorFlow dataset from the Reverb replay buffer. We will pass this to the Learner to sample experiences for training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ba7bilizt_qW"
      },
      "outputs": [],
      "source": [
        "dataset = reverb_replay.as_dataset(\n",
        "      sample_batch_size=batch_size, num_steps=2).prefetch(50)\n",
        "experience_dataset_fn = lambda: dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "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:\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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yq7JE8IwFe0E"
      },
      "outputs": [],
      "source": [
        "tf_eval_policy = tf_agent.policy\n",
        "eval_policy = py_tf_eager_policy.PyTFEagerPolicy(\n",
        "  tf_eval_policy, use_tf_function=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f_A4rZveEQzW"
      },
      "outputs": [],
      "source": [
        "tf_collect_policy = tf_agent.collect_policy\n",
        "collect_policy = py_tf_eager_policy.PyTFEagerPolicy(\n",
        "  tf_collect_policy, use_tf_function=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "azkJZ8oaF8uc"
      },
      "source": [
        "Policies can be created independently of agents. For example, use `tf_agents.policies.random_py_policy` to create a policy which will randomly select an action for each time_step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BwY7StuMkuV4"
      },
      "outputs": [],
      "source": [
        "random_policy = random_py_policy.RandomPyPolicy(\n",
        "  collect_env.time_step_spec(), collect_env.action_spec())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l1LMqw60Kuso"
      },
      "source": [
        "## Actors\n",
        "The actor manages interactions between a policy and an environment.\n",
        "  * The Actor components contain an instance of the environment (as `py_environment`) and a copy of the policy variables.\n",
        "  * Each Actor worker runs a sequence of data collection steps given the local values of the policy variables.\n",
        "  * Variable updates are done explicitly using the variable container client instance in the training script before calling `actor.run()`.\n",
        "  * The observed experience is written into the replay buffer in each data collection step."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XjE59ct9fU7W"
      },
      "source": [
        "As the Actors run data collection steps, they pass trajectories of (state, action, reward) to the observer, which caches and writes them to the Reverb replay system. \n",
        "\n",
        "We're storing trajectories for frames [(t0,t1) (t1,t2) (t2,t3), ...] because `stride_length=1`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HbyGmdiNfNDc"
      },
      "outputs": [],
      "source": [
        "rb_observer = reverb_utils.ReverbAddTrajectoryObserver(\n",
        "  reverb_replay.py_client,\n",
        "  table_name,\n",
        "  sequence_length=2,\n",
        "  stride_length=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6yaVVC22fOcF"
      },
      "source": [
        "We create an Actor with the random policy and collect experiences to seed the replay buffer with."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZGq3SY0kKwsa"
      },
      "outputs": [],
      "source": [
        "initial_collect_actor = actor.Actor(\n",
        "  collect_env,\n",
        "  random_policy,\n",
        "  train_step,\n",
        "  steps_per_run=initial_collect_steps,\n",
        "  observers=[rb_observer])\n",
        "initial_collect_actor.run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Pkg-0vZP_Ya"
      },
      "source": [
        "Instantiate an Actor with the collect policy to gather more experiences during training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A6ooXyk0FZ5j"
      },
      "outputs": [],
      "source": [
        "env_step_metric = py_metrics.EnvironmentSteps()\n",
        "collect_actor = actor.Actor(\n",
        "  collect_env,\n",
        "  collect_policy,\n",
        "  train_step,\n",
        "  steps_per_run=1,\n",
        "  metrics=actor.collect_metrics(10),\n",
        "  summary_dir=os.path.join(tempdir, learner.TRAIN_DIR),\n",
        "  observers=[rb_observer, env_step_metric])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FR9CZ-jfPN2T"
      },
      "source": [
        "Create an Actor which will be used to evaluate the policy during training. We pass in `actor.eval_metrics(num_eval_episodes)` to log metrics later."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vHY2BT5lFhgL"
      },
      "outputs": [],
      "source": [
        "eval_actor = actor.Actor(\n",
        "  eval_env,\n",
        "  eval_policy,\n",
        "  train_step,\n",
        "  episodes_per_run=num_eval_episodes,\n",
        "  metrics=actor.eval_metrics(num_eval_episodes),\n",
        "  summary_dir=os.path.join(tempdir, 'eval'),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y6eBGSYiOf83"
      },
      "source": [
        "## Learners\n",
        "The Learner component contains the agent and performs gradient step updates to the policy variables using experience data from the replay buffer. After one or more training steps, the Learner can push a new set of variable values to the variable container."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gi37YicSFTfF"
      },
      "outputs": [],
      "source": [
        "saved_model_dir = os.path.join(tempdir, learner.POLICY_SAVED_MODEL_DIR)\n",
        "\n",
        "# Triggers to save the agent's policy checkpoints.\n",
        "learning_triggers = [\n",
        "    triggers.PolicySavedModelTrigger(\n",
        "        saved_model_dir,\n",
        "        tf_agent,\n",
        "        train_step,\n",
        "        interval=policy_save_interval),\n",
        "    triggers.StepPerSecondLogTrigger(train_step, interval=1000),\n",
        "]\n",
        "\n",
        "agent_learner = learner.Learner(\n",
        "  tempdir,\n",
        "  train_step,\n",
        "  tf_agent,\n",
        "  experience_dataset_fn,\n",
        "  triggers=learning_triggers,\n",
        "  strategy=strategy)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "94rCXQtbUbXv"
      },
      "source": [
        "## Metrics and Evaluation\n",
        "\n",
        "We instantiated the eval Actor with `actor.eval_metrics` above, which creates most commonly used metrics during policy evaluation:\n",
        "* 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.\n",
        "* Average episode length.\n",
        "\n",
        "We run the Actor to generate these metrics."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "83iMSHUC71RG"
      },
      "outputs": [],
      "source": [
        "def get_eval_metrics():\n",
        "  eval_actor.run()\n",
        "  results = {}\n",
        "  for metric in eval_actor.metrics:\n",
        "    results[metric.name] = metric.result()\n",
        "  return results\n",
        "\n",
        "metrics = get_eval_metrics()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jnOMvX_eZvOW"
      },
      "outputs": [],
      "source": [
        "def log_eval_metrics(step, metrics):\n",
        "  eval_results = (', ').join(\n",
        "      '{} = {:.6f}'.format(name, result) for name, result in metrics.items())\n",
        "  print('step = {0}: {1}'.format(step, eval_results))\n",
        "\n",
        "log_eval_metrics(0, metrics)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hWWURm_rXG-f"
      },
      "source": [
        "Check out the [metrics module](https://github.com/tensorflow/agents/blob/master/tf_agents/metrics/tf_metrics.py) for other standard implementations of different metrics."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0pTbJ3PeyF-u"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "try:\n",
        "  %%time\n",
        "except:\n",
        "  pass\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 = get_eval_metrics()[\"AverageReturn\"]\n",
        "returns = [avg_return]\n",
        "\n",
        "for _ in range(num_iterations):\n",
        "  # Training.\n",
        "  collect_actor.run()\n",
        "  loss_info = agent_learner.run(iterations=1)\n",
        "\n",
        "  # Evaluating.\n",
        "  step = agent_learner.train_step_numpy\n",
        "\n",
        "  if eval_interval and step % eval_interval == 0:\n",
        "    metrics = get_eval_metrics()\n",
        "    log_eval_metrics(step, metrics)\n",
        "    returns.append(metrics[\"AverageReturn\"])\n",
        "\n",
        "  if log_interval and step % log_interval == 0:\n",
        "    print('step = {0}: loss = {1}'.format(step, loss_info.loss.numpy()))\n",
        "\n",
        "rb_observer.close()\n",
        "reverb_server.stop()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "68jNcA_TiJDq"
      },
      "source": [
        "## Visualization\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aO-LWCdbbOIC"
      },
      "source": [
        "### Plots\n",
        "\n",
        "We can plot average return vs global steps to see the performance of our agent. In `Minitaur`, the reward function is based on how far the minitaur walks in 1000 steps and penalizes the energy expenditure."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rXKzyGt72HS8"
      },
      "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()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M7-XpPP99Cy7"
      },
      "source": [
        "### Videos"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "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": null,
      "metadata": {
        "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": {
        "id": "9c_PH-pX4Pr5"
      },
      "source": [
        "The following code visualizes the agent's policy for a few episodes:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PSgaQN1nXT-h"
      },
      "outputs": [],
      "source": [
        "num_episodes = 3\n",
        "video_filename = 'sac_minitaur.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_env.render())\n",
        "    while not time_step.is_last():\n",
        "      action_step = eval_actor.policy.action(time_step)\n",
        "      time_step = eval_env.step(action_step.action)\n",
        "      video.append_data(eval_env.render())\n",
        "\n",
        "embed_mp4(video_filename)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "7_SAC_minitaur_tutorial.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}