diff --git a/.github/workflows/build-docker.yaml b/.github/workflows/build-docker.yaml
deleted file mode 100644
index 5c5504bdb..000000000
--- a/.github/workflows/build-docker.yaml
+++ /dev/null
@@ -1,46 +0,0 @@
-name: build and push Docker image
-
-on:
-  push:
-    tags:
-      - "release_[0-9]+"
-
-jobs:
-  build:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v2
-      - uses: google-github-actions/setup-gcloud@v0
-        with:
-          version: '270.0.0'
-          service_account_email: ${{ secrets.GCP_STG_EMAIL }}
-          service_account_key: ${{ secrets.GCP_STG_PRIVATE_KEY }}
-          export_default_credentials: true
-      - name: Configure GCR Docker auth (stg)
-        run: gcloud auth configure-docker
-      - name: Build docker image
-        run: docker build --tag ml-agents:${GITHUB_SHA} --build-arg SHA=${GITHUB_SHA} .
-      - name: Tag docker image (stg)
-        run: |
-          docker tag ml-agents:${GITHUB_SHA} gcr.io/unity-ai-ml-agents-stg/ml-agents:${GITHUB_SHA}
-          docker tag ml-agents:${GITHUB_SHA} gcr.io/unity-ai-ml-agents-stg/ml-agents:latest
-      - name: Push docker image (stg)
-        run: |
-          docker push gcr.io/unity-ai-ml-agents-stg/ml-agents:${GITHUB_SHA}
-          docker push gcr.io/unity-ai-ml-agents-stg/ml-agents:latest
-      - uses: google-github-actions/setup-gcloud@v0
-        with:
-          version: '270.0.0'
-          service_account_email: ${{ secrets.GCP_PRD_EMAIL }}
-          service_account_key: ${{ secrets.GCP_PRD_PRIVATE_KEY }}
-          export_default_credentials: true
-      - name: Configure GCR Docker auth (prd)
-        run: gcloud auth configure-docker
-      - name: Tag docker image (prd)
-        run: |
-          docker tag ml-agents:${GITHUB_SHA} gcr.io/unity-ai-ml-agents-prd/ml-agents:${GITHUB_SHA}
-          docker tag ml-agents:${GITHUB_SHA} gcr.io/unity-ai-ml-agents-prd/ml-agents:latest
-      - name: Push docker image (prd)
-        run: |
-          docker push gcr.io/unity-ai-ml-agents-prd/ml-agents:${GITHUB_SHA}
-          docker push gcr.io/unity-ai-ml-agents-prd/ml-agents:latest
diff --git a/.github/workflows/colab.yml b/.github/workflows/colab.yml
index de1547f95..fc4850e3e 100644
--- a/.github/workflows/colab.yml
+++ b/.github/workflows/colab.yml
@@ -9,16 +9,20 @@ on:
     - '.github/workflows/colab.yml'
   push:
     branches: [main]
+  workflow_dispatch:
 
 jobs:
   colab:
-    runs-on: ubuntu-latest
+    # Only run on GitHub Enterprise
+    if: ${{ github.server_url == 'https://github.cds.internal.unity3d.com' }}
+    runs-on: [self-hosted, Linux, X64]
     env:
         COLAB_ALWAYS_INSTALL_XVFB: 1
         QLEARNING_NUM_TRAINING_STEPS: 5
         QLEARNING_NUM_NEW_EXP: 64
         QLEARNING_BUFFER_SIZE: 64
     strategy:
+      fail-fast: false
       matrix:
         notebook_path: [Colab_UnityEnvironment_1_Run.ipynb, Colab_UnityEnvironment_2_Train.ipynb, Colab_UnityEnvironment_3_SideChannel.ipynb]
     steps:
@@ -29,17 +33,19 @@ jobs:
         python-version: 3.8.x
     - uses: actions/setup-dotnet@v1
       with:
-        dotnet-version: '3.1.x'
-    - name: Cache pip
-      uses: actions/cache@v2
-      with:
-        # This path is specific to Ubuntu
-        path: ~/.cache/pip
-        # Look to see if there is a cache hit for the corresponding requirements file
-        key: ${{ runner.os }}-pip-${{ hashFiles('ml-agents/setup.py', 'ml-agents-envs/setup.py', 'gym-unity/setup.py', 'colab_requirements.txt') }}
-        restore-keys: |
-          ${{ runner.os }}-pip-
-          ${{ runner.os }}-
+        dotnet-version: '6.0.x'
+# Caching not supported on GitHub Enterprise
+# See https://github.com/actions/cache/issues/505
+#    - name: Cache pip
+#      uses: actions/cache@v2
+#      with:
+#        # This path is specific to Ubuntu
+#        path: ~/.cache/pip
+#        # Look to see if there is a cache hit for the corresponding requirements file
+#        key: ${{ runner.os }}-pip-${{ hashFiles('ml-agents/setup.py', 'ml-agents-envs/setup.py', 'gym-unity/setup.py', 'colab_requirements.txt') }}
+#        restore-keys: |
+#          ${{ runner.os }}-pip-
+#          ${{ runner.os }}-
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
@@ -49,7 +55,7 @@ jobs:
         python -m pip install --progress-bar=off -e ./ml-agents
         python -m pip install --progress-bar=off -r colab_requirements.txt
     - name: Execute notebook
-      run: jupyter nbconvert --to notebook --execute --log-level=INFO --ExecutePreprocessor.kernel_name=python3 --output output-${{ matrix.notebook_path }} colab/${{ matrix.notebook_path }}
+      run: jupyter nbconvert --to notebook --execute --log-level=DEBUG --ExecutePreprocessor.kernel_name=python3 --output output-${{ matrix.notebook_path }} colab/${{ matrix.notebook_path }}
     - name: Upload colab results
       uses: actions/upload-artifact@v2
       with:
diff --git a/.github/workflows/lock.yml b/.github/workflows/lock.yml
index bae04df51..19f11f59c 100644
--- a/.github/workflows/lock.yml
+++ b/.github/workflows/lock.yml
@@ -3,12 +3,15 @@ name: 'Lock Threads'
 on:
   schedule:
     - cron: '0 0/4 * * *'
+  workflow_dispatch:
 
 jobs:
   lock:
     runs-on: ubuntu-latest
     steps:
       - uses: dessant/lock-threads@v2
+        # Only lock threads on public GitHub mirror, not enterprise.
+        if: ${{ github.server_url == 'https://github.com' }}
         with:
           github-token: ${{ github.token }}
           issue-lock-inactive-days: '30'
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index 8d1d80edd..134ea6183 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -3,17 +3,24 @@ name: nightly
 on:
   schedule:
     - cron: '0 7 * * *' # run at 7 AM UTC == midnight PST
+  workflow_dispatch:
 
 jobs:
   markdown-link-check-full:
-    runs-on: ubuntu-latest
+    # Only run on GitHub Enterprise
+    if: ${{ github.server_url == 'https://github.cds.internal.unity3d.com' }}
+    runs-on: [self-hosted, Linux, X64]
     steps:
     - uses: actions/checkout@v2
-    - uses: actions/setup-python@v1
-    - uses: actions/setup-node@v2-beta
+    - uses: actions/setup-python@v2
+    - uses: actions/setup-node@v2
       with:
         node-version: '12'
-    - run: sudo npm install -g markdown-link-check
-    - uses: pre-commit/action@v2.0.0
-      with:
-          extra_args: --hook-stage manual markdown-link-check-full --all-files
+    - name: Install manual dependencies
+      run: |
+        sudo npm install -g markdown-link-check
+        python -m pip install pre-commit
+        pre-commit install
+    - name: Run markdown checker
+      run: |
+        pre-commit run --hook-stage manual markdown-link-check-full --all-files
diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
index ebd636aca..6f0f29250 100644
--- a/.github/workflows/pre-commit.yml
+++ b/.github/workflows/pre-commit.yml
@@ -4,40 +4,56 @@ on:
   pull_request:
   push:
     branches: [main]
+  workflow_dispatch:
 
 jobs:
   pre-commit:
-    runs-on: ubuntu-latest
+    if: ${{ github.server_url == 'https://github.cds.internal.unity3d.com' }}
+    runs-on: [self-hosted, Linux, X64]
     steps:
     - uses: actions/checkout@v2
     - uses: actions/setup-python@v2
       with:
         python-version: 3.7.x
     - uses: actions/setup-ruby@v1
+      env:
+        ImageOS: ubuntu20
       with:
         ruby-version: '2.6'
     - uses: actions/setup-dotnet@v1
       with:
         dotnet-version: '6.0.x'
         include-prerelease: true
-    - uses: pre-commit/action@v2.0.0
-
+    - name: Install manual dependencies
+      run: |
+        python -m pip install pre-commit
+        pre-commit install
+    - name: Run pre-commit
+      run: |
+        pre-commit run --all-files
   markdown-link-check:
-    runs-on: ubuntu-latest
+    # Only run on GitHub Enterprise
+    if: ${{ github.server_url == 'https://github.cds.internal.unity3d.com' }}
+    runs-on: [self-hosted, Linux, X64]
     steps:
     - uses: actions/checkout@v2
-    - uses: actions/setup-python@v1
-    - uses: actions/setup-node@v2-beta
+    - uses: actions/setup-python@v2
+    - uses: actions/setup-node@v2
       with:
         node-version: '12'
-    - run: sudo npm install -g markdown-link-check
-    - uses: pre-commit/action@v2.0.0
-      with:
-          extra_args: --hook-stage manual markdown-link-check --all-files
-
+    - name: Install manual dependencies
+      run: |
+        sudo npm install -g markdown-link-check
+        python -m pip install pre-commit
+        pre-commit install
+    - name: Run markdown checker
+      run: |
+        pre-commit run --hook-stage manual markdown-link-check --all-files
   validate-meta-files:
-    runs-on: ubuntu-latest
+    # Only run on GitHub Enterprise
+    if: ${{ github.server_url == 'https://github.cds.internal.unity3d.com' }}
+    runs-on: [self-hosted, Linux, X64]
     steps:
     - uses: actions/checkout@v2
-    - uses: actions/setup-python@v1
+    - uses: actions/setup-python@v2
     - run: python utils/validate_meta_files.py
diff --git a/.github/workflows/publish_pypi.yaml b/.github/workflows/publish_pypi.yaml
index 9bcbbee4b..e39b4e243 100644
--- a/.github/workflows/publish_pypi.yaml
+++ b/.github/workflows/publish_pypi.yaml
@@ -1,16 +1,19 @@
 # Adapted from https://packaging.python.org/guides/publishing-package-distribution-releases-using-github-actions-ci-cd-workflows/
-name: publish to PyPI (or TestPyPI)
+name: Publish to PyPI (or TestPyPI)
 
 on:
   push:
     tags:
       - "release_[0-9]+_test[0-9]+"
       - "release_[0-9]+"
+  workflow_dispatch:
 
 jobs:
   build-and-publish:
+    # Only run on GitHub Enterprise
+    if: ${{ github.server_url == 'https://github.cds.internal.unity3d.com' }}
     name: publish to PyPI (or TestPyPI)
-    runs-on: ubuntu-18.04
+    runs-on: [self-hosted, Linux, X64]
     strategy:
         matrix:
             package-path: [ml-agents, ml-agents-envs, gym-unity]
@@ -34,14 +37,14 @@ jobs:
         python setup.py bdist_wheel
     - name: Publish distribution 📦 to Test PyPI
       if: startsWith(github.ref, 'refs/tags') && contains(github.ref, 'test')
-      uses: pypa/gh-action-pypi-publish@master
+      uses: actions/gh-action-pypi-publish@717ba43cfbb0387f6ce311b169a825772f54d295
       with:
-        password: ${{ secrets.test_pypi_password }}
+        password: ${{ secrets.TEST_PYPI_PASSWORD }}
         repository_url: https://test.pypi.org/legacy/
         packages_dir: ${{ matrix.package-path }}/dist/
     - name: Publish distribution 📦 to Production PyPI
       if: startsWith(github.ref, 'refs/tags') && !contains(github.ref, 'test')
-      uses: pypa/gh-action-pypi-publish@master
+      uses: actions/gh-action-pypi-publish@717ba43cfbb0387f6ce311b169a825772f54d295
       with:
-        password: ${{ secrets.pypi_password }}
+        password: ${{ secrets.PYPI_PASSWORD }}
         packages_dir: ${{ matrix.package-path }}/dist/
diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index ec43218be..cae1bdfde 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -11,15 +11,22 @@ on:
     - '.github/workflows/pytest.yml'
   push:
     branches: [main]
+  workflow_dispatch:
 
 jobs:
   pytest:
-    runs-on: ubuntu-latest
+    # Only run on GitHub Enterprise
+    if: ${{ github.server_url == 'https://github.cds.internal.unity3d.com' }}
+    runs-on: [self-hosted, Linux, X64]
     env:
         TEST_ENFORCE_BUFFER_KEY_TYPES: 1
     strategy:
+      # If one test in the matrix fails we still want to run the others.
+      fail-fast: false
       matrix:
-        python-version: [3.6.x, 3.7.x, 3.8.x, 3.9.x]
+        # NOTE: 3.9.10 was failing but is hard to test since it JUST came out. Delaying debugging for now.
+        # TODO: Set back to 3.9.x once we can test with 3.9.10 via pyenv.
+        python-version: [3.6.x, 3.7.x, 3.8.x, 3.9.9]
         include:
           - python-version: 3.6.x
             pip_constraints: test_constraints_min_version.txt
@@ -27,7 +34,7 @@ jobs:
             pip_constraints: test_constraints_mid_version.txt
           - python-version: 3.8.x
             pip_constraints: test_constraints_mid_version.txt
-          - python-version: 3.9.x
+          - python-version: 3.9.9
             pip_constraints: test_constraints_max_version.txt
     steps:
     - uses: actions/checkout@v2
@@ -35,16 +42,18 @@ jobs:
       uses: actions/setup-python@v2
       with:
         python-version: ${{ matrix.python-version }}
-    - name: Cache pip
-      uses: actions/cache@v2
-      with:
-        # This path is specific to Ubuntu
-        path: ~/.cache/pip
-        # Look to see if there is a cache hit for the corresponding requirements file
-        key: ${{ runner.os }}-pip-${{ hashFiles('ml-agents/setup.py', 'ml-agents-envs/setup.py', 'gym-unity/setup.py', 'test_requirements.txt', matrix.pip_constraints) }}
-        restore-keys: |
-          ${{ runner.os }}-pip-
-          ${{ runner.os }}-
+# Caching not supported on GitHub Enterprise
+# See https://github.com/actions/cache/issues/505
+#    - name: Cache pip
+#      uses: actions/cache@v2
+#      with:
+#        # This path is specific to Ubuntu
+#        path: ~/.cache/pip
+#        # Look to see if there is a cache hit for the corresponding requirements file
+#        key: ${{ runner.os }}-pip-${{ hashFiles('ml-agents/setup.py', 'ml-agents-envs/setup.py', 'gym-unity/setup.py', 'test_requirements.txt', matrix.pip_constraints) }}
+#        restore-keys: |
+#          ${{ runner.os }}-pip-
+#          ${{ runner.os }}-
     - name: Display Python version
       run: python -c "import sys; print(sys.version)"
     - name: Install dependencies
diff --git a/colab/Colab_UnityEnvironment_1_Run.ipynb b/colab/Colab_UnityEnvironment_1_Run.ipynb
index 70f4abfc0..f27fc3ff5 100644
--- a/colab/Colab_UnityEnvironment_1_Run.ipynb
+++ b/colab/Colab_UnityEnvironment_1_Run.ipynb
@@ -103,18 +103,17 @@
     "esac\n",
     "exit 0\n",
     "    \"\"\")\n",
-    "  pro_bar.update(progress(5, 100))\n",
-    "  !apt-get install daemon >/dev/null 2>&1\n",
+    "  !sudo apt-get update\n",
     "  pro_bar.update(progress(10, 100))\n",
-    "  !apt-get install wget >/dev/null 2>&1\n",
+    "  !sudo DEBIAN_FRONTEND=noninteractive apt install -y daemon wget gdebi-core build-essential libfontenc1 libfreetype6 xorg-dev xorg\n",
     "  pro_bar.update(progress(20, 100))\n",
-    "  !wget http://security.ubuntu.com/ubuntu/pool/main/libx/libxfont/libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb >/dev/null 2>&1\n",
+    "  !wget http://security.ubuntu.com/ubuntu/pool/main/libx/libxfont/libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb 2>&1\n",
     "  pro_bar.update(progress(30, 100))\n",
-    "  !wget --output-document xvfb.deb http://security.ubuntu.com/ubuntu/pool/universe/x/xorg-server/xvfb_1.18.4-0ubuntu0.12_amd64.deb >/dev/null 2>&1\n",
+    "  !wget --output-document xvfb.deb http://security.ubuntu.com/ubuntu/pool/universe/x/xorg-server/xvfb_1.18.4-0ubuntu0.12_amd64.deb 2>&1\n",
     "  pro_bar.update(progress(40, 100))\n",
-    "  !dpkg -i libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb >/dev/null 2>&1\n",
+    "  !sudo dpkg -i libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb 2>&1\n",
     "  pro_bar.update(progress(50, 100))\n",
-    "  !dpkg -i xvfb.deb >/dev/null 2>&1\n",
+    "  !sudo dpkg -i xvfb.deb 2>&1\n",
     "  pro_bar.update(progress(70, 100))\n",
     "  !rm libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb\n",
     "  pro_bar.update(progress(80, 100))\n",
@@ -501,4 +500,4 @@
    "outputs": []
   }
  ]
-}
+}
\ No newline at end of file
diff --git a/colab/Colab_UnityEnvironment_2_Train.ipynb b/colab/Colab_UnityEnvironment_2_Train.ipynb
index 1d85f578c..2c6cb798a 100644
--- a/colab/Colab_UnityEnvironment_2_Train.ipynb
+++ b/colab/Colab_UnityEnvironment_2_Train.ipynb
@@ -1,700 +1,689 @@
 {
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "name": "Colab-UnityEnvironment-2-Train.ipynb",
-      "private_outputs": true,
-      "provenance": [],
-      "collapsed_sections": []
-    },
-    "kernelspec": {
-      "name": "python3",
-      "language": "python",
-      "display_name": "Python 3"
-    },
-    "pycharm": {
-      "stem_cell": {
-        "cell_type": "raw",
-        "source": [],
-        "metadata": {
-          "collapsed": false
-        }
-      }
-    }
+ "nbformat": 4,
+ "nbformat_minor": 0,
+ "metadata": {
+  "colab": {
+   "name": "Colab-UnityEnvironment-2-Train.ipynb",
+   "private_outputs": true,
+   "provenance": [],
+   "collapsed_sections": []
   },
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "pbVXrmEsLXDt"
-      },
-      "source": [
-        "# ML-Agents Q-Learning with GridWorld\n",
-        "<img src=\"https://github.com/Unity-Technologies/ml-agents/blob/release_19_docs/docs/images/gridworld.png?raw=true\" align=\"middle\" width=\"435\"/>"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "WNKTwHU3d2-l"
-      },
-      "source": [
-        "## Setup"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "htb-p1hSNX7D"
-      },
-      "source": [
-        "#@title Install Rendering Dependencies { display-mode: \"form\" }\n",
-        "#@markdown (You only need to run this code when using Colab's hosted runtime)\n",
-        "\n",
-        "import os\n",
-        "from IPython.display import HTML, display\n",
-        "\n",
-        "def progress(value, max=100):\n",
-        "    return HTML(\"\"\"\n",
-        "        <progress\n",
-        "            value='{value}'\n",
-        "            max='{max}',\n",
-        "            style='width: 100%'\n",
-        "        >\n",
-        "            {value}\n",
-        "        </progress>\n",
-        "    \"\"\".format(value=value, max=max))\n",
-        "\n",
-        "pro_bar = display(progress(0, 100), display_id=True)\n",
-        "\n",
-        "try:\n",
-        "  import google.colab\n",
-        "  INSTALL_XVFB = True\n",
-        "except ImportError:\n",
-        "  INSTALL_XVFB = 'COLAB_ALWAYS_INSTALL_XVFB' in os.environ\n",
-        "\n",
-        "if INSTALL_XVFB:\n",
-        "  with open('frame-buffer', 'w') as writefile:\n",
-        "    writefile.write(\"\"\"#taken from https://gist.github.com/jterrace/2911875\n",
-        "XVFB=/usr/bin/Xvfb\n",
-        "XVFBARGS=\":1 -screen 0 1024x768x24 -ac +extension GLX +render -noreset\"\n",
-        "PIDFILE=./frame-buffer.pid\n",
-        "case \"$1\" in\n",
-        "  start)\n",
-        "    echo -n \"Starting virtual X frame buffer: Xvfb\"\n",
-        "    /sbin/start-stop-daemon --start --quiet --pidfile $PIDFILE --make-pidfile --background --exec $XVFB -- $XVFBARGS\n",
-        "    echo \".\"\n",
-        "    ;;\n",
-        "  stop)\n",
-        "    echo -n \"Stopping virtual X frame buffer: Xvfb\"\n",
-        "    /sbin/start-stop-daemon --stop --quiet --pidfile $PIDFILE\n",
-        "    rm $PIDFILE\n",
-        "    echo \".\"\n",
-        "    ;;\n",
-        "  restart)\n",
-        "    $0 stop\n",
-        "    $0 start\n",
-        "    ;;\n",
-        "  *)\n",
-        "        echo \"Usage: /etc/init.d/xvfb {start|stop|restart}\"\n",
-        "        exit 1\n",
-        "esac\n",
-        "exit 0\n",
-        "    \"\"\")\n",
-        "  pro_bar.update(progress(5, 100))\n",
-        "  !apt-get install daemon >/dev/null 2>&1\n",
-        "  pro_bar.update(progress(10, 100))\n",
-        "  !apt-get install wget >/dev/null 2>&1\n",
-        "  pro_bar.update(progress(20, 100))\n",
-        "  !wget http://security.ubuntu.com/ubuntu/pool/main/libx/libxfont/libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb >/dev/null 2>&1\n",
-        "  pro_bar.update(progress(30, 100))\n",
-        "  !wget --output-document xvfb.deb http://security.ubuntu.com/ubuntu/pool/universe/x/xorg-server/xvfb_1.18.4-0ubuntu0.12_amd64.deb >/dev/null 2>&1\n",
-        "  pro_bar.update(progress(40, 100))\n",
-        "  !dpkg -i libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb >/dev/null 2>&1\n",
-        "  pro_bar.update(progress(50, 100))\n",
-        "  !dpkg -i xvfb.deb >/dev/null 2>&1\n",
-        "  pro_bar.update(progress(70, 100))\n",
-        "  !rm libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb\n",
-        "  pro_bar.update(progress(80, 100))\n",
-        "  !rm xvfb.deb\n",
-        "  pro_bar.update(progress(90, 100))\n",
-        "  !bash frame-buffer start\n",
-        "  os.environ[\"DISPLAY\"] = \":1\"\n",
-        "pro_bar.update(progress(100, 100))"
-      ],
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Pzj7wgapAcDs"
-      },
-      "source": [
-        "### Installing ml-agents"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "N8yfQqkbebQ5"
-      },
-      "source": [
-        "try:\n",
-        "  import mlagents\n",
-        "  print(\"ml-agents already installed\")\n",
-        "except ImportError:\n",
-        "  !python -m pip install -q mlagents==0.28.0\n",
-        "  print(\"Installed ml-agents\")"
-      ],
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "jz81TWAkbuFY"
-      },
-      "source": [
-        "## Train the GridWorld Environment with Q-Learning"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "29n3dt1Zx5ty"
-      },
-      "source": [
-        "### What is the GridWorld Environment"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "pZhVRfdoyPmv"
-      },
-      "source": [
-        "The [GridWorld](https://github.com/Unity-Technologies/ml-agents/blob/release_19_docs/docs/Learning-Environment-Examples.md#gridworld) Environment is a simple Unity visual environment. The Agent is a blue square in a 3x3 grid that is trying to reach a green __`+`__ while avoiding a red __`x`__.\n",
-        "\n",
-        "The observation is an image obtained by a camera on top of the grid.\n",
-        "\n",
-        "The Action can be one of 5 :\n",
-        " - Do not move\n",
-        " - Move up\n",
-        " - Move down\n",
-        " - Move right\n",
-        " - Move left\n",
-        "\n",
-        "The Agent receives a reward of _1.0_ if it reaches the green __`+`__, a penalty of _-1.0_ if it touches the red __`x`__ and a penalty of `-0.01` at every step (to force the Agent to solve the task as fast as possible)\n",
-        "\n",
-        "__Note__ There are 9 Agents, each in their own grid, at once in the environment. This alows for faster data collection.\n",
-        "\n",
-        "\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "4Gt-ZydJyJWD"
-      },
-      "source": [
-        "### The Q-Learning Algorithm\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "KA1qOgfq0Xdv"
-      },
-      "source": [
-        "In this Notebook, we will implement a very simple Q-Learning algorithm. We will use [pytorch](https://pytorch.org/) to do so.\n",
-        "\n",
-        "Below is the code to create the neural network we will use in the Notebook."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "q79rUp_Sx6A_"
-      },
-      "source": [
-        "import torch\n",
-        "from typing import Tuple\n",
-        "from math import floor\n",
-        "from torch.nn import Parameter\n",
-        "\n",
-        "\n",
-        "class VisualQNetwork(torch.nn.Module):\n",
-        "  def __init__(\n",
-        "    self,\n",
-        "    input_shape: Tuple[int, int, int],\n",
-        "    encoding_size: int,\n",
-        "    output_size: int\n",
-        "  ):\n",
-        "    \"\"\"\n",
-        "    Creates a neural network that takes as input a batch of images (3\n",
-        "    dimensional tensors) and outputs a batch of outputs (1 dimensional\n",
-        "    tensors)\n",
-        "    \"\"\"\n",
-        "    super(VisualQNetwork, self).__init__()\n",
-        "    height = input_shape[1]\n",
-        "    width = input_shape[2]\n",
-        "    initial_channels = input_shape[0]\n",
-        "    conv_1_hw = self.conv_output_shape((height, width), 8, 4)\n",
-        "    conv_2_hw = self.conv_output_shape(conv_1_hw, 4, 2)\n",
-        "    self.final_flat = conv_2_hw[0] * conv_2_hw[1] * 32\n",
-        "    self.conv1 = torch.nn.Conv2d(initial_channels, 16, [8, 8], [4, 4])\n",
-        "    self.conv2 = torch.nn.Conv2d(16, 32, [4, 4], [2, 2])\n",
-        "    self.dense1 = torch.nn.Linear(self.final_flat, encoding_size)\n",
-        "    self.dense2 = torch.nn.Linear(encoding_size, output_size)\n",
-        "\n",
-        "  def forward(self, visual_obs: torch.tensor):\n",
-        "    conv_1 = torch.relu(self.conv1(visual_obs))\n",
-        "    conv_2 = torch.relu(self.conv2(conv_1))\n",
-        "    hidden = self.dense1(conv_2.reshape([-1, self.final_flat]))\n",
-        "    hidden = torch.relu(hidden)\n",
-        "    hidden = self.dense2(hidden)\n",
-        "    return hidden\n",
-        "\n",
-        "  @staticmethod\n",
-        "  def conv_output_shape(\n",
-        "    h_w: Tuple[int, int],\n",
-        "    kernel_size: int = 1,\n",
-        "    stride: int = 1,\n",
-        "    pad: int = 0,\n",
-        "    dilation: int = 1,\n",
-        "  ):\n",
-        "    \"\"\"\n",
-        "    Computes the height and width of the output of a convolution layer.\n",
-        "    \"\"\"\n",
-        "    h = floor(\n",
-        "      ((h_w[0] + (2 * pad) - (dilation * (kernel_size - 1)) - 1) / stride) + 1\n",
-        "    )\n",
-        "    w = floor(\n",
-        "      ((h_w[1] + (2 * pad) - (dilation * (kernel_size - 1)) - 1) / stride) + 1\n",
-        "    )\n",
-        "    return h, w\n"
-      ],
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "EZoaEBAo2L0F"
-      },
-      "source": [
-        "We will now create a few classes to help us store the data we will use to train the Q-Learning algorithm."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "L772fe2q39DO"
-      },
-      "source": [
-        "import numpy as np\n",
-        "from typing import NamedTuple, List\n",
-        "\n",
-        "\n",
-        "class Experience(NamedTuple):\n",
-        "  \"\"\"\n",
-        "  An experience contains the data of one Agent transition.\n",
-        "  - Observation\n",
-        "  - Action\n",
-        "  - Reward\n",
-        "  - Done flag\n",
-        "  - Next Observation\n",
-        "  \"\"\"\n",
-        "\n",
-        "  obs: np.ndarray\n",
-        "  action: np.ndarray\n",
-        "  reward: float\n",
-        "  done: bool\n",
-        "  next_obs: np.ndarray\n",
-        "\n",
-        "# A Trajectory is an ordered sequence of Experiences\n",
-        "Trajectory = List[Experience]\n",
-        "\n",
-        "# A Buffer is an unordered list of Experiences from multiple Trajectories\n",
-        "Buffer = List[Experience]"
-      ],
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6HsM1d5I3_Tj"
-      },
-      "source": [
-        "Now, we can create our trainer class. The role of this trainer is to collect data from the Environment according to a Policy, and then train the Q-Network with that data."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "KkzBoRJCb18t"
-      },
-      "source": [
-        "from mlagents_envs.environment import ActionTuple, BaseEnv\n",
-        "from typing import Dict\n",
-        "import random\n",
-        "\n",
-        "\n",
-        "class Trainer:\n",
-        "  @staticmethod\n",
-        "  def generate_trajectories(\n",
-        "    env: BaseEnv, q_net: VisualQNetwork, buffer_size: int, epsilon: float\n",
-        "  ):\n",
-        "    \"\"\"\n",
-        "    Given a Unity Environment and a Q-Network, this method will generate a\n",
-        "    buffer of Experiences obtained by running the Environment with the Policy\n",
-        "    derived from the Q-Network.\n",
-        "    :param BaseEnv: The UnityEnvironment used.\n",
-        "    :param q_net: The Q-Network used to collect the data.\n",
-        "    :param buffer_size: The minimum size of the buffer this method will return.\n",
-        "    :param epsilon: Will add a random normal variable with standard deviation.\n",
-        "    epsilon to the value heads of the Q-Network to encourage exploration.\n",
-        "    :returns: a Tuple containing the created buffer and the average cumulative\n",
-        "    the Agents obtained.\n",
-        "    \"\"\"\n",
-        "    # Create an empty Buffer\n",
-        "    buffer: Buffer = []\n",
-        "\n",
-        "    # Reset the environment\n",
-        "    env.reset()\n",
-        "    # Read and store the Behavior Name of the Environment\n",
-        "    behavior_name = list(env.behavior_specs)[0]\n",
-        "    # Read and store the Behavior Specs of the Environment\n",
-        "    spec = env.behavior_specs[behavior_name]\n",
-        "\n",
-        "    # Create a Mapping from AgentId to Trajectories. This will help us create\n",
-        "    # trajectories for each Agents\n",
-        "    dict_trajectories_from_agent: Dict[int, Trajectory] = {}\n",
-        "    # Create a Mapping from AgentId to the last observation of the Agent\n",
-        "    dict_last_obs_from_agent: Dict[int, np.ndarray] = {}\n",
-        "    # Create a Mapping from AgentId to the last observation of the Agent\n",
-        "    dict_last_action_from_agent: Dict[int, np.ndarray] = {}\n",
-        "    # Create a Mapping from AgentId to cumulative reward (Only for reporting)\n",
-        "    dict_cumulative_reward_from_agent: Dict[int, float] = {}\n",
-        "    # Create a list to store the cumulative rewards obtained so far\n",
-        "    cumulative_rewards: List[float] = []\n",
-        "\n",
-        "    while len(buffer) < buffer_size:  # While not enough data in the buffer\n",
-        "      # Get the Decision Steps and Terminal Steps of the Agents\n",
-        "      decision_steps, terminal_steps = env.get_steps(behavior_name)\n",
-        "\n",
-        "      # permute the tensor to go from NHWC to NCHW\n",
-        "      order = (0, 3, 1, 2)\n",
-        "      decision_steps.obs = [np.transpose(obs, order) for obs in decision_steps.obs]\n",
-        "      terminal_steps.obs = [np.transpose(obs, order) for obs in terminal_steps.obs]\n",
-        "\n",
-        "      # For all Agents with a Terminal Step:\n",
-        "      for agent_id_terminated in terminal_steps:\n",
-        "        # Create its last experience (is last because the Agent terminated)\n",
-        "        last_experience = Experience(\n",
-        "          obs=dict_last_obs_from_agent[agent_id_terminated].copy(),\n",
-        "          reward=terminal_steps[agent_id_terminated].reward,\n",
-        "          done=not terminal_steps[agent_id_terminated].interrupted,\n",
-        "          action=dict_last_action_from_agent[agent_id_terminated].copy(),\n",
-        "          next_obs=terminal_steps[agent_id_terminated].obs[0],\n",
-        "        )\n",
-        "        # Clear its last observation and action (Since the trajectory is over)\n",
-        "        dict_last_obs_from_agent.pop(agent_id_terminated)\n",
-        "        dict_last_action_from_agent.pop(agent_id_terminated)\n",
-        "        # Report the cumulative reward\n",
-        "        cumulative_reward = (\n",
-        "          dict_cumulative_reward_from_agent.pop(agent_id_terminated)\n",
-        "          + terminal_steps[agent_id_terminated].reward\n",
-        "        )\n",
-        "        cumulative_rewards.append(cumulative_reward)\n",
-        "        # Add the Trajectory and the last experience to the buffer\n",
-        "        buffer.extend(dict_trajectories_from_agent.pop(agent_id_terminated))\n",
-        "        buffer.append(last_experience)\n",
-        "\n",
-        "      # For all Agents with a Decision Step:\n",
-        "      for agent_id_decisions in decision_steps:\n",
-        "        # If the Agent does not have a Trajectory, create an empty one\n",
-        "        if agent_id_decisions not in dict_trajectories_from_agent:\n",
-        "          dict_trajectories_from_agent[agent_id_decisions] = []\n",
-        "          dict_cumulative_reward_from_agent[agent_id_decisions] = 0\n",
-        "\n",
-        "        # If the Agent requesting a decision has a \"last observation\"\n",
-        "        if agent_id_decisions in dict_last_obs_from_agent:\n",
-        "          # Create an Experience from the last observation and the Decision Step\n",
-        "          exp = Experience(\n",
-        "            obs=dict_last_obs_from_agent[agent_id_decisions].copy(),\n",
-        "            reward=decision_steps[agent_id_decisions].reward,\n",
-        "            done=False,\n",
-        "            action=dict_last_action_from_agent[agent_id_decisions].copy(),\n",
-        "            next_obs=decision_steps[agent_id_decisions].obs[0],\n",
-        "          )\n",
-        "          # Update the Trajectory of the Agent and its cumulative reward\n",
-        "          dict_trajectories_from_agent[agent_id_decisions].append(exp)\n",
-        "          dict_cumulative_reward_from_agent[agent_id_decisions] += (\n",
-        "            decision_steps[agent_id_decisions].reward\n",
-        "          )\n",
-        "        # Store the observation as the new \"last observation\"\n",
-        "        dict_last_obs_from_agent[agent_id_decisions] = (\n",
-        "          decision_steps[agent_id_decisions].obs[0]\n",
-        "        )\n",
-        "\n",
-        "      # Generate an action for all the Agents that requested a decision\n",
-        "      # Compute the values for each action given the observation\n",
-        "      actions_values = (\n",
-        "        q_net(torch.from_numpy(decision_steps.obs[0])).detach().numpy()\n",
-        "      )\n",
-        "      # Add some noise with epsilon to the values\n",
-        "      actions_values += epsilon * (\n",
-        "        np.random.randn(actions_values.shape[0], actions_values.shape[1])\n",
-        "      ).astype(np.float32)\n",
-        "      # Pick the best action using argmax\n",
-        "      actions = np.argmax(actions_values, axis=1)\n",
-        "      actions.resize((len(decision_steps), 1))\n",
-        "      # Store the action that was picked, it will be put in the trajectory later\n",
-        "      for agent_index, agent_id in enumerate(decision_steps.agent_id):\n",
-        "        dict_last_action_from_agent[agent_id] = actions[agent_index]\n",
-        "\n",
-        "      # Set the actions in the environment\n",
-        "      # Unity Environments expect ActionTuple instances.\n",
-        "      action_tuple = ActionTuple()\n",
-        "      action_tuple.add_discrete(actions)\n",
-        "      env.set_actions(behavior_name, action_tuple)\n",
-        "      # Perform a step in the simulation\n",
-        "      env.step()\n",
-        "    return buffer, np.mean(cumulative_rewards)\n",
-        "\n",
-        "  @staticmethod\n",
-        "  def update_q_net(\n",
-        "    q_net: VisualQNetwork,\n",
-        "    optimizer: torch.optim,\n",
-        "    buffer: Buffer,\n",
-        "    action_size: int\n",
-        "  ):\n",
-        "    \"\"\"\n",
-        "    Performs an update of the Q-Network using the provided optimizer and buffer\n",
-        "    \"\"\"\n",
-        "    BATCH_SIZE = 1000\n",
-        "    NUM_EPOCH = 3\n",
-        "    GAMMA = 0.9\n",
-        "    batch_size = min(len(buffer), BATCH_SIZE)\n",
-        "    random.shuffle(buffer)\n",
-        "    # Split the buffer into batches\n",
-        "    batches = [\n",
-        "      buffer[batch_size * start : batch_size * (start + 1)]\n",
-        "      for start in range(int(len(buffer) / batch_size))\n",
-        "    ]\n",
-        "    for _ in range(NUM_EPOCH):\n",
-        "      for batch in batches:\n",
-        "        # Create the Tensors that will be fed in the network\n",
-        "        obs = torch.from_numpy(np.stack([ex.obs for ex in batch]))\n",
-        "        reward = torch.from_numpy(\n",
-        "          np.array([ex.reward for ex in batch], dtype=np.float32).reshape(-1, 1)\n",
-        "        )\n",
-        "        done = torch.from_numpy(\n",
-        "          np.array([ex.done for ex in batch], dtype=np.float32).reshape(-1, 1)\n",
-        "        )\n",
-        "        action = torch.from_numpy(np.stack([ex.action for ex in batch]))\n",
-        "        next_obs = torch.from_numpy(np.stack([ex.next_obs for ex in batch]))\n",
-        "\n",
-        "        # Use the Bellman equation to update the Q-Network\n",
-        "        target = (\n",
-        "          reward\n",
-        "          + (1.0 - done)\n",
-        "          * GAMMA\n",
-        "          * torch.max(q_net(next_obs).detach(), dim=1, keepdim=True).values\n",
-        "        )\n",
-        "        mask = torch.zeros((len(batch), action_size))\n",
-        "        mask.scatter_(1, action, 1)\n",
-        "        prediction = torch.sum(q_net(obs) * mask, dim=1, keepdim=True)\n",
-        "        criterion = torch.nn.MSELoss()\n",
-        "        loss = criterion(prediction, target)\n",
-        "\n",
-        "        # Perform the backpropagation\n",
-        "        optimizer.zero_grad()\n",
-        "        loss.backward()\n",
-        "        optimizer.step()\n"
-      ],
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "vcU4ZMAEWCvX"
-      },
-      "source": [
-        "### Run Training"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "_lIHijQfbYjh"
-      },
-      "source": [
-        "# -----------------\n",
-        "# This code is used to close an env that might not have been closed before\n",
-        "try:\n",
-        "  env.close()\n",
-        "except:\n",
-        "  pass\n",
-        "# -----------------\n",
-        "\n",
-        "from mlagents_envs.registry import default_registry\n",
-        "from mlagents_envs.environment import UnityEnvironment\n",
-        "import matplotlib.pyplot as plt\n",
-        "%matplotlib inline\n",
-        "\n",
-        "# Create the GridWorld Environment from the registry\n",
-        "env = default_registry[\"GridWorld\"].make()\n",
-        "print(\"GridWorld environment created.\")\n",
-        "\n",
-        "num_actions = 5\n",
-        "\n",
-        "try:\n",
-        "  # Create a new Q-Network.\n",
-        "  qnet = VisualQNetwork((3, 64, 84), 126, num_actions)\n",
-        "\n",
-        "  experiences: Buffer = []\n",
-        "  optim = torch.optim.Adam(qnet.parameters(), lr= 0.001)\n",
-        "\n",
-        "  cumulative_rewards: List[float] = []\n",
-        "\n",
-        "  # The number of training steps that will be performed\n",
-        "  NUM_TRAINING_STEPS = int(os.getenv('QLEARNING_NUM_TRAINING_STEPS', 70))\n",
-        "  # The number of experiences to collect per training step\n",
-        "  NUM_NEW_EXP = int(os.getenv('QLEARNING_NUM_NEW_EXP', 1000))\n",
-        "  # The maximum size of the Buffer\n",
-        "  BUFFER_SIZE = int(os.getenv('QLEARNING_BUFFER_SIZE', 10000))\n",
-        "\n",
-        "  for n in range(NUM_TRAINING_STEPS):\n",
-        "    new_exp,_ = Trainer.generate_trajectories(env, qnet, NUM_NEW_EXP, epsilon=0.1)\n",
-        "    random.shuffle(experiences)\n",
-        "    if len(experiences) > BUFFER_SIZE:\n",
-        "      experiences = experiences[:BUFFER_SIZE]\n",
-        "    experiences.extend(new_exp)\n",
-        "    Trainer.update_q_net(qnet, optim, experiences, num_actions)\n",
-        "    _, rewards = Trainer.generate_trajectories(env, qnet, 100, epsilon=0)\n",
-        "    cumulative_rewards.append(rewards)\n",
-        "    print(\"Training step \", n+1, \"\\treward \", rewards)\n",
-        "except KeyboardInterrupt:\n",
-        "  print(\"\\nTraining interrupted, continue to next cell to save to save the model.\")\n",
-        "finally:\n",
-        "  env.close()\n",
-        "\n",
-        "# Show the training graph\n",
-        "try:\n",
-        "  plt.plot(range(NUM_TRAINING_STEPS), cumulative_rewards)\n",
-        "except ValueError:\n",
-        "  print(\"\\nPlot failed on interrupted training.\")\n",
-        " "
-      ],
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "### Export PyTorch Model to ONNX"
-      ],
-      "metadata": {
-        "id": "pFEP9JT5KuO-"
-      }
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "The following cell provides an example of some of the extra tensors a model needs to work for ML-Agents inference with Barracuda. The GridWorldColab scene is configured to work with this ONNX file.\n",
-        "Only policy models need to be exported for inference and they need the following additional tensors:  \n",
-        "\n",
-        "*   All models need version_number\n",
-        "*   All models need memory_size\n",
-        "*   Models with continuous outputs need continuous_action_output_shape\n",
-	    "*   Models with discrete outputs need discrete_action_output_shape and an additional mask input that matches the shape of the discrete outputs\n",
-        "*   The mask input must be connected to the outputs or it will be pruned on export, if mask values aren't being set they will be 1, so multiplying the discrete outputs by the mask will have no effect"
-      ],
-      "metadata": {
-        "id": "V0V7bH4mOnUF"
-      }
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "class WrapperNet(torch.nn.Module):\n",
-        "    def __init__(\n",
-        "            self,\n",
-        "            qnet: VisualQNetwork,\n",
-        "            discrete_output_sizes: List[int],\n",
-        "    ):\n",
-        "        \"\"\"\n",
-        "        Wraps the VisualQNetwork adding extra constants and dummy mask inputs\n",
-        "        required by runtime inference with Barracuda. \n",
-        "\n",
-        "        For environment continuous actions outputs would need to add them\n",
-        "        similarly to how discrete action outputs work, both in the wrapper\n",
-        "        and in the ONNX output_names / dynamic_axes.\n",
-        "        \"\"\"\n",
-        "        super(WrapperNet, self).__init__()\n",
-        "        self.qnet = qnet\n",
-        "\n",
-        "        # version_number\n",
-        "        #   MLAgents1_0 = 2   (not covered by this example)\n",
-        "        #   MLAgents2_0 = 3\n",
-        "        version_number = torch.Tensor([3])\n",
-        "        self.version_number = Parameter(version_number, requires_grad=False)\n",
-        "\n",
-        "        # memory_size\n",
-        "        # TODO: document case where memory is not zero.\n",
-        "        memory_size = torch.Tensor([0])\n",
-        "        self.memory_size = Parameter(memory_size, requires_grad=False)\n",
-        "\n",
-        "        # discrete_action_output_shape\n",
-        "        output_shape=torch.Tensor([discrete_output_sizes])\n",
-        "        self.discrete_shape = Parameter(output_shape, requires_grad=False)\n",
-        "\n",
-        "\n",
-        "    # if you have discrete actions ML-agents expects corresponding a mask\n",
-        "    # tensor with the same shape to exist as input\n",
-        "    def forward(self, visual_obs: torch.tensor, mask: torch.tensor):\n",
-        "        qnet_result = self.qnet(visual_obs)\n",
-        "        # Connect mask to keep it from getting pruned\n",
-        "        # Mask values will be 1 if you never call SetActionMask() in\n",
-        "        # WriteDiscreteActionMask() \n",
-        "        qnet_result = torch.mul(qnet_result, mask)\n",
-        "        action = torch.argmax(qnet_result, dim=1, keepdim=True)\n",
-        "        return [action], self.discrete_shape, self.version_number, self.memory_size\n",
-        "\n",
-        "\n",
-        "torch.onnx.export(\n",
-        "    WrapperNet(qnet, [num_actions]),\n",
-        "    # A tuple with an example of the input tensors\n",
-        "    (torch.tensor([experiences[0].obs]), torch.ones(1, num_actions)),\n",
-        "    'GridWorldColab.onnx',\n",
-        "    opset_version=9,\n",
-        "    # input_names must correspond to the WrapperNet forward parameters\n",
-        "    # obs will be obs_0, obs_1, etc.\n",
-        "    input_names=[\"obs_0\", \"action_masks\"],\n",
-        "    # output_names must correspond to the return tuple of the WrapperNet\n",
-        "    # forward function.\n",
-        "    output_names=[\"discrete_actions\", \"discrete_action_output_shape\",\n",
-        "                  \"version_number\", \"memory_size\"],\n",
-        "    # All inputs and outputs should have their 0th dimension be designated\n",
-        "    # as 'batch'\n",
-        "    dynamic_axes={'obs_0': {0: 'batch'},\n",
-        "                  'action_masks': {0: 'batch'},\n",
-        "                  'discrete_actions': {0: 'batch'},\n",
-        "                  'discrete_action_output_shape': {0: 'batch'}\n",
-        "                 }\n",
-        "    )"
-      ],
-      "metadata": {
-        "id": "T0zzg8fWpTYO"
-      },
-      "execution_count": null,
-      "outputs": []
-    }
-  ]
-}
+  "kernelspec": {
+   "name": "python3",
+   "language": "python",
+   "display_name": "Python 3"
+  }
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "pbVXrmEsLXDt"
+   },
+   "source": [
+    "# ML-Agents Q-Learning with GridWorld\n",
+    "<img src=\"https://github.com/Unity-Technologies/ml-agents/blob/release_19_docs/docs/images/gridworld.png?raw=true\" align=\"middle\" width=\"435\"/>"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "WNKTwHU3d2-l"
+   },
+   "source": [
+    "## Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "htb-p1hSNX7D"
+   },
+   "source": [
+    "#@title Install Rendering Dependencies { display-mode: \"form\" }\n",
+    "#@markdown (You only need to run this code when using Colab's hosted runtime)\n",
+    "\n",
+    "import os\n",
+    "from IPython.display import HTML, display\n",
+    "\n",
+    "def progress(value, max=100):\n",
+    "    return HTML(\"\"\"\n",
+    "        <progress\n",
+    "            value='{value}'\n",
+    "            max='{max}',\n",
+    "            style='width: 100%'\n",
+    "        >\n",
+    "            {value}\n",
+    "        </progress>\n",
+    "    \"\"\".format(value=value, max=max))\n",
+    "\n",
+    "pro_bar = display(progress(0, 100), display_id=True)\n",
+    "\n",
+    "try:\n",
+    "  import google.colab\n",
+    "  INSTALL_XVFB = True\n",
+    "except ImportError:\n",
+    "  INSTALL_XVFB = 'COLAB_ALWAYS_INSTALL_XVFB' in os.environ\n",
+    "\n",
+    "if INSTALL_XVFB:\n",
+    "  with open('frame-buffer', 'w') as writefile:\n",
+    "    writefile.write(\"\"\"#taken from https://gist.github.com/jterrace/2911875\n",
+    "XVFB=/usr/bin/Xvfb\n",
+    "XVFBARGS=\":1 -screen 0 1024x768x24 -ac +extension GLX +render -noreset\"\n",
+    "PIDFILE=./frame-buffer.pid\n",
+    "case \"$1\" in\n",
+    "  start)\n",
+    "    echo -n \"Starting virtual X frame buffer: Xvfb\"\n",
+    "    /sbin/start-stop-daemon --start --quiet --pidfile $PIDFILE --make-pidfile --background --exec $XVFB -- $XVFBARGS\n",
+    "    echo \".\"\n",
+    "    ;;\n",
+    "  stop)\n",
+    "    echo -n \"Stopping virtual X frame buffer: Xvfb\"\n",
+    "    /sbin/start-stop-daemon --stop --quiet --pidfile $PIDFILE\n",
+    "    rm $PIDFILE\n",
+    "    echo \".\"\n",
+    "    ;;\n",
+    "  restart)\n",
+    "    $0 stop\n",
+    "    $0 start\n",
+    "    ;;\n",
+    "  *)\n",
+    "        echo \"Usage: /etc/init.d/xvfb {start|stop|restart}\"\n",
+    "        exit 1\n",
+    "esac\n",
+    "exit 0\n",
+    "    \"\"\")\n",
+    "  !sudo apt-get update\n",
+    "  pro_bar.update(progress(10, 100))\n",
+    "  !sudo DEBIAN_FRONTEND=noninteractive apt install -y daemon wget gdebi-core build-essential libfontenc1 libfreetype6 xorg-dev xorg\n",
+    "  pro_bar.update(progress(20, 100))\n",
+    "  !wget http://security.ubuntu.com/ubuntu/pool/main/libx/libxfont/libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb 2>&1\n",
+    "  pro_bar.update(progress(30, 100))\n",
+    "  !wget --output-document xvfb.deb http://security.ubuntu.com/ubuntu/pool/universe/x/xorg-server/xvfb_1.18.4-0ubuntu0.12_amd64.deb 2>&1\n",
+    "  pro_bar.update(progress(40, 100))\n",
+    "  !sudo dpkg -i libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb 2>&1\n",
+    "  pro_bar.update(progress(50, 100))\n",
+    "  !sudo dpkg -i xvfb.deb 2>&1\n",
+    "  pro_bar.update(progress(70, 100))\n",
+    "  !rm libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb\n",
+    "  pro_bar.update(progress(80, 100))\n",
+    "  !rm xvfb.deb\n",
+    "  pro_bar.update(progress(90, 100))\n",
+    "  !bash frame-buffer start\n",
+    "  os.environ[\"DISPLAY\"] = \":1\"\n",
+    "pro_bar.update(progress(100, 100))"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Pzj7wgapAcDs"
+   },
+   "source": [
+    "### Installing ml-agents"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "N8yfQqkbebQ5"
+   },
+   "source": [
+    "try:\n",
+    "  import mlagents\n",
+    "  print(\"ml-agents already installed\")\n",
+    "except ImportError:\n",
+    "  !python -m pip install -q mlagents==0.28.0\n",
+    "  print(\"Installed ml-agents\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "jz81TWAkbuFY"
+   },
+   "source": [
+    "## Train the GridWorld Environment with Q-Learning"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "29n3dt1Zx5ty"
+   },
+   "source": [
+    "### What is the GridWorld Environment"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "pZhVRfdoyPmv"
+   },
+   "source": [
+    "The [GridWorld](https://github.com/Unity-Technologies/ml-agents/blob/release_19_docs/docs/Learning-Environment-Examples.md#gridworld) Environment is a simple Unity visual environment. The Agent is a blue square in a 3x3 grid that is trying to reach a green __`+`__ while avoiding a red __`x`__.\n",
+    "\n",
+    "The observation is an image obtained by a camera on top of the grid.\n",
+    "\n",
+    "The Action can be one of 5 :\n",
+    " - Do not move\n",
+    " - Move up\n",
+    " - Move down\n",
+    " - Move right\n",
+    " - Move left\n",
+    "\n",
+    "The Agent receives a reward of _1.0_ if it reaches the green __`+`__, a penalty of _-1.0_ if it touches the red __`x`__ and a penalty of `-0.01` at every step (to force the Agent to solve the task as fast as possible)\n",
+    "\n",
+    "__Note__ There are 9 Agents, each in their own grid, at once in the environment. This alows for faster data collection.\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "4Gt-ZydJyJWD"
+   },
+   "source": [
+    "### The Q-Learning Algorithm\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "KA1qOgfq0Xdv"
+   },
+   "source": [
+    "In this Notebook, we will implement a very simple Q-Learning algorithm. We will use [pytorch](https://pytorch.org/) to do so.\n",
+    "\n",
+    "Below is the code to create the neural network we will use in the Notebook."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "q79rUp_Sx6A_"
+   },
+   "source": [
+    "import torch\n",
+    "from typing import Tuple\n",
+    "from math import floor\n",
+    "from torch.nn import Parameter\n",
+    "\n",
+    "\n",
+    "class VisualQNetwork(torch.nn.Module):\n",
+    "  def __init__(\n",
+    "    self,\n",
+    "    input_shape: Tuple[int, int, int],\n",
+    "    encoding_size: int,\n",
+    "    output_size: int\n",
+    "  ):\n",
+    "    \"\"\"\n",
+    "    Creates a neural network that takes as input a batch of images (3\n",
+    "    dimensional tensors) and outputs a batch of outputs (1 dimensional\n",
+    "    tensors)\n",
+    "    \"\"\"\n",
+    "    super(VisualQNetwork, self).__init__()\n",
+    "    height = input_shape[1]\n",
+    "    width = input_shape[2]\n",
+    "    initial_channels = input_shape[0]\n",
+    "    conv_1_hw = self.conv_output_shape((height, width), 8, 4)\n",
+    "    conv_2_hw = self.conv_output_shape(conv_1_hw, 4, 2)\n",
+    "    self.final_flat = conv_2_hw[0] * conv_2_hw[1] * 32\n",
+    "    self.conv1 = torch.nn.Conv2d(initial_channels, 16, [8, 8], [4, 4])\n",
+    "    self.conv2 = torch.nn.Conv2d(16, 32, [4, 4], [2, 2])\n",
+    "    self.dense1 = torch.nn.Linear(self.final_flat, encoding_size)\n",
+    "    self.dense2 = torch.nn.Linear(encoding_size, output_size)\n",
+    "\n",
+    "  def forward(self, visual_obs: torch.tensor):\n",
+    "    conv_1 = torch.relu(self.conv1(visual_obs))\n",
+    "    conv_2 = torch.relu(self.conv2(conv_1))\n",
+    "    hidden = self.dense1(conv_2.reshape([-1, self.final_flat]))\n",
+    "    hidden = torch.relu(hidden)\n",
+    "    hidden = self.dense2(hidden)\n",
+    "    return hidden\n",
+    "\n",
+    "  @staticmethod\n",
+    "  def conv_output_shape(\n",
+    "    h_w: Tuple[int, int],\n",
+    "    kernel_size: int = 1,\n",
+    "    stride: int = 1,\n",
+    "    pad: int = 0,\n",
+    "    dilation: int = 1,\n",
+    "  ):\n",
+    "    \"\"\"\n",
+    "    Computes the height and width of the output of a convolution layer.\n",
+    "    \"\"\"\n",
+    "    h = floor(\n",
+    "      ((h_w[0] + (2 * pad) - (dilation * (kernel_size - 1)) - 1) / stride) + 1\n",
+    "    )\n",
+    "    w = floor(\n",
+    "      ((h_w[1] + (2 * pad) - (dilation * (kernel_size - 1)) - 1) / stride) + 1\n",
+    "    )\n",
+    "    return h, w\n"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "EZoaEBAo2L0F"
+   },
+   "source": [
+    "We will now create a few classes to help us store the data we will use to train the Q-Learning algorithm."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "L772fe2q39DO"
+   },
+   "source": [
+    "import numpy as np\n",
+    "from typing import NamedTuple, List\n",
+    "\n",
+    "\n",
+    "class Experience(NamedTuple):\n",
+    "  \"\"\"\n",
+    "  An experience contains the data of one Agent transition.\n",
+    "  - Observation\n",
+    "  - Action\n",
+    "  - Reward\n",
+    "  - Done flag\n",
+    "  - Next Observation\n",
+    "  \"\"\"\n",
+    "\n",
+    "  obs: np.ndarray\n",
+    "  action: np.ndarray\n",
+    "  reward: float\n",
+    "  done: bool\n",
+    "  next_obs: np.ndarray\n",
+    "\n",
+    "# A Trajectory is an ordered sequence of Experiences\n",
+    "Trajectory = List[Experience]\n",
+    "\n",
+    "# A Buffer is an unordered list of Experiences from multiple Trajectories\n",
+    "Buffer = List[Experience]"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "6HsM1d5I3_Tj"
+   },
+   "source": [
+    "Now, we can create our trainer class. The role of this trainer is to collect data from the Environment according to a Policy, and then train the Q-Network with that data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "KkzBoRJCb18t"
+   },
+   "source": [
+    "from mlagents_envs.environment import ActionTuple, BaseEnv\n",
+    "from typing import Dict\n",
+    "import random\n",
+    "\n",
+    "\n",
+    "class Trainer:\n",
+    "  @staticmethod\n",
+    "  def generate_trajectories(\n",
+    "    env: BaseEnv, q_net: VisualQNetwork, buffer_size: int, epsilon: float\n",
+    "  ):\n",
+    "    \"\"\"\n",
+    "    Given a Unity Environment and a Q-Network, this method will generate a\n",
+    "    buffer of Experiences obtained by running the Environment with the Policy\n",
+    "    derived from the Q-Network.\n",
+    "    :param BaseEnv: The UnityEnvironment used.\n",
+    "    :param q_net: The Q-Network used to collect the data.\n",
+    "    :param buffer_size: The minimum size of the buffer this method will return.\n",
+    "    :param epsilon: Will add a random normal variable with standard deviation.\n",
+    "    epsilon to the value heads of the Q-Network to encourage exploration.\n",
+    "    :returns: a Tuple containing the created buffer and the average cumulative\n",
+    "    the Agents obtained.\n",
+    "    \"\"\"\n",
+    "    # Create an empty Buffer\n",
+    "    buffer: Buffer = []\n",
+    "\n",
+    "    # Reset the environment\n",
+    "    env.reset()\n",
+    "    # Read and store the Behavior Name of the Environment\n",
+    "    behavior_name = list(env.behavior_specs)[0]\n",
+    "    # Read and store the Behavior Specs of the Environment\n",
+    "    spec = env.behavior_specs[behavior_name]\n",
+    "\n",
+    "    # Create a Mapping from AgentId to Trajectories. This will help us create\n",
+    "    # trajectories for each Agents\n",
+    "    dict_trajectories_from_agent: Dict[int, Trajectory] = {}\n",
+    "    # Create a Mapping from AgentId to the last observation of the Agent\n",
+    "    dict_last_obs_from_agent: Dict[int, np.ndarray] = {}\n",
+    "    # Create a Mapping from AgentId to the last observation of the Agent\n",
+    "    dict_last_action_from_agent: Dict[int, np.ndarray] = {}\n",
+    "    # Create a Mapping from AgentId to cumulative reward (Only for reporting)\n",
+    "    dict_cumulative_reward_from_agent: Dict[int, float] = {}\n",
+    "    # Create a list to store the cumulative rewards obtained so far\n",
+    "    cumulative_rewards: List[float] = []\n",
+    "\n",
+    "    while len(buffer) < buffer_size:  # While not enough data in the buffer\n",
+    "      # Get the Decision Steps and Terminal Steps of the Agents\n",
+    "      decision_steps, terminal_steps = env.get_steps(behavior_name)\n",
+    "\n",
+    "      # permute the tensor to go from NHWC to NCHW\n",
+    "      order = (0, 3, 1, 2)\n",
+    "      decision_steps.obs = [np.transpose(obs, order) for obs in decision_steps.obs]\n",
+    "      terminal_steps.obs = [np.transpose(obs, order) for obs in terminal_steps.obs]\n",
+    "\n",
+    "      # For all Agents with a Terminal Step:\n",
+    "      for agent_id_terminated in terminal_steps:\n",
+    "        # Create its last experience (is last because the Agent terminated)\n",
+    "        last_experience = Experience(\n",
+    "          obs=dict_last_obs_from_agent[agent_id_terminated].copy(),\n",
+    "          reward=terminal_steps[agent_id_terminated].reward,\n",
+    "          done=not terminal_steps[agent_id_terminated].interrupted,\n",
+    "          action=dict_last_action_from_agent[agent_id_terminated].copy(),\n",
+    "          next_obs=terminal_steps[agent_id_terminated].obs[0],\n",
+    "        )\n",
+    "        # Clear its last observation and action (Since the trajectory is over)\n",
+    "        dict_last_obs_from_agent.pop(agent_id_terminated)\n",
+    "        dict_last_action_from_agent.pop(agent_id_terminated)\n",
+    "        # Report the cumulative reward\n",
+    "        cumulative_reward = (\n",
+    "          dict_cumulative_reward_from_agent.pop(agent_id_terminated)\n",
+    "          + terminal_steps[agent_id_terminated].reward\n",
+    "        )\n",
+    "        cumulative_rewards.append(cumulative_reward)\n",
+    "        # Add the Trajectory and the last experience to the buffer\n",
+    "        buffer.extend(dict_trajectories_from_agent.pop(agent_id_terminated))\n",
+    "        buffer.append(last_experience)\n",
+    "\n",
+    "      # For all Agents with a Decision Step:\n",
+    "      for agent_id_decisions in decision_steps:\n",
+    "        # If the Agent does not have a Trajectory, create an empty one\n",
+    "        if agent_id_decisions not in dict_trajectories_from_agent:\n",
+    "          dict_trajectories_from_agent[agent_id_decisions] = []\n",
+    "          dict_cumulative_reward_from_agent[agent_id_decisions] = 0\n",
+    "\n",
+    "        # If the Agent requesting a decision has a \"last observation\"\n",
+    "        if agent_id_decisions in dict_last_obs_from_agent:\n",
+    "          # Create an Experience from the last observation and the Decision Step\n",
+    "          exp = Experience(\n",
+    "            obs=dict_last_obs_from_agent[agent_id_decisions].copy(),\n",
+    "            reward=decision_steps[agent_id_decisions].reward,\n",
+    "            done=False,\n",
+    "            action=dict_last_action_from_agent[agent_id_decisions].copy(),\n",
+    "            next_obs=decision_steps[agent_id_decisions].obs[0],\n",
+    "          )\n",
+    "          # Update the Trajectory of the Agent and its cumulative reward\n",
+    "          dict_trajectories_from_agent[agent_id_decisions].append(exp)\n",
+    "          dict_cumulative_reward_from_agent[agent_id_decisions] += (\n",
+    "            decision_steps[agent_id_decisions].reward\n",
+    "          )\n",
+    "        # Store the observation as the new \"last observation\"\n",
+    "        dict_last_obs_from_agent[agent_id_decisions] = (\n",
+    "          decision_steps[agent_id_decisions].obs[0]\n",
+    "        )\n",
+    "\n",
+    "      # Generate an action for all the Agents that requested a decision\n",
+    "      # Compute the values for each action given the observation\n",
+    "      actions_values = (\n",
+    "        q_net(torch.from_numpy(decision_steps.obs[0])).detach().numpy()\n",
+    "      )\n",
+    "      # Add some noise with epsilon to the values\n",
+    "      actions_values += epsilon * (\n",
+    "        np.random.randn(actions_values.shape[0], actions_values.shape[1])\n",
+    "      ).astype(np.float32)\n",
+    "      # Pick the best action using argmax\n",
+    "      actions = np.argmax(actions_values, axis=1)\n",
+    "      actions.resize((len(decision_steps), 1))\n",
+    "      # Store the action that was picked, it will be put in the trajectory later\n",
+    "      for agent_index, agent_id in enumerate(decision_steps.agent_id):\n",
+    "        dict_last_action_from_agent[agent_id] = actions[agent_index]\n",
+    "\n",
+    "      # Set the actions in the environment\n",
+    "      # Unity Environments expect ActionTuple instances.\n",
+    "      action_tuple = ActionTuple()\n",
+    "      action_tuple.add_discrete(actions)\n",
+    "      env.set_actions(behavior_name, action_tuple)\n",
+    "      # Perform a step in the simulation\n",
+    "      env.step()\n",
+    "    return buffer, np.mean(cumulative_rewards)\n",
+    "\n",
+    "  @staticmethod\n",
+    "  def update_q_net(\n",
+    "    q_net: VisualQNetwork,\n",
+    "    optimizer: torch.optim,\n",
+    "    buffer: Buffer,\n",
+    "    action_size: int\n",
+    "  ):\n",
+    "    \"\"\"\n",
+    "    Performs an update of the Q-Network using the provided optimizer and buffer\n",
+    "    \"\"\"\n",
+    "    BATCH_SIZE = 1000\n",
+    "    NUM_EPOCH = 3\n",
+    "    GAMMA = 0.9\n",
+    "    batch_size = min(len(buffer), BATCH_SIZE)\n",
+    "    random.shuffle(buffer)\n",
+    "    # Split the buffer into batches\n",
+    "    batches = [\n",
+    "      buffer[batch_size * start : batch_size * (start + 1)]\n",
+    "      for start in range(int(len(buffer) / batch_size))\n",
+    "    ]\n",
+    "    for _ in range(NUM_EPOCH):\n",
+    "      for batch in batches:\n",
+    "        # Create the Tensors that will be fed in the network\n",
+    "        obs = torch.from_numpy(np.stack([ex.obs for ex in batch]))\n",
+    "        reward = torch.from_numpy(\n",
+    "          np.array([ex.reward for ex in batch], dtype=np.float32).reshape(-1, 1)\n",
+    "        )\n",
+    "        done = torch.from_numpy(\n",
+    "          np.array([ex.done for ex in batch], dtype=np.float32).reshape(-1, 1)\n",
+    "        )\n",
+    "        action = torch.from_numpy(np.stack([ex.action for ex in batch]))\n",
+    "        next_obs = torch.from_numpy(np.stack([ex.next_obs for ex in batch]))\n",
+    "\n",
+    "        # Use the Bellman equation to update the Q-Network\n",
+    "        target = (\n",
+    "          reward\n",
+    "          + (1.0 - done)\n",
+    "          * GAMMA\n",
+    "          * torch.max(q_net(next_obs).detach(), dim=1, keepdim=True).values\n",
+    "        )\n",
+    "        mask = torch.zeros((len(batch), action_size))\n",
+    "        mask.scatter_(1, action, 1)\n",
+    "        prediction = torch.sum(q_net(obs) * mask, dim=1, keepdim=True)\n",
+    "        criterion = torch.nn.MSELoss()\n",
+    "        loss = criterion(prediction, target)\n",
+    "\n",
+    "        # Perform the backpropagation\n",
+    "        optimizer.zero_grad()\n",
+    "        loss.backward()\n",
+    "        optimizer.step()\n"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "vcU4ZMAEWCvX"
+   },
+   "source": [
+    "### Run Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "_lIHijQfbYjh"
+   },
+   "source": [
+    "# -----------------\n",
+    "# This code is used to close an env that might not have been closed before\n",
+    "try:\n",
+    "  env.close()\n",
+    "except:\n",
+    "  pass\n",
+    "# -----------------\n",
+    "\n",
+    "from mlagents_envs.registry import default_registry\n",
+    "from mlagents_envs.environment import UnityEnvironment\n",
+    "import matplotlib.pyplot as plt\n",
+    "%matplotlib inline\n",
+    "\n",
+    "# Create the GridWorld Environment from the registry\n",
+    "env = default_registry[\"GridWorld\"].make()\n",
+    "print(\"GridWorld environment created.\")\n",
+    "\n",
+    "num_actions = 5\n",
+    "\n",
+    "try:\n",
+    "  # Create a new Q-Network.\n",
+    "  qnet = VisualQNetwork((3, 64, 84), 126, num_actions)\n",
+    "\n",
+    "  experiences: Buffer = []\n",
+    "  optim = torch.optim.Adam(qnet.parameters(), lr= 0.001)\n",
+    "\n",
+    "  cumulative_rewards: List[float] = []\n",
+    "\n",
+    "  # The number of training steps that will be performed\n",
+    "  NUM_TRAINING_STEPS = int(os.getenv('QLEARNING_NUM_TRAINING_STEPS', 70))\n",
+    "  # The number of experiences to collect per training step\n",
+    "  NUM_NEW_EXP = int(os.getenv('QLEARNING_NUM_NEW_EXP', 1000))\n",
+    "  # The maximum size of the Buffer\n",
+    "  BUFFER_SIZE = int(os.getenv('QLEARNING_BUFFER_SIZE', 10000))\n",
+    "\n",
+    "  for n in range(NUM_TRAINING_STEPS):\n",
+    "    new_exp,_ = Trainer.generate_trajectories(env, qnet, NUM_NEW_EXP, epsilon=0.1)\n",
+    "    random.shuffle(experiences)\n",
+    "    if len(experiences) > BUFFER_SIZE:\n",
+    "      experiences = experiences[:BUFFER_SIZE]\n",
+    "    experiences.extend(new_exp)\n",
+    "    Trainer.update_q_net(qnet, optim, experiences, num_actions)\n",
+    "    _, rewards = Trainer.generate_trajectories(env, qnet, 100, epsilon=0)\n",
+    "    cumulative_rewards.append(rewards)\n",
+    "    print(\"Training step \", n+1, \"\\treward \", rewards)\n",
+    "except KeyboardInterrupt:\n",
+    "  print(\"\\nTraining interrupted, continue to next cell to save to save the model.\")\n",
+    "finally:\n",
+    "  env.close()\n",
+    "\n",
+    "# Show the training graph\n",
+    "try:\n",
+    "  plt.plot(range(NUM_TRAINING_STEPS), cumulative_rewards)\n",
+    "except ValueError:\n",
+    "  print(\"\\nPlot failed on interrupted training.\")\n"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Export PyTorch Model to ONNX"
+   ],
+   "metadata": {
+    "id": "pFEP9JT5KuO-"
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "The following cell provides an example of some of the extra tensors a model needs to work for ML-Agents inference with Barracuda. The GridWorldColab scene is configured to work with this ONNX file.\n",
+    "Only policy models need to be exported for inference and they need the following additional tensors:\n",
+    "\n",
+    "*   All models need version_number\n",
+    "*   All models need memory_size\n",
+    "*   Models with continuous outputs need continuous_action_output_shape\n",
+    "*   Models with discrete outputs need discrete_action_output_shape and an additional mask input that matches the shape of the discrete outputs\n",
+    "*   The mask input must be connected to the outputs or it will be pruned on export, if mask values aren't being set they will be 1, so multiplying the discrete outputs by the mask will have no effect"
+   ],
+   "metadata": {
+    "id": "V0V7bH4mOnUF"
+   }
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "class WrapperNet(torch.nn.Module):\n",
+    "    def __init__(\n",
+    "            self,\n",
+    "            qnet: VisualQNetwork,\n",
+    "            discrete_output_sizes: List[int],\n",
+    "    ):\n",
+    "        \"\"\"\n",
+    "        Wraps the VisualQNetwork adding extra constants and dummy mask inputs\n",
+    "        required by runtime inference with Barracuda.\n",
+    "\n",
+    "        For environment continuous actions outputs would need to add them\n",
+    "        similarly to how discrete action outputs work, both in the wrapper\n",
+    "        and in the ONNX output_names / dynamic_axes.\n",
+    "        \"\"\"\n",
+    "        super(WrapperNet, self).__init__()\n",
+    "        self.qnet = qnet\n",
+    "\n",
+    "        # version_number\n",
+    "        #   MLAgents1_0 = 2   (not covered by this example)\n",
+    "        #   MLAgents2_0 = 3\n",
+    "        version_number = torch.Tensor([3])\n",
+    "        self.version_number = Parameter(version_number, requires_grad=False)\n",
+    "\n",
+    "        # memory_size\n",
+    "        # TODO: document case where memory is not zero.\n",
+    "        memory_size = torch.Tensor([0])\n",
+    "        self.memory_size = Parameter(memory_size, requires_grad=False)\n",
+    "\n",
+    "        # discrete_action_output_shape\n",
+    "        output_shape=torch.Tensor([discrete_output_sizes])\n",
+    "        self.discrete_shape = Parameter(output_shape, requires_grad=False)\n",
+    "\n",
+    "\n",
+    "    # if you have discrete actions ML-agents expects corresponding a mask\n",
+    "    # tensor with the same shape to exist as input\n",
+    "    def forward(self, visual_obs: torch.tensor, mask: torch.tensor):\n",
+    "        qnet_result = self.qnet(visual_obs)\n",
+    "        # Connect mask to keep it from getting pruned\n",
+    "        # Mask values will be 1 if you never call SetActionMask() in\n",
+    "        # WriteDiscreteActionMask()\n",
+    "        qnet_result = torch.mul(qnet_result, mask)\n",
+    "        action = torch.argmax(qnet_result, dim=1, keepdim=True)\n",
+    "        return [action], self.discrete_shape, self.version_number, self.memory_size\n",
+    "\n",
+    "\n",
+    "torch.onnx.export(\n",
+    "    WrapperNet(qnet, [num_actions]),\n",
+    "    # A tuple with an example of the input tensors\n",
+    "    (torch.tensor([experiences[0].obs]), torch.ones(1, num_actions)),\n",
+    "    'GridWorldColab.onnx',\n",
+    "    opset_version=9,\n",
+    "    # input_names must correspond to the WrapperNet forward parameters\n",
+    "    # obs will be obs_0, obs_1, etc.\n",
+    "    input_names=[\"obs_0\", \"action_masks\"],\n",
+    "    # output_names must correspond to the return tuple of the WrapperNet\n",
+    "    # forward function.\n",
+    "    output_names=[\"discrete_actions\", \"discrete_action_output_shape\",\n",
+    "                  \"version_number\", \"memory_size\"],\n",
+    "    # All inputs and outputs should have their 0th dimension be designated\n",
+    "    # as 'batch'\n",
+    "    dynamic_axes={'obs_0': {0: 'batch'},\n",
+    "                  'action_masks': {0: 'batch'},\n",
+    "                  'discrete_actions': {0: 'batch'},\n",
+    "                  'discrete_action_output_shape': {0: 'batch'}\n",
+    "                 }\n",
+    "    )"
+   ],
+   "metadata": {
+    "id": "T0zzg8fWpTYO"
+   },
+   "execution_count": null,
+   "outputs": []
+  }
+ ]
+}
\ No newline at end of file
diff --git a/colab/Colab_UnityEnvironment_3_SideChannel.ipynb b/colab/Colab_UnityEnvironment_3_SideChannel.ipynb
index ede4adfc4..735a23380 100644
--- a/colab/Colab_UnityEnvironment_3_SideChannel.ipynb
+++ b/colab/Colab_UnityEnvironment_3_SideChannel.ipynb
@@ -13,15 +13,6 @@
    "name": "python3",
    "language": "python",
    "display_name": "Python 3"
-  },
-  "pycharm": {
-   "stem_cell": {
-    "cell_type": "raw",
-    "source": [],
-    "metadata": {
-     "collapsed": false
-    }
-   }
   }
  },
  "cells": [
@@ -103,18 +94,17 @@
     "esac\n",
     "exit 0\n",
     "    \"\"\")\n",
-    "  pro_bar.update(progress(5, 100))\n",
-    "  !apt-get install daemon >/dev/null 2>&1\n",
+    "  !sudo apt-get update\n",
     "  pro_bar.update(progress(10, 100))\n",
-    "  !apt-get install wget >/dev/null 2>&1\n",
+    "  !sudo DEBIAN_FRONTEND=noninteractive apt install -y daemon wget gdebi-core build-essential libfontenc1 libfreetype6 xorg-dev xorg\n",
     "  pro_bar.update(progress(20, 100))\n",
-    "  !wget http://security.ubuntu.com/ubuntu/pool/main/libx/libxfont/libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb >/dev/null 2>&1\n",
+    "  !wget http://security.ubuntu.com/ubuntu/pool/main/libx/libxfont/libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb 2>&1\n",
     "  pro_bar.update(progress(30, 100))\n",
-    "  !wget --output-document xvfb.deb http://security.ubuntu.com/ubuntu/pool/universe/x/xorg-server/xvfb_1.18.4-0ubuntu0.12_amd64.deb >/dev/null 2>&1\n",
+    "  !wget --output-document xvfb.deb http://security.ubuntu.com/ubuntu/pool/universe/x/xorg-server/xvfb_1.18.4-0ubuntu0.12_amd64.deb 2>&1\n",
     "  pro_bar.update(progress(40, 100))\n",
-    "  !dpkg -i libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb >/dev/null 2>&1\n",
+    "  !sudo dpkg -i libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb 2>&1\n",
     "  pro_bar.update(progress(50, 100))\n",
-    "  !dpkg -i xvfb.deb >/dev/null 2>&1\n",
+    "  !sudo dpkg -i xvfb.deb 2>&1\n",
     "  pro_bar.update(progress(70, 100))\n",
     "  !rm libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb\n",
     "  pro_bar.update(progress(80, 100))\n",
@@ -300,4 +290,4 @@
    ]
   }
  ]
-}
+}
\ No newline at end of file
diff --git a/test_requirements.txt b/test_requirements.txt
index 1a3d424ec..bdd4649be 100644
--- a/test_requirements.txt
+++ b/test_requirements.txt
@@ -1,5 +1,5 @@
 # Test-only dependencies should go here, not in setup.py
-pytest>4.0.0,<6.0.0
+pytest>=6.2.5
 pytest-cov==2.6.1
 pytest-xdist==1.34.0