PACE_IMPALA (1).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "PACE IMPALA",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ie11y0Fjlgtv"
      },
      "source": [
        "# Installations"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SdvYZ9QBU__4",
        "outputId": "907464e7-7b2e-4ce1-94c9-37343d282f61"
      },
      "source": [
        "#install necessary libraries\n",
        "!pip install wandb\n",
        "!pip install dm_env\n",
        "!pip install git+https://github.com/deepmind/dm-haiku\n",
        "!pip install git+git://github.com/deepmind/bsuite.git\n",
        "!pip install git+git://github.com/deepmind/optax.git\n",
        "!pip install git+git://github.com/deepmind/rlax.git\n",
        "!pip install dm-tree\n",
        "!pip install packaging\n",
        "!pip install tensorflow-datasets\n",
        "!pip install tensorflow\n",
        "!pip install pyvirtualdisplay\n",
        "!apt-get install xvfb\n",
        "!pip install git+https://github.com/tensorflow/docs\n",
        "!pip install tqdm \n",
        "!pip install chex"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting wandb\n",
            "  Downloading wandb-0.12.0-py2.py3-none-any.whl (1.6 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.6 MB 14.8 MB/s \n",
            "\u001b[?25hCollecting GitPython>=1.0.0\n",
            "  Downloading GitPython-3.1.18-py3-none-any.whl (170 kB)\n",
            "\u001b[K     |████████████████████████████████| 170 kB 67.0 MB/s \n",
            "\u001b[?25hRequirement already satisfied: six>=1.13.0 in /usr/local/lib/python3.7/dist-packages (from wandb) (1.15.0)\n",
            "Requirement already satisfied: protobuf>=3.12.0 in /usr/local/lib/python3.7/dist-packages (from wandb) (3.17.3)\n",
            "Collecting sentry-sdk>=1.0.0\n",
            "  Downloading sentry_sdk-1.3.1-py2.py3-none-any.whl (133 kB)\n",
            "\u001b[K     |████████████████████████████████| 133 kB 68.7 MB/s \n",
            "\u001b[?25hRequirement already satisfied: promise<3,>=2.0 in /usr/local/lib/python3.7/dist-packages (from wandb) (2.3)\n",
            "Collecting docker-pycreds>=0.4.0\n",
            "  Downloading docker_pycreds-0.4.0-py2.py3-none-any.whl (9.0 kB)\n",
            "Collecting subprocess32>=3.5.3\n",
            "  Downloading subprocess32-3.5.4.tar.gz (97 kB)\n",
            "\u001b[K     |████████████████████████████████| 97 kB 7.5 MB/s \n",
            "\u001b[?25hRequirement already satisfied: python-dateutil>=2.6.1 in /usr/local/lib/python3.7/dist-packages (from wandb) (2.8.1)\n",
            "Requirement already satisfied: PyYAML in /usr/local/lib/python3.7/dist-packages (from wandb) (3.13)\n",
            "Requirement already satisfied: Click!=8.0.0,>=7.0 in /usr/local/lib/python3.7/dist-packages (from wandb) (7.1.2)\n",
            "Requirement already satisfied: requests<3,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from wandb) (2.23.0)\n",
            "Collecting configparser>=3.8.1\n",
            "  Downloading configparser-5.0.2-py3-none-any.whl (19 kB)\n",
            "Collecting shortuuid>=0.5.0\n",
            "  Downloading shortuuid-1.0.1-py3-none-any.whl (7.5 kB)\n",
            "Collecting pathtools\n",
            "  Downloading pathtools-0.1.2.tar.gz (11 kB)\n",
            "Requirement already satisfied: psutil>=5.0.0 in /usr/local/lib/python3.7/dist-packages (from wandb) (5.4.8)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.0 in /usr/local/lib/python3.7/dist-packages (from GitPython>=1.0.0->wandb) (3.7.4.3)\n",
            "Collecting gitdb<5,>=4.0.1\n",
            "  Downloading gitdb-4.0.7-py3-none-any.whl (63 kB)\n",
            "\u001b[K     |████████████████████████████████| 63 kB 2.2 MB/s \n",
            "\u001b[?25hCollecting smmap<5,>=3.0.1\n",
            "  Downloading smmap-4.0.0-py2.py3-none-any.whl (24 kB)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.0.0->wandb) (2021.5.30)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.0.0->wandb) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.0.0->wandb) (2.10)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.0.0->wandb) (3.0.4)\n",
            "Building wheels for collected packages: subprocess32, pathtools\n",
            "  Building wheel for subprocess32 (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for subprocess32: filename=subprocess32-3.5.4-py3-none-any.whl size=6502 sha256=d7e3f01eac5f75eb1b4c8a3521afb1d9131917e636b6ac2794d7cfad8509699f\n",
            "  Stored in directory: /root/.cache/pip/wheels/50/ca/fa/8fca8d246e64f19488d07567547ddec8eb084e8c0d7a59226a\n",
            "  Building wheel for pathtools (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pathtools: filename=pathtools-0.1.2-py3-none-any.whl size=8806 sha256=765704333b8a0bfefad46ba4f663ffadbb7e500c02b1b66ee5c486552c7dac2b\n",
            "  Stored in directory: /root/.cache/pip/wheels/3e/31/09/fa59cef12cdcfecc627b3d24273699f390e71828921b2cbba2\n",
            "Successfully built subprocess32 pathtools\n",
            "Installing collected packages: smmap, gitdb, subprocess32, shortuuid, sentry-sdk, pathtools, GitPython, docker-pycreds, configparser, wandb\n",
            "Successfully installed GitPython-3.1.18 configparser-5.0.2 docker-pycreds-0.4.0 gitdb-4.0.7 pathtools-0.1.2 sentry-sdk-1.3.1 shortuuid-1.0.1 smmap-4.0.0 subprocess32-3.5.4 wandb-0.12.0\n",
            "Collecting dm_env\n",
            "  Downloading dm_env-1.5-py3-none-any.whl (26 kB)\n",
            "Requirement already satisfied: dm-tree in /usr/local/lib/python3.7/dist-packages (from dm_env) (0.1.6)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from dm_env) (1.19.5)\n",
            "Requirement already satisfied: absl-py in /usr/local/lib/python3.7/dist-packages (from dm_env) (0.12.0)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from absl-py->dm_env) (1.15.0)\n",
            "Installing collected packages: dm-env\n",
            "Successfully installed dm-env-1.5\n",
            "Collecting git+https://github.com/deepmind/dm-haiku\n",
            "  Cloning https://github.com/deepmind/dm-haiku to /tmp/pip-req-build-imlgdpn0\n",
            "  Running command git clone -q https://github.com/deepmind/dm-haiku /tmp/pip-req-build-imlgdpn0\n",
            "Requirement already satisfied: absl-py>=0.7.1 in /usr/local/lib/python3.7/dist-packages (from dm-haiku==0.0.5.dev0) (0.12.0)\n",
            "Collecting jmp>=0.0.2\n",
            "  Downloading jmp-0.0.2-py3-none-any.whl (16 kB)\n",
            "Requirement already satisfied: numpy>=1.18.0 in /usr/local/lib/python3.7/dist-packages (from dm-haiku==0.0.5.dev0) (1.19.5)\n",
            "Requirement already satisfied: tabulate>=0.8.9 in /usr/local/lib/python3.7/dist-packages (from dm-haiku==0.0.5.dev0) (0.8.9)\n",
            "Requirement already satisfied: typing_extensions in /usr/local/lib/python3.7/dist-packages (from dm-haiku==0.0.5.dev0) (3.7.4.3)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from absl-py>=0.7.1->dm-haiku==0.0.5.dev0) (1.15.0)\n",
            "Building wheels for collected packages: dm-haiku\n",
            "  Building wheel for dm-haiku (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for dm-haiku: filename=dm_haiku-0.0.5.dev0-py3-none-any.whl size=532207 sha256=881089eb77da4c604c4d9aa4fe88cb702f74d69267a3176673f98270e5edbd18\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-e_rowyki/wheels/06/28/69/ebaac5b2435641427299f29d88d005fb4e2627f4a108f0bdbc\n",
            "Successfully built dm-haiku\n",
            "Installing collected packages: jmp, dm-haiku\n",
            "Successfully installed dm-haiku-0.0.5.dev0 jmp-0.0.2\n",
            "Collecting git+git://github.com/deepmind/bsuite.git\n",
            "  Cloning git://github.com/deepmind/bsuite.git to /tmp/pip-req-build-sz5mp0ye\n",
            "  Running command git clone -q git://github.com/deepmind/bsuite.git /tmp/pip-req-build-sz5mp0ye\n",
            "Requirement already satisfied: absl-py in /usr/local/lib/python3.7/dist-packages (from bsuite==0.3.5) (0.12.0)\n",
            "Requirement already satisfied: dm_env in /usr/local/lib/python3.7/dist-packages (from bsuite==0.3.5) (1.5)\n",
            "Collecting immutabledict\n",
            "  Downloading immutabledict-2.1.0-py3-none-any.whl (4.0 kB)\n",
            "Requirement already satisfied: gym in /usr/local/lib/python3.7/dist-packages (from bsuite==0.3.5) (0.17.3)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.7/dist-packages (from bsuite==0.3.5) (3.2.2)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from bsuite==0.3.5) (1.19.5)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.7/dist-packages (from bsuite==0.3.5) (1.1.5)\n",
            "Requirement already satisfied: plotnine in /usr/local/lib/python3.7/dist-packages (from bsuite==0.3.5) (0.6.0)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from bsuite==0.3.5) (1.4.1)\n",
            "Requirement already satisfied: scikit-image in /usr/local/lib/python3.7/dist-packages (from bsuite==0.3.5) (0.16.2)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from bsuite==0.3.5) (1.15.0)\n",
            "Requirement already satisfied: termcolor in /usr/local/lib/python3.7/dist-packages (from bsuite==0.3.5) (1.1.0)\n",
            "Requirement already satisfied: dm-tree in /usr/local/lib/python3.7/dist-packages (from dm_env->bsuite==0.3.5) (0.1.6)\n",
            "Requirement already satisfied: pyglet<=1.5.0,>=1.4.0 in /usr/local/lib/python3.7/dist-packages (from gym->bsuite==0.3.5) (1.5.0)\n",
            "Requirement already satisfied: cloudpickle<1.7.0,>=1.2.0 in /usr/local/lib/python3.7/dist-packages (from gym->bsuite==0.3.5) (1.3.0)\n",
            "Requirement already satisfied: future in /usr/local/lib/python3.7/dist-packages (from pyglet<=1.5.0,>=1.4.0->gym->bsuite==0.3.5) (0.16.0)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib->bsuite==0.3.5) (0.10.0)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->bsuite==0.3.5) (2.8.1)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->bsuite==0.3.5) (2.4.7)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->bsuite==0.3.5) (1.3.1)\n",
            "Requirement already satisfied: pytz>=2017.2 in /usr/local/lib/python3.7/dist-packages (from pandas->bsuite==0.3.5) (2018.9)\n",
            "Requirement already satisfied: descartes>=1.1.0 in /usr/local/lib/python3.7/dist-packages (from plotnine->bsuite==0.3.5) (1.1.0)\n",
            "Requirement already satisfied: patsy>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from plotnine->bsuite==0.3.5) (0.5.1)\n",
            "Requirement already satisfied: mizani>=0.6.0 in /usr/local/lib/python3.7/dist-packages (from plotnine->bsuite==0.3.5) (0.6.0)\n",
            "Requirement already satisfied: statsmodels>=0.9.0 in /usr/local/lib/python3.7/dist-packages (from plotnine->bsuite==0.3.5) (0.10.2)\n",
            "Requirement already satisfied: palettable in /usr/local/lib/python3.7/dist-packages (from mizani>=0.6.0->plotnine->bsuite==0.3.5) (3.3.0)\n",
            "Requirement already satisfied: pillow>=4.3.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image->bsuite==0.3.5) (7.1.2)\n",
            "Requirement already satisfied: imageio>=2.3.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image->bsuite==0.3.5) (2.4.1)\n",
            "Requirement already satisfied: PyWavelets>=0.4.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image->bsuite==0.3.5) (1.1.1)\n",
            "Requirement already satisfied: networkx>=2.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image->bsuite==0.3.5) (2.5.1)\n",
            "Requirement already satisfied: decorator<5,>=4.3 in /usr/local/lib/python3.7/dist-packages (from networkx>=2.0->scikit-image->bsuite==0.3.5) (4.4.2)\n",
            "Building wheels for collected packages: bsuite\n",
            "  Building wheel for bsuite (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for bsuite: filename=bsuite-0.3.5-py3-none-any.whl size=250260 sha256=0886bb33f05f0952c121f9c53700135f1b24447f4d094c8586d0a4581cae609a\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-e_ot_amq/wheels/8e/cd/02/51ff02c07cffd573fb0a0095e53898313e70ebdc541e351d99\n",
            "Successfully built bsuite\n",
            "Installing collected packages: immutabledict, bsuite\n",
            "Successfully installed bsuite-0.3.5 immutabledict-2.1.0\n",
            "Collecting git+git://github.com/deepmind/optax.git\n",
            "  Cloning git://github.com/deepmind/optax.git to /tmp/pip-req-build-4_zjukh5\n",
            "  Running command git clone -q git://github.com/deepmind/optax.git /tmp/pip-req-build-4_zjukh5\n",
            "Requirement already satisfied: absl-py>=0.7.1 in /usr/local/lib/python3.7/dist-packages (from optax==0.0.9) (0.12.0)\n",
            "Collecting chex>=0.0.4\n",
            "  Downloading chex-0.0.8-py3-none-any.whl (57 kB)\n",
            "\u001b[K     |████████████████████████████████| 57 kB 5.3 MB/s \n",
            "\u001b[?25hRequirement already satisfied: jax>=0.1.55 in /usr/local/lib/python3.7/dist-packages (from optax==0.0.9) (0.2.17)\n",
            "Requirement already satisfied: jaxlib>=0.1.37 in /usr/local/lib/python3.7/dist-packages (from optax==0.0.9) (0.1.69+cuda110)\n",
            "Requirement already satisfied: numpy>=1.18.0 in /usr/local/lib/python3.7/dist-packages (from optax==0.0.9) (1.19.5)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from absl-py>=0.7.1->optax==0.0.9) (1.15.0)\n",
            "Requirement already satisfied: toolz>=0.9.0 in /usr/local/lib/python3.7/dist-packages (from chex>=0.0.4->optax==0.0.9) (0.11.1)\n",
            "Requirement already satisfied: dm-tree>=0.1.5 in /usr/local/lib/python3.7/dist-packages (from chex>=0.0.4->optax==0.0.9) (0.1.6)\n",
            "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.7/dist-packages (from jax>=0.1.55->optax==0.0.9) (3.3.0)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from jaxlib>=0.1.37->optax==0.0.9) (1.4.1)\n",
            "Requirement already satisfied: flatbuffers<3.0,>=1.12 in /usr/local/lib/python3.7/dist-packages (from jaxlib>=0.1.37->optax==0.0.9) (1.12)\n",
            "Building wheels for collected packages: optax\n",
            "  Building wheel for optax (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for optax: filename=optax-0.0.9-py3-none-any.whl size=120437 sha256=5a4274a4f8503987d29cab6ed31879ee21eae40dc13ed92474b06e4f1c0f103e\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-l3ve02ho/wheels/99/a8/ea/844576cb39aaa6929daf91bd459d92d5d2e0c52abcc31d5b8c\n",
            "Successfully built optax\n",
            "Installing collected packages: chex, optax\n",
            "Successfully installed chex-0.0.8 optax-0.0.9\n",
            "Collecting git+git://github.com/deepmind/rlax.git\n",
            "  Cloning git://github.com/deepmind/rlax.git to /tmp/pip-req-build-pqewwh4t\n",
            "  Running command git clone -q git://github.com/deepmind/rlax.git /tmp/pip-req-build-pqewwh4t\n",
            "Requirement already satisfied: absl-py>=0.9.0 in /usr/local/lib/python3.7/dist-packages (from rlax==0.0.4) (0.12.0)\n",
            "Requirement already satisfied: chex>=0.0.8 in /usr/local/lib/python3.7/dist-packages (from rlax==0.0.4) (0.0.8)\n",
            "Requirement already satisfied: jax>=0.1.55 in /usr/local/lib/python3.7/dist-packages (from rlax==0.0.4) (0.2.17)\n",
            "Requirement already satisfied: jaxlib>=0.1.37 in /usr/local/lib/python3.7/dist-packages (from rlax==0.0.4) (0.1.69+cuda110)\n",
            "Requirement already satisfied: numpy>=1.18.0 in /usr/local/lib/python3.7/dist-packages (from rlax==0.0.4) (1.19.5)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from absl-py>=0.9.0->rlax==0.0.4) (1.15.0)\n",
            "Requirement already satisfied: dm-tree>=0.1.5 in /usr/local/lib/python3.7/dist-packages (from chex>=0.0.8->rlax==0.0.4) (0.1.6)\n",
            "Requirement already satisfied: toolz>=0.9.0 in /usr/local/lib/python3.7/dist-packages (from chex>=0.0.8->rlax==0.0.4) (0.11.1)\n",
            "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.7/dist-packages (from jax>=0.1.55->rlax==0.0.4) (3.3.0)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from jaxlib>=0.1.37->rlax==0.0.4) (1.4.1)\n",
            "Requirement already satisfied: flatbuffers<3.0,>=1.12 in /usr/local/lib/python3.7/dist-packages (from jaxlib>=0.1.37->rlax==0.0.4) (1.12)\n",
            "Building wheels for collected packages: rlax\n",
            "  Building wheel for rlax (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for rlax: filename=rlax-0.0.4-py3-none-any.whl size=112310 sha256=dc633171b839bec694749e3e33ea8f10d04c02f060e46c1049bf18cfd830935f\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-5lm_teys/wheels/e8/39/c2/aaadc821fd338957de135c0b4a068a698bdbc49cb7f0c2257a\n",
            "Successfully built rlax\n",
            "Installing collected packages: rlax\n",
            "Successfully installed rlax-0.0.4\n",
            "Requirement already satisfied: dm-tree in /usr/local/lib/python3.7/dist-packages (0.1.6)\n",
            "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.7/dist-packages (from dm-tree) (1.15.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (21.0)\n",
            "Requirement already satisfied: pyparsing>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging) (2.4.7)\n",
            "Requirement already satisfied: tensorflow-datasets in /usr/local/lib/python3.7/dist-packages (4.0.1)\n",
            "Requirement already satisfied: termcolor in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets) (1.1.0)\n",
            "Requirement already satisfied: future in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets) (0.16.0)\n",
            "Requirement already satisfied: dill in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets) (0.3.4)\n",
            "Requirement already satisfied: promise in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets) (2.3)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets) (4.41.1)\n",
            "Requirement already satisfied: dm-tree in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets) (0.1.6)\n",
            "Requirement already satisfied: importlib-resources in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets) (5.2.0)\n",
            "Requirement already satisfied: tensorflow-metadata in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets) (1.1.0)\n",
            "Requirement already satisfied: absl-py in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets) (0.12.0)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets) (1.19.5)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets) (1.15.0)\n",
            "Requirement already satisfied: attrs>=18.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets) (21.2.0)\n",
            "Requirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets) (2.23.0)\n",
            "Requirement already satisfied: protobuf>=3.6.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets) (3.17.3)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests>=2.19.0->tensorflow-datasets) (3.0.4)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests>=2.19.0->tensorflow-datasets) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests>=2.19.0->tensorflow-datasets) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests>=2.19.0->tensorflow-datasets) (2021.5.30)\n",
            "Requirement already satisfied: zipp>=3.1.0 in /usr/local/lib/python3.7/dist-packages (from importlib-resources->tensorflow-datasets) (3.5.0)\n",
            "Requirement already satisfied: googleapis-common-protos<2,>=1.52.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-metadata->tensorflow-datasets) (1.53.0)\n",
            "Requirement already satisfied: tensorflow in /usr/local/lib/python3.7/dist-packages (2.5.0)\n",
            "Requirement already satisfied: tensorboard~=2.5 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (2.5.0)\n",
            "Requirement already satisfied: flatbuffers~=1.12.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.12)\n",
            "Requirement already satisfied: wheel~=0.35 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (0.36.2)\n",
            "Requirement already satisfied: wrapt~=1.12.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.12.1)\n",
            "Requirement already satisfied: typing-extensions~=3.7.4 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (3.7.4.3)\n",
            "Requirement already satisfied: google-pasta~=0.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (0.2.0)\n",
            "Requirement already satisfied: grpcio~=1.34.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.34.1)\n",
            "Requirement already satisfied: keras-nightly~=2.5.0.dev in /usr/local/lib/python3.7/dist-packages (from tensorflow) (2.5.0.dev2021032900)\n",
            "Requirement already satisfied: absl-py~=0.10 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (0.12.0)\n",
            "Requirement already satisfied: numpy~=1.19.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.19.5)\n",
            "Requirement already satisfied: termcolor~=1.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.1.0)\n",
            "Requirement already satisfied: gast==0.4.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (0.4.0)\n",
            "Requirement already satisfied: keras-preprocessing~=1.1.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.1.2)\n",
            "Requirement already satisfied: astunparse~=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.6.3)\n",
            "Requirement already satisfied: tensorflow-estimator<2.6.0,>=2.5.0rc0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (2.5.0)\n",
            "Requirement already satisfied: protobuf>=3.9.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (3.17.3)\n",
            "Requirement already satisfied: h5py~=3.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (3.1.0)\n",
            "Requirement already satisfied: opt-einsum~=3.3.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (3.3.0)\n",
            "Requirement already satisfied: six~=1.15.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.15.0)\n",
            "Requirement already satisfied: cached-property in /usr/local/lib/python3.7/dist-packages (from h5py~=3.1.0->tensorflow) (1.5.2)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow) (1.8.0)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow) (3.3.4)\n",
            "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow) (1.0.1)\n",
            "Requirement already satisfied: google-auth<2,>=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow) (1.32.1)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow) (0.4.4)\n",
            "Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow) (57.2.0)\n",
            "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow) (0.6.1)\n",
            "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow) (2.23.0)\n",
            "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard~=2.5->tensorflow) (4.2.2)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard~=2.5->tensorflow) (4.7.2)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard~=2.5->tensorflow) (0.2.8)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard~=2.5->tensorflow) (1.3.0)\n",
            "Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from markdown>=2.6.8->tensorboard~=2.5->tensorflow) (4.6.1)\n",
            "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.7/dist-packages (from pyasn1-modules>=0.2.1->google-auth<2,>=1.6.3->tensorboard~=2.5->tensorflow) (0.4.8)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard~=2.5->tensorflow) (3.0.4)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard~=2.5->tensorflow) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard~=2.5->tensorflow) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard~=2.5->tensorflow) (2021.5.30)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.7/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard~=2.5->tensorflow) (3.1.1)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->markdown>=2.6.8->tensorboard~=2.5->tensorflow) (3.5.0)\n",
            "Collecting pyvirtualdisplay\n",
            "  Downloading PyVirtualDisplay-2.2-py3-none-any.whl (15 kB)\n",
            "Collecting EasyProcess\n",
            "  Downloading EasyProcess-0.3-py2.py3-none-any.whl (7.9 kB)\n",
            "Installing collected packages: EasyProcess, pyvirtualdisplay\n",
            "Successfully installed EasyProcess-0.3 pyvirtualdisplay-2.2\n",
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "The following NEW packages will be installed:\n",
            "  xvfb\n",
            "0 upgraded, 1 newly installed, 0 to remove and 40 not upgraded.\n",
            "Need to get 784 kB of archives.\n",
            "After this operation, 2,270 kB of additional disk space will be used.\n",
            "Get:1 http://archive.ubuntu.com/ubuntu bionic-updates/universe amd64 xvfb amd64 2:1.19.6-1ubuntu4.9 [784 kB]\n",
            "Fetched 784 kB in 0s (8,766 kB/s)\n",
            "Selecting previously unselected package xvfb.\n",
            "(Reading database ... 160837 files and directories currently installed.)\n",
            "Preparing to unpack .../xvfb_2%3a1.19.6-1ubuntu4.9_amd64.deb ...\n",
            "Unpacking xvfb (2:1.19.6-1ubuntu4.9) ...\n",
            "Setting up xvfb (2:1.19.6-1ubuntu4.9) ...\n",
            "Processing triggers for man-db (2.8.3-2ubuntu0.1) ...\n",
            "Collecting git+https://github.com/tensorflow/docs\n",
            "  Cloning https://github.com/tensorflow/docs to /tmp/pip-req-build-e5d3s68r\n",
            "  Running command git clone -q https://github.com/tensorflow/docs /tmp/pip-req-build-e5d3s68r\n",
            "Requirement already satisfied: astor in /usr/local/lib/python3.7/dist-packages (from tensorflow-docs===0.0.02b3cacb3742ac64d4a0b20e3f8d19bf62cf47730-) (0.8.1)\n",
            "Requirement already satisfied: absl-py in /usr/local/lib/python3.7/dist-packages (from tensorflow-docs===0.0.02b3cacb3742ac64d4a0b20e3f8d19bf62cf47730-) (0.12.0)\n",
            "Requirement already satisfied: protobuf>=3.14 in /usr/local/lib/python3.7/dist-packages (from tensorflow-docs===0.0.02b3cacb3742ac64d4a0b20e3f8d19bf62cf47730-) (3.17.3)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.7/dist-packages (from tensorflow-docs===0.0.02b3cacb3742ac64d4a0b20e3f8d19bf62cf47730-) (3.13)\n",
            "Requirement already satisfied: six>=1.9 in /usr/local/lib/python3.7/dist-packages (from protobuf>=3.14->tensorflow-docs===0.0.02b3cacb3742ac64d4a0b20e3f8d19bf62cf47730-) (1.15.0)\n",
            "Building wheels for collected packages: tensorflow-docs\n",
            "  Building wheel for tensorflow-docs (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for tensorflow-docs: filename=tensorflow_docs-0.0.02b3cacb3742ac64d4a0b20e3f8d19bf62cf47730_-py3-none-any.whl size=154265 sha256=b022c22edf0ef953ff6c9d91261748c658d4bfd210ec18cf31129693eefef91f\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-09wm5ogd/wheels/cc/c4/d8/5341e93b6376c5c929c49469fce21155eb69cef1a4da4ce32c\n",
            "\u001b[33m  WARNING: Built wheel for tensorflow-docs is invalid: Metadata 1.2 mandates PEP 440 version, but '0.0.02b3cacb3742ac64d4a0b20e3f8d19bf62cf47730-' is not\u001b[0m\n",
            "Failed to build tensorflow-docs\n",
            "Installing collected packages: tensorflow-docs\n",
            "    Running setup.py install for tensorflow-docs ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[33m  DEPRECATION: tensorflow-docs was installed using the legacy 'setup.py install' method, because a wheel could not be built for it. A possible replacement is to fix the wheel build issue reported above. You can find discussion regarding this at https://github.com/pypa/pip/issues/8368.\u001b[0m\n",
            "Successfully installed tensorflow-docs-0.0.02b3cacb3742ac64d4a0b20e3f8d19bf62cf47730-\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (4.41.1)\n",
            "Requirement already satisfied: chex in /usr/local/lib/python3.7/dist-packages (0.0.8)\n",
            "Requirement already satisfied: numpy>=1.18.0 in /usr/local/lib/python3.7/dist-packages (from chex) (1.19.5)\n",
            "Requirement already satisfied: dm-tree>=0.1.5 in /usr/local/lib/python3.7/dist-packages (from chex) (0.1.6)\n",
            "Requirement already satisfied: toolz>=0.9.0 in /usr/local/lib/python3.7/dist-packages (from chex) (0.11.1)\n",
            "Requirement already satisfied: jax>=0.1.55 in /usr/local/lib/python3.7/dist-packages (from chex) (0.2.17)\n",
            "Requirement already satisfied: absl-py>=0.9.0 in /usr/local/lib/python3.7/dist-packages (from chex) (0.12.0)\n",
            "Requirement already satisfied: jaxlib>=0.1.37 in /usr/local/lib/python3.7/dist-packages (from chex) (0.1.69+cuda110)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from absl-py>=0.9.0->chex) (1.15.0)\n",
            "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.7/dist-packages (from jax>=0.1.55->chex) (3.3.0)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from jaxlib>=0.1.37->chex) (1.4.1)\n",
            "Requirement already satisfied: flatbuffers<3.0,>=1.12 in /usr/local/lib/python3.7/dist-packages (from jaxlib>=0.1.37->chex) (1.12)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5tTbMKAo02k2"
      },
      "source": [
        "#import libraries\n",
        "import jax\n",
        "import collections\n",
        "import functools\n",
        "from typing import Any, Callable, Optional, Tuple, Dict\n",
        "import dm_env\n",
        "import haiku as hk\n",
        "from examples.impala import util\n",
        "import jax.numpy as jnp\n",
        "import numpy as np\n",
        "import chex\n",
        "import itertools\n",
        "import queue\n",
        "import threading\n",
        "import warnings\n",
        "from examples.impala import util\n",
        "from jax.experimental import optimizers\n",
        "import optax\n",
        "import rlax\n",
        "from tqdm import tqdm  \n",
        "from rlax._src import base\n",
        "import haiku as hk\n",
        "from examples.impala import util\n",
        "\n",
        "from PIL import Image\n",
        "import threading\n",
        "from typing import List\n",
        "import wandb\n",
        "from absl import app\n",
        "from bsuite.environments import cartpole\n",
        "from examples.impala import util\n",
        "import jax\n",
        "import optax\n",
        "import collections\n",
        "import gym\n",
        "import numpy as np\n",
        "import statistics\n",
        "import tensorflow as tf\n",
        "import re\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import collections\n",
        "\n",
        "import dm_env\n",
        "import haiku as hk\n",
        "import jax.nn\n",
        "import jax.numpy as jnp"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ALCe54IQvxqg"
      },
      "source": [
        "# Agent\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rx_Lqd99wYS8"
      },
      "source": [
        "#create var for NN output\n",
        "AgentOutput = collections.namedtuple(\"AgentOutput\",\n",
        "                                     [\"policy_logits\", \"values\", \"action\"])\n",
        "\n",
        "Action = int\n",
        "Nest = Any\n",
        "NetFactory = Callable[[int], hk.RNNCore]\n",
        "\n",
        "#agent interface\n",
        "class Agent:\n",
        "  def __init__(self, num_actions: int, obs_spec: Nest,\n",
        "               net_factory: NetFactory):\n",
        "    \"\"\"\n",
        "    The interface for the agent.\n",
        "    Args:\n",
        "      num_actions: Number of possible actions for the agent. Assumes a flat,\n",
        "        discrete, 0-indexed action space.\n",
        "      obs_spec: The observation spec of the environment.\n",
        "      net_factory: A function from num_actions to a Haiku module representing\n",
        "        the agent. This module should have an initial_state() function and an\n",
        "        unroll function.\n",
        "    \"\"\"\n",
        "    #set observation spec\n",
        "    self._obs_spec = obs_spec\n",
        "\n",
        "    #set nn being used- for non-CatchNet use second version\n",
        "    net_factory = functools.partial(net_factory, num_actions)\n",
        "    #net_factory = functools.partial(net_factory, num_actions)\n",
        "    \n",
        "    #start the nn in initial starting space\n",
        "    _, self._initial_state_apply_fn = hk.without_apply_rng(\n",
        "        hk.transform(\n",
        "            lambda batch_size: net_factory().initial_state(batch_size)))\n",
        "\n",
        "    #create a quick way to call nn on data later\n",
        "    self._init_fn, self._apply_fn = hk.without_apply_rng(\n",
        "        hk.transform(lambda obs, state: net_factory().unroll(obs, state)))\n",
        "    \n",
        "\n",
        "  #initialize the agent\n",
        "  @functools.partial(jax.jit, static_argnums=0)\n",
        "  def initial_params(self, rng_key):\n",
        "    \"\"\"\n",
        "    Initializes the agent params.\n",
        "    Args:\n",
        "      rng_key: initial state of nn\n",
        "    Returns: initial state function\n",
        "    \"\"\"\n",
        "    #create framework to feed data into nn later\n",
        "    dummy_inputs = jax.tree_map(lambda t: np.zeros(t.shape, t.dtype),\n",
        "                                self._obs_spec)\n",
        "    dummy_inputs = util.preprocess_step(dm_env.restart(dummy_inputs))\n",
        "    dummy_inputs = jax.tree_map(lambda t: t[None, None, ...], dummy_inputs)\n",
        "    return self._init_fn(rng_key, dummy_inputs, self.initial_state(1))\n",
        "\n",
        "  #return initial state\n",
        "  @functools.partial(jax.jit, static_argnums=(0, 1))\n",
        "  def initial_state(self, batch_size: Optional[int]):\n",
        "    \"\"\"\n",
        "    Returns agent initial state\n",
        "    Args:\n",
        "      batch_size: batch size, int\n",
        "    Returns:\n",
        "      function to generate initial state\"\"\"\n",
        "    # We expect that generating the initial_state does not require parameters.\n",
        "    return self._initial_state_apply_fn(None, batch_size)\n",
        "\n",
        "  @functools.partial(jax.jit, static_argnums=(0,))\n",
        "  def step(\n",
        "      self,\n",
        "      rng_key,\n",
        "      params: hk.Params,\n",
        "      timestep: dm_env.TimeStep,\n",
        "      state: Nest,\n",
        "  ) -> Tuple[AgentOutput, Nest]:\n",
        "    \"\"\"For a given single-step, output the chosen action\"\"\"\n",
        "    # Pad timestep, state to be [T, B, ...] and [B, ...] respectively.\n",
        "    timestep = jax.tree_map(lambda t: t[None, None, ...], timestep)\n",
        "    state = jax.tree_map(lambda t: t[None, ...], state)\n",
        "\n",
        "    net_out, next_state = self._apply_fn(params, timestep, state)\n",
        "    # Remove the padding from above.\n",
        "    net_out = jax.tree_map(lambda t: jnp.squeeze(t, axis=(0, 1)), net_out)\n",
        "    next_state = jax.tree_map(lambda t: jnp.squeeze(t, axis=0), next_state)\n",
        "    # Sample an action and return.\n",
        "    action = hk.multinomial(rng_key, net_out.policy_logits, num_samples=1)\n",
        "    action = jnp.squeeze(action, axis=-1)\n",
        "    return AgentOutput(net_out.policy_logits, net_out.value, action), next_state\n",
        "\n",
        "  def unroll(\n",
        "      self,\n",
        "      params: hk.Params,\n",
        "      trajectory: dm_env.TimeStep,\n",
        "      state: Nest,\n",
        "  ) -> AgentOutput:\n",
        "    \"\"\"Unroll the agent along trajectory.\"\"\"\n",
        "    net_out, _ = self._apply_fn(params, trajectory, state)\n",
        "    return AgentOutput(net_out.policy_logits, net_out.value, action=[])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DlZ7YS35ga7K"
      },
      "source": [
        "# VTrace\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YcDvhaj5Rb7P"
      },
      "source": [
        "#create vtrace output array\n",
        "Array = chex.Array\n",
        "VTraceOutput = collections.namedtuple(\n",
        "    'vtrace_output', ['errors', 'pg_advantage', 'q_estimate'])\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tKyujOWAgVas"
      },
      "source": [
        "def vtrace_td_error_and_advantage(\n",
        "    c_help: Array,\n",
        "    v_tm1: Array,\n",
        "    v_t: Array,\n",
        "    r_t: Array,\n",
        "    discount_t: Array,\n",
        "    rho_tm1: Array,\n",
        "    lambda_: float = 1.0,\n",
        "    clip_rho_threshold: float = 1.0,\n",
        "    clip_pg_rho_threshold: float = 1.0,\n",
        "    stop_target_gradients: bool = True,\n",
        ") -> VTraceOutput:\n",
        "  \"\"\"Calculates V-Trace errors and PG advantage from importance weights.\n",
        "  This functions computes the TD-errors and policy gradient Advantage terms\n",
        "  as used by the IMPALA distributed actor-critic agent.\n",
        "  See \"IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor\n",
        "  Learner Architectures\" by Espeholt et al. (https://arxiv.org/abs/1802.01561)\n",
        "  Args:\n",
        "    c_help: target policy values\n",
        "    v_tm1: values at time t-1.\n",
        "    v_t: values at time t.\n",
        "    r_t: reward at time t.\n",
        "    discount_t: discount at time t.\n",
        "    rho_tm1: importance weights at time t.\n",
        "    lambda_: scalar mixing parameter lambda.\n",
        "    clip_rho_threshold: clip threshold for importance ratios.\n",
        "    clip_pg_rho_threshold: clip threshold for policy gradient importance ratios.\n",
        "    stop_target_gradients: whether or not to apply stop gradient to targets.\n",
        "  Returns:\n",
        "    a tuple of V-Trace error, policy gradient advantage, and estimated Q-values.\n",
        "  \"\"\"\n",
        "  #check shapes are correct\n",
        "  chex.assert_rank([v_tm1, v_t, r_t, discount_t, rho_tm1], 1)\n",
        "  chex.assert_type([v_tm1, v_t, r_t, discount_t, rho_tm1], float)\n",
        "  chex.assert_equal_shape([v_tm1, v_t, r_t, discount_t, rho_tm1])\n",
        "\n",
        "  #calculate TD error\n",
        "  errors = vtrace(\n",
        "      c_help, v_tm1, v_t, r_t, discount_t, rho_tm1,\n",
        "      lambda_, clip_rho_threshold, stop_target_gradients)\n",
        "  targets_tm1 = errors + v_tm1\n",
        "  #calculate q bootstrap value\n",
        "  q_bootstrap = jnp.concatenate([\n",
        "      lambda_ * targets_tm1[1:] + (1 - lambda_) * v_tm1[1:],\n",
        "      v_t[-1:],\n",
        "  ], axis=0)\n",
        "  #estimate updated q value\n",
        "  q_estimate = r_t + discount_t * q_bootstrap\n",
        "  #estimate pg advantage for loss function\n",
        "  clipped_pg_rho_tm1 = jnp.minimum(clip_pg_rho_threshold, rho_tm1)\n",
        "  pg_advantages = clipped_pg_rho_tm1 * (q_estimate - v_tm1)\n",
        "  #return vtrace outputs\n",
        "  return VTraceOutput(\n",
        "      errors=errors, pg_advantage=pg_advantages, q_estimate=q_estimate)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "li9UC6-tgCfA"
      },
      "source": [
        "def vtrace(\n",
        "    c_help: Array,\n",
        "    v_tm1: Array,\n",
        "    v_t: Array,\n",
        "    r_t: Array,\n",
        "    discount_t: Array,\n",
        "    rho_tm1: Array,\n",
        "    lambda_: float = 1.0,\n",
        "    clip_rho_threshold: float = 1.0,\n",
        "    stop_target_gradients: bool = True,\n",
        ") -> Array:\n",
        "  \"\"\"Calculates V-Trace errors from importance weights.\n",
        "  V-trace computes TD-errors from multistep trajectories by applying\n",
        "  off-policy corrections based on clipped importance sampling ratios.\n",
        "  See \"IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor\n",
        "  Learner Architectures\" by Espeholt et al. (https://arxiv.org/abs/1802.01561).\n",
        "  Args:\n",
        "    c_help: target policy values \n",
        "    v_tm1: values at time t-1.\n",
        "    v_t: values at time t.\n",
        "    r_t: reward at time t.\n",
        "    discount_t: discount at time t.\n",
        "    rho_tm1: importance sampling ratios.\n",
        "    lambda_: scalar mixing parameter lambda.\n",
        "    clip_rho_threshold: clip threshold for importance weights.\n",
        "    stop_target_gradients: whether or not to apply stop gradient to targets.\n",
        "  Returns:\n",
        "    V-Trace error.\n",
        "  \"\"\"\n",
        "  #check shapes of everything\n",
        "  chex.assert_rank([v_tm1, v_t, r_t, discount_t, rho_tm1], [1, 1, 1, 1, 1])\n",
        "  chex.assert_type([v_tm1, v_t, r_t, discount_t, rho_tm1],\n",
        "                   [float, float, float, float, float])\n",
        "  chex.assert_equal_shape([v_tm1, v_t, r_t, discount_t, rho_tm1])\n",
        "\n",
        "  # Clip importance sampling ratios if needed-- change depending on IS algo\n",
        "  clipped_rhos = rho_tm1 \n",
        "  c_t = rho_tm1\n",
        "\n",
        "  #lambda num-- change depending on IS algo\n",
        "  lambda_num = 1.\n",
        "\n",
        "  # Compute the temporal difference errors.\n",
        "  td_errors = clipped_rhos * (r_t + discount_t * v_t - v_tm1)\n",
        "\n",
        "  # Work backwards computing the td-errors.\n",
        "  err = 0.0\n",
        "  errors = []\n",
        "  for i in jnp.arange(v_t.shape[0] - 1, -1, -1):\n",
        "    err = td_errors[i] + discount_t[i] * lambda_num * err\n",
        "    errors.insert(0, err)\n",
        "\n",
        "  # Return errors, maybe disabling gradient flow through bootstrap targets.\n",
        "  return jax.lax.select(\n",
        "      stop_target_gradients,\n",
        "      jax.lax.stop_gradient(jnp.array(errors) + v_tm1) - v_tm1,\n",
        "      jnp.array(errors))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y0y3cm_wxRsn"
      },
      "source": [
        "# Learner"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sOJGN1fQlQZd"
      },
      "source": [
        "\n",
        "def policy_gradient_loss(logits, *args):\n",
        "  \"\"\"calculate policy gradient loss\n",
        "  Args:\n",
        "    logits: policy probabilities\n",
        "  Returns:\n",
        "    policy gradient loss\"\"\"\n",
        "  # calculate mean for batch\n",
        "  mean_per_batch = jax.vmap(rlax.policy_gradient_loss, in_axes=1)(logits, *args)\n",
        "  # multiply according to probabilities\n",
        "  total_loss_per_batch = mean_per_batch * logits.shape[0]\n",
        "  # return sum of those means\n",
        "  return jnp.sum(total_loss_per_batch)\n",
        "\n",
        "\n",
        "def entropy_loss(logits, *args):\n",
        "  \"\"\"calculate entropy loss\n",
        "  Args:\n",
        "    logits: policy probabilities\n",
        "  Return: \n",
        "    policy gradient loss\"\"\"\n",
        "  # calculate mean for batch\n",
        "  mean_per_batch = jax.vmap(rlax.entropy_loss, in_axes=1)(logits, *args)\n",
        "  # multiply according to probabilities\n",
        "  total_loss_per_batch = mean_per_batch * logits.shape[0]\n",
        "  # return sum of those means\n",
        "  return jnp.sum(total_loss_per_batch)\n",
        "\n",
        "\n",
        "class Learner:\n",
        "  \"\"\"Manages state and performs updates for IMPALA learner.\"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      agent: Agent,\n",
        "      rng_key,\n",
        "      opt: optax.GradientTransformation,\n",
        "      batch_size: int,\n",
        "      discount_factor: float,\n",
        "      frames_per_iter: int,\n",
        "      learnersActor,\n",
        "      max_abs_reward: float = 0,\n",
        "      logger=None,\n",
        "  ):\n",
        "    if jax.device_count() > 1:\n",
        "      warnings.warn('Note: the impala example will only take advantage of a '\n",
        "                    'single accelerator.')\n",
        "\n",
        "    # initialize vars\n",
        "    self._agent = agent\n",
        "    self._opt = opt\n",
        "    self._batch_size = batch_size\n",
        "    self._discount_factor = discount_factor\n",
        "    self._frames_per_iter = frames_per_iter\n",
        "    self._max_abs_reward = max_abs_reward\n",
        "    self._learners_actor = learnersActor\n",
        "\n",
        "    self._all_total_loss = {\"total_loss\":[], \"PG_loss\":[], \"baseline_loss\":[], \"entropy_loss\":[],\n",
        "                            \"grad_norm_unclipped\":[], \"weight_norm\":[],\n",
        "                            \"num_frames\":[], \"epoch\":[], \"error\":[], \"q\":[],\n",
        "                            \"error_length\":[], \"sum\":[]}\n",
        "\n",
        "\n",
        "\n",
        "    # Data pipeline objects.\n",
        "    self._done = False\n",
        "    self._host_q = queue.Queue(maxsize=self._batch_size)\n",
        "    self._device_q = queue.Queue(maxsize=1)\n",
        "    self.logit_log = None\n",
        "\n",
        "    # Prepare the parameters to be served to actors.\n",
        "    params = agent.initial_params(rng_key)\n",
        "    self._params_for_actor = (0, jax.device_get(params))\n",
        "\n",
        "    # keep track of data-collection vars\n",
        "    self._best_return = 0\n",
        "    self._best_visual_return = 0\n",
        "    self._tracker = 0\n",
        "    self._discount_log = 0\n",
        "\n",
        "    # Set up logging.\n",
        "    if logger is None:\n",
        "      logger = util.AbslLogger()\n",
        "    self._logger = logger\n",
        "\n",
        "  def _loss(\n",
        "      self,\n",
        "      theta: hk.Params,\n",
        "      trajectories: util.Transition,\n",
        "  ) -> Tuple[jnp.ndarray, Dict[str, jnp.ndarray]]:\n",
        "    \"\"\"Compute vtrace-based actor-critic loss.\n",
        "    Args:\n",
        "      theta: current params\n",
        "      trajectories: batch collected by actors\n",
        "    Returns:\n",
        "      total_loss: loss for all the updates\n",
        "      logs: data logs\"\"\"\n",
        "    # feed batched actor info into nn for learner outputs\n",
        "    initial_state = jax.tree_map(lambda t: t[0], trajectories.agent_state)\n",
        "    learner_outputs = self._agent.unroll(theta, trajectories.timestep,\n",
        "                                         initial_state)\n",
        "    v_t = learner_outputs.values[1:]\n",
        "    # Remove bootstrap timestep from non-timesteps.\n",
        "    _, actor_out, _ = jax.tree_map(lambda t: t[:-1], trajectories)\n",
        "    learner_outputs = jax.tree_map(lambda t: t[:-1], learner_outputs)\n",
        "    v_tm1 = learner_outputs.values\n",
        "\n",
        "    # Get the discount, reward, step_type from the next timestep.\n",
        "    timestep = jax.tree_map(lambda t: t[1:], trajectories.timestep)\n",
        "    discounts = timestep.discount * self._discount_factor\n",
        "    self._discount_log = timestep.discount\n",
        "    rewards = timestep.reward\n",
        "    if self._max_abs_reward > 0:\n",
        "      rewards = jnp.clip(rewards, -self._max_abs_reward, self._max_abs_reward)\n",
        "\n",
        "    # check to see if we've gone back to first step (failed and restarted)\n",
        "    mask = jnp.not_equal(timestep.step_type, int(dm_env.StepType.FIRST))\n",
        "    mask = mask.astype(jnp.float32)\n",
        "\n",
        "    #calculate rhos and target policy probabilities for importance sampling\n",
        "    rhos = rlax.categorical_importance_sampling_ratios(\n",
        "        learner_outputs.policy_logits, actor_out.policy_logits,\n",
        "        actor_out.action)\n",
        "    c_help = base.batched_index(jax.nn.log_softmax(learner_outputs.policy_logits), actor_out.action)\n",
        "    \n",
        "    # calculate vtrace\n",
        "    vtrace_td_error_and_advantage_map = jax.vmap(\n",
        "        vtrace_td_error_and_advantage, in_axes=1, out_axes=1)\n",
        "    vtrace_returns = vtrace_td_error_and_advantage_map(\n",
        "        c_help, v_tm1, v_t, rewards, discounts, rhos)\n",
        "    \n",
        "    #use vtrace outputs to calculate losses\n",
        "    error = vtrace_returns.errors\n",
        "    q = vtrace_returns.q_estimate\n",
        "    pg_advs = vtrace_returns.pg_advantage\n",
        "    pg_loss = policy_gradient_loss(learner_outputs.policy_logits,\n",
        "                                   actor_out.action, pg_advs, mask)\n",
        "    baseline_loss = 0.5 * jnp.sum(jnp.square(vtrace_returns.errors) * mask)\n",
        "    ent_loss = entropy_loss(learner_outputs.policy_logits, mask)\n",
        "\n",
        "    #add together for overall loss\n",
        "    total_loss = pg_loss\n",
        "    total_loss += 0.5 * baseline_loss\n",
        "    total_loss += 0.01 * ent_loss\n",
        "\n",
        "    #data logs\n",
        "    logs = {}\n",
        "    logs['PG_loss'] = pg_loss\n",
        "    logs['baseline_loss'] = baseline_loss\n",
        "    logs['entropy_loss'] = ent_loss\n",
        "    logs['total_loss'] = total_loss\n",
        "    logs['error'] = vtrace_returns.errors\n",
        "    logs['q'] = q\n",
        "    logs['error_length'] = jnp.sum(vtrace_returns.errors) / len(vtrace_returns.errors)\n",
        "    logs['sum'] = jnp.sum(vtrace_returns.errors)\n",
        "\n",
        "    return total_loss, logs\n",
        "\n",
        "  @functools.partial(jax.jit, static_argnums=0)\n",
        "  def update(self, params, opt_state, batch: util.Transition):\n",
        "    \"\"\"The actual update function.\n",
        "    Args:\n",
        "      params: current policy params\n",
        "      opt_state: current optimal state\n",
        "      batch: batch from actors\n",
        "    Returns:\n",
        "      new params, updated optimal state, data logs\"\"\"\n",
        "\n",
        "    # calculate total loss\n",
        "    (loss_val, logs), grads = jax.value_and_grad(\n",
        "        self._loss, has_aux=True)(params, batch)\n",
        "\n",
        "    # use total loss to update policy \n",
        "    grad_norm_unclipped = optimizers.l2_norm(grads)\n",
        "    updates, updated_opt_state = self._opt.update(grads, opt_state)\n",
        "    params = optax.apply_updates(params, updates)\n",
        "    weight_norm = optimizers.l2_norm(params)\n",
        "\n",
        "    #data logs\n",
        "    logs.update({\n",
        "        'grad_norm_unclipped': grad_norm_unclipped,\n",
        "        'weight_norm': weight_norm,\n",
        "    })\n",
        "    \n",
        "    return params, updated_opt_state, logs\n",
        "\n",
        "  def enqueue_traj(self, traj: util.Transition):\n",
        "    \"\"\"Enqueue trajectory from actor\n",
        "    Args:\n",
        "      atraj: actor trajectory\"\"\"\n",
        "    self._host_q.put(traj)\n",
        "\n",
        "  def best_return(self, actor_return):\n",
        "    \"\"\"update best actor return if needed\n",
        "    Args:\n",
        "      actor_return: best return from actor trajectories\"\"\"\n",
        "    # if this actor produced the best return, update it\n",
        "    if(actor_return > self._best_return):\n",
        "      self._best_return = actor_return\n",
        "\n",
        "  def best_visual_return(self, actor_return):\n",
        "     \"\"\"update best visual actor return if needed\n",
        "    Args:\n",
        "      actor_return: best return from the visual actor trajectories\"\"\"\n",
        "    # if this actor produced the best return out of the visual actors, update it\n",
        "    if(actor_return > self._best_visual_return):\n",
        "      self._best_visual_return = actor_return\n",
        "\n",
        "  def params_for_actor(self) -> Tuple[int, hk.Params]:\n",
        "    \"\"\"return current actor params\"\"\"\n",
        "    return self._params_for_actor\n",
        "\n",
        "  def host_to_device_worker(self):\n",
        "    \"\"\"Elementary data pipeline.\"\"\"\n",
        "    batch = []\n",
        "    while not self._done:\n",
        "      # Try to get a batch. Skip the iteration if we couldn't.\n",
        "      try:\n",
        "        for _ in range(len(batch), self._batch_size):\n",
        "          # As long as possible while keeping learner_test time reasonable.\n",
        "          batch.append(self._host_q.get(timeout=10))\n",
        "      except queue.Empty:\n",
        "        continue\n",
        "\n",
        "      assert len(batch) == self._batch_size\n",
        "      # Prepare for consumption, then put batch onto device.\n",
        "      stacked_batch = jax.tree_multimap(lambda *xs: np.stack(xs, axis=1),\n",
        "                                        *batch)\n",
        "      self._device_q.put(jax.device_put(stacked_batch))\n",
        "\n",
        "      # Clean out the built-up batch.\n",
        "      batch = []\n",
        "\n",
        "  def run(self, max_iterations: int = -1):\n",
        "    \"\"\"Runs the learner for max_iterations updates\n",
        "    Args:\n",
        "      max_iterations: how many times to run learner update\n",
        "    \"\"\"\n",
        "    # Start host-to-device transfer worker.\n",
        "    transfer_thread = threading.Thread(target=self.host_to_device_worker)\n",
        "    transfer_thread.start()\n",
        "\n",
        "    # get current params and optimal state\n",
        "    (num_frames, params) = self._params_for_actor\n",
        "    opt_state = self._opt.init(params)\n",
        "\n",
        "    # iterate as many times as needed\n",
        "    steps = range(max_iterations) if max_iterations != -1 else itertools.count()\n",
        "    for epoch in tqdm(steps):\n",
        "      # get actor trajectories\n",
        "      batch = self._device_q.get()\n",
        "\n",
        "      # update actor params and optimal state\n",
        "      params, opt_state, logs = self.update(params, opt_state, batch)\n",
        "\n",
        "      # move time tracker to next frame sequence\n",
        "      num_frames += self._frames_per_iter\n",
        "\n",
        "      # Collect parameters to distribute to downstream actors\n",
        "      self._params_for_actor = (num_frames, jax.device_get(params))\n",
        "\n",
        "      # Collect and write logs out\n",
        "      logs = jax.device_get(logs)\n",
        "      element_names = [\"total_loss\", \"PG_loss\", \"baseline_loss\", \"entropy_loss\", \"grad_norm_unclipped\", \"weight_norm\",\n",
        "       \"num_frames\", \"epoch\", \"error\", \"q\", \"error_length\", \"sum\"]\n",
        "      \n",
        "      logs.update({\n",
        "          'num_frames': num_frames,\n",
        "          'epoch': epoch,\n",
        "      })\n",
        "\n",
        "      for element in element_names:\n",
        "        new_element = {element: np.append(self._all_total_loss[element], logs[element])}\n",
        "        self._all_total_loss.update(new_element)\n",
        "\n",
        "      self._logger.write(logs)\n",
        "\n",
        "    # Shut down.\n",
        "    self._done = True\n",
        "    self._logger.close()\n",
        "    transfer_thread.join()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2R9cxaU31Gnt"
      },
      "source": [
        "# Actor"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5yPQoAHujRzl"
      },
      "source": [
        "class SeparateActor:\n",
        "  \"\"\"Manages the state of a single agent/environment interaction loop to keep track of actor returns\"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      agent: Agent,\n",
        "      env: dm_env.Environment,\n",
        "      unroll_length: int,\n",
        "      rng_seed: int = 42,\n",
        "      logger=None,\n",
        "  ):\n",
        "    # initialize vars\n",
        "    self._agent = agent\n",
        "    self._env = env\n",
        "    self._unroll_length = unroll_length\n",
        "    self._timestep = env.reset()\n",
        "    self._agent_state = agent.initial_state(None)\n",
        "    self._traj = []\n",
        "    self._rng_key = jax.random.PRNGKey(rng_seed)\n",
        "    self._best_return = 0 \n",
        "    self._all_returns = []\n",
        "    self._all_frames = []\n",
        "    self._num_runs = 0\n",
        "\n",
        "    if logger is None:\n",
        "      logger = util.AbslLogger()\n",
        "    self._logger = logger\n",
        "\n",
        "    self._episode_return = 0.\n",
        "    self._tracker = 0\n",
        "\n",
        "  def unroll(self, rng_key, frame_count: int, params: hk.Params,\n",
        "             unroll_length: int) -> util.Transition:\n",
        "    \"\"\"Run unroll_length agent/environment steps, returning the trajectory\n",
        "    Args:\n",
        "      rng_key: rng_key\n",
        "      frame_count: what frame you're on\n",
        "      params: current actor params\n",
        "      unroll_length: how far to unroll\"\"\"\n",
        "    \n",
        "    # set actor to current time and state\n",
        "    timestep = self._timestep\n",
        "    agent_state = self._agent_state\n",
        "    # Unroll one longer if trajectory is empty.\n",
        "    num_interactions = unroll_length + int(not self._traj)\n",
        "    subkeys = jax.random.split(rng_key, num_interactions)\n",
        "    # data logs\n",
        "    self._best_return = 0\n",
        "    self._tracker += 1\n",
        "    # loop through unroll length collecting data\n",
        "    for i in range(num_interactions):\n",
        "      # load timestep\n",
        "      timestep = util.preprocess_step(timestep)\n",
        "      # get next state\n",
        "      agent_out, next_state = self._agent.step(subkeys[i], params, timestep,\n",
        "                                               agent_state)\n",
        "      # transtion to that state\n",
        "      transition = util.Transition(\n",
        "          timestep=timestep,\n",
        "          agent_out=agent_out,\n",
        "          agent_state=agent_state)\n",
        "      # keep track of that data\n",
        "      self._traj.append(transition)\n",
        "      # move to next state + timestep\n",
        "      agent_state = next_state\n",
        "      timestep = self._env.step(agent_out.action)\n",
        "\n",
        "      # if it's the last timestep, save rewards and logs\n",
        "      if timestep.last():\n",
        "        # save reward\n",
        "        self._episode_return += timestep.reward\n",
        "        self._num_runs = self._num_runs + 1\n",
        "        # data logs\n",
        "        self._logger.write({\n",
        "           'num_frames': frame_count,\n",
        "           'episode_return': self._episode_return,\n",
        "        })\n",
        "        # if best return from this actor, save it\n",
        "        if(self._episode_return > self._best_return):\n",
        "          self._best_return = self._episode_return\n",
        "        # set reward to 0 to start again\n",
        "        self._episode_return = 0.\n",
        "      else:\n",
        "        # add 1 for staying upright, or 0 for falling\n",
        "        self._episode_return += timestep.reward or 0.\n",
        "      \n",
        "      # keep track of rewards and times\n",
        "      self._all_returns.append(self._episode_return)\n",
        "      self._all_frames.append(frame_count)\n",
        "\n",
        "\n",
        "  def unroll_without_push(self, frame_count, params):\n",
        "    \"\"\"Run one unroll and send trajectory to learner\n",
        "    Args:\n",
        "      frame_count: which frame/time actor is in from learner\n",
        "      params: actor params from learner\n",
        "      \"\"\"\n",
        "    \n",
        "    # keep track of new params\n",
        "    params = jax.device_put(params)\n",
        "    self._rng_key, subkey = jax.random.split(self._rng_key)\n",
        "    # create 1 set of data\n",
        "    act_out = self.unroll(\n",
        "        rng_key=subkey,\n",
        "        frame_count=frame_count,\n",
        "        params=params,\n",
        "        unroll_length=self._unroll_length)\n",
        "\n",
        "  def pull_params(self):\n",
        "    \"\"\"pull the new params from the learner\"\"\"\n",
        "    return self._learner.params_for_actor()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U3FEGCeOSw08"
      },
      "source": [
        "import dm_env\n",
        "import haiku as hk\n",
        "from examples.impala import agent as agent_lib\n",
        "from examples.impala import learner as learner_lib\n",
        "from examples.impala import util\n",
        "import jax\n",
        "import numpy as np\n",
        "\n",
        "\n",
        "class Actor:\n",
        "  \"\"\"Manages the state of a single agent/environment interaction loop. This is the basic actor.\"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      agent: agent_lib.Agent,\n",
        "      env: dm_env.Environment,\n",
        "      unroll_length: int,\n",
        "      learner: learner_lib.Learner,\n",
        "      rng_seed: int = 42,\n",
        "      logger=None,\n",
        "  ):\n",
        "    #set vars\n",
        "    self._agent = agent\n",
        "    self._env = env\n",
        "    self._unroll_length = unroll_length\n",
        "    self._learner = learner\n",
        "    self._timestep = env.reset()\n",
        "    self._agent_state = agent.initial_state(None)\n",
        "    self._traj = []\n",
        "    self._rng_key = jax.random.PRNGKey(rng_seed)\n",
        "    self._episode_return = 0.\n",
        "\n",
        "    #create logger\n",
        "    if logger is None:\n",
        "      logger = util.NullLogger()\n",
        "    self._logger = logger\n",
        "\n",
        "  def unroll(self, rng_key, frame_count: int, params: hk.Params,\n",
        "             unroll_length: int) -> util.Transition:\n",
        "    \"\"\"Run unroll_length agent/environment steps, returning the trajectory\n",
        "    Args:\n",
        "      rng_key: rng_key\n",
        "      frame_count: what frame actor is on\n",
        "      params: params pulled from learner\n",
        "      unroll_length: how many data samples to collect\n",
        "    Returns\n",
        "      trajectory: data from actor\"\"\"\n",
        "    # set to current timestep and state\n",
        "    timestep = self._timestep\n",
        "    agent_state = self._agent_state\n",
        "    # Unroll one longer if trajectory is empty.\n",
        "    num_interactions = unroll_length + int(not self._traj)\n",
        "    subkeys = jax.random.split(rng_key, num_interactions)\n",
        "    # loop through number of data collections needed\n",
        "    for i in range(num_interactions):\n",
        "      # load next timestep\n",
        "      timestep = util.preprocess_step(timestep)\n",
        "      # get next state\n",
        "      agent_out, next_state = self._agent.step(subkeys[i], params, timestep,\n",
        "                                               agent_state)\n",
        "      # transition to next state\n",
        "      transition = util.Transition(\n",
        "          timestep=timestep,\n",
        "          agent_out=agent_out,\n",
        "          agent_state=agent_state)\n",
        "      # keep track of that transition info\n",
        "      self._traj.append(transition)\n",
        "      # move to next state/timestep\n",
        "      agent_state = next_state\n",
        "      timestep = self._env.step(agent_out.action)\n",
        "\n",
        "      # if last step, return total reward\n",
        "      if timestep.last():\n",
        "        # add last timestep reward\n",
        "        self._episode_return += timestep.reward\n",
        "        # log data\n",
        "        self._logger.write({\n",
        "           'num_frames': frame_count,\n",
        "           'episode_return': self._episode_return,\n",
        "        })\n",
        "        # set to 0 to begin cycle again\n",
        "        self._episode_return = 0.\n",
        "      else:\n",
        "        # add reward for current timestep: +1 for staying upright\n",
        "        self._episode_return += timestep.reward or 0.\n",
        "\n",
        "    # organize data for learner\n",
        "    trajectory = jax.device_get(self._traj)\n",
        "    trajectory = jax.tree_multimap(lambda *xs: np.stack(xs), *trajectory)\n",
        "    self._timestep = timestep\n",
        "    self._agent_state = agent_state\n",
        "    # Keep the bootstrap timestep for next trajectory.\n",
        "    self._traj = self._traj[-1:]\n",
        "    return trajectory\n",
        "\n",
        "  def unroll_and_push(self, frame_count: int, params: hk.Params):\n",
        "    \"\"\"Run one unroll and send trajectory to learner.\n",
        "    Args:\n",
        "      frame_count: current frame count from learner\n",
        "      params: current params from learner\"\"\"\n",
        "    # keep track of current params\n",
        "    params = jax.device_put(params)\n",
        "    self._rng_key, subkey = jax.random.split(self._rng_key)\n",
        "    # collect necessary data\n",
        "    act_out = self.unroll(\n",
        "        rng_key=subkey,\n",
        "        frame_count=frame_count,\n",
        "        params=params,\n",
        "        unroll_length=self._unroll_length)\n",
        "    # send learner the data\n",
        "    self._learner.enqueue_traj(act_out)\n",
        "\n",
        "  def pull_params(self):\n",
        "    \"\"\"pull new params from learner\"\"\"\n",
        "    return self._learner.params_for_actor()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "evNkFk6GeXqN"
      },
      "source": [
        "class ActorVisual:\n",
        "  \"\"\"Manages the state of a single agent/environment interaction loop and produces visual for CartPole\"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      agent: Agent,\n",
        "      env: dm_env.Environment,\n",
        "      unroll_length: int,\n",
        "      learner: Learner,\n",
        "      rng_seed: int = 42,\n",
        "      logger=None,\n",
        "  ):\n",
        "    #set necessary vars\n",
        "    self._agent = agent\n",
        "    self._env = env\n",
        "    self._unroll_length = unroll_length\n",
        "    self._learner = learner\n",
        "    self._timestep = env.reset()\n",
        "    self._agent_state = agent.initial_state(None)\n",
        "    self._traj = []\n",
        "    self._rng_key = jax.random.PRNGKey(rng_seed)\n",
        "\n",
        "    self._images = None\n",
        "    self._viewer = None\n",
        "\n",
        "    # data log\n",
        "    if logger is None:\n",
        "      logger = util.AbslLogger()\n",
        "    self._logger = logger\n",
        "\n",
        "    self._episode_return = 0.\n",
        "    self.best_return = 0\n",
        "\n",
        "  def unroll(self, rng_key, frame_count: int, params: hk.Params,\n",
        "             unroll_length: int) -> util.Transition:\n",
        "    \"\"\"Run unroll_length agent/environment steps, returning the trajectory\n",
        "    Args:\n",
        "      rng_key: rng_key\n",
        "      frame_count: pulled from learner, frame we're on\n",
        "      params: params pulled from learner\n",
        "      unroll_length: how much data to collect\n",
        "    Returns:\n",
        "      collected data\"\"\"\n",
        "\n",
        "    # start screen rendering\n",
        "    screen = self.render(mode='rgb_array')\n",
        "    im = Image.fromarray(screen)\n",
        "    temp_images = [im]\n",
        "\n",
        "    # set current time step and state\n",
        "    timestep = self._timestep\n",
        "    agent_state = self._agent_state\n",
        "    # Unroll one longer if trajectory is empty.\n",
        "    num_interactions = unroll_length + int(not self._traj)\n",
        "    subkeys = jax.random.split(rng_key, num_interactions)\n",
        "    self._best_return = 0\n",
        "    # loop through amount of data needed\n",
        "    for i in range(num_interactions):\n",
        "      # load timestep\n",
        "      timestep = util.preprocess_step(timestep)\n",
        "      #get next state\n",
        "      agent_out, next_state = self._agent.step(subkeys[i], params, timestep,\n",
        "                                               agent_state)\n",
        "      # transition to next state\n",
        "      transition = util.Transition(\n",
        "          timestep=timestep,\n",
        "          agent_out=agent_out,\n",
        "          agent_state=agent_state)\n",
        "      # keep track of transition data\n",
        "      self._traj.append(transition)\n",
        "      # move to next state + timestep\n",
        "      agent_state = next_state\n",
        "      timestep = self._env.step(agent_out.action)\n",
        "\n",
        "      # every few loops save an image for the gif\n",
        "      if i % 10 == 0:\n",
        "        screen = self.render(mode='rgb_array')\n",
        "        temp_images.append(Image.fromarray(screen))\n",
        "\n",
        "      # if last timestep, save reward and start over\n",
        "      if timestep.last():\n",
        "        # add reward for timestep\n",
        "        self._episode_return += timestep.reward\n",
        "        # log data\n",
        "        self._logger.write({\n",
        "           'num_frames': frame_count,\n",
        "           'episode_return': self._episode_return,\n",
        "        })\n",
        "        # if best return for this actor, save it and its gif\n",
        "        if (self._episode_return > self._best_return):\n",
        "          self._best_return = self._episode_return\n",
        "          self._images = temp_images\n",
        "        # restart\n",
        "        self._episode_return = 0.\n",
        "      else:\n",
        "        # add reward if it kept upright\n",
        "        self._episode_return += timestep.reward or 0.\n",
        "\n",
        "    # Pack the trajectory and reset parent state.\n",
        "    trajectory = jax.device_get(self._traj)\n",
        "    trajectory = jax.tree_multimap(lambda *xs: np.stack(xs), *trajectory)\n",
        "    self._timestep = timestep\n",
        "    self._agent_state = agent_state\n",
        "    # Keep the bootstrap timestep for next trajectory.\n",
        "    self._traj = self._traj[-1:]\n",
        "    return trajectory\n",
        "\n",
        "  def unroll_and_push(self, frame_count: int, params: hk.Params):\n",
        "    \"\"\"Run one unroll and send trajectory to learner\n",
        "    Args:\n",
        "      frame_count: frame we're on\n",
        "      params: params pulled from learner\"\"\"\n",
        "    # save params\n",
        "    params = jax.device_put(params)\n",
        "    self._rng_key, subkey = jax.random.split(self._rng_key)\n",
        "    # collect data\n",
        "    act_out = self.unroll(\n",
        "        rng_key=subkey,\n",
        "        frame_count=frame_count,\n",
        "        params=params,\n",
        "        unroll_length=self._unroll_length)\n",
        "    # send data to learner\n",
        "    self._learner.enqueue_traj(act_out)\n",
        "    self._learner.best_visual_return(self._best_return)\n",
        "\n",
        "  def pull_params(self):\n",
        "    \"\"\"pull params from learner\"\"\"\n",
        "    return self._learner.params_for_actor()\n",
        "\n",
        "  def render(self, mode='human'):\n",
        "    \"\"\"render the cartpole gif\"\"\"\n",
        "    # set image size\n",
        "      screen_width = 600\n",
        "      screen_height = 400\n",
        "\n",
        "      # draw parts of cartpole\n",
        "      world_width = self._env._x_threshold  * 2\n",
        "      scale = screen_width/world_width\n",
        "      carty = 100  # TOP OF CART\n",
        "      polewidth = 10.0\n",
        "      polelen = scale * (2 * 0.5)\n",
        "      cartwidth = 50.0\n",
        "      cartheight = 30.0\n",
        "\n",
        "      test = self._viewer\n",
        "      if self._viewer is None:\n",
        "        # add a new render if it doesn't exist-- render according to data cartpole keeps track of\n",
        "        from gym.envs.classic_control import rendering\n",
        "        self._viewer = rendering.Viewer(screen_width, screen_height)\n",
        "        l, r, t, b = -cartwidth / 2, cartwidth / 2, cartheight / 2, -cartheight / 2\n",
        "        axleoffset = cartheight / 4.0\n",
        "        cart = rendering.FilledPolygon([(l, b), (l, t), (r, t), (r, b)])\n",
        "        self._env.carttrans = rendering.Transform()\n",
        "        cart.add_attr(self._env.carttrans)\n",
        "        self._viewer.add_geom(cart)\n",
        "        l, r, t, b = -polewidth / 2, polewidth / 2, polelen - polewidth / 2, -polewidth / 2\n",
        "        pole = rendering.FilledPolygon([(l, b), (l, t), (r, t), (r, b)])\n",
        "        pole.set_color(.8, .6, .4)\n",
        "        self._env.poletrans = rendering.Transform(translation=(0, axleoffset))\n",
        "        pole.add_attr(self._env.poletrans)\n",
        "        pole.add_attr(self._env.carttrans)\n",
        "        self._viewer.add_geom(pole)\n",
        "        self._env.axle = rendering.make_circle(polewidth/2)\n",
        "        self._env.axle.add_attr(self._env.poletrans)\n",
        "        self._env.axle.add_attr(self._env.carttrans)\n",
        "        self._env.axle.set_color(.5, .5, .8)\n",
        "        self._viewer.add_geom(self._env.axle)\n",
        "        self._env.track = rendering.Line((0, carty), (screen_width, carty))\n",
        "        self._env.track.set_color(0, 0, 0)\n",
        "        self._viewer.add_geom(self._env.track)\n",
        "\n",
        "        self._env._pole_geom = pole\n",
        "\n",
        "      if self._env._state is None:\n",
        "          return None\n",
        "\n",
        "      # Edit the pole polygon vertex\n",
        "      pole = self._env._pole_geom\n",
        "      l, r, t, b = -polewidth / 2, polewidth / 2, polelen - polewidth / 2, -polewidth / 2\n",
        "      pole.v = [(l, b), (l, t), (r, t), (r, b)]\n",
        "\n",
        "      x = self._env._state\n",
        "      cartx = x[0] * scale + screen_width / 2.0  # MIDDLE OF CART\n",
        "      self._env.carttrans.set_translation(cartx, carty)\n",
        "      self._env.poletrans.set_rotation(-x[2])\n",
        "\n",
        "      return self._viewer.render(return_rgb_array=mode == 'rgb_array')\n",
        "\n",
        "  def close(self):\n",
        "    \"\"\"close render\"\"\"\n",
        "      if self._viewer:\n",
        "          self._viewer.close()\n",
        "          self._viewer = Non\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WFPfK8kYRlKw"
      },
      "source": [
        "class ActorGrapher:\n",
        "  \"\"\"Manages the state of a single agent/environment interaction loop.\"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      agent: Agent,\n",
        "      env: dm_env.Environment,\n",
        "      unroll_length: int,\n",
        "      learner: Learner,\n",
        "      rng_seed: int = 42,\n",
        "      logger=None,\n",
        "  ):\n",
        "    # create necessary vars\n",
        "    self._agent = agent\n",
        "    self._env = env\n",
        "    self._unroll_length = unroll_length\n",
        "    self._learner = learner\n",
        "    self._timestep = env.reset()\n",
        "    self._agent_state = agent.initial_state(None)\n",
        "    self._traj = []\n",
        "    self._rng_key = jax.random.PRNGKey(rng_seed)\n",
        "    self._best_return = 0 \n",
        "    self._all_returns = []\n",
        "    self._all_frames = []\n",
        "    self._num_runs = 0\n",
        "\n",
        "    # create logger\n",
        "    if logger is None:\n",
        "      logger = util.AbslLogger()\n",
        "    self._logger = logger\n",
        "\n",
        "    self._episode_return = 0.\n",
        "\n",
        "  def unroll(self, rng_key, frame_count: int, params: hk.Params,\n",
        "             unroll_length: int) -> util.Transition:\n",
        "    \"\"\"Run unroll_length agent/environment steps, returning the trajectory.\"\"\"\n",
        "    self._num_runs = 0\n",
        "    timestep = self._timestep\n",
        "    agent_state = self._agent_state\n",
        "    # Unroll one longer if trajectory is empty.\n",
        "    num_interactions = unroll_length + int(not self._traj)\n",
        "    subkeys = jax.random.split(rng_key, num_interactions)\n",
        "    self._best_return = 0\n",
        "    # run through and collect data\n",
        "    for i in range(num_interactions):\n",
        "      # load timestep\n",
        "      timestep = util.preprocess_step(timestep)\n",
        "      # figure out next state\n",
        "      agent_out, next_state = self._agent.step(subkeys[i], params, timestep,\n",
        "                                               agent_state)\n",
        "      # transition to next state\n",
        "      transition = util.Transition(\n",
        "          timestep=timestep,\n",
        "          agent_out=agent_out,\n",
        "          agent_state=agent_state)\n",
        "      self._traj.append(transition)\n",
        "      agent_state = next_state\n",
        "      timestep = self._env.step(agent_out.action)\n",
        "\n",
        "      # if last timestep, save episode return, if not continue on\n",
        "      if timestep.last():\n",
        "        self._episode_return += timestep.reward\n",
        "        self._all_returns.append(self._episode_return)\n",
        "        self._all_frames.append(frame_count)\n",
        "        self._num_runs = self._num_runs + 1\n",
        "\n",
        "        self._logger.write({\n",
        "           'num_frames': frame_count,\n",
        "           'episode_return': self._episode_return,\n",
        "        })\n",
        "        if(self._episode_return > self._best_return):\n",
        "          self._best_return = self._episode_return\n",
        "        self._episode_return = 0.\n",
        "      else:\n",
        "        self._episode_return += timestep.reward or 0.\n",
        "\n",
        "    # Pack the trajectory and reset parent state.\n",
        "    trajectory = jax.device_get(self._traj)\n",
        "    trajectory = jax.tree_multimap(lambda *xs: np.stack(xs), *trajectory)\n",
        "    self._timestep = timestep\n",
        "    self._agent_state = agent_state\n",
        "    # Keep the bootstrap timestep for next trajectory.\n",
        "    self._traj = self._traj[-1:]\n",
        "    return trajectory\n",
        "\n",
        "  def unroll_and_push(self, frame_count: int, params: hk.Params):\n",
        "    \"\"\"Run one unroll and send trajectory to learner.\"\"\"\n",
        "    params = jax.device_put(params)\n",
        "    self._rng_key, subkey = jax.random.split(self._rng_key)\n",
        "    # collect needed amount of data\n",
        "    act_out = self.unroll(\n",
        "        rng_key=subkey,\n",
        "        frame_count=frame_count,\n",
        "        params=params,\n",
        "        unroll_length=self._unroll_length)\n",
        "    self._learner.enqueue_traj(act_out)\n",
        "\n",
        "  def pull_params(self):\n",
        "    # pull data from learner\n",
        "    return self._learner.params_for_actor()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BoWUiZ6D1Om4"
      },
      "source": [
        "# Common Networks"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hc0UfOmgGLON"
      },
      "source": [
        "# keep track of nn variables\n",
        "NetOutput = collections.namedtuple('NetOutput', ['policy_logits', 'value'])\n",
        "\n",
        "class CatchNet(hk.RNNCore):\n",
        "  \"\"\"The easiest IMPALA nn\"\"\"\n",
        "\n",
        "  def __init__(self, num_actions, name=None):\n",
        "    \"\"\"basic init function\n",
        "    Args:\n",
        "    num_actions: number of actions that can be taken in an env\"\"\"\n",
        "    super(CatchNet, self).__init__(name=name)\n",
        "    self._num_actions = num_actions\n",
        "\n",
        "  def initial_state(self, batch_size):\n",
        "    \"\"\"set up for nn, just the shape\n",
        "    Args:\n",
        "      batch_size: learner batch size\"\"\"\n",
        "    if batch_size is None:\n",
        "      shape = []\n",
        "    else:\n",
        "      shape = [batch_size]\n",
        "    return jnp.zeros(shape)  \n",
        "\n",
        "  def __call__(self, x: dm_env.TimeStep, state):\n",
        "    \"\"\"one loop of the NN\n",
        "    Args:\n",
        "    x: timestep class for current time\n",
        "    state: the current state\"\"\"\n",
        "    # feed through linear function\n",
        "    torso_net = hk.Sequential(\n",
        "        [hk.Flatten(),\n",
        "         hk.Linear(128), jax.nn.relu,\n",
        "         hk.Linear(64), jax.nn.relu])\n",
        "    torso_output = torso_net(x.observation)\n",
        "    policy_logits = hk.Linear(self._num_actions)(torso_output)\n",
        "    value = hk.Linear(1)(torso_output)\n",
        "    value = jnp.squeeze(value, axis=-1)\n",
        "    # output values\n",
        "    return NetOutput(policy_logits=policy_logits, value=value), state\n",
        "\n",
        "  def unroll(self, x, state):\n",
        "    \"\"\"Apply nn to all data\n",
        "    Args:\n",
        "      x: timesteps\n",
        "      state: states\"\"\"\n",
        "    out, _ = hk.BatchApply(self)(x, None)\n",
        "    return out, state\n",
        "\n",
        "\n",
        "class AtariShallowTorso(hk.Module):\n",
        "  \"\"\"Shallow torso for Atari, from the DQN paper.\"\"\"\n",
        "\n",
        "  def __init__(self, name=None):\n",
        "    super(AtariShallowTorso, self).__init__(name=name)\n",
        "\n",
        "  def __call__(self, x):\n",
        "    # put data through nn\n",
        "    torso_net = hk.Sequential([\n",
        "        lambda x: x / 255.,\n",
        "        hk.Conv2D(32, kernel_shape=[8, 8], stride=[4, 4], padding='VALID'),\n",
        "        jax.nn.relu,\n",
        "        hk.Conv2D(64, kernel_shape=[4, 4], stride=[2, 2], padding='VALID'),\n",
        "        jax.nn.relu,\n",
        "        hk.Conv2D(64, kernel_shape=[3, 3], stride=[1, 1], padding='VALID'),\n",
        "        jax.nn.relu,\n",
        "        hk.Flatten(),\n",
        "        hk.Linear(512),\n",
        "        jax.nn.relu,\n",
        "    ])\n",
        "    return torso_net(x)\n",
        "\n",
        "\n",
        "class ResidualBlock(hk.Module):\n",
        "  \"\"\"Residual block.\"\"\"\n",
        "\n",
        "  def __init__(self, num_channels, name=None):\n",
        "    super(ResidualBlock, self).__init__(name=name)\n",
        "    self._num_channels = num_channels\n",
        "\n",
        "  def __call__(self, x):\n",
        "    # feed through Residual Block\n",
        "    main_branch = hk.Sequential([\n",
        "        jax.nn.relu,\n",
        "        hk.Conv2D(\n",
        "            self._num_channels,\n",
        "            kernel_shape=[3, 3],\n",
        "            stride=[1, 1],\n",
        "            padding='SAME'),\n",
        "        jax.nn.relu,\n",
        "        hk.Conv2D(\n",
        "            self._num_channels,\n",
        "            kernel_shape=[3, 3],\n",
        "            stride=[1, 1],\n",
        "            padding='SAME'),\n",
        "    ])\n",
        "    return main_branch(x) + x\n",
        "\n",
        "\n",
        "class AtariDeepTorso(hk.Module):\n",
        "  \"\"\"Deep torso for Atari, from the IMPALA paper.\"\"\"\n",
        "\n",
        "  def __init__(self, name=None):\n",
        "    super(AtariDeepTorso, self).__init__(name=name)\n",
        "\n",
        "  def __call__(self, x):\n",
        "    # go through convs, res. blocks and relus\n",
        "    torso_out = x / 255.\n",
        "    for i, (num_channels, num_blocks) in enumerate([(16, 2), (32, 2), (32, 2)]):\n",
        "      conv = hk.Conv2D(\n",
        "          num_channels, kernel_shape=[3, 3], stride=[1, 1], padding='SAME')\n",
        "      torso_out = conv(torso_out)\n",
        "      torso_out = hk.max_pool(\n",
        "          torso_out,\n",
        "          window_shape=[1, 3, 1],\n",
        "          strides=[1, 2, 1],\n",
        "          padding='SAME',\n",
        "      )\n",
        "      for j in range(num_blocks):\n",
        "        block = ResidualBlock(num_channels, name='residual_{}_{}'.format(i, j))\n",
        "        torso_out = block(torso_out)\n",
        "\n",
        "    torso_out = jax.nn.relu(torso_out)\n",
        "    torso_out = hk.Flatten()(torso_out)\n",
        "    torso_out = hk.Linear(256)(torso_out)\n",
        "    torso_out = jax.nn.relu(torso_out)\n",
        "    return torso_out\n",
        "\n",
        "\n",
        "class AtariNet(hk.RNNCore):\n",
        "  \"\"\"Network for Atari.\"\"\"\n",
        "\n",
        "  def __init__(self, num_actions, use_resnet, use_lstm, name=None):\n",
        "    super(AtariNet, self).__init__(name=name)\n",
        "    self._num_actions = num_actions\n",
        "    self._use_resnet = use_resnet\n",
        "    self._use_lstm = use_lstm\n",
        "    self._core = hk.ResetCore(hk.LSTM(256))\n",
        "\n",
        "  def initial_state(self, batch_size):\n",
        "    return self._core.initial_state(batch_size)\n",
        "\n",
        "  def __call__(self, x: dm_env.TimeStep, state):\n",
        "    x = jax.tree_map(lambda t: t[None, ...], x)\n",
        "    return self.unroll(x, state)\n",
        "\n",
        "  def unroll(self, x, state):\n",
        "    \"\"\"Unrolls more efficiently than dynamic_unroll.\"\"\"\n",
        "    if self._use_resnet:\n",
        "      torso = AtariDeepTorso()\n",
        "    else:\n",
        "      torso = AtariShallowTorso()\n",
        "\n",
        "    torso_output = hk.BatchApply(torso)(x.observation)\n",
        "    if self._use_lstm:\n",
        "      should_reset = jnp.equal(x.step_type, int(dm_env.StepType.FIRST))\n",
        "      core_input = (torso_output, should_reset)\n",
        "      core_output, state = hk.dynamic_unroll(self._core, core_input, state)\n",
        "    else:\n",
        "      core_output = torso_output\n",
        "      # state passes through.\n",
        "\n",
        "    return hk.BatchApply(self._head)(core_output), state\n",
        "\n",
        "  def _head(self, core_output):\n",
        "    policy_logits = hk.Linear(self._num_actions)(core_output)\n",
        "    value = hk.Linear(1)(core_output)\n",
        "    value = jnp.squeeze(value, axis=-1)\n",
        "    return NetOutput(policy_logits=policy_logits, value=value)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fFKze95f1cq6"
      },
      "source": [
        "# Main"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "background_save": true,
          "base_uri": "https://localhost:8080/"
        },
        "id": "P5JfZgLCBBVz",
        "outputId": "8fd55dce-0f9b-4367-b32a-4d1a380c5d57"
      },
      "source": [
        "#set necessary vars\n",
        "ACTION_REPEAT = 4\n",
        "BATCH_SIZE = 32\n",
        "DISCOUNT_FACTOR = 0.99\n",
        "NUM_ACTORS = 3\n",
        "UNROLL_LENGTH = 20\n",
        "FRAMES_PER_ITER = ACTION_REPEAT * BATCH_SIZE * UNROLL_LENGTH\n",
        "MAX_ENV_FRAMES = FRAMES_PER_ITER * 400\n",
        "MAX_TIME = 10\n",
        "\n",
        "\n",
        "\n",
        "def run_actor(actor, stop_signal: List[bool]):\n",
        "  \"\"\"Runs an actor to produce num_trajectories trajectories\n",
        "  Args:\n",
        "    actor: an actor to collect samples\n",
        "    stop_signal: whether the actor should keep running\"\"\"\n",
        "  # actor collects data until told otherwise\n",
        "  while not stop_signal[0]:\n",
        "    # actor updates its frame count and params from learner\n",
        "    frame_count, params = actor.pull_params()\n",
        "    # actor collects data and sends back to learner\n",
        "    actor.unroll_and_push(frame_count, params)\n",
        "\n",
        "# build cartpole environment\n",
        "build_env = cartpole.Cartpole\n",
        "\n",
        "# construct agent\n",
        "env_for_spec = build_env(max_time = MAX_TIME)\n",
        "num_actions = env_for_spec.action_spec().num_values\n",
        "agent = Agent(num_actions, env_for_spec.observation_spec(),\n",
        "                          CatchNet)\n",
        "\n",
        "# Calculate number of updates learner will do\n",
        "max_updates = MAX_ENV_FRAMES / FRAMES_PER_ITER\n",
        "print(\"Running \", max_updates, \" iterations for learner\")\n",
        "# Construct optimizer\n",
        "opt = optax.rmsprop(5e-3, decay=0.99, eps=1e-7)\n",
        "\n",
        "# create actor to keep track of learner changes\n",
        "learnersActor = SeparateActor(\n",
        "      agent,\n",
        "      build_env(max_time = MAX_TIME),\n",
        "      UNROLL_LENGTH,\n",
        "      rng_seed=0,\n",
        "      logger=util.AbslLogger(),  \n",
        "  )\n",
        "\n",
        "# Construct the learner\n",
        "print(\"making learner now...\")\n",
        "learner = Learner(\n",
        "    agent,\n",
        "    jax.random.PRNGKey(428),\n",
        "    opt,\n",
        "    BATCH_SIZE,\n",
        "    DISCOUNT_FACTOR,\n",
        "    FRAMES_PER_ITER,\n",
        "    learnersActor,\n",
        "    max_abs_reward=1.,\n",
        "    logger=util.AbslLogger(),  # Provide your own logger here.\n",
        ")\n",
        "\n",
        "# Construct the actors on different threads\n",
        "print(\"making actors now...\")\n",
        "actor_threads = []\n",
        "\n",
        "stop_signal = [False]\n",
        "\n",
        "# create the actor with image capabilities\n",
        "actorVisual = ActorGrapher(\n",
        "      agent,\n",
        "      build_env(max_time = MAX_TIME),\n",
        "      UNROLL_LENGTH,\n",
        "      learner,\n",
        "      rng_seed=0,\n",
        "      logger=util.AbslLogger(),  \n",
        "  )\n",
        "args = (actorVisual, stop_signal)\n",
        "actor_threads.append(threading.Thread(target=run_actor, args=args))\n",
        "\n",
        "# create the regular actors\n",
        "for i in range(1, NUM_ACTORS):\n",
        "  actor = Actor(\n",
        "      agent,\n",
        "      build_env(max_time = MAX_TIME),\n",
        "      UNROLL_LENGTH,\n",
        "      learner,\n",
        "      rng_seed=i,\n",
        "      logger=util.AbslLogger(),  # Provide your own logger here.\n",
        "  )\n",
        "  args = (actor, stop_signal)\n",
        "  actor_threads.append(threading.Thread(target=run_actor, args=args))\n",
        "\n",
        "# Start the actors and learner\n",
        "print(\"starting the learning!\")\n",
        "for t in actor_threads:\n",
        "  t.start()\n",
        "learner.run(int(max_updates))\n",
        "\n",
        "\n",
        "  # Stop.\n",
        "stop_signal[0] = True\n",
        "for t in actor_threads:\n",
        " t.join()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Running  400.0  iterations for learner\n",
            "making learner now...\n",
            "making actors now...\n",
            "starting the learning!\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 400/400 [1:03:19<00:00,  9.50s/it]\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "w0ePL9e8cvHu",
        "outputId": "4d18d605-9250-4667-a84e-e39d4d3939bf"
      },
      "source": [
        "# initialize necessary vars\n",
        "current_frame = 0\n",
        "current_num = 0.0\n",
        "\n",
        "all_averages = []\n",
        "all_frames = []\n",
        "\n",
        "# whenever the tracking actor restarts an episode, log its episode return\n",
        "for frame_num in range(len(learnersActor._all_frames)):\n",
        "  if (learnersActor._all_returns[frame_num] == 0.0):\n",
        "    current_num = 0.0\n",
        "    if (learnersActor._all_returns[frame_num-1] != 0.0):\n",
        "      # if episode just ended, look back 1 step to see highest value\n",
        "      all_averages.append(learnersActor._all_returns[frame_num-1])\n",
        "      all_frames.append(learnersActor._all_frames[frame_num-1])\n",
        "  else:\n",
        "    current_num +=1\n",
        "\n",
        "# print episode returns and corresponding frame numbers\n",
        "print(\"Episode Returns: \")\n",
        "print(all_averages)\n",
        "print(\"Frame numbers: \")\n",
        "print(all_frames)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Averages: \n",
            "[72.0, 27.0, 30.0, 29.0, 30.0, 29.0, 27.0, 29.0, 27.0, 28.0, 28.0, 29.0, 43.0, 30.0, 28.0, 27.0, 30.0, 29.0, 28.0, 30.0, 29.0, 29.0, 29.0, 27.0, 29.0, 29.0, 27.0, 28.0, 30.0, 29.0, 28.0, 28.0, 29.0, 29.0, 28.0, 29.0, 29.0, 30.0, 29.0, 27.0, 30.0, 29.0, 28.0, 28.0, 29.0, 28.0, 29.0, 30.0, 28.0, 31.0, 31.0, 32.0, 34.0, 65.0, 28.0, 30.0, 28.0, 27.0, 28.0, 28.0, 30.0, 28.0, 30.0, 30.0, 29.0, 30.0, 28.0, 28.0, 30.0, 29.0, 30.0, 28.0, 28.0, 30.0, 29.0, 33.0, 84.0, 44.0, 72.0, 118.0, 76.0, 56.0, 61.0, 74.0, 57.0, 66.0, 113.0, 149.0, 156.0, 171.0, 172.0, 192.0, 148.0, 150.0, 185.0, 351.0, 245.0, 211.0, 217.0, 234.0, 447.0, 270.0, 309.0, 279.0, 403.0, 328.0]\n",
            "Frame numbers: \n",
            "[7680, 10240, 15360, 20480, 23040, 28160, 30720, 35840, 38400, 43520, 46080, 51200, 56320, 61440, 64000, 69120, 71680, 76800, 79360, 84480, 87040, 92160, 94720, 99840, 102400, 107520, 110080, 115200, 120320, 122880, 128000, 130560, 135680, 138240, 143360, 145920, 151040, 153600, 158720, 161280, 166400, 171520, 174080, 179200, 181760, 186880, 189440, 194560, 197120, 202240, 207360, 209920, 215040, 222720, 227840, 232960, 235520, 238080, 243200, 245760, 250880, 256000, 258560, 263680, 266240, 271360, 273920, 279040, 284160, 286720, 291840, 294400, 299520, 302080, 307200, 312320, 322560, 327680, 337920, 353280, 363520, 371200, 378880, 389120, 396800, 404480, 419840, 437760, 458240, 481280, 504320, 527360, 547840, 565760, 591360, 634880, 668160, 693760, 721920, 752640, 811520, 844800, 885760, 921600, 972800, 1013760]\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}