Final_Project.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wub_l5mXARiV"
      },
      "source": [
        "# Multi-Agent Reinforcement Learning.\n",
        "\n",
        "Multi-agent reinforcement learning (MARL) is a deep learning paradigm that uses two or more agents to train reinforcement learning models. MARL esssentially uses several standard reinforcement learning agents with connected reward functions in order to train each model. \n",
        "\n",
        "MARL is primarily used in cooperative and competitive environments such as chess, go, tic-tac-toe etc. \n",
        "\n",
        "In this project, we are going to train a MARL model, based off of the Deep Q-Network models we made in class, to play chess using the environment provided by `PettingZoo` an similar environnment to the `gym` environments we used in class, but with support for multiplayer games.\n",
        "\n",
        "This project will be based on [this tutorial](https://pettingzoo.farama.org/tutorials/tianshou/advanced/) from PettingZoo. \n",
        "\n",
        "We will also be using the `Tianshou` framework to create and train our model, this makes creating our model architecture and training functions easier, and allows us to focus on how the multiple agents interact with each other."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-ps07tPJCmag"
      },
      "source": [
        "# Setup\n",
        "\n",
        "We need to install PettingZoo, and Tianshou along with some other frameworks to get started. We will also import a lot of libraries as usual."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nP6FYAYkS-yJ",
        "outputId": "caa80047-4eb1-47af-e9f4-349ee96597b2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting pettingzoo[classic]\n",
            "  Downloading PettingZoo-1.22.3-py3-none-any.whl (816 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m816.1/816.1 KB\u001b[0m \u001b[31m27.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.18.0 in /usr/local/lib/python3.8/dist-packages (from pettingzoo[classic]) (1.21.6)\n",
            "Collecting gymnasium>=0.26.0\n",
            "  Downloading gymnasium-0.27.1-py3-none-any.whl (883 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m883.7/883.7 KB\u001b[0m \u001b[31m45.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting rlcard==1.0.5\n",
            "  Downloading rlcard-1.0.5.tar.gz (251 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m251.1/251.1 KB\u001b[0m \u001b[31m15.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting hanabi-learning-environment==0.0.4\n",
            "  Downloading hanabi_learning_environment-0.0.4-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (115 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m115.5/115.5 KB\u001b[0m \u001b[31m5.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting pygame==2.1.3.dev8\n",
            "  Downloading pygame-2.1.3.dev8-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (13.7 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m13.7/13.7 MB\u001b[0m \u001b[31m25.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting chess==1.7.0\n",
            "  Downloading chess-1.7.0-py3-none-any.whl (147 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m147.1/147.1 KB\u001b[0m \u001b[31m10.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: cffi in /usr/local/lib/python3.8/dist-packages (from hanabi-learning-environment==0.0.4->pettingzoo[classic]) (1.15.1)\n",
            "Requirement already satisfied: termcolor in /usr/local/lib/python3.8/dist-packages (from rlcard==1.0.5->pettingzoo[classic]) (2.2.0)\n",
            "Collecting gymnasium-notices>=0.0.1\n",
            "  Downloading gymnasium_notices-0.0.1-py3-none-any.whl (2.8 kB)\n",
            "Requirement already satisfied: importlib-metadata>=4.8.0 in /usr/local/lib/python3.8/dist-packages (from gymnasium>=0.26.0->pettingzoo[classic]) (6.0.0)\n",
            "Collecting jax-jumpy>=0.2.0\n",
            "  Downloading jax_jumpy-0.2.0-py3-none-any.whl (11 kB)\n",
            "Requirement already satisfied: cloudpickle>=1.2.0 in /usr/local/lib/python3.8/dist-packages (from gymnasium>=0.26.0->pettingzoo[classic]) (2.2.0)\n",
            "Requirement already satisfied: typing-extensions>=4.3.0 in /usr/local/lib/python3.8/dist-packages (from gymnasium>=0.26.0->pettingzoo[classic]) (4.4.0)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.8/dist-packages (from importlib-metadata>=4.8.0->gymnasium>=0.26.0->pettingzoo[classic]) (3.11.0)\n",
            "Requirement already satisfied: pycparser in /usr/local/lib/python3.8/dist-packages (from cffi->hanabi-learning-environment==0.0.4->pettingzoo[classic]) (2.21)\n",
            "Building wheels for collected packages: rlcard\n",
            "  Building wheel for rlcard (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for rlcard: filename=rlcard-1.0.5-py3-none-any.whl size=307114 sha256=49f0a6f789df2de5e553371c109fe44fd0ea678fdb0762589100c037566036f5\n",
            "  Stored in directory: /root/.cache/pip/wheels/c2/d5/af/3a3274185e5eca0e51acad1c1c26daad0b341b2ea717dc0d29\n",
            "Successfully built rlcard\n",
            "Installing collected packages: gymnasium-notices, rlcard, pygame, jax-jumpy, chess, hanabi-learning-environment, gymnasium, pettingzoo\n",
            "Successfully installed chess-1.7.0 gymnasium-0.27.1 gymnasium-notices-0.0.1 hanabi-learning-environment-0.0.4 jax-jumpy-0.2.0 pettingzoo-1.22.3 pygame-2.1.3.dev8 rlcard-1.0.5\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting stable-baselines3[extra]\n",
            "  Downloading stable_baselines3-1.7.0-py3-none-any.whl (171 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m171.8/171.8 KB\u001b[0m \u001b[31m6.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.8/dist-packages (from stable-baselines3[extra]) (1.21.6)\n",
            "Requirement already satisfied: cloudpickle in /usr/local/lib/python3.8/dist-packages (from stable-baselines3[extra]) (2.2.0)\n",
            "Requirement already satisfied: torch>=1.11 in /usr/local/lib/python3.8/dist-packages (from stable-baselines3[extra]) (1.13.1+cu116)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.8/dist-packages (from stable-baselines3[extra]) (3.2.2)\n",
            "Collecting gym==0.21\n",
            "  Downloading gym-0.21.0.tar.gz (1.5 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.5/1.5 MB\u001b[0m \u001b[31m50.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting importlib-metadata~=4.13\n",
            "  Downloading importlib_metadata-4.13.0-py3-none-any.whl (23 kB)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.8/dist-packages (from stable-baselines3[extra]) (1.3.5)\n",
            "Requirement already satisfied: tensorboard>=2.9.1 in /usr/local/lib/python3.8/dist-packages (from stable-baselines3[extra]) (2.9.1)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.8/dist-packages (from stable-baselines3[extra]) (5.4.8)\n",
            "Requirement already satisfied: pillow in /usr/local/lib/python3.8/dist-packages (from stable-baselines3[extra]) (7.1.2)\n",
            "Collecting rich\n",
            "  Downloading rich-13.3.1-py3-none-any.whl (239 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m239.0/239.0 KB\u001b[0m \u001b[31m22.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting autorom[accept-rom-license]~=0.4.2\n",
            "  Downloading AutoROM-0.4.2-py3-none-any.whl (16 kB)\n",
            "Collecting ale-py==0.7.4\n",
            "  Downloading ale_py-0.7.4-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.6 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.6/1.6 MB\u001b[0m \u001b[31m59.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: opencv-python in /usr/local/lib/python3.8/dist-packages (from stable-baselines3[extra]) (4.6.0.66)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.8/dist-packages (from stable-baselines3[extra]) (4.64.1)\n",
            "Requirement already satisfied: importlib-resources in /usr/local/lib/python3.8/dist-packages (from ale-py==0.7.4->stable-baselines3[extra]) (5.10.2)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.8/dist-packages (from autorom[accept-rom-license]~=0.4.2->stable-baselines3[extra]) (7.1.2)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.8/dist-packages (from autorom[accept-rom-license]~=0.4.2->stable-baselines3[extra]) (2.25.1)\n",
            "Collecting AutoROM.accept-rom-license\n",
            "  Downloading AutoROM.accept-rom-license-0.5.4.tar.gz (12 kB)\n",
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.8/dist-packages (from importlib-metadata~=4.13->stable-baselines3[extra]) (3.11.0)\n",
            "Requirement already satisfied: werkzeug>=1.0.1 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.9.1->stable-baselines3[extra]) (1.0.1)\n",
            "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.9.1->stable-baselines3[extra]) (0.6.1)\n",
            "Requirement already satisfied: wheel>=0.26 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.9.1->stable-baselines3[extra]) (0.38.4)\n",
            "Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.9.1->stable-baselines3[extra]) (57.4.0)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.9.1->stable-baselines3[extra]) (0.4.6)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.9.1->stable-baselines3[extra]) (3.4.1)\n",
            "Requirement already satisfied: protobuf<3.20,>=3.9.2 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.9.1->stable-baselines3[extra]) (3.19.6)\n",
            "Requirement already satisfied: grpcio>=1.24.3 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.9.1->stable-baselines3[extra]) (1.51.1)\n",
            "Requirement already satisfied: absl-py>=0.4 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.9.1->stable-baselines3[extra]) (1.3.0)\n",
            "Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.9.1->stable-baselines3[extra]) (2.16.0)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.9.1->stable-baselines3[extra]) (1.8.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.8/dist-packages (from torch>=1.11->stable-baselines3[extra]) (4.4.0)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.8/dist-packages (from matplotlib->stable-baselines3[extra]) (0.11.0)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.8/dist-packages (from matplotlib->stable-baselines3[extra]) (2.8.2)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.8/dist-packages (from matplotlib->stable-baselines3[extra]) (3.0.9)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.8/dist-packages (from matplotlib->stable-baselines3[extra]) (1.4.4)\n",
            "Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.8/dist-packages (from pandas->stable-baselines3[extra]) (2022.7)\n",
            "Collecting markdown-it-py<3.0.0,>=2.1.0\n",
            "  Downloading markdown_it_py-2.1.0-py3-none-any.whl (84 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m84.5/84.5 KB\u001b[0m \u001b[31m11.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting pygments<3.0.0,>=2.14.0\n",
            "  Downloading Pygments-2.14.0-py3-none-any.whl (1.1 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.1/1.1 MB\u001b[0m \u001b[31m54.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->stable-baselines3[extra]) (4.9)\n",
            "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->stable-baselines3[extra]) (5.2.1)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->stable-baselines3[extra]) (0.2.8)\n",
            "Requirement already satisfied: six>=1.9.0 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->stable-baselines3[extra]) (1.15.0)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.8/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.9.1->stable-baselines3[extra]) (1.3.1)\n",
            "Collecting mdurl~=0.1\n",
            "  Downloading mdurl-0.1.2-py3-none-any.whl (10.0 kB)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests->autorom[accept-rom-license]~=0.4.2->stable-baselines3[extra]) (2022.12.7)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests->autorom[accept-rom-license]~=0.4.2->stable-baselines3[extra]) (2.10)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests->autorom[accept-rom-license]~=0.4.2->stable-baselines3[extra]) (1.24.3)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests->autorom[accept-rom-license]~=0.4.2->stable-baselines3[extra]) (4.0.0)\n",
            "Collecting libtorrent\n",
            "  Using cached libtorrent-2.0.7-cp38-cp38-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (8.6 MB)\n",
            "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.8/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard>=2.9.1->stable-baselines3[extra]) (0.4.8)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.8/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.9.1->stable-baselines3[extra]) (3.2.2)\n",
            "Building wheels for collected packages: gym, AutoROM.accept-rom-license\n",
            "  Building wheel for gym (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for gym: filename=gym-0.21.0-py3-none-any.whl size=1616824 sha256=a70288935f2323b06565a8d48e399099c2691d74a0456654969c61d5d4939d46\n",
            "  Stored in directory: /root/.cache/pip/wheels/27/6d/b3/a3a6e10704795c9b9000f1ab2dc480dfe7bed42f5972806e73\n",
            "  Building wheel for AutoROM.accept-rom-license (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for AutoROM.accept-rom-license: filename=AutoROM.accept_rom_license-0.5.4-py3-none-any.whl size=441148 sha256=fe2b56ed5e184370a6a014f09bee769902d5b37e89852d46b815f3cdf1f0fa13\n",
            "  Stored in directory: /root/.cache/pip/wheels/64/60/90/db006a24f232de90641041430b5913a601345c9efc4cb883ea\n",
            "Successfully built gym AutoROM.accept-rom-license\n",
            "Installing collected packages: libtorrent, pygments, mdurl, importlib-metadata, gym, markdown-it-py, AutoROM.accept-rom-license, autorom, ale-py, stable-baselines3, rich\n",
            "  Attempting uninstall: pygments\n",
            "    Found existing installation: Pygments 2.6.1\n",
            "    Uninstalling Pygments-2.6.1:\n",
            "      Successfully uninstalled Pygments-2.6.1\n",
            "  Attempting uninstall: importlib-metadata\n",
            "    Found existing installation: importlib-metadata 6.0.0\n",
            "    Uninstalling importlib-metadata-6.0.0:\n",
            "      Successfully uninstalled importlib-metadata-6.0.0\n",
            "  Attempting uninstall: gym\n",
            "    Found existing installation: gym 0.25.2\n",
            "    Uninstalling gym-0.25.2:\n",
            "      Successfully uninstalled gym-0.25.2\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "ipython 7.9.0 requires jedi>=0.10, which is not installed.\u001b[0m\u001b[31m\n",
            "\u001b[0mSuccessfully installed AutoROM.accept-rom-license-0.5.4 ale-py-0.7.4 autorom-0.4.2 gym-0.21.0 importlib-metadata-4.13.0 libtorrent-2.0.7 markdown-it-py-2.1.0 mdurl-0.1.2 pygments-2.14.0 rich-13.3.1 stable-baselines3-1.7.0\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting supersuit\n",
            "  Downloading SuperSuit-3.7.1-py3-none-any.whl (66 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m66.7/66.7 KB\u001b[0m \u001b[31m7.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.19.0 in /usr/local/lib/python3.8/dist-packages (from supersuit) (1.21.6)\n",
            "Requirement already satisfied: gymnasium>=0.26.0 in /usr/local/lib/python3.8/dist-packages (from supersuit) (0.27.1)\n",
            "Requirement already satisfied: jax-jumpy>=0.2.0 in /usr/local/lib/python3.8/dist-packages (from gymnasium>=0.26.0->supersuit) (0.2.0)\n",
            "Requirement already satisfied: typing-extensions>=4.3.0 in /usr/local/lib/python3.8/dist-packages (from gymnasium>=0.26.0->supersuit) (4.4.0)\n",
            "Requirement already satisfied: importlib-metadata>=4.8.0 in /usr/local/lib/python3.8/dist-packages (from gymnasium>=0.26.0->supersuit) (4.13.0)\n",
            "Requirement already satisfied: cloudpickle>=1.2.0 in /usr/local/lib/python3.8/dist-packages (from gymnasium>=0.26.0->supersuit) (2.2.0)\n",
            "Requirement already satisfied: gymnasium-notices>=0.0.1 in /usr/local/lib/python3.8/dist-packages (from gymnasium>=0.26.0->supersuit) (0.0.1)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.8/dist-packages (from importlib-metadata>=4.8.0->gymnasium>=0.26.0->supersuit) (3.11.0)\n",
            "Installing collected packages: supersuit\n",
            "Successfully installed supersuit-3.7.1\n",
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "The following package was automatically installed and is no longer required:\n",
            "  libnvidia-common-510\n",
            "Use 'apt autoremove' to remove it.\n",
            "The following additional packages will be installed:\n",
            "  freeglut3 libpython2-stdlib python2 python2-minimal\n",
            "Suggested packages:\n",
            "  python-tk python-numpy libgle3 python2-doc\n",
            "The following NEW packages will be installed:\n",
            "  freeglut3 libpython2-stdlib python-opengl python2 python2-minimal\n",
            "0 upgraded, 5 newly installed, 0 to remove and 27 not upgraded.\n",
            "Need to get 621 kB of archives.\n",
            "After this operation, 6,059 kB of additional disk space will be used.\n",
            "Get:1 http://archive.ubuntu.com/ubuntu focal/universe amd64 python2-minimal amd64 2.7.17-2ubuntu4 [27.5 kB]\n",
            "Get:2 http://archive.ubuntu.com/ubuntu focal/universe amd64 libpython2-stdlib amd64 2.7.17-2ubuntu4 [7,072 B]\n",
            "Get:3 http://archive.ubuntu.com/ubuntu focal/universe amd64 python2 amd64 2.7.17-2ubuntu4 [26.5 kB]\n",
            "Get:4 http://archive.ubuntu.com/ubuntu focal/universe amd64 freeglut3 amd64 2.8.1-3 [73.6 kB]\n",
            "Get:5 http://archive.ubuntu.com/ubuntu focal/universe amd64 python-opengl all 3.1.0+dfsg-2build1 [486 kB]\n",
            "Fetched 621 kB in 1s (923 kB/s)\n",
            "Selecting previously unselected package python2-minimal.\n",
            "(Reading database ... 129496 files and directories currently installed.)\n",
            "Preparing to unpack .../python2-minimal_2.7.17-2ubuntu4_amd64.deb ...\n",
            "Unpacking python2-minimal (2.7.17-2ubuntu4) ...\n",
            "Selecting previously unselected package libpython2-stdlib:amd64.\n",
            "Preparing to unpack .../libpython2-stdlib_2.7.17-2ubuntu4_amd64.deb ...\n",
            "Unpacking libpython2-stdlib:amd64 (2.7.17-2ubuntu4) ...\n",
            "Setting up python2-minimal (2.7.17-2ubuntu4) ...\n",
            "Selecting previously unselected package python2.\n",
            "(Reading database ... 129525 files and directories currently installed.)\n",
            "Preparing to unpack .../python2_2.7.17-2ubuntu4_amd64.deb ...\n",
            "Unpacking python2 (2.7.17-2ubuntu4) ...\n",
            "Selecting previously unselected package freeglut3:amd64.\n",
            "Preparing to unpack .../freeglut3_2.8.1-3_amd64.deb ...\n",
            "Unpacking freeglut3:amd64 (2.8.1-3) ...\n",
            "Selecting previously unselected package python-opengl.\n",
            "Preparing to unpack .../python-opengl_3.1.0+dfsg-2build1_all.deb ...\n",
            "Unpacking python-opengl (3.1.0+dfsg-2build1) ...\n",
            "Setting up freeglut3:amd64 (2.8.1-3) ...\n",
            "Setting up libpython2-stdlib:amd64 (2.7.17-2ubuntu4) ...\n",
            "Setting up python2 (2.7.17-2ubuntu4) ...\n",
            "Setting up python-opengl (3.1.0+dfsg-2build1) ...\n",
            "Processing triggers for man-db (2.9.1-1) ...\n",
            "Processing triggers for libc-bin (2.31-0ubuntu9.9) ...\n",
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "ffmpeg is already the newest version (7:4.2.7-0ubuntu0.1).\n",
            "The following package was automatically installed and is no longer required:\n",
            "  libnvidia-common-510\n",
            "Use 'apt autoremove' to remove it.\n",
            "0 upgraded, 0 newly installed, 0 to remove and 27 not upgraded.\n",
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "The following package was automatically installed and is no longer required:\n",
            "  libnvidia-common-510\n",
            "Use 'apt autoremove' to remove it.\n",
            "The following NEW packages will be installed:\n",
            "  xvfb\n",
            "0 upgraded, 1 newly installed, 0 to remove and 27 not upgraded.\n",
            "Need to get 780 kB of archives.\n",
            "After this operation, 2,271 kB of additional disk space will be used.\n",
            "Get:1 http://archive.ubuntu.com/ubuntu focal-updates/universe amd64 xvfb amd64 2:1.20.13-1ubuntu1~20.04.5 [780 kB]\n",
            "Fetched 780 kB in 1s (1,232 kB/s)\n",
            "Selecting previously unselected package xvfb.\n",
            "(Reading database ... 131913 files and directories currently installed.)\n",
            "Preparing to unpack .../xvfb_2%3a1.20.13-1ubuntu1~20.04.5_amd64.deb ...\n",
            "Unpacking xvfb (2:1.20.13-1ubuntu1~20.04.5) ...\n",
            "Setting up xvfb (2:1.20.13-1ubuntu1~20.04.5) ...\n",
            "Processing triggers for man-db (2.9.1-1) ...\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting pyvirtualdisplay\n",
            "  Downloading PyVirtualDisplay-3.0-py3-none-any.whl (15 kB)\n",
            "Installing collected packages: pyvirtualdisplay\n",
            "Successfully installed pyvirtualdisplay-3.0\n",
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "The following packages will be REMOVED:\n",
            "  libnvidia-common-510\n",
            "0 upgraded, 0 newly installed, 1 to remove and 27 not upgraded.\n",
            "After this operation, 35.8 kB disk space will be freed.\n",
            "(Reading database ... 131920 files and directories currently installed.)\n",
            "Removing libnvidia-common-510 (510.108.03-0ubuntu1) ...\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting tianshou\n",
            "  Downloading tianshou-0.4.11-py3-none-any.whl (163 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m163.7/163.7 KB\u001b[0m \u001b[31m12.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy>1.16.0 in /usr/local/lib/python3.8/dist-packages (from tianshou) (1.21.6)\n",
            "Requirement already satisfied: protobuf~=3.19.0 in /usr/local/lib/python3.8/dist-packages (from tianshou) (3.19.6)\n",
            "Requirement already satisfied: h5py>=2.10.0 in /usr/local/lib/python3.8/dist-packages (from tianshou) (3.1.0)\n",
            "Requirement already satisfied: tensorboard>=2.5.0 in /usr/local/lib/python3.8/dist-packages (from tianshou) (2.9.1)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.8/dist-packages (from tianshou) (21.3)\n",
            "Collecting gym>=0.23.1\n",
            "  Downloading gym-0.26.2.tar.gz (721 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m721.7/721.7 KB\u001b[0m \u001b[31m54.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.8/dist-packages (from tianshou) (4.64.1)\n",
            "Requirement already satisfied: torch>=1.4.0 in /usr/local/lib/python3.8/dist-packages (from tianshou) (1.13.1+cu116)\n",
            "Requirement already satisfied: numba>=0.51.0 in /usr/local/lib/python3.8/dist-packages (from tianshou) (0.56.4)\n",
            "Requirement already satisfied: importlib-metadata>=4.8.0 in /usr/local/lib/python3.8/dist-packages (from gym>=0.23.1->tianshou) (4.13.0)\n",
            "Requirement already satisfied: cloudpickle>=1.2.0 in /usr/local/lib/python3.8/dist-packages (from gym>=0.23.1->tianshou) (2.2.0)\n",
            "Requirement already satisfied: gym-notices>=0.0.4 in /usr/local/lib/python3.8/dist-packages (from gym>=0.23.1->tianshou) (0.0.8)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.8/dist-packages (from numba>=0.51.0->tianshou) (57.4.0)\n",
            "Requirement already satisfied: llvmlite<0.40,>=0.39.0dev0 in /usr/local/lib/python3.8/dist-packages (from numba>=0.51.0->tianshou) (0.39.1)\n",
            "Requirement already satisfied: grpcio>=1.24.3 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.5.0->tianshou) (1.51.1)\n",
            "Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.5.0->tianshou) (2.16.0)\n",
            "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.5.0->tianshou) (2.25.1)\n",
            "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.5.0->tianshou) (0.6.1)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.5.0->tianshou) (0.4.6)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.5.0->tianshou) (3.4.1)\n",
            "Requirement already satisfied: absl-py>=0.4 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.5.0->tianshou) (1.3.0)\n",
            "Requirement already satisfied: wheel>=0.26 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.5.0->tianshou) (0.38.4)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.5.0->tianshou) (1.8.1)\n",
            "Requirement already satisfied: werkzeug>=1.0.1 in /usr/local/lib/python3.8/dist-packages (from tensorboard>=2.5.0->tianshou) (1.0.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.8/dist-packages (from torch>=1.4.0->tianshou) (4.4.0)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.8/dist-packages (from packaging->tianshou) (3.0.9)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.5.0->tianshou) (0.2.8)\n",
            "Requirement already satisfied: six>=1.9.0 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.5.0->tianshou) (1.15.0)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.5.0->tianshou) (4.9)\n",
            "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.5.0->tianshou) (5.2.1)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.8/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.5.0->tianshou) (1.3.1)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.8/dist-packages (from importlib-metadata>=4.8.0->gym>=0.23.1->tianshou) (3.11.0)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.21.0->tensorboard>=2.5.0->tianshou) (2022.12.7)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.21.0->tensorboard>=2.5.0->tianshou) (2.10)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.21.0->tensorboard>=2.5.0->tianshou) (4.0.0)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.21.0->tensorboard>=2.5.0->tianshou) (1.24.3)\n",
            "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.8/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard>=2.5.0->tianshou) (0.4.8)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.8/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.5.0->tianshou) (3.2.2)\n",
            "Building wheels for collected packages: gym\n",
            "  Building wheel for gym (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for gym: filename=gym-0.26.2-py3-none-any.whl size=827650 sha256=22e8988a5850aaa50702206160e093c52c498e95f68b0754a8cbde8e61a3ef64\n",
            "  Stored in directory: /root/.cache/pip/wheels/17/79/65/7afedc162d858b02708a3b8f7a6dd5b1000dcd5b0f894f7cc1\n",
            "Successfully built gym\n",
            "Installing collected packages: gym, tianshou\n",
            "  Attempting uninstall: gym\n",
            "    Found existing installation: gym 0.21.0\n",
            "    Uninstalling gym-0.21.0:\n",
            "      Successfully uninstalled gym-0.21.0\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "stable-baselines3 1.7.0 requires gym==0.21, but you have gym 0.26.2 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0mSuccessfully installed gym-0.26.2 tianshou-0.4.11\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: pygame in /usr/local/lib/python3.8/dist-packages (2.1.3.dev8)\n"
          ]
        }
      ],
      "source": [
        "import sys\n",
        "IN_COLAB = \"google.colab\" in sys.modules\n",
        "\n",
        "if IN_COLAB:\n",
        "  !pip install pettingzoo[classic]\n",
        "  !pip install stable-baselines3[extra]\n",
        "  !pip install supersuit\n",
        "  !apt install python-opengl\n",
        "  !apt install ffmpeg\n",
        "  !apt install xvfb\n",
        "  !pip install pyvirtualdisplay\n",
        "  !apt autoremove\n",
        "  !pip install tianshou\n",
        "  !pip install pygame"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "zEkw1gj4TUDR"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/Users/rohandatar/Documents/Middlebury Classes/Winter 2023/Deep Learning/Project/Code/mlenv/lib/python3.10/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
            "  from .autonotebook import tqdm as notebook_tqdm\n"
          ]
        }
      ],
      "source": [
        "import argparse\n",
        "import os\n",
        "from copy import deepcopy\n",
        "from typing import Optional, Tuple\n",
        "\n",
        "import gym\n",
        "import numpy as np\n",
        "import torch\n",
        "import pygame\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from tianshou.data import Collector, VectorReplayBuffer\n",
        "from tianshou.env import DummyVectorEnv\n",
        "from tianshou.env.pettingzoo_env import PettingZooEnv\n",
        "from tianshou.policy import BasePolicy, DQNPolicy, MultiAgentPolicyManager, RandomPolicy\n",
        "from tianshou.trainer import offpolicy_trainer\n",
        "from tianshou.utils import TensorboardLogger\n",
        "from tianshou.utils.net.common import Net\n",
        "from torch.utils.tensorboard.writer import SummaryWriter\n",
        "\n",
        "from pettingzoo.classic import chess_v5"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gQTUYNRpC48K"
      },
      "source": [
        "We will use the `argparse` library to format and collect the parameters we need for our model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "id": "_lcPmXZDgv1J"
      },
      "outputs": [
        {
          "ename": "SyntaxError",
          "evalue": "'(' was never closed (1053855554.py, line 18)",
          "output_type": "error",
          "traceback": [
            "\u001b[0;36m  Cell \u001b[0;32mIn[33], line 18\u001b[0;36m\u001b[0m\n\u001b[0;31m    parser.add_argument(\u001b[0m\n\u001b[0m                       ^\u001b[0m\n\u001b[0;31mSyntaxError\u001b[0m\u001b[0;31m:\u001b[0m '(' was never closed\n"
          ]
        }
      ],
      "source": [
        "def get_parser() -> argparse.ArgumentParser:\n",
        "    parser = argparse.ArgumentParser()\n",
        "    parser.add_argument(\"--seed\", type=int, default=1626)\n",
        "    parser.add_argument(\"--eps-test\", type=float, default=0.05)\n",
        "    parser.add_argument(\"--eps-train\", type=float, default=0.1)\n",
        "    parser.add_argument(\"--buffer-size\", type=int, default=20000)\n",
        "    parser.add_argument(\"--lr\", type=float, default=1e-4)\n",
        "    parser.add_argument(\n",
        "        \"--gamma\", type=float, default=0.9, help=\"a smaller gamma favors earlier win\"\n",
        "    )\n",
        "    parser.add_argument(\"--n-step\", type=int, default=3)\n",
        "    parser.add_argument(\"--target-update-freq\", type=int, default=320)\n",
        "    parser.add_argument(\"--epoch\", type=int, default=10) #we only need to train for a few epochs\n",
        "    parser.add_argument(\"--step-per-epoch\", type=int, default=1000)\n",
        "    parser.add_argument(\"--step-per-collect\", type=int, default=10)\n",
        "    parser.add_argument(\"--update-per-step\", type=float, default=0.1)\n",
        "    parser.add_argument(\"--batch-size\", type=int, default=64)\n",
        "    parser.add_argument(\n",
        "        \"--hidden-sizes\", type=int, nargs=\"*\", default=[128, 128, 128, 128] \n",
        "    parser.add_argument(\"--training-num\", type=int, default=10)\n",
        "    parser.add_argument(\"--test-num\", type=int, default=10)\n",
        "    parser.add_argument(\"--logdir\", type=str, default=\"log\")\n",
        "    parser.add_argument(\"--render\", type=float, default=0.1)\n",
        "    parser.add_argument(\n",
        "        \"--win-rate\",\n",
        "        type=float,\n",
        "        default=0.6,\n",
        "        help=\"the expected winning rate: Optimal policy can get 0.7\",\n",
        "    )\n",
        "    parser.add_argument(\n",
        "        \"--watch\",\n",
        "        default=False,\n",
        "        action=\"store_true\",\n",
        "        help=\"no training, \" \"watch the play of pre-trained models\",\n",
        "    )\n",
        "    parser.add_argument(\n",
        "        \"--agent-id\",\n",
        "        type=int,\n",
        "        default=2,\n",
        "        help=\"the learned agent plays as the\"\n",
        "        \" agent_id-th player. Choices are 1 and 2.\",\n",
        "    )\n",
        "    parser.add_argument(\n",
        "        \"--resume-path\",\n",
        "        type=str,\n",
        "        default=\"\",\n",
        "        help=\"the path of agent pth file \" \"for resuming from a pre-trained agent\",\n",
        "    )\n",
        "    parser.add_argument(\n",
        "        \"--opponent-path\",\n",
        "        type=str,\n",
        "        default=\"\",\n",
        "        help=\"the path of opponent agent pth file \"\n",
        "        \"for resuming from a pre-trained agent\",\n",
        "    )\n",
        "    parser.add_argument(\n",
        "        \"--device\", type=str, default=\"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "    )\n",
        "    return parser"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "K7gHlwt5g3Mp"
      },
      "outputs": [],
      "source": [
        "def get_args() -> argparse.Namespace:\n",
        "    parser = get_parser()\n",
        "    return parser.parse_known_args()[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "wIVZdzKqg9xw"
      },
      "outputs": [],
      "source": [
        "def get_env(render_mode=None):\n",
        "    return PettingZooEnv(chess_v5.env(render_mode=render_mode))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mlFbYZXDFvqF"
      },
      "source": [
        "# Architecture\n",
        "\n",
        "Each agent in our model will be a Deep Q-Network."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "7Ew5GRrkXxcU"
      },
      "outputs": [],
      "source": [
        "# This is a rewrite of Tianshou's Net class\n",
        "# Original source at: https://tianshou.readthedocs.io/en/master/_modules/tianshou/utils/net/common.html\n",
        "class Network(nn.Module):\n",
        "  def __init__(self, state_shape, action_shape, hidden_sizes, device=\"cpu\", num_atoms=1)->None:\n",
        "    super(Network, self).__init__()\n",
        "    #get input and output dimensions from state and action space\n",
        "    input_dim = int(np.prod(state_shape))\n",
        "    output_dim = int(np.prod(action_shape)) * num_atoms\n",
        "    # the size of our Linear layers is determined by the hidden_sizes param\n",
        "    # We use ReLU activations as usual\n",
        "    hidden_sizes = [input_dim] + list(hidden_sizes)\n",
        "    model = []\n",
        "    for in_dim, out_dim in zip(hidden_sizes[:-1], hidden_sizes[1:]):\n",
        "        model += [nn.Linear(in_dim, out_dim)]\n",
        "        model += [nn.ReLU()]\n",
        "    model += [nn.Linear(hidden_sizes[-1], output_dim)]\n",
        "    self.layers = nn.Sequential(*model)\n",
        "    self.device = device\n",
        "        \n",
        "\n",
        "  def forward(self, x, state=None, info={}):\n",
        "    if self.device is not None:\n",
        "      x = torch.as_tensor(x, device=self.device, dtype=torch.float32)\n",
        "    x = x.flatten(1)\n",
        "    logits = self.layers(x)\n",
        "    return logits, state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "VS26nBXhg329"
      },
      "outputs": [],
      "source": [
        "def get_agents(\n",
        "    args: argparse.Namespace = get_args(),\n",
        "    agent_learn: Optional[BasePolicy] = None,\n",
        "    agent_opponent: Optional[BasePolicy] = None,\n",
        "    optim: Optional[torch.optim.Optimizer] = None,\n",
        ") -> Tuple[BasePolicy, torch.optim.Optimizer, list]:\n",
        "    env = get_env()\n",
        "    observation_space = (\n",
        "        env.observation_space[\"observation\"]\n",
        "        if isinstance(env.observation_space, gym.spaces.Dict)\n",
        "        else env.observation_space\n",
        "    )\n",
        "    args.state_shape = (\n",
        "        observation_space[\"observation\"].shape or observation_space[\"observation\"].n\n",
        "    )\n",
        "    args.action_shape = env.action_space.shape or env.action_space.n\n",
        "    if agent_learn is None:\n",
        "        # model\n",
        "        net = Network(\n",
        "            args.state_shape,\n",
        "            args.action_shape,\n",
        "            hidden_sizes=args.hidden_sizes,\n",
        "            device=args.device,\n",
        "        ).to(args.device)\n",
        "        if optim is None:\n",
        "            optim = torch.optim.Adam(net.parameters(), lr=args.lr)\n",
        "        agent_learn = DQNPolicy(\n",
        "            net,\n",
        "            optim,\n",
        "            args.gamma,\n",
        "            args.n_step,\n",
        "            target_update_freq=args.target_update_freq,\n",
        "        )\n",
        "        if args.resume_path:\n",
        "            if args.device == \"cpu\":\n",
        "                agent_learn.load_state_dict(torch.load(args.resume_path, map_location=\"cpu\"))\n",
        "            else:\n",
        "                agent_learn.load_state_dict(torch.load(args.resume_path))\n",
        "\n",
        "    if agent_opponent is None:\n",
        "        if args.opponent_path:\n",
        "            agent_opponent = deepcopy(agent_learn)\n",
        "            if args.device == \"cpu\":\n",
        "                agent_learn.load_state_dict(torch.load(args.opponent_path, map_location=\"cpu\"))\n",
        "            else:\n",
        "                agent_learn.load_state_dict(torch.load(args.opponent_path))\n",
        "        else:\n",
        "            agent_opponent = RandomPolicy() #change between RandomPolicy() and agent_learn to train against random agent or self play\n",
        "\n",
        "    if args.agent_id == 1:\n",
        "        agents = [agent_learn, agent_opponent]\n",
        "    else:\n",
        "        agents = [agent_opponent, agent_learn]\n",
        "    policy = MultiAgentPolicyManager(agents, env)\n",
        "    return policy, optim, env.agents"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "jIeVdpwaGAO6"
      },
      "source": [
        "# Training\n",
        "\n",
        "We will use Adam for optimization as usual."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "yI5U_yu5hAlQ"
      },
      "outputs": [],
      "source": [
        "def train_agent(\n",
        "    args: argparse.Namespace = get_args(),\n",
        "    agent_learn: Optional[BasePolicy] = None,\n",
        "    agent_opponent: Optional[BasePolicy] = None,\n",
        "    optim: Optional[torch.optim.Optimizer] = None,\n",
        ") -> Tuple[dict, BasePolicy]:\n",
        "    # ======== environment setup =========\n",
        "    train_envs = DummyVectorEnv([get_env for _ in range(args.training_num)])\n",
        "    test_envs = DummyVectorEnv([get_env for _ in range(args.test_num)])\n",
        "    # seed\n",
        "    np.random.seed(args.seed)\n",
        "    torch.manual_seed(args.seed)\n",
        "    train_envs.seed(args.seed)\n",
        "    test_envs.seed(args.seed)\n",
        "\n",
        "    # ======== agent setup =========\n",
        "    policy, optim, agents = get_agents(\n",
        "        args, agent_learn=agent_learn, agent_opponent=agent_opponent, optim=optim\n",
        "    )\n",
        "\n",
        "    # ======== collector setup =========\n",
        "    train_collector = Collector(\n",
        "        policy,\n",
        "        train_envs,\n",
        "        VectorReplayBuffer(args.buffer_size, len(train_envs)),\n",
        "        exploration_noise=True,\n",
        "    )\n",
        "    test_collector = Collector(policy, test_envs, exploration_noise=True)\n",
        "    # policy.set_eps(1)\n",
        "    train_collector.collect(n_step=args.batch_size * args.training_num)\n",
        "\n",
        "    # ======== tensorboard logging setup =========\n",
        "    log_path = os.path.join(args.logdir, \"chess\", \"dqn\")\n",
        "    writer = SummaryWriter(log_path)\n",
        "    writer.add_text(\"args\", str(args))\n",
        "    logger = TensorboardLogger(writer)\n",
        "\n",
        "    # ======== callback functions used during training =========\n",
        "    def save_best_fn(policy):\n",
        "        if hasattr(args, \"model_save_path\"):\n",
        "            model_save_path = args.model_save_path\n",
        "        else:\n",
        "            model_save_path = os.path.join(\n",
        "                args.logdir, \"chess\", \"dqn\", \"policy.pth\"\n",
        "            )\n",
        "        torch.save(\n",
        "            policy.policies[agents[args.agent_id - 1]].state_dict(), model_save_path\n",
        "        )\n",
        "\n",
        "    def stop_fn(mean_rewards):\n",
        "        return mean_rewards >= args.win_rate\n",
        "\n",
        "    def train_fn(epoch, env_step):\n",
        "        policy.policies[agents[args.agent_id - 1]].set_eps(args.eps_train)\n",
        "\n",
        "    def test_fn(epoch, env_step):\n",
        "        policy.policies[agents[args.agent_id - 1]].set_eps(args.eps_test)\n",
        "\n",
        "    def reward_metric(rews):\n",
        "        return rews[:, args.agent_id - 1]\n",
        "\n",
        "    # trainer\n",
        "    result = offpolicy_trainer(\n",
        "        policy,\n",
        "        train_collector,\n",
        "        test_collector,\n",
        "        args.epoch,\n",
        "        args.step_per_epoch,\n",
        "        args.step_per_collect,\n",
        "        args.test_num,\n",
        "        args.batch_size,\n",
        "        train_fn=train_fn,\n",
        "        test_fn=test_fn,\n",
        "        stop_fn=stop_fn,\n",
        "        save_best_fn=save_best_fn,\n",
        "        update_per_step=args.update_per_step,\n",
        "        logger=logger,\n",
        "        test_in_train=False,\n",
        "        reward_metric=reward_metric,\n",
        "    )\n",
        "\n",
        "    return result, policy.policies[agents[args.agent_id - 1]]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "gsbG2qPUhFSu"
      },
      "outputs": [],
      "source": [
        "# ======== a test function that tests a pre-trained agent ======\n",
        "def watch(\n",
        "    args: argparse.Namespace = get_args(),\n",
        "    agent_learn: Optional[BasePolicy] = None,\n",
        "    agent_opponent: Optional[BasePolicy] = None,\n",
        ") -> None:\n",
        "    env = DummyVectorEnv([lambda: get_env(render_mode=\"human\")])\n",
        "    policy, optim, agents = get_agents(\n",
        "        args, agent_learn=agent_learn, agent_opponent=agent_opponent\n",
        "    )\n",
        "    policy.eval()\n",
        "    policy.policies[agents[args.agent_id - 1]].set_eps(args.eps_test)\n",
        "    collector = Collector(policy, env, exploration_noise=True)\n",
        "    result = collector.collect(n_episode=1, render=args.render)\n",
        "    rews, lens = result[\"rews\"], result[\"lens\"]\n",
        "    print(f\"Final reward: {rews[:, args.agent_id - 1].mean()}, length: {lens.mean()}\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now we can actually train and run our model. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {},
      "outputs": [],
      "source": [
        "args = get_args()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1lMHxJK1hIWj",
        "outputId": "bbb2afe6-2ec6-4dba-83e5-e778f6735aac"
      },
      "outputs": [],
      "source": [
        "# basic trainer\n",
        "# result, agent = train_agent(args)\n",
        "\n",
        "# mixed random and self play trainer\n",
        "self_epochs = 10 #number of epochs to self play\n",
        "rand_epochs = 5 #number of epochs to play against a random player\n",
        "rand_agent = RandomPolicy()\n",
        "\n",
        "args.epoch = rand_epochs\n",
        "result, agent = train_agent(args, agent_opponent=rand_agent)\n",
        "\n",
        "for i in range(5):\n",
        "  #self play\n",
        "  args.epoch = self_epochs\n",
        "  result, agent = train_agent(args, agent_learn=agent, agent_opponent=agent)\n",
        "\n",
        "  #random agent\n",
        "  args.epoch = rand_epochs\n",
        "  result, agent = train_agent(args, agent_learn=agent, agent_opponent=rand_agent)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Watching a Game\n",
        "\n",
        "From this cell, we can either run our existing model against itself or a random agent, or input our own models to play against."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mZnnWVT5Evar",
        "outputId": "6f9e3fd9-8c57-455d-9de2-e06d26d8b15b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Final reward: 0.0, length: 131.0\n"
          ]
        }
      ],
      "source": [
        "#use this line to test against preselected opponent in get_agents\n",
        "# watch(args, agent)\n",
        "\n",
        "# uncomment the following lines to test pretrained agents\n",
        "# args.hidden_sizes = [128, 256, 256, 128] # this line is needed to load mix_train.pth and mix_train_adv.pth\n",
        "# to load other pretrained agents, rerun cell above to reset args\n",
        "args.resume_path = \"self_play.pth\"\n",
        "args.opponent_path = \"Random.pth\"\n",
        "args.agent_id = 2 # 1 for white, 2 for black\n",
        "watch(args)\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "mlenv",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.8"
    },
    "vscode": {
      "interpreter": {
        "hash": "565c47540da7e65ea615b44c43371586a51f28685fd543ee76cf0c99bf32f5e1"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}