PPO playing MicrortsDefeatCoacAIShaped-v3 from https://github.com/sgoodfriend/rl-algo-impls/tree/f7c6f26745a35b21529f65cf3c71dfd6bbf33919

.gitignore

@@ -0,0 +1,147 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# Logging into tensorboard and wandb
+runs/*
+wandb
+
+# macOS
+.DS_STORE
+
+# Local scratch work
+scratch/*
+
+# vscode
+.vscode/
+
+# Don't bother tracking saved_models or videos
+saved_models/*
+downloaded_models/*
+videos/*

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Scott Goodfriend
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,207 @@
+---
+library_name: rl-algo-impls
+tags:
+- MicrortsDefeatCoacAIShaped-v3
+- ppo
+- deep-reinforcement-learning
+- reinforcement-learning
+model-index:
+- name: ppo
+  results:
+  - metrics:
+    - type: mean_reward
+      value: 0.85 +/- 0.53
+      name: mean_reward
+    task:
+      type: reinforcement-learning
+      name: reinforcement-learning
+    dataset:
+      name: MicrortsDefeatCoacAIShaped-v3
+      type: MicrortsDefeatCoacAIShaped-v3
+---
+# **PPO** Agent playing **MicrortsDefeatCoacAIShaped-v3**
+
+This is a trained model of a **PPO** agent playing **MicrortsDefeatCoacAIShaped-v3** using the [/sgoodfriend/rl-algo-impls](https://github.com/sgoodfriend/rl-algo-impls) repo.
+
+All models trained at this commit can be found at https://api.wandb.ai/links/sgoodfriend/sjo3qukl.
+
+## Training Results
+
+This model was trained from 3 trainings of **PPO** agents using different initial seeds. These agents were trained by checking out [f7c6f26](https://github.com/sgoodfriend/rl-algo-impls/tree/f7c6f26745a35b21529f65cf3c71dfd6bbf33919). The best and last models were kept from each training. This submission has loaded the best models from each training, reevaluates them, and selects the best model from these latest evaluations (mean - std).
+
+| algo   | env                           |   seed |   reward_mean |   reward_std |   eval_episodes | best   | wandb_url                                                                    |
+|:-------|:------------------------------|-------:|--------------:|-------------:|----------------:|:-------|:-----------------------------------------------------------------------------|
+| ppo    | MicrortsDefeatCoacAIShaped-v3 |      1 |      0.846154 |     0.532939 |              26 | *      | [wandb](https://wandb.ai/sgoodfriend/rl-algo-impls-benchmarks/runs/coezd1h3) |
+| ppo    | MicrortsDefeatCoacAIShaped-v3 |      2 |      0.615385 |     0.788227 |              26 |        | [wandb](https://wandb.ai/sgoodfriend/rl-algo-impls-benchmarks/runs/leapictl) |
+| ppo    | MicrortsDefeatCoacAIShaped-v3 |      3 |      0.461538 |     0.88712  |              26 |        | [wandb](https://wandb.ai/sgoodfriend/rl-algo-impls-benchmarks/runs/eznziz1s) |
+
+
+### Prerequisites: Weights & Biases (WandB)
+Training and benchmarking assumes you have a Weights & Biases project to upload runs to.
+By default training goes to a rl-algo-impls project while benchmarks go to
+rl-algo-impls-benchmarks. During training and benchmarking runs, videos of the best
+models and the model weights are uploaded to WandB.
+
+Before doing anything below, you'll need to create a wandb account and run `wandb
+login`.
+
+
+
+## Usage
+/sgoodfriend/rl-algo-impls: https://github.com/sgoodfriend/rl-algo-impls
+
+Note: While the model state dictionary and hyperaparameters are saved, the latest
+implementation could be sufficiently different to not be able to reproduce similar
+results. You might need to checkout the commit the agent was trained on:
+[f7c6f26](https://github.com/sgoodfriend/rl-algo-impls/tree/f7c6f26745a35b21529f65cf3c71dfd6bbf33919).
+```
+# Downloads the model, sets hyperparameters, and runs agent for 3 episodes
+python enjoy.py --wandb-run-path=sgoodfriend/rl-algo-impls-benchmarks/coezd1h3
+```
+
+Setup hasn't been completely worked out yet, so you might be best served by using Google
+Colab starting from the
+[colab_enjoy.ipynb](https://github.com/sgoodfriend/rl-algo-impls/blob/main/colab_enjoy.ipynb)
+notebook.
+
+
+
+## Training
+If you want the highest chance to reproduce these results, you'll want to checkout the
+commit the agent was trained on: [f7c6f26](https://github.com/sgoodfriend/rl-algo-impls/tree/f7c6f26745a35b21529f65cf3c71dfd6bbf33919). While
+training is deterministic, different hardware will give different results.
+
+```
+python train.py --algo ppo --env MicrortsDefeatCoacAIShaped-v3 --seed 1
+```
+
+Setup hasn't been completely worked out yet, so you might be best served by using Google
+Colab starting from the
+[colab_train.ipynb](https://github.com/sgoodfriend/rl-algo-impls/blob/main/colab_train.ipynb)
+notebook.
+
+
+
+## Benchmarking (with Lambda Labs instance)
+This and other models from https://api.wandb.ai/links/sgoodfriend/sjo3qukl were generated by running a script on a Lambda
+Labs instance. In a Lambda Labs instance terminal:
+```
+git clone git@github.com:sgoodfriend/rl-algo-impls.git
+cd rl-algo-impls
+bash ./lambda_labs/setup.sh
+wandb login
+bash ./lambda_labs/benchmark.sh [-a {"ppo a2c dqn vpg"}] [-e ENVS] [-j {6}] [-p {rl-algo-impls-benchmarks}] [-s {"1 2 3"}]
+```
+
+### Alternative: Google Colab Pro+
+As an alternative,
+[colab_benchmark.ipynb](https://github.com/sgoodfriend/rl-algo-impls/tree/main/benchmarks#:~:text=colab_benchmark.ipynb),
+can be used. However, this requires a Google Colab Pro+ subscription and running across
+4 separate instances because otherwise running all jobs will exceed the 24-hour limit.
+
+
+
+## Hyperparameters
+This isn't exactly the format of hyperparams in hyperparams/ppo.yml, but instead the Wandb Run Config. However, it's very
+close and has some additional data:
+```
+additional_keys_to_log:
+- microrts_stats
+algo: ppo
+algo_hyperparams:
+  batch_size: 3072
+  clip_range: 0.1
+  clip_range_decay: none
+  clip_range_vf: 0.1
+  ent_coef: 0.01
+  gamma_end: 0.999
+  learning_rate: 0.00025
+  learning_rate_decay: spike
+  max_grad_norm: 0.5
+  n_epochs: 4
+  n_steps: 512
+  ppo2_vf_coef_halving: true
+  vf_coef: 0.5
+device: auto
+env: Microrts-selfplay-unet-decay
+env_hyperparams:
+  env_type: microrts
+  make_kwargs:
+    map_paths:
+    - maps/16x16/basesWorkers16x16.xml
+    max_steps: 2000
+    num_selfplay_envs: 36
+    render_theme: 2
+    reward_weight:
+    - 10
+    - 1
+    - 1
+    - 0.2
+    - 1
+    - 4
+  n_envs: 24
+  self_play_kwargs:
+    num_old_policies: 12
+    save_steps: 200000
+    swap_steps: 10000
+    swap_window_size: 4
+    window: 25
+env_id: MicrortsDefeatCoacAIShaped-v3
+eval_hyperparams:
+  deterministic: false
+  env_overrides:
+    bots:
+      coacAI: 2
+      droplet: 2
+      guidedRojoA3N: 2
+      izanagi: 2
+      lightRushAI: 2
+      mixedBot: 2
+      naiveMCTSAI: 2
+      passiveAI: 2
+      randomAI: 2
+      randomBiasedAI: 2
+      rojo: 2
+      tiamat: 2
+      workerRushAI: 2
+    make_kwargs:
+      map_paths:
+      - maps/16x16/basesWorkers16x16.xml
+      max_steps: 4000
+      num_selfplay_envs: 0
+      render_theme: 2
+      reward_weight:
+      - 1
+      - 0
+      - 0
+      - 0
+      - 0
+      - 0
+    n_envs: 26
+    self_play_kwargs: {}
+  max_video_length: 4000
+  n_episodes: 26
+  score_function: mean
+  step_freq: 1000000
+microrts_reward_decay_callback: true
+n_timesteps: 300000000
+policy_hyperparams:
+  activation_fn: relu
+  actor_head_style: unet
+  cnn_flatten_dim: 256
+  cnn_style: microrts
+  v_hidden_sizes:
+  - 256
+  - 128
+seed: 1
+use_deterministic_algorithms: true
+wandb_entity: null
+wandb_group: null
+wandb_project_name: rl-algo-impls-benchmarks
+wandb_tags:
+- benchmark_f7c6f26
+- host_192-9-151-120
+- branch_selfplay
+- v0.0.9
+
+```

benchmark_publish.py

@@ -0,0 +1,4 @@
+from rl_algo_impls.benchmark_publish import benchmark_publish
+
+if __name__ == "__main__":
+    benchmark_publish()

colab/colab_atari1.sh

@@ -0,0 +1,4 @@
+ALGO="ppo"
+ENVS="PongNoFrameskip-v4 BreakoutNoFrameskip-v4"
+BENCHMARK_MAX_PROCS="${BENCHMARK_MAX_PROCS:-3}"
+bash scripts/train_loop.sh -a $ALGO -e "$ENVS" | xargs -I CMD -P $BENCHMARK_MAX_PROCS bash -c CMD

colab/colab_atari2.sh

@@ -0,0 +1,4 @@
+ALGO="ppo"
+ENVS="SpaceInvadersNoFrameskip-v4 QbertNoFrameskip-v4"
+BENCHMARK_MAX_PROCS="${BENCHMARK_MAX_PROCS:-3}"
+bash scripts/train_loop.sh -a $ALGO -e "$ENVS" | xargs -I CMD -P $BENCHMARK_MAX_PROCS bash -c CMD

colab/colab_basic.sh

@@ -0,0 +1,4 @@
+ALGO="ppo"
+ENVS="CartPole-v1 MountainCar-v0 MountainCarContinuous-v0 Acrobot-v1 LunarLander-v2"
+BENCHMARK_MAX_PROCS="${BENCHMARK_MAX_PROCS:-3}"
+bash scripts/train_loop.sh -a $ALGO -e "$ENVS" | xargs -I CMD -P $BENCHMARK_MAX_PROCS bash -c CMD

colab/colab_benchmark.ipynb

@@ -0,0 +1,195 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "machine_shape": "hm",
+      "authorship_tag": "ABX9TyOGIH7rqgasim3Sz7b1rpoE",
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "gpuClass": "standard",
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sgoodfriend/rl-algo-impls/blob/main/benchmarks/colab_benchmark.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# [sgoodfriend/rl-algo-impls](https://github.com/sgoodfriend/rl-algo-impls) in Google Colaboratory\n",
+        "## Parameters\n",
+        "\n",
+        "\n",
+        "1.   Wandb\n",
+        "\n"
+      ],
+      "metadata": {
+        "id": "S-tXDWP8WTLc"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from getpass import getpass\n",
+        "import os\n",
+        "os.environ[\"WANDB_API_KEY\"] = getpass(\"Wandb API key to upload metrics, videos, and models: \")"
+      ],
+      "metadata": {
+        "id": "1ZtdYgxWNGwZ"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Setup\n",
+        "Clone [sgoodfriend/rl-algo-impls](https://github.com/sgoodfriend/rl-algo-impls) "
+      ],
+      "metadata": {
+        "id": "bsG35Io0hmKG"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%%capture\n",
+        "!git clone https://github.com/sgoodfriend/rl-algo-impls.git"
+      ],
+      "metadata": {
+        "id": "k5ynTV25hdAf"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "Installing the correct packages:\n",
+        "\n",
+        "While conda and poetry are generally used for package management, the mismatch in Python versions (3.10 in the project file vs 3.8 in Colab) makes using the package yml files difficult to use. For now, instead I'm going to specify the list of requirements manually below:"
+      ],
+      "metadata": {
+        "id": "jKxGok-ElYQ7"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%%capture\n",
+        "!apt install python-opengl\n",
+        "!apt install ffmpeg\n",
+        "!apt install xvfb\n",
+        "!apt install swig"
+      ],
+      "metadata": {
+        "id": "nn6EETTc2Ewf"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%%capture\n",
+        "%cd /content/rl-algo-impls\n",
+        "python -m pip install ."
+      ],
+      "metadata": {
+        "id": "AfZh9rH3yQii"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Run Once Per Runtime"
+      ],
+      "metadata": {
+        "id": "4o5HOLjc4wq7"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import wandb\n",
+        "wandb.login()"
+      ],
+      "metadata": {
+        "id": "PCXa5tdS2qFX"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Restart Session beteween runs"
+      ],
+      "metadata": {
+        "id": "AZBZfSUV43JQ"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%%capture\n",
+        "from pyvirtualdisplay import Display\n",
+        "\n",
+        "virtual_display = Display(visible=0, size=(1400, 900))\n",
+        "virtual_display.start()"
+      ],
+      "metadata": {
+        "id": "VzemeQJP2NO9"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "The below 5 bash scripts train agents on environments with 3 seeds each:\n",
+        "- colab_basic.sh and colab_pybullet.sh test on a set of basic gym environments and 4 PyBullet environments. Running both together will likely take about 18 hours. This is likely to run into runtime limits for free Colab and Colab Pro, but is fine for Colab Pro+.\n",
+        "- colab_carracing.sh only trains 3 seeds on CarRacing-v0, which takes almost 22 hours on Colab Pro+ on high-RAM, standard GPU.\n",
+        "- colab_atari1.sh and colab_atari2.sh likely need to be run separately because each takes about 19 hours on high-RAM, standard GPU."
+      ],
+      "metadata": {
+        "id": "nSHfna0hLlO1"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%cd /content/rl-algo-impls\n",
+        "os.environ[\"BENCHMARK_MAX_PROCS\"] = str(1) # Can't reliably raise this to 2+, but would make it faster.\n",
+        "!./benchmarks/colab_basic.sh\n",
+        "!./benchmarks/colab_pybullet.sh\n",
+        "# !./benchmarks/colab_carracing.sh\n",
+        "# !./benchmarks/colab_atari1.sh\n",
+        "# !./benchmarks/colab_atari2.sh"
+      ],
+      "metadata": {
+        "id": "07aHYFH1zfXa"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}

colab/colab_carracing.sh

@@ -0,0 +1,4 @@
+ALGO="ppo"
+ENVS="CarRacing-v0"
+BENCHMARK_MAX_PROCS="${BENCHMARK_MAX_PROCS:-3}"
+bash scripts/train_loop.sh -a $ALGO -e "$ENVS" | xargs -I CMD -P $BENCHMARK_MAX_PROCS bash -c CMD

colab/colab_enjoy.ipynb

@@ -0,0 +1,198 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "machine_shape": "hm",
+      "authorship_tag": "ABX9TyN6S7kyJKrM5x0OOiN+CgTc",
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "gpuClass": "standard",
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sgoodfriend/rl-algo-impls/blob/main/colab_enjoy.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# [sgoodfriend/rl-algo-impls](https://github.com/sgoodfriend/rl-algo-impls) in Google Colaboratory\n",
+        "## Parameters\n",
+        "\n",
+        "\n",
+        "1.   Wandb\n",
+        "\n"
+      ],
+      "metadata": {
+        "id": "S-tXDWP8WTLc"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from getpass import getpass\n",
+        "import os\n",
+        "os.environ[\"WANDB_API_KEY\"] = getpass(\"Wandb API key to upload metrics, videos, and models: \")"
+      ],
+      "metadata": {
+        "id": "1ZtdYgxWNGwZ"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "2. enjoy.py parameters"
+      ],
+      "metadata": {
+        "id": "ao0nAh3MOdN7"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "WANDB_RUN_PATH=\"sgoodfriend/rl-algo-impls-benchmarks/rd0lisee\""
+      ],
+      "metadata": {
+        "id": "jKL_NFhVOjSc"
+      },
+      "execution_count": 2,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Setup\n",
+        "Clone [sgoodfriend/rl-algo-impls](https://github.com/sgoodfriend/rl-algo-impls) "
+      ],
+      "metadata": {
+        "id": "bsG35Io0hmKG"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%%capture\n",
+        "!git clone https://github.com/sgoodfriend/rl-algo-impls.git"
+      ],
+      "metadata": {
+        "id": "k5ynTV25hdAf"
+      },
+      "execution_count": 3,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "Installing the correct packages:\n",
+        "\n",
+        "While conda and poetry are generally used for package management, the mismatch in Python versions (3.10 in the project file vs 3.8 in Colab) makes using the package yml files difficult to use. For now, instead I'm going to specify the list of requirements manually below:"
+      ],
+      "metadata": {
+        "id": "jKxGok-ElYQ7"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%%capture\n",
+        "!apt install python-opengl\n",
+        "!apt install ffmpeg\n",
+        "!apt install xvfb\n",
+        "!apt install swig"
+      ],
+      "metadata": {
+        "id": "nn6EETTc2Ewf"
+      },
+      "execution_count": 4,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%%capture\n",
+        "%cd /content/rl-algo-impls\n",
+        "python -m pip install ."
+      ],
+      "metadata": {
+        "id": "AfZh9rH3yQii"
+      },
+      "execution_count": 5,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Run Once Per Runtime"
+      ],
+      "metadata": {
+        "id": "4o5HOLjc4wq7"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import wandb\n",
+        "wandb.login()"
+      ],
+      "metadata": {
+        "id": "PCXa5tdS2qFX"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Restart Session beteween runs"
+      ],
+      "metadata": {
+        "id": "AZBZfSUV43JQ"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%%capture\n",
+        "from pyvirtualdisplay import Display\n",
+        "\n",
+        "virtual_display = Display(visible=0, size=(1400, 900))\n",
+        "virtual_display.start()"
+      ],
+      "metadata": {
+        "id": "VzemeQJP2NO9"
+      },
+      "execution_count": 7,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%cd /content/rl-algo-impls\n",
+        "!python enjoy.py --wandb-run-path={WANDB_RUN_PATH}"
+      ],
+      "metadata": {
+        "id": "07aHYFH1zfXa"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}

colab/colab_pybullet.sh

@@ -0,0 +1,4 @@
+ALGO="ppo"
+ENVS="HalfCheetahBulletEnv-v0 AntBulletEnv-v0 HopperBulletEnv-v0 Walker2DBulletEnv-v0"
+BENCHMARK_MAX_PROCS="${BENCHMARK_MAX_PROCS:-3}"
+bash scripts/train_loop.sh -a $ALGO -e "$ENVS" | xargs -I CMD -P $BENCHMARK_MAX_PROCS bash -c CMD

colab/colab_train.ipynb

@@ -0,0 +1,200 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "machine_shape": "hm",
+      "authorship_tag": "ABX9TyMmemQnx6G7GOnn6XBdjgxY",
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "gpuClass": "standard",
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sgoodfriend/rl-algo-impls/blob/main/colab_train.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# [sgoodfriend/rl-algo-impls](https://github.com/sgoodfriend/rl-algo-impls) in Google Colaboratory\n",
+        "## Parameters\n",
+        "\n",
+        "\n",
+        "1.   Wandb\n",
+        "\n"
+      ],
+      "metadata": {
+        "id": "S-tXDWP8WTLc"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from getpass import getpass\n",
+        "import os\n",
+        "os.environ[\"WANDB_API_KEY\"] = getpass(\"Wandb API key to upload metrics, videos, and models: \")"
+      ],
+      "metadata": {
+        "id": "1ZtdYgxWNGwZ"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "2. train run parameters"
+      ],
+      "metadata": {
+        "id": "ao0nAh3MOdN7"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "ALGO = \"ppo\"\n",
+        "ENV = \"CartPole-v1\"\n",
+        "SEED = 1"
+      ],
+      "metadata": {
+        "id": "jKL_NFhVOjSc"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Setup\n",
+        "Clone [sgoodfriend/rl-algo-impls](https://github.com/sgoodfriend/rl-algo-impls) "
+      ],
+      "metadata": {
+        "id": "bsG35Io0hmKG"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%%capture\n",
+        "!git clone https://github.com/sgoodfriend/rl-algo-impls.git"
+      ],
+      "metadata": {
+        "id": "k5ynTV25hdAf"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "Installing the correct packages:\n",
+        "\n",
+        "While conda and poetry are generally used for package management, the mismatch in Python versions (3.10 in the project file vs 3.8 in Colab) makes using the package yml files difficult to use. For now, instead I'm going to specify the list of requirements manually below:"
+      ],
+      "metadata": {
+        "id": "jKxGok-ElYQ7"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%%capture\n",
+        "!apt install python-opengl\n",
+        "!apt install ffmpeg\n",
+        "!apt install xvfb\n",
+        "!apt install swig"
+      ],
+      "metadata": {
+        "id": "nn6EETTc2Ewf"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%%capture\n",
+        "%cd /content/rl-algo-impls\n",
+        "python -m pip install ."
+      ],
+      "metadata": {
+        "id": "AfZh9rH3yQii"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Run Once Per Runtime"
+      ],
+      "metadata": {
+        "id": "4o5HOLjc4wq7"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import wandb\n",
+        "wandb.login()"
+      ],
+      "metadata": {
+        "id": "PCXa5tdS2qFX"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Restart Session beteween runs"
+      ],
+      "metadata": {
+        "id": "AZBZfSUV43JQ"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%%capture\n",
+        "from pyvirtualdisplay import Display\n",
+        "\n",
+        "virtual_display = Display(visible=0, size=(1400, 900))\n",
+        "virtual_display.start()"
+      ],
+      "metadata": {
+        "id": "VzemeQJP2NO9"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%cd /content/rl-algo-impls\n",
+        "!python train.py --algo {ALGO} --env {ENV} --seed {SEED}"
+      ],
+      "metadata": {
+        "id": "07aHYFH1zfXa"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}

compare_runs.py

@@ -0,0 +1,4 @@
+from rl_algo_impls.compare_runs import compare_runs
+
+if __name__ == "__main__":
+    compare_runs()

enjoy.py

@@ -0,0 +1,4 @@
+from rl_algo_impls.enjoy import enjoy
+
+if __name__ == "__main__":
+    enjoy()

environment.yml

@@ -0,0 +1,12 @@
+name: rl_algo_impls
+channels:
+  - pytorch
+  - conda-forge
+  - nodefaults
+dependencies:
+  - python>=3.8, <3.10
+  - mamba
+  - pip
+  - pytorch
+  - torchvision
+  - torchaudio

huggingface_publish.py

@@ -0,0 +1,4 @@
+from rl_algo_impls.huggingface_publish import huggingface_publish
+
+if __name__ == "__main__":
+    huggingface_publish()

optimize.py

@@ -0,0 +1,4 @@
+from rl_algo_impls.optimize import optimize
+
+if __name__ == "__main__":
+    optimize()

pyproject.toml

@@ -0,0 +1,88 @@
+[project]
+name = "rl_algo_impls"
+version = "0.0.9"
+description = "Implementations of reinforcement learning algorithms"
+authors = [
+    {name = "Scott Goodfriend", email = "goodfriend.scott@gmail.com"},
+]
+license = {file = "LICENSE"}
+readme = "README.md"
+requires-python = ">= 3.8"
+classifiers = [
+    "License :: OSI Approved :: MIT License",
+    "Development Status :: 3 - Alpha",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+]
+dependencies = [
+    "cmake",
+    "swig",
+    "scipy",
+    "torch",
+    "torchvision",
+    "tensorboard >= 2.11.2, < 2.12",
+    "AutoROM.accept-rom-license >= 0.4.2, < 0.5",
+    "stable-baselines3[extra] >= 1.7.0, < 1.8",
+    "gym[box2d] >= 0.21.0, < 0.22",
+    "pyglet == 1.5.27",
+    "wandb",
+    "pyvirtualdisplay",
+    "pybullet",
+    "tabulate",
+    "huggingface-hub",
+    "optuna",
+    "dash",
+    "kaleido",
+    "PyYAML",
+    "scikit-learn",
+]
+
+[tool.setuptools]
+packages = ["rl_algo_impls"]
+
+[project.optional-dependencies]
+test = [
+    "pytest",
+    "black",
+    "mypy",
+    "flake8",
+    "flake8-bugbear",
+    "isort",
+]
+procgen = [
+    "numexpr >= 2.8.4",
+    "gym3",
+    "glfw >= 1.12.0, < 1.13",
+    "procgen; platform_machine=='x86_64'",
+]
+microrts-ppo = [
+    "numpy < 1.24.0", # Support for gym-microrts < 0.6.0
+    "gym-microrts == 0.2.0", # Match ppo-implementation-details
+]
+microrts-paper = [
+    "numpy < 1.24.0", # Support for gym-microrts < 0.6.0
+    "gym-microrts == 0.3.2",
+]
+microrts = [
+    "gym-microrts",
+]
+jupyter = [
+    "jupyter",
+    "notebook"
+]
+all = [
+    "rl-algo-impls[test]",
+    "rl-algo-impls[procgen]",
+    "rl-algo-impls[microrts]",
+]
+
+[project.urls]
+"Homepage" = "https://github.com/sgoodfriend/rl-algo-impls"
+
+[build-system]
+requires = ["setuptools==65.5.0", "setuptools-scm"]
+build-backend = "setuptools.build_meta"
+
+[tool.isort]
+profile = "black" 

replay.meta.json

@@ -0,0 +1 @@
+{"content_type": "video/mp4", "encoder_version": {"backend": "ffmpeg", "version": "b'ffmpeg version 4.2.7-0ubuntu0.1 Copyright (c) 2000-2022 the FFmpeg developers\\nbuilt with gcc 9 (Ubuntu 9.4.0-1ubuntu1~20.04.1)\\nconfiguration: --prefix=/usr --extra-version=0ubuntu0.1 --toolchain=hardened --libdir=/usr/lib/x86_64-linux-gnu --incdir=/usr/include/x86_64-linux-gnu --arch=amd64 --enable-gpl --disable-stripping --enable-avresample --disable-filter=resample --enable-avisynth --enable-gnutls --enable-ladspa --enable-libaom --enable-libass --enable-libbluray --enable-libbs2b --enable-libcaca --enable-libcdio --enable-libcodec2 --enable-libflite --enable-libfontconfig --enable-libfreetype --enable-libfribidi --enable-libgme --enable-libgsm --enable-libjack --enable-libmp3lame --enable-libmysofa --enable-libopenjpeg --enable-libopenmpt --enable-libopus --enable-libpulse --enable-librsvg --enable-librubberband --enable-libshine --enable-libsnappy --enable-libsoxr --enable-libspeex --enable-libssh --enable-libtheora --enable-libtwolame --enable-libvidstab --enable-libvorbis --enable-libvpx --enable-libwavpack --enable-libwebp --enable-libx265 --enable-libxml2 --enable-libxvid --enable-libzmq --enable-libzvbi --enable-lv2 --enable-omx --enable-openal --enable-opencl --enable-opengl --enable-sdl2 --enable-libdc1394 --enable-libdrm --enable-libiec61883 --enable-nvenc --enable-chromaprint --enable-frei0r --enable-libx264 --enable-shared\\nlibavutil      56. 31.100 / 56. 31.100\\nlibavcodec     58. 54.100 / 58. 54.100\\nlibavformat    58. 29.100 / 58. 29.100\\nlibavdevice    58.  8.100 / 58.  8.100\\nlibavfilter     7. 57.100 /  7. 57.100\\nlibavresample   4.  0.  0 /  4.  0.  0\\nlibswscale      5.  5.100 /  5.  5.100\\nlibswresample   3.  5.100 /  3.  5.100\\nlibpostproc    55.  5.100 / 55.  5.100\\n'", "cmdline": ["ffmpeg", "-nostats", "-loglevel", "error", "-y", "-f", "rawvideo", "-s:v", "640x640", "-pix_fmt", "rgb24", "-framerate", "150", "-i", "-", "-vf", "scale=trunc(iw/2)*2:trunc(ih/2)*2", "-vcodec", "libx264", "-pix_fmt", "yuv420p", "-r", "150", "/tmp/tmp2qhg_um5/ppo-Microrts-selfplay-unet-decay/replay.mp4"]}, "episode": {"r": 1.0, "l": 801, "t": 10.68842}}

rl_algo_impls/a2c/a2c.py

@@ -0,0 +1,205 @@
+import logging
+from time import perf_counter
+from typing import List, Optional, TypeVar
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.tensorboard.writer import SummaryWriter
+
+from rl_algo_impls.shared.algorithm import Algorithm
+from rl_algo_impls.shared.callbacks import Callback
+from rl_algo_impls.shared.gae import compute_advantages
+from rl_algo_impls.shared.policy.actor_critic import ActorCritic
+from rl_algo_impls.shared.schedule import schedule, update_learning_rate
+from rl_algo_impls.shared.stats import log_scalars
+from rl_algo_impls.wrappers.vectorable_wrapper import (
+    VecEnv,
+    single_action_space,
+    single_observation_space,
+)
+
+A2CSelf = TypeVar("A2CSelf", bound="A2C")
+
+
+class A2C(Algorithm):
+    def __init__(
+        self,
+        policy: ActorCritic,
+        env: VecEnv,
+        device: torch.device,
+        tb_writer: SummaryWriter,
+        learning_rate: float = 7e-4,
+        learning_rate_decay: str = "none",
+        n_steps: int = 5,
+        gamma: float = 0.99,
+        gae_lambda: float = 1.0,
+        ent_coef: float = 0.0,
+        ent_coef_decay: str = "none",
+        vf_coef: float = 0.5,
+        max_grad_norm: float = 0.5,
+        rms_prop_eps: float = 1e-5,
+        use_rms_prop: bool = True,
+        sde_sample_freq: int = -1,
+        normalize_advantage: bool = False,
+    ) -> None:
+        super().__init__(policy, env, device, tb_writer)
+        self.policy = policy
+
+        self.lr_schedule = schedule(learning_rate_decay, learning_rate)
+        if use_rms_prop:
+            self.optimizer = torch.optim.RMSprop(
+                policy.parameters(), lr=learning_rate, eps=rms_prop_eps
+            )
+        else:
+            self.optimizer = torch.optim.Adam(policy.parameters(), lr=learning_rate)
+
+        self.n_steps = n_steps
+
+        self.gamma = gamma
+        self.gae_lambda = gae_lambda
+
+        self.vf_coef = vf_coef
+        self.ent_coef_schedule = schedule(ent_coef_decay, ent_coef)
+        self.max_grad_norm = max_grad_norm
+
+        self.sde_sample_freq = sde_sample_freq
+        self.normalize_advantage = normalize_advantage
+
+    def learn(
+        self: A2CSelf,
+        train_timesteps: int,
+        callbacks: Optional[List[Callback]] = None,
+        total_timesteps: Optional[int] = None,
+        start_timesteps: int = 0,
+    ) -> A2CSelf:
+        if total_timesteps is None:
+            total_timesteps = train_timesteps
+        assert start_timesteps + train_timesteps <= total_timesteps
+        epoch_dim = (self.n_steps, self.env.num_envs)
+        step_dim = (self.env.num_envs,)
+        obs_space = single_observation_space(self.env)
+        act_space = single_action_space(self.env)
+
+        obs = np.zeros(epoch_dim + obs_space.shape, dtype=obs_space.dtype)
+        actions = np.zeros(epoch_dim + act_space.shape, dtype=act_space.dtype)
+        rewards = np.zeros(epoch_dim, dtype=np.float32)
+        episode_starts = np.zeros(epoch_dim, dtype=np.bool8)
+        values = np.zeros(epoch_dim, dtype=np.float32)
+        logprobs = np.zeros(epoch_dim, dtype=np.float32)
+
+        next_obs = self.env.reset()
+        next_episode_starts = np.full(step_dim, True, dtype=np.bool8)
+
+        timesteps_elapsed = start_timesteps
+        while timesteps_elapsed < start_timesteps + train_timesteps:
+            start_time = perf_counter()
+
+            progress = timesteps_elapsed / total_timesteps
+            ent_coef = self.ent_coef_schedule(progress)
+            learning_rate = self.lr_schedule(progress)
+            update_learning_rate(self.optimizer, learning_rate)
+            log_scalars(
+                self.tb_writer,
+                "charts",
+                {
+                    "ent_coef": ent_coef,
+                    "learning_rate": learning_rate,
+                },
+                timesteps_elapsed,
+            )
+
+            self.policy.eval()
+            self.policy.reset_noise()
+            for s in range(self.n_steps):
+                timesteps_elapsed += self.env.num_envs
+                if self.sde_sample_freq > 0 and s > 0 and s % self.sde_sample_freq == 0:
+                    self.policy.reset_noise()
+
+                obs[s] = next_obs
+                episode_starts[s] = next_episode_starts
+
+                actions[s], values[s], logprobs[s], clamped_action = self.policy.step(
+                    next_obs
+                )
+                next_obs, rewards[s], next_episode_starts, _ = self.env.step(
+                    clamped_action
+                )
+
+            advantages = compute_advantages(
+                rewards,
+                values,
+                episode_starts,
+                next_episode_starts,
+                next_obs,
+                self.policy,
+                self.gamma,
+                self.gae_lambda,
+            )
+            returns = advantages + values
+
+            b_obs = torch.tensor(obs.reshape((-1,) + obs_space.shape)).to(self.device)
+            b_actions = torch.tensor(actions.reshape((-1,) + act_space.shape)).to(
+                self.device
+            )
+            b_advantages = torch.tensor(advantages.reshape(-1)).to(self.device)
+            b_returns = torch.tensor(returns.reshape(-1)).to(self.device)
+
+            if self.normalize_advantage:
+                b_advantages = (b_advantages - b_advantages.mean()) / (
+                    b_advantages.std() + 1e-8
+                )
+
+            self.policy.train()
+            logp_a, entropy, v = self.policy(b_obs, b_actions)
+
+            pi_loss = -(b_advantages * logp_a).mean()
+            value_loss = F.mse_loss(b_returns, v)
+            entropy_loss = -entropy.mean()
+
+            loss = pi_loss + self.vf_coef * value_loss + ent_coef * entropy_loss
+
+            self.optimizer.zero_grad()
+            loss.backward()
+            nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
+            self.optimizer.step()
+
+            y_pred = values.reshape(-1)
+            y_true = returns.reshape(-1)
+            var_y = np.var(y_true).item()
+            explained_var = (
+                np.nan if var_y == 0 else 1 - np.var(y_true - y_pred).item() / var_y
+            )
+
+            end_time = perf_counter()
+            rollout_steps = self.n_steps * self.env.num_envs
+            self.tb_writer.add_scalar(
+                "train/steps_per_second",
+                (rollout_steps) / (end_time - start_time),
+                timesteps_elapsed,
+            )
+
+            log_scalars(
+                self.tb_writer,
+                "losses",
+                {
+                    "loss": loss.item(),
+                    "pi_loss": pi_loss.item(),
+                    "v_loss": value_loss.item(),
+                    "entropy_loss": entropy_loss.item(),
+                    "explained_var": explained_var,
+                },
+                timesteps_elapsed,
+            )
+
+            if callbacks:
+                if not all(
+                    c.on_step(timesteps_elapsed=rollout_steps) for c in callbacks
+                ):
+                    logging.info(
+                        f"Callback terminated training at {timesteps_elapsed} timesteps"
+                    )
+                    break
+
+        return self

rl_algo_impls/a2c/optimize.py

@@ -0,0 +1,77 @@
+import optuna
+
+from copy import deepcopy
+
+from rl_algo_impls.runner.config import Config, Hyperparams, EnvHyperparams
+from rl_algo_impls.shared.vec_env import make_eval_env
+from rl_algo_impls.shared.policy.optimize_on_policy import sample_on_policy_hyperparams
+from rl_algo_impls.tuning.optimize_env import sample_env_hyperparams
+
+
+def sample_params(
+    trial: optuna.Trial,
+    base_hyperparams: Hyperparams,
+    base_config: Config,
+) -> Hyperparams:
+    hyperparams = deepcopy(base_hyperparams)
+
+    base_env_hyperparams = EnvHyperparams(**hyperparams.env_hyperparams)
+    env = make_eval_env(base_config, base_env_hyperparams, override_n_envs=1)
+
+    # env_hyperparams
+    env_hyperparams = sample_env_hyperparams(trial, hyperparams.env_hyperparams, env)
+
+    # policy_hyperparams
+    policy_hyperparams = sample_on_policy_hyperparams(
+        trial, hyperparams.policy_hyperparams, env
+    )
+
+    # algo_hyperparams
+    algo_hyperparams = hyperparams.algo_hyperparams
+
+    learning_rate = trial.suggest_float("learning_rate", 1e-5, 2e-3, log=True)
+    learning_rate_decay = trial.suggest_categorical(
+        "learning_rate_decay", ["none", "linear"]
+    )
+    n_steps_exp = trial.suggest_int("n_steps_exp", 1, 10)
+    n_steps = 2**n_steps_exp
+    trial.set_user_attr("n_steps", n_steps)
+    gamma = 1.0 - trial.suggest_float("gamma_om", 1e-4, 1e-1, log=True)
+    trial.set_user_attr("gamma", gamma)
+    gae_lambda = 1 - trial.suggest_float("gae_lambda_om", 1e-4, 1e-1)
+    trial.set_user_attr("gae_lambda", gae_lambda)
+    ent_coef = trial.suggest_float("ent_coef", 1e-8, 2.5e-2, log=True)
+    ent_coef_decay = trial.suggest_categorical("ent_coef_decay", ["none", "linear"])
+    vf_coef = trial.suggest_float("vf_coef", 0.1, 0.7)
+    max_grad_norm = trial.suggest_float("max_grad_norm", 1e-1, 1e1, log=True)
+    use_rms_prop = trial.suggest_categorical("use_rms_prop", [True, False])
+    normalize_advantage = trial.suggest_categorical(
+        "normalize_advantage", [True, False]
+    )
+
+    algo_hyperparams.update(
+        {
+            "learning_rate": learning_rate,
+            "learning_rate_decay": learning_rate_decay,
+            "n_steps": n_steps,
+            "gamma": gamma,
+            "gae_lambda": gae_lambda,
+            "ent_coef": ent_coef,
+            "ent_coef_decay": ent_coef_decay,
+            "vf_coef": vf_coef,
+            "max_grad_norm": max_grad_norm,
+            "use_rms_prop": use_rms_prop,
+            "normalize_advantage": normalize_advantage,
+        }
+    )
+
+    if policy_hyperparams.get("use_sde", False):
+        sde_sample_freq = 2 ** trial.suggest_int("sde_sample_freq_exp", 0, n_steps_exp)
+        trial.set_user_attr("sde_sample_freq", sde_sample_freq)
+        algo_hyperparams["sde_sample_freq"] = sde_sample_freq
+    elif "sde_sample_freq" in algo_hyperparams:
+        del algo_hyperparams["sde_sample_freq"]
+
+    env.close()
+
+    return hyperparams

rl_algo_impls/benchmark_publish.py

@@ -0,0 +1,111 @@
+import argparse
+import subprocess
+import wandb
+import wandb.apis.public
+
+from collections import defaultdict
+from multiprocessing.pool import ThreadPool
+from typing import List, NamedTuple
+
+
+class RunGroup(NamedTuple):
+    algo: str
+    env_id: str
+
+
+def benchmark_publish() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--wandb-project-name",
+        type=str,
+        default="rl-algo-impls-benchmarks",
+        help="WandB project name to load runs from",
+    )
+    parser.add_argument(
+        "--wandb-entity",
+        type=str,
+        default=None,
+        help="WandB team of project. None uses default entity",
+    )
+    parser.add_argument("--wandb-tags", type=str, nargs="+", help="WandB tags")
+    parser.add_argument("--wandb-report-url", type=str, help="Link to WandB report")
+    parser.add_argument(
+        "--envs", type=str, nargs="*", help="Optional filter down to these envs"
+    )
+    parser.add_argument(
+        "--exclude-envs",
+        type=str,
+        nargs="*",
+        help="Environments to exclude from publishing",
+    )
+    parser.add_argument(
+        "--huggingface-user",
+        type=str,
+        default=None,
+        help="Huggingface user or team to upload model cards. Defaults to huggingface-cli login user",
+    )
+    parser.add_argument(
+        "--pool-size",
+        type=int,
+        default=3,
+        help="How many publish jobs can run in parallel",
+    )
+    parser.add_argument(
+        "--virtual-display", action="store_true", help="Use headless virtual display"
+    )
+    # parser.set_defaults(
+    #     wandb_tags=["benchmark_e47a44c", "host_129-146-2-230"],
+    #     wandb_report_url="https://api.wandb.ai/links/sgoodfriend/v4wd7cp5",
+    #     envs=[],
+    #     exclude_envs=[],
+    # )
+    args = parser.parse_args()
+    print(args)
+
+    api = wandb.Api()
+    all_runs = api.runs(
+        f"{args.wandb_entity or api.default_entity}/{args.wandb_project_name}"
+    )
+
+    required_tags = set(args.wandb_tags)
+    runs: List[wandb.apis.public.Run] = [
+        r
+        for r in all_runs
+        if required_tags.issubset(set(r.config.get("wandb_tags", [])))
+    ]
+
+    runs_paths_by_group = defaultdict(list)
+    for r in runs:
+        if r.state != "finished":
+            continue
+        algo = r.config["algo"]
+        env = r.config["env"]
+        if args.envs and env not in args.envs:
+            continue
+        if args.exclude_envs and env in args.exclude_envs:
+            continue
+        run_group = RunGroup(algo, env)
+        runs_paths_by_group[run_group].append("/".join(r.path))
+
+    def run(run_paths: List[str]) -> None:
+        publish_args = ["python", "huggingface_publish.py"]
+        publish_args.append("--wandb-run-paths")
+        publish_args.extend(run_paths)
+        publish_args.append("--wandb-report-url")
+        publish_args.append(args.wandb_report_url)
+        if args.huggingface_user:
+            publish_args.append("--huggingface-user")
+            publish_args.append(args.huggingface_user)
+        if args.virtual_display:
+            publish_args.append("--virtual-display")
+        subprocess.run(publish_args)
+
+    tp = ThreadPool(args.pool_size)
+    for run_paths in runs_paths_by_group.values():
+        tp.apply_async(run, (run_paths,))
+    tp.close()
+    tp.join()
+
+
+if __name__ == "__main__":
+    benchmark_publish()

rl_algo_impls/compare_runs.py

@@ -0,0 +1,199 @@
+import argparse
+import itertools
+import numpy as np
+import pandas as pd
+import wandb
+import wandb.apis.public
+
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Dict, Iterable, List, TypeVar
+
+from rl_algo_impls.benchmark_publish import RunGroup
+
+
+@dataclass
+class Comparison:
+    control_values: List[float]
+    experiment_values: List[float]
+
+    def mean_diff_percentage(self) -> float:
+        return self._diff_percentage(
+            np.mean(self.control_values).item(), np.mean(self.experiment_values).item()
+        )
+
+    def median_diff_percentage(self) -> float:
+        return self._diff_percentage(
+            np.median(self.control_values).item(),
+            np.median(self.experiment_values).item(),
+        )
+
+    def _diff_percentage(self, c: float, e: float) -> float:
+        if c == e:
+            return 0
+        elif c == 0:
+            return float("inf") if e > 0 else float("-inf")
+        return 100 * (e - c) / c
+
+    def score(self) -> float:
+        return (
+            np.sum(
+                np.sign((self.mean_diff_percentage(), self.median_diff_percentage()))
+            ).item()
+            / 2
+        )
+
+
+RunGroupRunsSelf = TypeVar("RunGroupRunsSelf", bound="RunGroupRuns")
+
+
+class RunGroupRuns:
+    def __init__(
+        self,
+        run_group: RunGroup,
+        control: List[str],
+        experiment: List[str],
+        summary_stats: List[str] = ["best_eval", "eval", "train_rolling"],
+        summary_metrics: List[str] = ["mean", "result"],
+    ) -> None:
+        self.algo = run_group.algo
+        self.env = run_group.env_id
+        self.control = set(control)
+        self.experiment = set(experiment)
+
+        self.summary_stats = summary_stats
+        self.summary_metrics = summary_metrics
+
+        self.control_runs = []
+        self.experiment_runs = []
+
+    def add_run(self, run: wandb.apis.public.Run) -> None:
+        wandb_tags = set(run.config.get("wandb_tags", []))
+        if self.control & wandb_tags:
+            self.control_runs.append(run)
+        elif self.experiment & wandb_tags:
+            self.experiment_runs.append(run)
+
+    def comparisons_by_metric(self) -> Dict[str, Comparison]:
+        c_by_m = {}
+        for metric in (
+            f"{s}/{m}"
+            for s, m in itertools.product(self.summary_stats, self.summary_metrics)
+        ):
+            c_by_m[metric] = Comparison(
+                [c.summary[metric] for c in self.control_runs],
+                [e.summary[metric] for e in self.experiment_runs],
+            )
+        return c_by_m
+
+    @staticmethod
+    def data_frame(rows: Iterable[RunGroupRunsSelf]) -> pd.DataFrame:
+        results = defaultdict(list)
+        for r in rows:
+            if not r.control_runs or not r.experiment_runs:
+                continue
+            results["algo"].append(r.algo)
+            results["env"].append(r.env)
+            results["control"].append(r.control)
+            results["expierment"].append(r.experiment)
+            c_by_m = r.comparisons_by_metric()
+            results["score"].append(
+                sum(m.score() for m in c_by_m.values()) / len(c_by_m)
+            )
+            for m, c in c_by_m.items():
+                results[f"{m}_mean"].append(c.mean_diff_percentage())
+                results[f"{m}_median"].append(c.median_diff_percentage())
+        return pd.DataFrame(results)
+
+
+def compare_runs() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-p",
+        "--wandb-project-name",
+        type=str,
+        default="rl-algo-impls-benchmarks",
+        help="WandB project name to load runs from",
+    )
+    parser.add_argument(
+        "--wandb-entity",
+        type=str,
+        default=None,
+        help="WandB team. None uses default entity",
+    )
+    parser.add_argument(
+        "-n",
+        "--wandb-hostname-tag",
+        type=str,
+        nargs="*",
+        help="WandB tags for hostname (i.e. host_192-9-145-26)",
+    )
+    parser.add_argument(
+        "-c",
+        "--wandb-control-tag",
+        type=str,
+        nargs="+",
+        help="WandB tag for control commit (i.e. benchmark_5598ebc)",
+    )
+    parser.add_argument(
+        "-e",
+        "--wandb-experiment-tag",
+        type=str,
+        nargs="+",
+        help="WandB tag for experiment commit (i.e. benchmark_5540e1f)",
+    )
+    parser.add_argument(
+        "--envs",
+        type=str,
+        nargs="*",
+        help="If specified, only compare these envs",
+    )
+    parser.add_argument(
+        "--exclude-envs",
+        type=str,
+        nargs="*",
+        help="Environments to exclude from comparison",
+    )
+    # parser.set_defaults(
+    #     wandb_hostname_tag=["host_150-230-44-105", "host_155-248-214-128"],
+    #     wandb_control_tag=["benchmark_fbc943f"],
+    #     wandb_experiment_tag=["benchmark_f59bf74"],
+    #     exclude_envs=[],
+    # )
+    args = parser.parse_args()
+    print(args)
+
+    api = wandb.Api()
+    all_runs = api.runs(
+        path=f"{args.wandb_entity or api.default_entity}/{args.wandb_project_name}",
+        order="+created_at",
+    )
+
+    runs_by_run_group: Dict[RunGroup, RunGroupRuns] = {}
+    wandb_hostname_tags = set(args.wandb_hostname_tag)
+    for r in all_runs:
+        if r.state != "finished":
+            continue
+        wandb_tags = set(r.config.get("wandb_tags", []))
+        if not wandb_tags or not wandb_hostname_tags & wandb_tags:
+            continue
+        rg = RunGroup(r.config["algo"], r.config.get("env_id") or r.config["env"])
+        if args.exclude_envs and rg.env_id in args.exclude_envs:
+            continue
+        if args.envs and rg.env_id not in args.envs:
+            continue
+        if rg not in runs_by_run_group:
+            runs_by_run_group[rg] = RunGroupRuns(
+                rg,
+                args.wandb_control_tag,
+                args.wandb_experiment_tag,
+            )
+        runs_by_run_group[rg].add_run(r)
+    df = RunGroupRuns.data_frame(runs_by_run_group.values()).round(decimals=2)
+    print(f"**Total Score: {sum(df.score)}**")
+    df.loc["mean"] = df.mean(numeric_only=True)
+    print(df.to_markdown())
+
+
+if __name__ == "__main__":
+    compare_runs()

rl_algo_impls/dqn/dqn.py

@@ -0,0 +1,189 @@
+import copy
+import logging
+import random
+from collections import deque
+from typing import List, NamedTuple, Optional, TypeVar
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.optim import Adam
+from torch.utils.tensorboard.writer import SummaryWriter
+
+from rl_algo_impls.dqn.policy import DQNPolicy
+from rl_algo_impls.shared.algorithm import Algorithm
+from rl_algo_impls.shared.callbacks import Callback
+from rl_algo_impls.shared.schedule import linear_schedule
+from rl_algo_impls.wrappers.vectorable_wrapper import VecEnv, VecEnvObs
+
+
+class Transition(NamedTuple):
+    obs: np.ndarray
+    action: np.ndarray
+    reward: float
+    done: bool
+    next_obs: np.ndarray
+
+
+class Batch(NamedTuple):
+    obs: np.ndarray
+    actions: np.ndarray
+    rewards: np.ndarray
+    dones: np.ndarray
+    next_obs: np.ndarray
+
+
+class ReplayBuffer:
+    def __init__(self, num_envs: int, maxlen: int) -> None:
+        self.num_envs = num_envs
+        self.buffer = deque(maxlen=maxlen)
+
+    def add(
+        self,
+        obs: VecEnvObs,
+        action: np.ndarray,
+        reward: np.ndarray,
+        done: np.ndarray,
+        next_obs: VecEnvObs,
+    ) -> None:
+        assert isinstance(obs, np.ndarray)
+        assert isinstance(next_obs, np.ndarray)
+        for i in range(self.num_envs):
+            self.buffer.append(
+                Transition(obs[i], action[i], reward[i], done[i], next_obs[i])
+            )
+
+    def sample(self, batch_size: int) -> Batch:
+        ts = random.sample(self.buffer, batch_size)
+        return Batch(
+            obs=np.array([t.obs for t in ts]),
+            actions=np.array([t.action for t in ts]),
+            rewards=np.array([t.reward for t in ts]),
+            dones=np.array([t.done for t in ts]),
+            next_obs=np.array([t.next_obs for t in ts]),
+        )
+
+    def __len__(self) -> int:
+        return len(self.buffer)
+
+
+DQNSelf = TypeVar("DQNSelf", bound="DQN")
+
+
+class DQN(Algorithm):
+    def __init__(
+        self,
+        policy: DQNPolicy,
+        env: VecEnv,
+        device: torch.device,
+        tb_writer: SummaryWriter,
+        learning_rate: float = 1e-4,
+        buffer_size: int = 1_000_000,
+        learning_starts: int = 50_000,
+        batch_size: int = 32,
+        tau: float = 1.0,
+        gamma: float = 0.99,
+        train_freq: int = 4,
+        gradient_steps: int = 1,
+        target_update_interval: int = 10_000,
+        exploration_fraction: float = 0.1,
+        exploration_initial_eps: float = 1.0,
+        exploration_final_eps: float = 0.05,
+        max_grad_norm: float = 10.0,
+    ) -> None:
+        super().__init__(policy, env, device, tb_writer)
+        self.policy = policy
+
+        self.optimizer = Adam(self.policy.q_net.parameters(), lr=learning_rate)
+
+        self.target_q_net = copy.deepcopy(self.policy.q_net).to(self.device)
+        self.target_q_net.train(False)
+        self.tau = tau
+        self.target_update_interval = target_update_interval
+
+        self.replay_buffer = ReplayBuffer(self.env.num_envs, buffer_size)
+        self.batch_size = batch_size
+
+        self.learning_starts = learning_starts
+        self.train_freq = train_freq
+        self.gradient_steps = gradient_steps
+
+        self.gamma = gamma
+        self.exploration_eps_schedule = linear_schedule(
+            exploration_initial_eps,
+            exploration_final_eps,
+            end_fraction=exploration_fraction,
+        )
+
+        self.max_grad_norm = max_grad_norm
+
+    def learn(
+        self: DQNSelf, total_timesteps: int, callbacks: Optional[List[Callback]] = None
+    ) -> DQNSelf:
+        self.policy.train(True)
+        obs = self.env.reset()
+        obs = self._collect_rollout(self.learning_starts, obs, 1)
+        learning_steps = total_timesteps - self.learning_starts
+        timesteps_elapsed = 0
+        steps_since_target_update = 0
+        while timesteps_elapsed < learning_steps:
+            progress = timesteps_elapsed / learning_steps
+            eps = self.exploration_eps_schedule(progress)
+            obs = self._collect_rollout(self.train_freq, obs, eps)
+            rollout_steps = self.train_freq
+            timesteps_elapsed += rollout_steps
+            for _ in range(
+                self.gradient_steps if self.gradient_steps > 0 else self.train_freq
+            ):
+                self.train()
+            steps_since_target_update += rollout_steps
+            if steps_since_target_update >= self.target_update_interval:
+                self._update_target()
+                steps_since_target_update = 0
+            if callbacks:
+                if not all(
+                    c.on_step(timesteps_elapsed=rollout_steps) for c in callbacks
+                ):
+                    logging.info(
+                        f"Callback terminated training at {timesteps_elapsed} timesteps"
+                    )
+                    break
+        return self
+
+    def train(self) -> None:
+        if len(self.replay_buffer) < self.batch_size:
+            return
+        o, a, r, d, next_o = self.replay_buffer.sample(self.batch_size)
+        o = torch.as_tensor(o, device=self.device)
+        a = torch.as_tensor(a, device=self.device).unsqueeze(1)
+        r = torch.as_tensor(r, dtype=torch.float32, device=self.device)
+        d = torch.as_tensor(d, dtype=torch.long, device=self.device)
+        next_o = torch.as_tensor(next_o, device=self.device)
+
+        with torch.no_grad():
+            target = r + (1 - d) * self.gamma * self.target_q_net(next_o).max(1).values
+        current = self.policy.q_net(o).gather(dim=1, index=a).squeeze(1)
+        loss = F.smooth_l1_loss(current, target)
+
+        self.optimizer.zero_grad()
+        loss.backward()
+        if self.max_grad_norm:
+            nn.utils.clip_grad_norm_(self.policy.q_net.parameters(), self.max_grad_norm)
+        self.optimizer.step()
+
+    def _collect_rollout(self, timesteps: int, obs: VecEnvObs, eps: float) -> VecEnvObs:
+        for _ in range(0, timesteps, self.env.num_envs):
+            action = self.policy.act(obs, eps, deterministic=False)
+            next_obs, reward, done, _ = self.env.step(action)
+            self.replay_buffer.add(obs, action, reward, done, next_obs)
+            obs = next_obs
+        return obs
+
+    def _update_target(self) -> None:
+        for target_param, param in zip(
+            self.target_q_net.parameters(), self.policy.q_net.parameters()
+        ):
+            target_param.data.copy_(
+                self.tau * param.data + (1 - self.tau) * target_param.data
+            )

rl_algo_impls/dqn/policy.py

@@ -0,0 +1,62 @@
+import os
+from typing import Optional, Sequence, TypeVar
+
+import numpy as np
+import torch
+
+from rl_algo_impls.dqn.q_net import QNetwork
+from rl_algo_impls.shared.policy.policy import Policy
+from rl_algo_impls.wrappers.vectorable_wrapper import (
+    VecEnv,
+    VecEnvObs,
+    single_action_space,
+    single_observation_space,
+)
+
+DQNPolicySelf = TypeVar("DQNPolicySelf", bound="DQNPolicy")
+
+
+class DQNPolicy(Policy):
+    def __init__(
+        self,
+        env: VecEnv,
+        hidden_sizes: Sequence[int] = [],
+        cnn_flatten_dim: int = 512,
+        cnn_style: str = "nature",
+        cnn_layers_init_orthogonal: Optional[bool] = None,
+        impala_channels: Sequence[int] = (16, 32, 32),
+        **kwargs,
+    ) -> None:
+        super().__init__(env, **kwargs)
+        self.q_net = QNetwork(
+            single_observation_space(env),
+            single_action_space(env),
+            hidden_sizes,
+            cnn_flatten_dim=cnn_flatten_dim,
+            cnn_style=cnn_style,
+            cnn_layers_init_orthogonal=cnn_layers_init_orthogonal,
+            impala_channels=impala_channels,
+        )
+
+    def act(
+        self,
+        obs: VecEnvObs,
+        eps: float = 0,
+        deterministic: bool = True,
+        action_masks: Optional[np.ndarray] = None,
+    ) -> np.ndarray:
+        assert eps == 0 if deterministic else eps >= 0
+        assert (
+            action_masks is None
+        ), f"action_masks not currently supported in {self.__class__.__name__}"
+        if not deterministic and np.random.random() < eps:
+            return np.array(
+                [
+                    single_action_space(self.env).sample()
+                    for _ in range(self.env.num_envs)
+                ]
+            )
+        else:
+            o = self._as_tensor(obs)
+            with torch.no_grad():
+                return self.q_net(o).argmax(axis=1).cpu().numpy()

rl_algo_impls/dqn/q_net.py

@@ -0,0 +1,41 @@
+from typing import Optional, Sequence, Type
+
+import gym
+import torch as th
+import torch.nn as nn
+from gym.spaces import Discrete
+
+from rl_algo_impls.shared.encoder import Encoder
+from rl_algo_impls.shared.module.utils import mlp
+
+
+class QNetwork(nn.Module):
+    def __init__(
+        self,
+        observation_space: gym.Space,
+        action_space: gym.Space,
+        hidden_sizes: Sequence[int] = [],
+        activation: Type[nn.Module] = nn.ReLU,  # Used by stable-baselines3
+        cnn_flatten_dim: int = 512,
+        cnn_style: str = "nature",
+        cnn_layers_init_orthogonal: Optional[bool] = None,
+        impala_channels: Sequence[int] = (16, 32, 32),
+    ) -> None:
+        super().__init__()
+        assert isinstance(action_space, Discrete)
+        self._feature_extractor = Encoder(
+            observation_space,
+            activation,
+            cnn_flatten_dim=cnn_flatten_dim,
+            cnn_style=cnn_style,
+            cnn_layers_init_orthogonal=cnn_layers_init_orthogonal,
+            impala_channels=impala_channels,
+        )
+        layer_sizes = (
+            (self._feature_extractor.out_dim,) + tuple(hidden_sizes) + (action_space.n,)
+        )
+        self._fc = mlp(layer_sizes, activation)
+
+    def forward(self, obs: th.Tensor) -> th.Tensor:
+        x = self._feature_extractor(obs)
+        return self._fc(x)

rl_algo_impls/enjoy.py

@@ -0,0 +1,35 @@
+# Support for PyTorch mps mode (https://pytorch.org/docs/stable/notes/mps.html)
+import os
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+from rl_algo_impls.runner.evaluate import EvalArgs, evaluate_model
+from rl_algo_impls.runner.running_utils import base_parser
+
+
+def enjoy() -> None:
+    parser = base_parser(multiple=False)
+    parser.add_argument("--render", default=True, type=bool)
+    parser.add_argument("--best", default=True, type=bool)
+    parser.add_argument("--n_envs", default=1, type=int)
+    parser.add_argument("--n_episodes", default=3, type=int)
+    parser.add_argument("--deterministic-eval", default=None, type=bool)
+    parser.add_argument(
+        "--no-print-returns", action="store_true", help="Limit printing"
+    )
+    # wandb-run-path overrides base RunArgs
+    parser.add_argument("--wandb-run-path", default=None, type=str)
+    parser.set_defaults(
+        algo=["ppo"],
+        wandb_run_path="sgoodfriend/rl-algo-impls/m5c1t7g5",
+    )
+    args = parser.parse_args()
+    args.algo = args.algo[0]
+    args.env = args.env[0]
+    args = EvalArgs(**vars(args))
+
+    evaluate_model(args, os.getcwd())
+
+
+if __name__ == "__main__":
+    enjoy()

rl_algo_impls/huggingface_publish.py

@@ -0,0 +1,193 @@
+import os
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+import argparse
+import shutil
+import subprocess
+import tempfile
+from typing import List, Optional
+
+import requests
+import wandb.apis.public
+from huggingface_hub.hf_api import HfApi, upload_folder
+from huggingface_hub.repocard import metadata_save
+from pyvirtualdisplay.display import Display
+
+import wandb
+from rl_algo_impls.publish.markdown_format import EvalTableData, model_card_text
+from rl_algo_impls.runner.config import EnvHyperparams
+from rl_algo_impls.runner.evaluate import EvalArgs, evaluate_model
+from rl_algo_impls.shared.callbacks.eval_callback import evaluate
+from rl_algo_impls.shared.vec_env import make_eval_env
+from rl_algo_impls.wrappers.vec_episode_recorder import VecEpisodeRecorder
+
+
+def publish(
+    wandb_run_paths: List[str],
+    wandb_report_url: str,
+    huggingface_user: Optional[str] = None,
+    huggingface_token: Optional[str] = None,
+    virtual_display: bool = False,
+) -> None:
+    if virtual_display:
+        display = Display(visible=False, size=(1400, 900))
+        display.start()
+
+    api = wandb.Api()
+    runs = [api.run(rp) for rp in wandb_run_paths]
+    algo = runs[0].config["algo"]
+    hyperparam_id = runs[0].config["env"]
+    evaluations = [
+        evaluate_model(
+            EvalArgs(
+                algo,
+                hyperparam_id,
+                seed=r.config.get("seed", None),
+                render=False,
+                best=True,
+                n_envs=None,
+                n_episodes=10,
+                no_print_returns=True,
+                wandb_run_path="/".join(r.path),
+            ),
+            os.getcwd(),
+        )
+        for r in runs
+    ]
+    run_metadata = requests.get(runs[0].file("wandb-metadata.json").url).json()
+    table_data = list(EvalTableData(r, e) for r, e in zip(runs, evaluations))
+    best_eval = sorted(
+        table_data, key=lambda d: d.evaluation.stats.score, reverse=True
+    )[0]
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        _, (policy, stats, config) = best_eval
+
+        repo_name = config.model_name(include_seed=False)
+        repo_dir_path = os.path.join(tmpdirname, repo_name)
+        # Locally clone this repo to a temp directory
+        subprocess.run(["git", "clone", ".", repo_dir_path])
+        shutil.rmtree(os.path.join(repo_dir_path, ".git"))
+        model_path = config.model_dir_path(best=True, downloaded=True)
+        shutil.copytree(
+            model_path,
+            os.path.join(
+                repo_dir_path, "saved_models", config.model_dir_name(best=True)
+            ),
+        )
+
+        github_url = "https://github.com/sgoodfriend/rl-algo-impls"
+        commit_hash = run_metadata.get("git", {}).get("commit", None)
+        env_id = runs[0].config.get("env_id") or runs[0].config["env"]
+        card_text = model_card_text(
+            algo,
+            env_id,
+            github_url,
+            commit_hash,
+            wandb_report_url,
+            table_data,
+            best_eval,
+        )
+        readme_filepath = os.path.join(repo_dir_path, "README.md")
+        os.remove(readme_filepath)
+        with open(readme_filepath, "w") as f:
+            f.write(card_text)
+
+        metadata = {
+            "library_name": "rl-algo-impls",
+            "tags": [
+                env_id,
+                algo,
+                "deep-reinforcement-learning",
+                "reinforcement-learning",
+            ],
+            "model-index": [
+                {
+                    "name": algo,
+                    "results": [
+                        {
+                            "metrics": [
+                                {
+                                    "type": "mean_reward",
+                                    "value": str(stats.score),
+                                    "name": "mean_reward",
+                                }
+                            ],
+                            "task": {
+                                "type": "reinforcement-learning",
+                                "name": "reinforcement-learning",
+                            },
+                            "dataset": {
+                                "name": env_id,
+                                "type": env_id,
+                            },
+                        }
+                    ],
+                }
+            ],
+        }
+        metadata_save(readme_filepath, metadata)
+
+        video_env = VecEpisodeRecorder(
+            make_eval_env(
+                config,
+                EnvHyperparams(**config.env_hyperparams),
+                override_n_envs=1,
+                normalize_load_path=model_path,
+            ),
+            os.path.join(repo_dir_path, "replay"),
+            max_video_length=3600,
+        )
+        evaluate(
+            video_env,
+            policy,
+            1,
+            deterministic=config.eval_hyperparams.get("deterministic", True),
+        )
+
+        api = HfApi()
+        huggingface_user = huggingface_user or api.whoami()["name"]
+        huggingface_repo = f"{huggingface_user}/{repo_name}"
+        api.create_repo(
+            token=huggingface_token,
+            repo_id=huggingface_repo,
+            private=False,
+            exist_ok=True,
+        )
+        repo_url = upload_folder(
+            repo_id=huggingface_repo,
+            folder_path=repo_dir_path,
+            path_in_repo="",
+            commit_message=f"{algo.upper()} playing {env_id} from {github_url}/tree/{commit_hash}",
+            token=huggingface_token,
+            delete_patterns="*",
+        )
+        print(f"Pushed model to the hub: {repo_url}")
+
+
+def huggingface_publish():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--wandb-run-paths",
+        type=str,
+        nargs="+",
+        help="Run paths of the form entity/project/run_id",
+    )
+    parser.add_argument("--wandb-report-url", type=str, help="Link to WandB report")
+    parser.add_argument(
+        "--huggingface-user",
+        type=str,
+        help="Huggingface user or team to upload model cards",
+        default=None,
+    )
+    parser.add_argument(
+        "--virtual-display", action="store_true", help="Use headless virtual display"
+    )
+    args = parser.parse_args()
+    print(args)
+    publish(**vars(args))
+
+
+if __name__ == "__main__":
+    huggingface_publish()

rl_algo_impls/hyperparams/a2c.yml

@@ -0,0 +1,142 @@
+CartPole-v1: &cartpole-defaults
+  n_timesteps: !!float 5e5
+  env_hyperparams:
+    n_envs: 8
+
+CartPole-v0:
+  <<: *cartpole-defaults
+
+MountainCar-v0:
+  n_timesteps: !!float 1e6
+  env_hyperparams:
+    n_envs: 16
+    normalize: true
+
+MountainCarContinuous-v0:
+  n_timesteps: !!float 1e5
+  env_hyperparams:
+    n_envs: 4
+    normalize: true
+  # policy_hyperparams:
+  #   use_sde: true
+  #   log_std_init: 0.0
+  #   init_layers_orthogonal: false
+  algo_hyperparams:
+    n_steps: 100
+    sde_sample_freq: 16
+
+Acrobot-v1:
+  n_timesteps: !!float 5e5
+  env_hyperparams:
+    normalize: true
+    n_envs: 16
+
+# Tuned
+LunarLander-v2:
+  device: cpu
+  n_timesteps: !!float 1e6
+  env_hyperparams:
+    n_envs: 4
+    normalize: true
+  algo_hyperparams:
+    n_steps: 2
+    gamma: 0.9955517404308908
+    gae_lambda: 0.9875340918797773
+    learning_rate: 0.0013814130817068916
+    learning_rate_decay: linear
+    ent_coef: !!float 3.388369146384422e-7
+    ent_coef_decay: none
+    max_grad_norm: 3.33982095073364
+    normalize_advantage: true
+    vf_coef: 0.1667838310548184
+
+BipedalWalker-v3:
+  n_timesteps: !!float 5e6
+  env_hyperparams:
+    n_envs: 16
+    normalize: true
+  policy_hyperparams:
+    use_sde: true
+    log_std_init: -2
+    init_layers_orthogonal: false
+  algo_hyperparams:
+    ent_coef: 0
+    max_grad_norm: 0.5
+    n_steps: 8
+    gae_lambda: 0.9
+    vf_coef: 0.4
+    gamma: 0.99
+    learning_rate: !!float 9.6e-4
+    learning_rate_decay: linear
+
+HalfCheetahBulletEnv-v0: &pybullet-defaults
+  n_timesteps: !!float 2e6
+  env_hyperparams:
+    n_envs: 4
+    normalize: true
+  policy_hyperparams:
+    use_sde: true
+    log_std_init: -2
+    init_layers_orthogonal: false
+  algo_hyperparams: &pybullet-algo-defaults
+    n_steps: 8
+    ent_coef: 0
+    max_grad_norm: 0.5
+    gae_lambda: 0.9
+    gamma: 0.99
+    vf_coef: 0.4
+    learning_rate: !!float 9.6e-4
+    learning_rate_decay: linear
+
+AntBulletEnv-v0:
+  <<: *pybullet-defaults
+
+Walker2DBulletEnv-v0:
+  <<: *pybullet-defaults
+
+HopperBulletEnv-v0:
+  <<: *pybullet-defaults
+
+# Tuned
+CarRacing-v0:
+  n_timesteps: !!float 4e6
+  env_hyperparams:
+    n_envs: 16
+    frame_stack: 4
+    normalize: true
+    normalize_kwargs:
+      norm_obs: false
+      norm_reward: true
+  policy_hyperparams:
+    use_sde: false
+    log_std_init: -1.3502584927786276
+    init_layers_orthogonal: true
+    activation_fn: tanh
+    share_features_extractor: false
+    cnn_flatten_dim: 256
+    hidden_sizes: [256]
+  algo_hyperparams:
+    n_steps: 16
+    learning_rate: 0.000025630993245026736
+    learning_rate_decay: linear
+    gamma: 0.99957617037542
+    gae_lambda: 0.949455676599436
+    ent_coef: !!float 1.707983205298309e-7
+    vf_coef: 0.10428178193833336
+    max_grad_norm: 0.5406643389792273
+    normalize_advantage: true
+    use_rms_prop: false
+
+_atari: &atari-defaults
+  n_timesteps: !!float 1e7
+  env_hyperparams: &atari-env-defaults
+    n_envs: 16
+    frame_stack: 4
+    no_reward_timeout_steps: 1000
+    no_reward_fire_steps: 500
+    vec_env_class: async
+  policy_hyperparams: &atari-policy-defaults
+    activation_fn: relu
+  algo_hyperparams:
+    ent_coef: 0.01
+    vf_coef: 0.25

rl_algo_impls/hyperparams/dqn.yml

@@ -0,0 +1,130 @@
+CartPole-v1: &cartpole-defaults
+  n_timesteps: !!float 5e4
+  env_hyperparams:
+    rolling_length: 50
+  policy_hyperparams:
+    hidden_sizes: [256, 256]
+  algo_hyperparams:
+    learning_rate: !!float 2.3e-3
+    batch_size: 64
+    buffer_size: 100000
+    learning_starts: 1000
+    gamma: 0.99
+    target_update_interval: 10
+    train_freq: 256
+    gradient_steps: 128
+    exploration_fraction: 0.16
+    exploration_final_eps: 0.04
+  eval_hyperparams:
+    step_freq: !!float 1e4
+
+CartPole-v0:
+  <<: *cartpole-defaults
+  n_timesteps: !!float 4e4
+
+MountainCar-v0:
+  n_timesteps: !!float 1.2e5
+  env_hyperparams:
+    rolling_length: 50
+  policy_hyperparams:
+    hidden_sizes: [256, 256]
+  algo_hyperparams:
+    learning_rate: !!float 4e-3
+    batch_size: 128
+    buffer_size: 10000
+    learning_starts: 1000
+    gamma: 0.98
+    target_update_interval: 600
+    train_freq: 16
+    gradient_steps: 8
+    exploration_fraction: 0.2
+    exploration_final_eps: 0.07
+
+Acrobot-v1:
+  n_timesteps: !!float 1e5
+  env_hyperparams:
+    rolling_length: 50
+  policy_hyperparams:
+    hidden_sizes: [256, 256]
+  algo_hyperparams:
+    learning_rate: !!float 6.3e-4
+    batch_size: 128
+    buffer_size: 50000
+    learning_starts: 0
+    gamma: 0.99
+    target_update_interval: 250
+    train_freq: 4
+    gradient_steps: -1
+    exploration_fraction: 0.12
+    exploration_final_eps: 0.1
+
+LunarLander-v2:
+  n_timesteps: !!float 5e5
+  env_hyperparams:
+    rolling_length: 50
+  policy_hyperparams:
+    hidden_sizes: [256, 256]
+  algo_hyperparams:
+    learning_rate: !!float 1e-4
+    batch_size: 256
+    buffer_size: 100000
+    learning_starts: 10000
+    gamma: 0.99
+    target_update_interval: 250
+    train_freq: 8
+    gradient_steps: -1
+    exploration_fraction: 0.12
+    exploration_final_eps: 0.1
+    max_grad_norm: 0.5
+  eval_hyperparams:
+    step_freq: 25_000
+
+_atari: &atari-defaults
+  n_timesteps: !!float 1e7
+  env_hyperparams:
+    frame_stack: 4
+    no_reward_timeout_steps: 1_000
+    no_reward_fire_steps: 500
+    n_envs: 8
+    vec_env_class: async
+  algo_hyperparams:
+    buffer_size: 100000
+    learning_rate: !!float 1e-4
+    batch_size: 32
+    learning_starts: 100000
+    target_update_interval: 1000
+    train_freq: 8
+    gradient_steps: 2
+    exploration_fraction: 0.1
+    exploration_final_eps: 0.01
+  eval_hyperparams:
+    deterministic: false
+
+PongNoFrameskip-v4:
+  <<: *atari-defaults
+  n_timesteps: !!float 2.5e6
+
+_impala-atari: &impala-atari-defaults
+  <<: *atari-defaults
+  policy_hyperparams:
+    cnn_style: impala
+    cnn_flatten_dim: 256
+    init_layers_orthogonal: true
+    cnn_layers_init_orthogonal: false
+
+impala-PongNoFrameskip-v4:
+  <<: *impala-atari-defaults
+  env_id: PongNoFrameskip-v4
+  n_timesteps: !!float 2.5e6
+
+impala-BreakoutNoFrameskip-v4:
+  <<: *impala-atari-defaults
+  env_id: BreakoutNoFrameskip-v4
+
+impala-SpaceInvadersNoFrameskip-v4:
+  <<: *impala-atari-defaults
+  env_id: SpaceInvadersNoFrameskip-v4
+
+impala-QbertNoFrameskip-v4:
+  <<: *impala-atari-defaults
+  env_id: QbertNoFrameskip-v4

rl_algo_impls/hyperparams/ppo.yml

@@ -0,0 +1,616 @@
+CartPole-v1: &cartpole-defaults
+  n_timesteps: !!float 1e5
+  env_hyperparams:
+    n_envs: 8
+  algo_hyperparams:
+    n_steps: 32
+    batch_size: 256
+    n_epochs: 20
+    gae_lambda: 0.8
+    gamma: 0.98
+    ent_coef: 0.0
+    learning_rate: 0.001
+    learning_rate_decay: linear
+    clip_range: 0.2
+    clip_range_decay: linear
+  eval_hyperparams:
+    step_freq: !!float 2.5e4
+
+CartPole-v0:
+  <<: *cartpole-defaults
+  n_timesteps: !!float 5e4
+
+MountainCar-v0:
+  n_timesteps: !!float 1e6
+  env_hyperparams:
+    normalize: true
+    n_envs: 16
+  algo_hyperparams:
+    n_steps: 16
+    n_epochs: 4
+    gae_lambda: 0.98
+    gamma: 0.99
+    ent_coef: 0.0
+
+MountainCarContinuous-v0:
+  n_timesteps: !!float 1e5
+  env_hyperparams:
+    normalize: true
+    n_envs: 4
+  # policy_hyperparams:
+  #   init_layers_orthogonal: false
+  #   log_std_init: -3.29
+  #   use_sde: true
+  algo_hyperparams:
+    n_steps: 512
+    batch_size: 256
+    n_epochs: 10
+    learning_rate: !!float 7.77e-5
+    ent_coef: 0.01 # 0.00429
+    ent_coef_decay: linear
+    clip_range: 0.1
+    gae_lambda: 0.9
+    max_grad_norm: 5
+    vf_coef: 0.19
+  eval_hyperparams:
+    step_freq: 5000
+
+Acrobot-v1:
+  n_timesteps: !!float 1e6
+  env_hyperparams:
+    n_envs: 16
+    normalize: true
+  algo_hyperparams:
+    n_steps: 256
+    n_epochs: 4
+    gae_lambda: 0.94
+    gamma: 0.99
+    ent_coef: 0.0
+
+LunarLander-v2:
+  n_timesteps: !!float 4e6
+  env_hyperparams:
+    n_envs: 16
+  algo_hyperparams:
+    n_steps: 1024
+    batch_size: 64
+    n_epochs: 4
+    gae_lambda: 0.98
+    gamma: 0.999
+    learning_rate: !!float 5e-4
+    learning_rate_decay: linear
+    clip_range: 0.2
+    clip_range_decay: linear
+    ent_coef: 0.01
+    normalize_advantage: false
+
+BipedalWalker-v3:
+  n_timesteps: !!float 10e6
+  env_hyperparams:
+    n_envs: 16
+    normalize: true
+  algo_hyperparams:
+    n_steps: 2048
+    batch_size: 64
+    gae_lambda: 0.95
+    gamma: 0.99
+    n_epochs: 10
+    ent_coef: 0.001
+    learning_rate: !!float 2.5e-4
+    learning_rate_decay: linear
+    clip_range: 0.2
+    clip_range_decay: linear
+
+CarRacing-v0: &carracing-defaults
+  n_timesteps: !!float 4e6
+  env_hyperparams:
+    n_envs: 8
+    frame_stack: 4
+  policy_hyperparams: &carracing-policy-defaults
+    use_sde: true
+    log_std_init: -2
+    init_layers_orthogonal: false
+    activation_fn: relu
+    share_features_extractor: false
+    cnn_flatten_dim: 256
+    hidden_sizes: [256]
+  algo_hyperparams:
+    n_steps: 512
+    batch_size: 128
+    n_epochs: 10
+    learning_rate: !!float 1e-4
+    learning_rate_decay: linear
+    gamma: 0.99
+    gae_lambda: 0.95
+    ent_coef: 0.0
+    sde_sample_freq: 4
+    max_grad_norm: 0.5
+    vf_coef: 0.5
+    clip_range: 0.2
+
+impala-CarRacing-v0:
+  <<: *carracing-defaults
+  env_id: CarRacing-v0
+  policy_hyperparams:
+    <<: *carracing-policy-defaults
+    cnn_style: impala
+    init_layers_orthogonal: true
+    cnn_layers_init_orthogonal: false
+    hidden_sizes: []
+
+# BreakoutNoFrameskip-v4
+# PongNoFrameskip-v4
+# SpaceInvadersNoFrameskip-v4
+# QbertNoFrameskip-v4
+_atari: &atari-defaults
+  n_timesteps: !!float 1e7
+  env_hyperparams: &atari-env-defaults
+    n_envs: 8
+    frame_stack: 4
+    no_reward_timeout_steps: 1000
+    no_reward_fire_steps: 500
+    vec_env_class: async
+  policy_hyperparams: &atari-policy-defaults
+    activation_fn: relu
+  algo_hyperparams: &atari-algo-defaults
+    n_steps: 128
+    batch_size: 256
+    n_epochs: 4
+    learning_rate: !!float 2.5e-4
+    learning_rate_decay: linear
+    clip_range: 0.1
+    clip_range_decay: linear
+    vf_coef: 0.5
+    ent_coef: 0.01
+  eval_hyperparams:
+    deterministic: false
+
+_norm-rewards-atari: &norm-rewards-atari-default
+  <<: *atari-defaults
+  env_hyperparams:
+    <<: *atari-env-defaults
+    clip_atari_rewards: false
+    normalize: true
+    normalize_kwargs:
+      norm_obs: false
+      norm_reward: true
+
+norm-rewards-BreakoutNoFrameskip-v4:
+  <<: *norm-rewards-atari-default
+  env_id: BreakoutNoFrameskip-v4
+
+debug-PongNoFrameskip-v4:
+  <<: *atari-defaults
+  device: cpu
+  env_id: PongNoFrameskip-v4
+  env_hyperparams:
+    <<: *atari-env-defaults
+    vec_env_class: sync
+
+_impala-atari: &impala-atari-defaults
+  <<: *atari-defaults
+  policy_hyperparams:
+    <<: *atari-policy-defaults
+    cnn_style: impala
+    cnn_flatten_dim: 256
+    init_layers_orthogonal: true
+    cnn_layers_init_orthogonal: false
+
+impala-PongNoFrameskip-v4:
+  <<: *impala-atari-defaults
+  env_id: PongNoFrameskip-v4
+
+impala-BreakoutNoFrameskip-v4:
+  <<: *impala-atari-defaults
+  env_id: BreakoutNoFrameskip-v4
+
+impala-SpaceInvadersNoFrameskip-v4:
+  <<: *impala-atari-defaults
+  env_id: SpaceInvadersNoFrameskip-v4
+
+impala-QbertNoFrameskip-v4:
+  <<: *impala-atari-defaults
+  env_id: QbertNoFrameskip-v4
+
+_microrts: &microrts-defaults
+  <<: *atari-defaults
+  n_timesteps: !!float 2e6
+  env_hyperparams: &microrts-env-defaults
+    n_envs: 8
+    vec_env_class: sync
+    mask_actions: true
+  policy_hyperparams: &microrts-policy-defaults
+    <<: *atari-policy-defaults
+    cnn_style: microrts
+    cnn_flatten_dim: 128
+  algo_hyperparams: &microrts-algo-defaults
+    <<: *atari-algo-defaults
+    clip_range_decay: none
+    clip_range_vf: 0.1
+    ppo2_vf_coef_halving: true
+  eval_hyperparams: &microrts-eval-defaults
+    deterministic: false # Good idea because MultiCategorical mode isn't great
+
+_no-mask-microrts: &no-mask-microrts-defaults
+  <<: *microrts-defaults
+  env_hyperparams:
+    <<: *microrts-env-defaults
+    mask_actions: false
+
+MicrortsMining-v1-NoMask:
+  <<: *no-mask-microrts-defaults
+  env_id: MicrortsMining-v1
+
+MicrortsAttackShapedReward-v1-NoMask:
+  <<: *no-mask-microrts-defaults
+  env_id: MicrortsAttackShapedReward-v1
+
+MicrortsRandomEnemyShapedReward3-v1-NoMask:
+  <<: *no-mask-microrts-defaults
+  env_id: MicrortsRandomEnemyShapedReward3-v1
+
+_microrts_ai: &microrts-ai-defaults
+  <<: *microrts-defaults
+  n_timesteps: !!float 100e6
+  additional_keys_to_log: ["microrts_stats"]
+  env_hyperparams: &microrts-ai-env-defaults
+    n_envs: 24
+    env_type: microrts
+    make_kwargs: &microrts-ai-env-make-kwargs-defaults
+      num_selfplay_envs: 0
+      max_steps: 2000
+      render_theme: 2
+      map_paths: [maps/16x16/basesWorkers16x16.xml]
+      reward_weight: [10.0, 1.0, 1.0, 0.2, 1.0, 4.0]
+  policy_hyperparams: &microrts-ai-policy-defaults
+    <<: *microrts-policy-defaults
+    cnn_flatten_dim: 256
+    actor_head_style: gridnet
+  algo_hyperparams: &microrts-ai-algo-defaults
+    <<: *microrts-algo-defaults
+    learning_rate: !!float 2.5e-4
+    learning_rate_decay: linear
+    n_steps: 512
+    batch_size: 3072
+    n_epochs: 4
+    ent_coef: 0.01
+    vf_coef: 0.5
+    max_grad_norm: 0.5
+    clip_range: 0.1
+    clip_range_vf: 0.1
+  eval_hyperparams: &microrts-ai-eval-defaults
+    <<: *microrts-eval-defaults
+    score_function: mean
+    max_video_length: 4000
+    env_overrides: &microrts-ai-eval-env-overrides
+      make_kwargs:
+        <<: *microrts-ai-env-make-kwargs-defaults
+        max_steps: 4000
+        reward_weight: [1.0, 0, 0, 0, 0, 0]
+
+MicrortsAttackPassiveEnemySparseReward-v3:
+  <<: *microrts-ai-defaults
+  n_timesteps: !!float 2e6
+  env_id: MicrortsAttackPassiveEnemySparseReward-v3 # Workaround to keep model name simple
+  env_hyperparams:
+    <<: *microrts-ai-env-defaults
+    bots:
+      passiveAI: 24
+
+MicrortsDefeatRandomEnemySparseReward-v3: &microrts-random-ai-defaults
+  <<: *microrts-ai-defaults
+  n_timesteps: !!float 2e6
+  env_id: MicrortsDefeatRandomEnemySparseReward-v3 # Workaround to keep model name simple
+  env_hyperparams:
+    <<: *microrts-ai-env-defaults
+    bots:
+      randomBiasedAI: 24
+
+enc-dec-MicrortsDefeatRandomEnemySparseReward-v3:
+  <<: *microrts-random-ai-defaults
+  policy_hyperparams:
+    <<: *microrts-ai-policy-defaults
+    cnn_style: gridnet_encoder
+    actor_head_style: gridnet_decoder
+    v_hidden_sizes: [128]
+
+unet-MicrortsDefeatRandomEnemySparseReward-v3:
+  <<: *microrts-random-ai-defaults
+  # device: cpu
+  policy_hyperparams:
+    <<: *microrts-ai-policy-defaults
+    actor_head_style: unet
+    v_hidden_sizes: [256, 128]
+  algo_hyperparams:
+    <<: *microrts-ai-algo-defaults
+    learning_rate: !!float 2.5e-4
+    learning_rate_decay: spike
+
+MicrortsDefeatCoacAIShaped-v3: &microrts-coacai-defaults
+  <<: *microrts-ai-defaults
+  env_id: MicrortsDefeatCoacAIShaped-v3 # Workaround to keep model name simple
+  n_timesteps: !!float 300e6
+  env_hyperparams: &microrts-coacai-env-defaults
+    <<: *microrts-ai-env-defaults
+    bots:
+      coacAI: 24
+  eval_hyperparams: &microrts-coacai-eval-defaults
+    <<: *microrts-ai-eval-defaults
+    step_freq: !!float 1e6
+    n_episodes: 26
+    env_overrides: &microrts-coacai-eval-env-overrides
+      <<: *microrts-ai-eval-env-overrides
+      n_envs: 26
+      bots:
+        coacAI: 2
+        randomBiasedAI: 2
+        randomAI: 2
+        passiveAI: 2
+        workerRushAI: 2
+        lightRushAI: 2
+        naiveMCTSAI: 2
+        mixedBot: 2
+        rojo: 2
+        izanagi: 2
+        tiamat: 2
+        droplet: 2
+        guidedRojoA3N: 2
+
+MicrortsDefeatCoacAIShaped-v3-diverseBots: &microrts-diverse-defaults
+  <<: *microrts-coacai-defaults
+  env_hyperparams:
+    <<: *microrts-coacai-env-defaults
+    bots:
+      coacAI: 18
+      randomBiasedAI: 2
+      lightRushAI: 2
+      workerRushAI: 2
+
+enc-dec-MicrortsDefeatCoacAIShaped-v3-diverseBots:
+  &microrts-env-dec-diverse-defaults
+  <<: *microrts-diverse-defaults
+  policy_hyperparams:
+    <<: *microrts-ai-policy-defaults
+    cnn_style: gridnet_encoder
+    actor_head_style: gridnet_decoder
+    v_hidden_sizes: [128]
+
+debug-enc-dec-MicrortsDefeatCoacAIShaped-v3-diverseBots:
+  <<: *microrts-env-dec-diverse-defaults
+  n_timesteps: !!float 1e6
+
+unet-MicrortsDefeatCoacAIShaped-v3-diverseBots: &microrts-unet-defaults
+  <<: *microrts-diverse-defaults
+  policy_hyperparams:
+    <<: *microrts-ai-policy-defaults
+    actor_head_style: unet
+    v_hidden_sizes: [256, 128]
+  algo_hyperparams: &microrts-unet-algo-defaults
+    <<: *microrts-ai-algo-defaults
+    learning_rate: !!float 2.5e-4
+    learning_rate_decay: spike
+
+Microrts-selfplay-unet: &microrts-selfplay-defaults
+  <<: *microrts-unet-defaults
+  env_hyperparams: &microrts-selfplay-env-defaults
+    <<: *microrts-ai-env-defaults
+    make_kwargs: &microrts-selfplay-env-make-kwargs-defaults
+      <<: *microrts-ai-env-make-kwargs-defaults
+      num_selfplay_envs: 36
+    self_play_kwargs:
+      num_old_policies: 12
+      save_steps: 200000
+      swap_steps: 10000
+      swap_window_size: 4
+      window: 25
+  eval_hyperparams: &microrts-selfplay-eval-defaults
+    <<: *microrts-coacai-eval-defaults
+    env_overrides: &microrts-selfplay-eval-env-overrides
+      <<: *microrts-coacai-eval-env-overrides
+      self_play_kwargs: {}
+
+Microrts-selfplay-unet-winloss: &microrts-selfplay-winloss-defaults
+  <<: *microrts-selfplay-defaults
+  env_hyperparams:
+    <<: *microrts-selfplay-env-defaults
+    make_kwargs:
+      <<: *microrts-selfplay-env-make-kwargs-defaults
+      reward_weight: [1.0, 0, 0, 0, 0, 0]
+  algo_hyperparams: &microrts-selfplay-winloss-algo-defaults
+    <<: *microrts-unet-algo-defaults
+    gamma: 0.999
+
+Microrts-selfplay-unet-decay: &microrts-selfplay-decay-defaults
+  <<: *microrts-selfplay-defaults
+  microrts_reward_decay_callback: true
+  algo_hyperparams:
+    <<: *microrts-unet-algo-defaults
+    gamma_end: 0.999
+
+Microrts-selfplay-unet-debug: &microrts-selfplay-debug-defaults
+  <<: *microrts-selfplay-decay-defaults
+  eval_hyperparams:
+    <<: *microrts-selfplay-eval-defaults
+    step_freq: !!float 1e5
+    env_overrides:
+      <<: *microrts-selfplay-eval-env-overrides
+      n_envs: 24
+      bots:
+        coacAI: 12
+        randomBiasedAI: 4
+        workerRushAI: 4
+        lightRushAI: 4
+
+Microrts-selfplay-unet-debug-mps:
+  <<: *microrts-selfplay-debug-defaults
+  device: mps
+
+HalfCheetahBulletEnv-v0: &pybullet-defaults
+  n_timesteps: !!float 2e6
+  env_hyperparams: &pybullet-env-defaults
+    n_envs: 16
+    normalize: true
+  policy_hyperparams: &pybullet-policy-defaults
+    pi_hidden_sizes: [256, 256]
+    v_hidden_sizes: [256, 256]
+    activation_fn: relu
+  algo_hyperparams: &pybullet-algo-defaults
+    n_steps: 512
+    batch_size: 128
+    n_epochs: 20
+    gamma: 0.99
+    gae_lambda: 0.9
+    ent_coef: 0.0
+    max_grad_norm: 0.5
+    vf_coef: 0.5
+    learning_rate: !!float 3e-5
+    clip_range: 0.4
+
+AntBulletEnv-v0:
+  <<: *pybullet-defaults
+  policy_hyperparams:
+    <<: *pybullet-policy-defaults
+  algo_hyperparams:
+    <<: *pybullet-algo-defaults
+
+Walker2DBulletEnv-v0:
+  <<: *pybullet-defaults
+  algo_hyperparams:
+    <<: *pybullet-algo-defaults
+    clip_range_decay: linear
+
+HopperBulletEnv-v0:
+  <<: *pybullet-defaults
+  algo_hyperparams:
+    <<: *pybullet-algo-defaults
+    clip_range_decay: linear
+
+HumanoidBulletEnv-v0:
+  <<: *pybullet-defaults
+  n_timesteps: !!float 1e7
+  env_hyperparams:
+    <<: *pybullet-env-defaults
+    n_envs: 8
+  policy_hyperparams:
+    <<: *pybullet-policy-defaults
+    # log_std_init: -1
+  algo_hyperparams:
+    <<: *pybullet-algo-defaults
+    n_steps: 2048
+    batch_size: 64
+    n_epochs: 10
+    gae_lambda: 0.95
+    learning_rate: !!float 2.5e-4
+    clip_range: 0.2
+
+_procgen: &procgen-defaults
+  env_hyperparams: &procgen-env-defaults
+    env_type: procgen
+    n_envs: 64
+    # grayscale: false
+    # frame_stack: 4
+    normalize: true # procgen only normalizes reward
+    make_kwargs: &procgen-make-kwargs-defaults
+      num_threads: 8
+  policy_hyperparams: &procgen-policy-defaults
+    activation_fn: relu
+    cnn_style: impala
+    cnn_flatten_dim: 256
+    init_layers_orthogonal: true
+    cnn_layers_init_orthogonal: false
+  algo_hyperparams: &procgen-algo-defaults
+    gamma: 0.999
+    gae_lambda: 0.95
+    n_steps: 256
+    batch_size: 2048
+    n_epochs: 3
+    ent_coef: 0.01
+    clip_range: 0.2
+    # clip_range_decay: linear
+    clip_range_vf: 0.2
+    learning_rate: !!float 5e-4
+    # learning_rate_decay: linear
+    vf_coef: 0.5
+  eval_hyperparams: &procgen-eval-defaults
+    ignore_first_episode: true
+    # deterministic: false
+    step_freq: !!float 1e5
+
+_procgen-easy: &procgen-easy-defaults
+  <<: *procgen-defaults
+  n_timesteps: !!float 25e6
+  env_hyperparams: &procgen-easy-env-defaults
+    <<: *procgen-env-defaults
+    make_kwargs:
+      <<: *procgen-make-kwargs-defaults
+      distribution_mode: easy
+
+procgen-coinrun-easy: &coinrun-easy-defaults
+  <<: *procgen-easy-defaults
+  env_id: coinrun
+
+debug-procgen-coinrun:
+  <<: *coinrun-easy-defaults
+  device: cpu
+
+procgen-starpilot-easy:
+  <<: *procgen-easy-defaults
+  env_id: starpilot
+
+procgen-bossfight-easy:
+  <<: *procgen-easy-defaults
+  env_id: bossfight
+
+procgen-bigfish-easy:
+  <<: *procgen-easy-defaults
+  env_id: bigfish
+
+_procgen-hard: &procgen-hard-defaults
+  <<: *procgen-defaults
+  n_timesteps: !!float 200e6
+  env_hyperparams: &procgen-hard-env-defaults
+    <<: *procgen-env-defaults
+    n_envs: 256
+    make_kwargs:
+      <<: *procgen-make-kwargs-defaults
+      distribution_mode: hard
+  algo_hyperparams: &procgen-hard-algo-defaults
+    <<: *procgen-algo-defaults
+    batch_size: 8192
+    clip_range_decay: linear
+    learning_rate_decay: linear
+  eval_hyperparams:
+    <<: *procgen-eval-defaults
+    step_freq: !!float 5e5
+
+procgen-starpilot-hard: &procgen-starpilot-hard-defaults
+  <<: *procgen-hard-defaults
+  env_id: starpilot
+
+procgen-starpilot-hard-2xIMPALA:
+  <<: *procgen-starpilot-hard-defaults
+  policy_hyperparams:
+    <<: *procgen-policy-defaults
+    impala_channels: [32, 64, 64]
+  algo_hyperparams:
+    <<: *procgen-hard-algo-defaults
+    learning_rate: !!float 3.3e-4
+
+procgen-starpilot-hard-2xIMPALA-fat:
+  <<: *procgen-starpilot-hard-defaults
+  policy_hyperparams:
+    <<: *procgen-policy-defaults
+    impala_channels: [32, 64, 64]
+    cnn_flatten_dim: 512
+  algo_hyperparams:
+    <<: *procgen-hard-algo-defaults
+    learning_rate: !!float 2.5e-4
+
+procgen-starpilot-hard-4xIMPALA:
+  <<: *procgen-starpilot-hard-defaults
+  policy_hyperparams:
+    <<: *procgen-policy-defaults
+    impala_channels: [64, 128, 128]
+  algo_hyperparams:
+    <<: *procgen-hard-algo-defaults
+    learning_rate: !!float 2.1e-4

rl_algo_impls/hyperparams/vpg.yml

@@ -0,0 +1,197 @@
+CartPole-v1: &cartpole-defaults
+  n_timesteps: !!float 4e5
+  algo_hyperparams:
+    n_steps: 4096
+    pi_lr: 0.01
+    gamma: 0.99
+    gae_lambda: 1
+    val_lr: 0.01
+    train_v_iters: 80
+  eval_hyperparams:
+    step_freq: !!float 2.5e4
+
+CartPole-v0:
+  <<: *cartpole-defaults
+  n_timesteps: !!float 1e5
+  algo_hyperparams:
+    n_steps: 1024
+    pi_lr: 0.01
+    gamma: 0.99
+    gae_lambda: 1
+    val_lr: 0.01
+    train_v_iters: 80
+
+MountainCar-v0:
+  n_timesteps: !!float 1e6
+  env_hyperparams:
+    normalize: true
+    n_envs: 16
+  algo_hyperparams:
+    n_steps: 200
+    pi_lr: 0.005
+    gamma: 0.99
+    gae_lambda: 0.97
+    val_lr: 0.01
+    train_v_iters: 80
+    max_grad_norm: 0.5
+
+MountainCarContinuous-v0:
+  n_timesteps: !!float 3e5
+  env_hyperparams:
+    normalize: true
+    n_envs: 4
+  # policy_hyperparams:
+  #   init_layers_orthogonal: false
+  #   log_std_init: -3.29
+  #   use_sde: true
+  algo_hyperparams:
+    n_steps: 1000
+    pi_lr: !!float 5e-4
+    gamma: 0.99
+    gae_lambda: 0.9
+    val_lr: !!float 1e-3
+    train_v_iters: 80
+    max_grad_norm: 5
+  eval_hyperparams:
+    step_freq: 5000
+
+Acrobot-v1:
+  n_timesteps: !!float 2e5
+  algo_hyperparams:
+    n_steps: 2048
+    pi_lr: 0.005
+    gamma: 0.99
+    gae_lambda: 0.97
+    val_lr: 0.01
+    train_v_iters: 80
+    max_grad_norm: 0.5
+
+LunarLander-v2:
+  n_timesteps: !!float 4e6
+  policy_hyperparams:
+    hidden_sizes: [256, 256]
+  algo_hyperparams:
+    n_steps: 2048
+    pi_lr: 0.0001
+    gamma: 0.999
+    gae_lambda: 0.97
+    val_lr: 0.0001
+    train_v_iters: 80
+    max_grad_norm: 0.5
+  eval_hyperparams:
+    deterministic: false
+
+BipedalWalker-v3:
+  n_timesteps: !!float 10e6
+  env_hyperparams:
+    n_envs: 16
+    normalize: true
+  policy_hyperparams:
+    hidden_sizes: [256, 256]
+  algo_hyperparams:
+    n_steps: 1600
+    gae_lambda: 0.95
+    gamma: 0.99
+    pi_lr: !!float 1e-4
+    val_lr: !!float 1e-4
+    train_v_iters: 80
+    max_grad_norm: 0.5
+  eval_hyperparams:
+    deterministic: false
+
+CarRacing-v0:
+  n_timesteps: !!float 4e6
+  env_hyperparams:
+    frame_stack: 4
+    n_envs: 4
+    vec_env_class: sync
+  policy_hyperparams:
+    use_sde: true
+    log_std_init: -2
+    init_layers_orthogonal: false
+    activation_fn: relu
+    cnn_flatten_dim: 256
+    hidden_sizes: [256]
+  algo_hyperparams:
+    n_steps: 1000
+    pi_lr: !!float 5e-5
+    gamma: 0.99
+    gae_lambda: 0.95
+    val_lr: !!float 1e-4
+    train_v_iters: 40
+    max_grad_norm: 0.5
+    sde_sample_freq: 4
+
+HalfCheetahBulletEnv-v0: &pybullet-defaults
+  n_timesteps: !!float 2e6
+  env_hyperparams: &pybullet-env-defaults
+    normalize: true
+  policy_hyperparams: &pybullet-policy-defaults
+    hidden_sizes: [256, 256]
+  algo_hyperparams: &pybullet-algo-defaults
+    n_steps: 4000
+    pi_lr: !!float 3e-4
+    gamma: 0.99
+    gae_lambda: 0.97
+    val_lr: !!float 1e-3
+    train_v_iters: 80
+    max_grad_norm: 0.5
+
+AntBulletEnv-v0:
+  <<: *pybullet-defaults
+  policy_hyperparams:
+    <<: *pybullet-policy-defaults
+    hidden_sizes: [400, 300]
+  algo_hyperparams:
+    <<: *pybullet-algo-defaults
+    pi_lr: !!float 7e-4
+    val_lr: !!float 7e-3
+
+HopperBulletEnv-v0:
+  <<: *pybullet-defaults
+
+Walker2DBulletEnv-v0:
+  <<: *pybullet-defaults
+
+FrozenLake-v1:
+  n_timesteps: !!float 8e5
+  env_params:
+    make_kwargs:
+      map_name: 8x8
+      is_slippery: true
+  policy_hyperparams:
+    hidden_sizes: [64]
+  algo_hyperparams:
+    n_steps: 2048
+    pi_lr: 0.01
+    gamma: 0.99
+    gae_lambda: 0.98
+    val_lr: 0.01
+    train_v_iters: 80
+    max_grad_norm: 0.5
+  eval_hyperparams:
+    step_freq: !!float 5e4
+    n_episodes: 10
+    save_best: true
+
+_atari: &atari-defaults
+  n_timesteps: !!float 10e6
+  env_hyperparams:
+    n_envs: 2
+    frame_stack: 4
+    no_reward_timeout_steps: 1000
+    no_reward_fire_steps: 500
+    vec_env_class: async
+  policy_hyperparams:
+    activation_fn: relu
+  algo_hyperparams:
+    n_steps: 3072
+    pi_lr: !!float 5e-5
+    gamma: 0.99
+    gae_lambda: 0.95
+    val_lr: !!float 1e-4
+    train_v_iters: 80
+    max_grad_norm: 0.5
+    ent_coef: 0.01
+  eval_hyperparams:
+    deterministic: false

rl_algo_impls/optimize.py

@@ -0,0 +1,475 @@
+import dataclasses
+import gc
+import inspect
+import logging
+import os
+from dataclasses import asdict, dataclass
+from typing import Callable, List, NamedTuple, Optional, Sequence, Union
+
+import numpy as np
+import optuna
+import torch
+from optuna.pruners import HyperbandPruner
+from optuna.samplers import TPESampler
+from optuna.visualization import plot_optimization_history, plot_param_importances
+from torch.utils.tensorboard.writer import SummaryWriter
+
+import wandb
+from rl_algo_impls.a2c.optimize import sample_params as a2c_sample_params
+from rl_algo_impls.runner.config import Config, EnvHyperparams, RunArgs
+from rl_algo_impls.runner.running_utils import (
+    ALGOS,
+    base_parser,
+    get_device,
+    hparam_dict,
+    load_hyperparams,
+    make_policy,
+    set_seeds,
+)
+from rl_algo_impls.shared.callbacks import Callback
+from rl_algo_impls.shared.callbacks.microrts_reward_decay_callback import (
+    MicrortsRewardDecayCallback,
+)
+from rl_algo_impls.shared.callbacks.optimize_callback import (
+    Evaluation,
+    OptimizeCallback,
+    evaluation,
+)
+from rl_algo_impls.shared.callbacks.self_play_callback import SelfPlayCallback
+from rl_algo_impls.shared.stats import EpisodesStats
+from rl_algo_impls.shared.vec_env import make_env, make_eval_env
+from rl_algo_impls.wrappers.self_play_wrapper import SelfPlayWrapper
+from rl_algo_impls.wrappers.vectorable_wrapper import find_wrapper
+
+
+@dataclass
+class StudyArgs:
+    load_study: bool
+    study_name: Optional[str] = None
+    storage_path: Optional[str] = None
+    n_trials: int = 100
+    n_jobs: int = 1
+    n_evaluations: int = 4
+    n_eval_envs: int = 8
+    n_eval_episodes: int = 16
+    timeout: Union[int, float, None] = None
+    wandb_project_name: Optional[str] = None
+    wandb_entity: Optional[str] = None
+    wandb_tags: Sequence[str] = dataclasses.field(default_factory=list)
+    wandb_group: Optional[str] = None
+    virtual_display: bool = False
+
+
+class Args(NamedTuple):
+    train_args: Sequence[RunArgs]
+    study_args: StudyArgs
+
+
+def parse_args() -> Args:
+    parser = base_parser()
+    parser.add_argument(
+        "--load-study",
+        action="store_true",
+        help="Load a preexisting study, useful for parallelization",
+    )
+    parser.add_argument("--study-name", type=str, help="Optuna study name")
+    parser.add_argument(
+        "--storage-path",
+        type=str,
+        help="Path of database for Optuna to persist to",
+    )
+    parser.add_argument(
+        "--wandb-project-name",
+        type=str,
+        default="rl-algo-impls-tuning",
+        help="WandB project name to upload tuning data to. If none, won't upload",
+    )
+    parser.add_argument(
+        "--wandb-entity",
+        type=str,
+        help="WandB team. None uses the default entity",
+    )
+    parser.add_argument(
+        "--wandb-tags", type=str, nargs="*", help="WandB tags to add to run"
+    )
+    parser.add_argument(
+        "--wandb-group", type=str, help="WandB group to group trials under"
+    )
+    parser.add_argument(
+        "--n-trials", type=int, default=100, help="Maximum number of trials"
+    )
+    parser.add_argument(
+        "--n-jobs", type=int, default=1, help="Number of jobs to run in parallel"
+    )
+    parser.add_argument(
+        "--n-evaluations",
+        type=int,
+        default=4,
+        help="Number of evaluations during the training",
+    )
+    parser.add_argument(
+        "--n-eval-envs",
+        type=int,
+        default=8,
+        help="Number of envs in vectorized eval environment",
+    )
+    parser.add_argument(
+        "--n-eval-episodes",
+        type=int,
+        default=16,
+        help="Number of episodes to complete for evaluation",
+    )
+    parser.add_argument("--timeout", type=int, help="Seconds to timeout optimization")
+    parser.add_argument(
+        "--virtual-display", action="store_true", help="Use headless virtual display"
+    )
+    # parser.set_defaults(
+    #     algo=["a2c"],
+    #     env=["CartPole-v1"],
+    #     seed=[100, 200, 300],
+    #     n_trials=5,
+    #     virtual_display=True,
+    # )
+    train_dict, study_dict = {}, {}
+    for k, v in vars(parser.parse_args()).items():
+        if k in inspect.signature(StudyArgs).parameters:
+            study_dict[k] = v
+        else:
+            train_dict[k] = v
+
+    study_args = StudyArgs(**study_dict)
+    # Hyperparameter tuning across algos and envs not supported
+    assert len(train_dict["algo"]) == 1
+    assert len(train_dict["env"]) == 1
+    train_args = RunArgs.expand_from_dict(train_dict)
+
+    if not all((study_args.study_name, study_args.storage_path)):
+        hyperparams = load_hyperparams(train_args[0].algo, train_args[0].env)
+        config = Config(train_args[0], hyperparams, os.getcwd())
+        if study_args.study_name is None:
+            study_args.study_name = config.run_name(include_seed=False)
+        if study_args.storage_path is None:
+            study_args.storage_path = (
+                f"sqlite:///{os.path.join(config.runs_dir, 'tuning.db')}"
+            )
+    # Default set group name to study name
+    study_args.wandb_group = study_args.wandb_group or study_args.study_name
+
+    return Args(train_args, study_args)
+
+
+def objective_fn(
+    args: Sequence[RunArgs], study_args: StudyArgs
+) -> Callable[[optuna.Trial], float]:
+    def objective(trial: optuna.Trial) -> float:
+        if len(args) == 1:
+            return simple_optimize(trial, args[0], study_args)
+        else:
+            return stepwise_optimize(trial, args, study_args)
+
+    return objective
+
+
+def simple_optimize(trial: optuna.Trial, args: RunArgs, study_args: StudyArgs) -> float:
+    base_hyperparams = load_hyperparams(args.algo, args.env)
+    base_config = Config(args, base_hyperparams, os.getcwd())
+    if args.algo == "a2c":
+        hyperparams = a2c_sample_params(trial, base_hyperparams, base_config)
+    else:
+        raise ValueError(f"Optimizing {args.algo} isn't supported")
+    config = Config(args, hyperparams, os.getcwd())
+
+    wandb_enabled = bool(study_args.wandb_project_name)
+    if wandb_enabled:
+        wandb.init(
+            project=study_args.wandb_project_name,
+            entity=study_args.wandb_entity,
+            config=asdict(hyperparams),
+            name=f"{config.model_name()}-{str(trial.number)}",
+            tags=study_args.wandb_tags,
+            group=study_args.wandb_group,
+            sync_tensorboard=True,
+            monitor_gym=True,
+            save_code=True,
+            reinit=True,
+        )
+        wandb.config.update(args)
+
+    tb_writer = SummaryWriter(config.tensorboard_summary_path)
+    set_seeds(args.seed, args.use_deterministic_algorithms)
+
+    env = make_env(
+        config, EnvHyperparams(**config.env_hyperparams), tb_writer=tb_writer
+    )
+    device = get_device(config, env)
+    policy_factory = lambda: make_policy(
+        args.algo, env, device, **config.policy_hyperparams
+    )
+    policy = policy_factory()
+    algo = ALGOS[args.algo](policy, env, device, tb_writer, **config.algo_hyperparams)
+
+    eval_env = make_eval_env(
+        config,
+        EnvHyperparams(**config.env_hyperparams),
+        override_n_envs=study_args.n_eval_envs,
+    )
+    optimize_callback = OptimizeCallback(
+        policy,
+        eval_env,
+        trial,
+        tb_writer,
+        step_freq=config.n_timesteps // study_args.n_evaluations,
+        n_episodes=study_args.n_eval_episodes,
+        deterministic=config.eval_hyperparams.get("deterministic", True),
+    )
+    callbacks: List[Callback] = [optimize_callback]
+    if config.hyperparams.microrts_reward_decay_callback:
+        callbacks.append(MicrortsRewardDecayCallback(config, env))
+    selfPlayWrapper = find_wrapper(env, SelfPlayWrapper)
+    if selfPlayWrapper:
+        callbacks.append(SelfPlayCallback(policy, policy_factory, selfPlayWrapper))
+    try:
+        algo.learn(config.n_timesteps, callbacks=callbacks)
+
+        if not optimize_callback.is_pruned:
+            optimize_callback.evaluate()
+            if not optimize_callback.is_pruned:
+                policy.save(config.model_dir_path(best=False))
+
+        eval_stat: EpisodesStats = callback.last_eval_stat  # type: ignore
+        train_stat: EpisodesStats = callback.last_train_stat  # type: ignore
+
+        tb_writer.add_hparams(
+            hparam_dict(hyperparams, vars(args)),
+            {
+                "hparam/last_mean": eval_stat.score.mean,
+                "hparam/last_result": eval_stat.score.mean - eval_stat.score.std,
+                "hparam/train_mean": train_stat.score.mean,
+                "hparam/train_result": train_stat.score.mean - train_stat.score.std,
+                "hparam/score": optimize_callback.last_score,
+                "hparam/is_pruned": optimize_callback.is_pruned,
+            },
+            None,
+            config.run_name(),
+        )
+        tb_writer.close()
+
+        if wandb_enabled:
+            wandb.run.summary["state"] = (  # type: ignore
+                "Pruned" if optimize_callback.is_pruned else "Complete"
+            )
+            wandb.finish(quiet=True)
+
+        if optimize_callback.is_pruned:
+            raise optuna.exceptions.TrialPruned()
+
+        return optimize_callback.last_score
+    except AssertionError as e:
+        logging.warning(e)
+        return np.nan
+    finally:
+        env.close()
+        eval_env.close()
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+def stepwise_optimize(
+    trial: optuna.Trial, args: Sequence[RunArgs], study_args: StudyArgs
+) -> float:
+    algo = args[0].algo
+    env_id = args[0].env
+    base_hyperparams = load_hyperparams(algo, env_id)
+    base_config = Config(args[0], base_hyperparams, os.getcwd())
+    if algo == "a2c":
+        hyperparams = a2c_sample_params(trial, base_hyperparams, base_config)
+    else:
+        raise ValueError(f"Optimizing {algo} isn't supported")
+
+    wandb_enabled = bool(study_args.wandb_project_name)
+    if wandb_enabled:
+        wandb.init(
+            project=study_args.wandb_project_name,
+            entity=study_args.wandb_entity,
+            config=asdict(hyperparams),
+            name=f"{str(trial.number)}-S{base_config.seed()}",
+            tags=study_args.wandb_tags,
+            group=study_args.wandb_group,
+            save_code=True,
+            reinit=True,
+        )
+
+    score = -np.inf
+
+    for i in range(study_args.n_evaluations):
+        evaluations: List[Evaluation] = []
+
+        for arg in args:
+            config = Config(arg, hyperparams, os.getcwd())
+
+            tb_writer = SummaryWriter(config.tensorboard_summary_path)
+            set_seeds(arg.seed, arg.use_deterministic_algorithms)
+
+            env = make_env(
+                config,
+                EnvHyperparams(**config.env_hyperparams),
+                normalize_load_path=config.model_dir_path() if i > 0 else None,
+                tb_writer=tb_writer,
+            )
+            device = get_device(config, env)
+            policy_factory = lambda: make_policy(
+                arg.algo, env, device, **config.policy_hyperparams
+            )
+            policy = policy_factory()
+            if i > 0:
+                policy.load(config.model_dir_path())
+            algo = ALGOS[arg.algo](
+                policy, env, device, tb_writer, **config.algo_hyperparams
+            )
+
+            eval_env = make_eval_env(
+                config,
+                EnvHyperparams(**config.env_hyperparams),
+                normalize_load_path=config.model_dir_path() if i > 0 else None,
+                override_n_envs=study_args.n_eval_envs,
+            )
+
+            start_timesteps = int(i * config.n_timesteps / study_args.n_evaluations)
+            train_timesteps = (
+                int((i + 1) * config.n_timesteps / study_args.n_evaluations)
+                - start_timesteps
+            )
+
+            callbacks = []
+            if config.hyperparams.microrts_reward_decay_callback:
+                callbacks.append(
+                    MicrortsRewardDecayCallback(
+                        config, env, start_timesteps=start_timesteps
+                    )
+                )
+            selfPlayWrapper = find_wrapper(env, SelfPlayWrapper)
+            if selfPlayWrapper:
+                callbacks.append(
+                    SelfPlayCallback(policy, policy_factory, selfPlayWrapper)
+                )
+            try:
+                algo.learn(
+                    train_timesteps,
+                    callbacks=callbacks,
+                    total_timesteps=config.n_timesteps,
+                    start_timesteps=start_timesteps,
+                )
+
+                evaluations.append(
+                    evaluation(
+                        policy,
+                        eval_env,
+                        tb_writer,
+                        study_args.n_eval_episodes,
+                        config.eval_hyperparams.get("deterministic", True),
+                        start_timesteps + train_timesteps,
+                    )
+                )
+
+                policy.save(config.model_dir_path())
+
+                tb_writer.close()
+
+            except AssertionError as e:
+                logging.warning(e)
+                if wandb_enabled:
+                    wandb_finish("Error")
+                return np.nan
+            finally:
+                env.close()
+                eval_env.close()
+                gc.collect()
+                torch.cuda.empty_cache()
+
+        d = {}
+        for idx, e in enumerate(evaluations):
+            d[f"{idx}/eval_mean"] = e.eval_stat.score.mean
+            d[f"{idx}/train_mean"] = e.train_stat.score.mean
+            d[f"{idx}/score"] = e.score
+        d["eval"] = np.mean([e.eval_stat.score.mean for e in evaluations]).item()
+        d["train"] = np.mean([e.train_stat.score.mean for e in evaluations]).item()
+        score = np.mean([e.score for e in evaluations]).item()
+        d["score"] = score
+
+        step = i + 1
+        wandb.log(d, step=step)
+
+        print(f"Trial #{trial.number} Step {step} Score: {round(score, 2)}")
+        trial.report(score, step)
+        if trial.should_prune():
+            if wandb_enabled:
+                wandb_finish("Pruned")
+            raise optuna.exceptions.TrialPruned()
+
+    if wandb_enabled:
+        wandb_finish("Complete")
+    return score
+
+
+def wandb_finish(state: str) -> None:
+    wandb.run.summary["state"] = state  # type: ignore
+    wandb.finish(quiet=True)
+
+
+def optimize() -> None:
+    from pyvirtualdisplay.display import Display
+
+    train_args, study_args = parse_args()
+    if study_args.virtual_display:
+        virtual_display = Display(visible=False, size=(1400, 900))
+        virtual_display.start()
+
+    sampler = TPESampler(**TPESampler.hyperopt_parameters())
+    pruner = HyperbandPruner()
+    if study_args.load_study:
+        assert study_args.study_name
+        assert study_args.storage_path
+        study = optuna.load_study(
+            study_name=study_args.study_name,
+            storage=study_args.storage_path,
+            sampler=sampler,
+            pruner=pruner,
+        )
+    else:
+        study = optuna.create_study(
+            study_name=study_args.study_name,
+            storage=study_args.storage_path,
+            sampler=sampler,
+            pruner=pruner,
+            direction="maximize",
+        )
+
+    try:
+        study.optimize(
+            objective_fn(train_args, study_args),
+            n_trials=study_args.n_trials,
+            n_jobs=study_args.n_jobs,
+            timeout=study_args.timeout,
+        )
+    except KeyboardInterrupt:
+        pass
+
+    best = study.best_trial
+    print(f"Best Trial Value: {best.value}")
+    print("Attributes:")
+    for key, value in list(best.params.items()) + list(best.user_attrs.items()):
+        print(f"  {key}: {value}")
+
+    df = study.trials_dataframe()
+    df = df[df.state == "COMPLETE"].sort_values(by=["value"], ascending=False)
+    print(df.to_markdown(index=False))
+
+    fig1 = plot_optimization_history(study)
+    fig1.write_image("opt_history.png")
+
+    fig2 = plot_param_importances(study)
+    fig2.write_image("param_importances.png")
+
+
+if __name__ == "__main__":
+    optimize()

rl_algo_impls/ppo/ppo.py

@@ -0,0 +1,385 @@
+import logging
+from dataclasses import asdict, dataclass
+from time import perf_counter
+from typing import List, NamedTuple, Optional, TypeVar
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.optim import Adam
+from torch.utils.tensorboard.writer import SummaryWriter
+
+from rl_algo_impls.shared.algorithm import Algorithm
+from rl_algo_impls.shared.callbacks import Callback
+from rl_algo_impls.shared.gae import compute_advantages
+from rl_algo_impls.shared.policy.actor_critic import ActorCritic
+from rl_algo_impls.shared.schedule import (
+    constant_schedule,
+    linear_schedule,
+    schedule,
+    update_learning_rate,
+)
+from rl_algo_impls.shared.stats import log_scalars
+from rl_algo_impls.wrappers.vectorable_wrapper import (
+    VecEnv,
+    single_action_space,
+    single_observation_space,
+)
+
+
+class TrainStepStats(NamedTuple):
+    loss: float
+    pi_loss: float
+    v_loss: float
+    entropy_loss: float
+    approx_kl: float
+    clipped_frac: float
+    val_clipped_frac: float
+
+
+@dataclass
+class TrainStats:
+    loss: float
+    pi_loss: float
+    v_loss: float
+    entropy_loss: float
+    approx_kl: float
+    clipped_frac: float
+    val_clipped_frac: float
+    explained_var: float
+
+    def __init__(self, step_stats: List[TrainStepStats], explained_var: float) -> None:
+        self.loss = np.mean([s.loss for s in step_stats]).item()
+        self.pi_loss = np.mean([s.pi_loss for s in step_stats]).item()
+        self.v_loss = np.mean([s.v_loss for s in step_stats]).item()
+        self.entropy_loss = np.mean([s.entropy_loss for s in step_stats]).item()
+        self.approx_kl = np.mean([s.approx_kl for s in step_stats]).item()
+        self.clipped_frac = np.mean([s.clipped_frac for s in step_stats]).item()
+        self.val_clipped_frac = np.mean([s.val_clipped_frac for s in step_stats]).item()
+        self.explained_var = explained_var
+
+    def write_to_tensorboard(self, tb_writer: SummaryWriter, global_step: int) -> None:
+        for name, value in asdict(self).items():
+            tb_writer.add_scalar(f"losses/{name}", value, global_step=global_step)
+
+    def __repr__(self) -> str:
+        return " | ".join(
+            [
+                f"Loss: {round(self.loss, 2)}",
+                f"Pi L: {round(self.pi_loss, 2)}",
+                f"V L: {round(self.v_loss, 2)}",
+                f"E L: {round(self.entropy_loss, 2)}",
+                f"Apx KL Div: {round(self.approx_kl, 2)}",
+                f"Clip Frac: {round(self.clipped_frac, 2)}",
+                f"Val Clip Frac: {round(self.val_clipped_frac, 2)}",
+            ]
+        )
+
+
+PPOSelf = TypeVar("PPOSelf", bound="PPO")
+
+
+class PPO(Algorithm):
+    def __init__(
+        self,
+        policy: ActorCritic,
+        env: VecEnv,
+        device: torch.device,
+        tb_writer: SummaryWriter,
+        learning_rate: float = 3e-4,
+        learning_rate_decay: str = "none",
+        n_steps: int = 2048,
+        batch_size: int = 64,
+        n_epochs: int = 10,
+        gamma: float = 0.99,
+        gae_lambda: float = 0.95,
+        clip_range: float = 0.2,
+        clip_range_decay: str = "none",
+        clip_range_vf: Optional[float] = None,
+        clip_range_vf_decay: str = "none",
+        normalize_advantage: bool = True,
+        ent_coef: float = 0.0,
+        ent_coef_decay: str = "none",
+        vf_coef: float = 0.5,
+        ppo2_vf_coef_halving: bool = False,
+        max_grad_norm: float = 0.5,
+        sde_sample_freq: int = -1,
+        update_advantage_between_epochs: bool = True,
+        update_returns_between_epochs: bool = False,
+        gamma_end: Optional[float] = None,
+    ) -> None:
+        super().__init__(policy, env, device, tb_writer)
+        self.policy = policy
+        self.get_action_mask = getattr(env, "get_action_mask")
+
+        self.gamma_schedule = (
+            linear_schedule(gamma, gamma_end)
+            if gamma_end is not None
+            else constant_schedule(gamma)
+        )
+        self.gae_lambda = gae_lambda
+        self.optimizer = Adam(self.policy.parameters(), lr=learning_rate, eps=1e-7)
+        self.lr_schedule = schedule(learning_rate_decay, learning_rate)
+        self.max_grad_norm = max_grad_norm
+        self.clip_range_schedule = schedule(clip_range_decay, clip_range)
+        self.clip_range_vf_schedule = None
+        if clip_range_vf:
+            self.clip_range_vf_schedule = schedule(clip_range_vf_decay, clip_range_vf)
+
+        if normalize_advantage:
+            assert (
+                env.num_envs * n_steps > 1 and batch_size > 1
+            ), f"Each minibatch must be larger than 1 to support normalization"
+        self.normalize_advantage = normalize_advantage
+
+        self.ent_coef_schedule = schedule(ent_coef_decay, ent_coef)
+        self.vf_coef = vf_coef
+        self.ppo2_vf_coef_halving = ppo2_vf_coef_halving
+
+        self.n_steps = n_steps
+        self.batch_size = batch_size
+        self.n_epochs = n_epochs
+        self.sde_sample_freq = sde_sample_freq
+
+        self.update_advantage_between_epochs = update_advantage_between_epochs
+        self.update_returns_between_epochs = update_returns_between_epochs
+
+    def learn(
+        self: PPOSelf,
+        train_timesteps: int,
+        callbacks: Optional[List[Callback]] = None,
+        total_timesteps: Optional[int] = None,
+        start_timesteps: int = 0,
+    ) -> PPOSelf:
+        if total_timesteps is None:
+            total_timesteps = train_timesteps
+        assert start_timesteps + train_timesteps <= total_timesteps
+
+        epoch_dim = (self.n_steps, self.env.num_envs)
+        step_dim = (self.env.num_envs,)
+        obs_space = single_observation_space(self.env)
+        act_space = single_action_space(self.env)
+        act_shape = self.policy.action_shape
+
+        next_obs = self.env.reset()
+        next_action_masks = self.get_action_mask() if self.get_action_mask else None
+        next_episode_starts = np.full(step_dim, True, dtype=np.bool_)
+
+        obs = np.zeros(epoch_dim + obs_space.shape, dtype=obs_space.dtype)  # type: ignore
+        actions = np.zeros(epoch_dim + act_shape, dtype=act_space.dtype)  # type: ignore
+        rewards = np.zeros(epoch_dim, dtype=np.float32)
+        episode_starts = np.zeros(epoch_dim, dtype=np.bool_)
+        values = np.zeros(epoch_dim, dtype=np.float32)
+        logprobs = np.zeros(epoch_dim, dtype=np.float32)
+        action_masks = (
+            np.zeros(
+                (self.n_steps,) + next_action_masks.shape, dtype=next_action_masks.dtype
+            )
+            if next_action_masks is not None
+            else None
+        )
+
+        timesteps_elapsed = start_timesteps
+        while timesteps_elapsed < start_timesteps + train_timesteps:
+            start_time = perf_counter()
+
+            progress = timesteps_elapsed / total_timesteps
+            ent_coef = self.ent_coef_schedule(progress)
+            learning_rate = self.lr_schedule(progress)
+            update_learning_rate(self.optimizer, learning_rate)
+            pi_clip = self.clip_range_schedule(progress)
+            gamma = self.gamma_schedule(progress)
+            chart_scalars = {
+                "learning_rate": self.optimizer.param_groups[0]["lr"],
+                "ent_coef": ent_coef,
+                "pi_clip": pi_clip,
+                "gamma": gamma,
+            }
+            if self.clip_range_vf_schedule:
+                v_clip = self.clip_range_vf_schedule(progress)
+                chart_scalars["v_clip"] = v_clip
+            else:
+                v_clip = None
+            log_scalars(self.tb_writer, "charts", chart_scalars, timesteps_elapsed)
+
+            self.policy.eval()
+            self.policy.reset_noise()
+            for s in range(self.n_steps):
+                timesteps_elapsed += self.env.num_envs
+                if self.sde_sample_freq > 0 and s > 0 and s % self.sde_sample_freq == 0:
+                    self.policy.reset_noise()
+
+                obs[s] = next_obs
+                episode_starts[s] = next_episode_starts
+                if action_masks is not None:
+                    action_masks[s] = next_action_masks
+
+                (
+                    actions[s],
+                    values[s],
+                    logprobs[s],
+                    clamped_action,
+                ) = self.policy.step(next_obs, action_masks=next_action_masks)
+                next_obs, rewards[s], next_episode_starts, _ = self.env.step(
+                    clamped_action
+                )
+                next_action_masks = (
+                    self.get_action_mask() if self.get_action_mask else None
+                )
+
+            self.policy.train()
+
+            b_obs = torch.tensor(obs.reshape((-1,) + obs_space.shape)).to(self.device)  # type: ignore
+            b_actions = torch.tensor(actions.reshape((-1,) + act_shape)).to(  # type: ignore
+                self.device
+            )
+            b_logprobs = torch.tensor(logprobs.reshape(-1)).to(self.device)
+            b_action_masks = (
+                torch.tensor(action_masks.reshape((-1,) + next_action_masks.shape[1:])).to(  # type: ignore
+                    self.device
+                )
+                if action_masks is not None
+                else None
+            )
+
+            y_pred = values.reshape(-1)
+            b_values = torch.tensor(y_pred).to(self.device)
+
+            step_stats = []
+            # Define variables that will definitely be set through the first epoch
+            advantages: np.ndarray = None  # type: ignore
+            b_advantages: torch.Tensor = None  # type: ignore
+            y_true: np.ndarray = None  # type: ignore
+            b_returns: torch.Tensor = None  # type: ignore
+            for e in range(self.n_epochs):
+                if e == 0 or self.update_advantage_between_epochs:
+                    advantages = compute_advantages(
+                        rewards,
+                        values,
+                        episode_starts,
+                        next_episode_starts,
+                        next_obs,
+                        self.policy,
+                        gamma,
+                        self.gae_lambda,
+                    )
+                    b_advantages = torch.tensor(advantages.reshape(-1)).to(self.device)
+                if e == 0 or self.update_returns_between_epochs:
+                    returns = advantages + values
+                    y_true = returns.reshape(-1)
+                    b_returns = torch.tensor(y_true).to(self.device)
+
+                b_idxs = torch.randperm(len(b_obs))
+                # Only record last epoch's stats
+                step_stats.clear()
+                for i in range(0, len(b_obs), self.batch_size):
+                    self.policy.reset_noise(self.batch_size)
+
+                    mb_idxs = b_idxs[i : i + self.batch_size]
+
+                    mb_obs = b_obs[mb_idxs]
+                    mb_actions = b_actions[mb_idxs]
+                    mb_values = b_values[mb_idxs]
+                    mb_logprobs = b_logprobs[mb_idxs]
+                    mb_action_masks = (
+                        b_action_masks[mb_idxs] if b_action_masks is not None else None
+                    )
+
+                    mb_adv = b_advantages[mb_idxs]
+                    if self.normalize_advantage:
+                        mb_adv = (mb_adv - mb_adv.mean()) / (mb_adv.std() + 1e-8)
+                    mb_returns = b_returns[mb_idxs]
+
+                    new_logprobs, entropy, new_values = self.policy(
+                        mb_obs, mb_actions, action_masks=mb_action_masks
+                    )
+
+                    logratio = new_logprobs - mb_logprobs
+                    ratio = torch.exp(logratio)
+                    clipped_ratio = torch.clamp(ratio, min=1 - pi_clip, max=1 + pi_clip)
+                    pi_loss = torch.max(-ratio * mb_adv, -clipped_ratio * mb_adv).mean()
+
+                    v_loss_unclipped = (new_values - mb_returns) ** 2
+                    if v_clip:
+                        v_loss_clipped = (
+                            mb_values
+                            + torch.clamp(new_values - mb_values, -v_clip, v_clip)
+                            - mb_returns
+                        ) ** 2
+                        v_loss = torch.max(v_loss_unclipped, v_loss_clipped).mean()
+                    else:
+                        v_loss = v_loss_unclipped.mean()
+
+                    if self.ppo2_vf_coef_halving:
+                        v_loss *= 0.5
+
+                    entropy_loss = -entropy.mean()
+
+                    loss = pi_loss + ent_coef * entropy_loss + self.vf_coef * v_loss
+
+                    self.optimizer.zero_grad()
+                    loss.backward()
+                    nn.utils.clip_grad_norm_(
+                        self.policy.parameters(), self.max_grad_norm
+                    )
+                    self.optimizer.step()
+
+                    with torch.no_grad():
+                        approx_kl = ((ratio - 1) - logratio).mean().cpu().numpy().item()
+                        clipped_frac = (
+                            ((ratio - 1).abs() > pi_clip)
+                            .float()
+                            .mean()
+                            .cpu()
+                            .numpy()
+                            .item()
+                        )
+                        val_clipped_frac = (
+                            ((new_values - mb_values).abs() > v_clip)
+                            .float()
+                            .mean()
+                            .cpu()
+                            .numpy()
+                            .item()
+                            if v_clip
+                            else 0
+                        )
+
+                    step_stats.append(
+                        TrainStepStats(
+                            loss.item(),
+                            pi_loss.item(),
+                            v_loss.item(),
+                            entropy_loss.item(),
+                            approx_kl,
+                            clipped_frac,
+                            val_clipped_frac,
+                        )
+                    )
+
+            var_y = np.var(y_true).item()
+            explained_var = (
+                np.nan if var_y == 0 else 1 - np.var(y_true - y_pred).item() / var_y
+            )
+            TrainStats(step_stats, explained_var).write_to_tensorboard(
+                self.tb_writer, timesteps_elapsed
+            )
+
+            end_time = perf_counter()
+            rollout_steps = self.n_steps * self.env.num_envs
+            self.tb_writer.add_scalar(
+                "train/steps_per_second",
+                rollout_steps / (end_time - start_time),
+                timesteps_elapsed,
+            )
+
+            if callbacks:
+                if not all(
+                    c.on_step(timesteps_elapsed=rollout_steps) for c in callbacks
+                ):
+                    logging.info(
+                        f"Callback terminated training at {timesteps_elapsed} timesteps"
+                    )
+                    break
+
+        return self

rl_algo_impls/publish/markdown_format.py

@@ -0,0 +1,210 @@
+import os
+import pandas as pd
+import wandb.apis.public
+import yaml
+
+from collections import defaultdict
+from dataclasses import dataclass, asdict
+from typing import Any, Dict, Iterable, List, NamedTuple, Optional, TypeVar
+from urllib.parse import urlparse
+
+from rl_algo_impls.runner.evaluate import Evaluation
+
+EvaluationRowSelf = TypeVar("EvaluationRowSelf", bound="EvaluationRow")
+
+
+@dataclass
+class EvaluationRow:
+    algo: str
+    env: str
+    seed: Optional[int]
+    reward_mean: float
+    reward_std: float
+    eval_episodes: int
+    best: str
+    wandb_url: str
+
+    @staticmethod
+    def data_frame(rows: List[EvaluationRowSelf]) -> pd.DataFrame:
+        results = defaultdict(list)
+        for r in rows:
+            for k, v in asdict(r).items():
+                results[k].append(v)
+        return pd.DataFrame(results)
+
+
+class EvalTableData(NamedTuple):
+    run: wandb.apis.public.Run
+    evaluation: Evaluation
+
+
+def evaluation_table(table_data: Iterable[EvalTableData]) -> str:
+    best_stats = sorted(
+        [d.evaluation.stats for d in table_data], key=lambda r: r.score, reverse=True
+    )[0]
+    table_data = sorted(table_data, key=lambda d: d.evaluation.config.seed() or 0)
+    rows = [
+        EvaluationRow(
+            config.algo,
+            config.env_id,
+            config.seed(),
+            stats.score.mean,
+            stats.score.std,
+            len(stats),
+            "*" if stats == best_stats else "",
+            f"[wandb]({r.url})",
+        )
+        for (r, (_, stats, config)) in table_data
+    ]
+    df = EvaluationRow.data_frame(rows)
+    return df.to_markdown(index=False)
+
+
+def github_project_link(github_url: str) -> str:
+    return f"[{urlparse(github_url).path}]({github_url})"
+
+
+def header_section(algo: str, env: str, github_url: str, wandb_report_url: str) -> str:
+    algo_caps = algo.upper()
+    lines = [
+        f"# **{algo_caps}** Agent playing **{env}**",
+        f"This is a trained model of a **{algo_caps}** agent playing **{env}** using "
+        f"the {github_project_link(github_url)} repo.",
+        f"All models trained at this commit can be found at {wandb_report_url}.",
+    ]
+    return "\n\n".join(lines)
+
+
+def github_tree_link(github_url: str, commit_hash: Optional[str]) -> str:
+    if not commit_hash:
+        return github_project_link(github_url)
+    return f"[{commit_hash[:7]}]({github_url}/tree/{commit_hash})"
+
+
+def results_section(
+    table_data: List[EvalTableData], algo: str, github_url: str, commit_hash: str
+) -> str:
+    # type: ignore
+    lines = [
+        "## Training Results",
+        f"This model was trained from {len(table_data)} trainings of **{algo.upper()}** "
+        + "agents using different initial seeds. "
+        + f"These agents were trained by checking out "
+        + f"{github_tree_link(github_url, commit_hash)}. "
+        + "The best and last models were kept from each training. "
+        + "This submission has loaded the best models from each training, reevaluates "
+        + "them, and selects the best model from these latest evaluations (mean - std).",
+    ]
+    lines.append(evaluation_table(table_data))
+    return "\n\n".join(lines)
+
+
+def prerequisites_section() -> str:
+    return """
+### Prerequisites: Weights & Biases (WandB)
+Training and benchmarking assumes you have a Weights & Biases project to upload runs to.
+By default training goes to a rl-algo-impls project while benchmarks go to
+rl-algo-impls-benchmarks. During training and benchmarking runs, videos of the best
+models and the model weights are uploaded to WandB.
+
+Before doing anything below, you'll need to create a wandb account and run `wandb
+login`.
+"""
+
+
+def usage_section(github_url: str, run_path: str, commit_hash: str) -> str:
+    return f"""
+## Usage
+{urlparse(github_url).path}: {github_url}
+
+Note: While the model state dictionary and hyperaparameters are saved, the latest
+implementation could be sufficiently different to not be able to reproduce similar
+results. You might need to checkout the commit the agent was trained on:
+{github_tree_link(github_url, commit_hash)}.
+```
+# Downloads the model, sets hyperparameters, and runs agent for 3 episodes
+python enjoy.py --wandb-run-path={run_path}
+```
+
+Setup hasn't been completely worked out yet, so you might be best served by using Google
+Colab starting from the
+[colab_enjoy.ipynb](https://github.com/sgoodfriend/rl-algo-impls/blob/main/colab_enjoy.ipynb)
+notebook.
+"""
+
+
+def training_setion(
+    github_url: str, commit_hash: str, algo: str, env: str, seed: Optional[int]
+) -> str:
+    return f"""
+## Training
+If you want the highest chance to reproduce these results, you'll want to checkout the
+commit the agent was trained on: {github_tree_link(github_url, commit_hash)}. While
+training is deterministic, different hardware will give different results.
+
+```
+python train.py --algo {algo} --env {env} {'--seed ' + str(seed) if seed is not None else ''}
+```
+
+Setup hasn't been completely worked out yet, so you might be best served by using Google
+Colab starting from the
+[colab_train.ipynb](https://github.com/sgoodfriend/rl-algo-impls/blob/main/colab_train.ipynb)
+notebook.
+"""
+
+
+def benchmarking_section(report_url: str) -> str:
+    return f"""
+## Benchmarking (with Lambda Labs instance)
+This and other models from {report_url} were generated by running a script on a Lambda
+Labs instance. In a Lambda Labs instance terminal:
+```
+git clone git@github.com:sgoodfriend/rl-algo-impls.git
+cd rl-algo-impls
+bash ./lambda_labs/setup.sh
+wandb login
+bash ./lambda_labs/benchmark.sh [-a {{"ppo a2c dqn vpg"}}] [-e ENVS] [-j {{6}}] [-p {{rl-algo-impls-benchmarks}}] [-s {{"1 2 3"}}]
+```
+
+### Alternative: Google Colab Pro+
+As an alternative,
+[colab_benchmark.ipynb](https://github.com/sgoodfriend/rl-algo-impls/tree/main/benchmarks#:~:text=colab_benchmark.ipynb),
+can be used. However, this requires a Google Colab Pro+ subscription and running across
+4 separate instances because otherwise running all jobs will exceed the 24-hour limit.
+"""
+
+
+def hyperparams_section(run_config: Dict[str, Any]) -> str:
+    return f"""
+## Hyperparameters
+This isn't exactly the format of hyperparams in {os.path.join("hyperparams",
+run_config["algo"] + ".yml")}, but instead the Wandb Run Config. However, it's very
+close and has some additional data:
+```
+{yaml.dump(run_config)}
+```
+"""
+
+
+def model_card_text(
+    algo: str,
+    env: str,
+    github_url: str,
+    commit_hash: str,
+    wandb_report_url: str,
+    table_data: List[EvalTableData],
+    best_eval: EvalTableData,
+) -> str:
+    run, (_, _, config) = best_eval
+    run_path = "/".join(run.path)
+    return "\n\n".join(
+        [
+            header_section(algo, env, github_url, wandb_report_url),
+            results_section(table_data, algo, github_url, commit_hash),
+            prerequisites_section(),
+            usage_section(github_url, run_path, commit_hash),
+            training_setion(github_url, commit_hash, algo, env, config.seed()),
+            benchmarking_section(wandb_report_url),
+            hyperparams_section(run.config),
+        ]
+    )

rl_algo_impls/runner/config.py

@@ -0,0 +1,203 @@
+import dataclasses
+import inspect
+import itertools
+import os
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Type, TypeVar, Union
+
+RunArgsSelf = TypeVar("RunArgsSelf", bound="RunArgs")
+
+
+@dataclass
+class RunArgs:
+    algo: str
+    env: str
+    seed: Optional[int] = None
+    use_deterministic_algorithms: bool = True
+
+    @classmethod
+    def expand_from_dict(
+        cls: Type[RunArgsSelf], d: Dict[str, Any]
+    ) -> List[RunArgsSelf]:
+        maybe_listify = lambda v: [v] if isinstance(v, str) or isinstance(v, int) else v
+        algos = maybe_listify(d["algo"])
+        envs = maybe_listify(d["env"])
+        seeds = maybe_listify(d["seed"])
+        args = []
+        for algo, env, seed in itertools.product(algos, envs, seeds):
+            _d = d.copy()
+            _d.update({"algo": algo, "env": env, "seed": seed})
+            args.append(cls(**_d))
+        return args
+
+
+@dataclass
+class EnvHyperparams:
+    env_type: str = "gymvec"
+    n_envs: int = 1
+    frame_stack: int = 1
+    make_kwargs: Optional[Dict[str, Any]] = None
+    no_reward_timeout_steps: Optional[int] = None
+    no_reward_fire_steps: Optional[int] = None
+    vec_env_class: str = "sync"
+    normalize: bool = False
+    normalize_kwargs: Optional[Dict[str, Any]] = None
+    rolling_length: int = 100
+    train_record_video: bool = False
+    video_step_interval: Union[int, float] = 1_000_000
+    initial_steps_to_truncate: Optional[int] = None
+    clip_atari_rewards: bool = True
+    normalize_type: Optional[str] = None
+    mask_actions: bool = False
+    bots: Optional[Dict[str, int]] = None
+    self_play_kwargs: Optional[Dict[str, Any]] = None
+
+
+HyperparamsSelf = TypeVar("HyperparamsSelf", bound="Hyperparams")
+
+
+@dataclass
+class Hyperparams:
+    device: str = "auto"
+    n_timesteps: Union[int, float] = 100_000
+    env_hyperparams: Dict[str, Any] = dataclasses.field(default_factory=dict)
+    policy_hyperparams: Dict[str, Any] = dataclasses.field(default_factory=dict)
+    algo_hyperparams: Dict[str, Any] = dataclasses.field(default_factory=dict)
+    eval_hyperparams: Dict[str, Any] = dataclasses.field(default_factory=dict)
+    env_id: Optional[str] = None
+    additional_keys_to_log: List[str] = dataclasses.field(default_factory=list)
+    microrts_reward_decay_callback: bool = False
+
+    @classmethod
+    def from_dict_with_extra_fields(
+        cls: Type[HyperparamsSelf], d: Dict[str, Any]
+    ) -> HyperparamsSelf:
+        return cls(
+            **{k: v for k, v in d.items() if k in inspect.signature(cls).parameters}
+        )
+
+
+@dataclass
+class Config:
+    args: RunArgs
+    hyperparams: Hyperparams
+    root_dir: str
+    run_id: str = datetime.now().isoformat()
+
+    def seed(self, training: bool = True) -> Optional[int]:
+        seed = self.args.seed
+        if training or seed is None:
+            return seed
+        return seed + self.env_hyperparams.get("n_envs", 1)
+
+    @property
+    def device(self) -> str:
+        return self.hyperparams.device
+
+    @property
+    def n_timesteps(self) -> int:
+        return int(self.hyperparams.n_timesteps)
+
+    @property
+    def env_hyperparams(self) -> Dict[str, Any]:
+        return self.hyperparams.env_hyperparams
+
+    @property
+    def policy_hyperparams(self) -> Dict[str, Any]:
+        return self.hyperparams.policy_hyperparams
+
+    @property
+    def algo_hyperparams(self) -> Dict[str, Any]:
+        return self.hyperparams.algo_hyperparams
+
+    @property
+    def eval_hyperparams(self) -> Dict[str, Any]:
+        return self.hyperparams.eval_hyperparams
+
+    def eval_callback_params(self) -> Dict[str, Any]:
+        eval_hyperparams = self.eval_hyperparams.copy()
+        if "env_overrides" in eval_hyperparams:
+            del eval_hyperparams["env_overrides"]
+        return eval_hyperparams
+
+    @property
+    def algo(self) -> str:
+        return self.args.algo
+
+    @property
+    def env_id(self) -> str:
+        return self.hyperparams.env_id or self.args.env
+
+    @property
+    def additional_keys_to_log(self) -> List[str]:
+        return self.hyperparams.additional_keys_to_log
+
+    def model_name(self, include_seed: bool = True) -> str:
+        # Use arg env name instead of environment name
+        parts = [self.algo, self.args.env]
+        if include_seed and self.args.seed is not None:
+            parts.append(f"S{self.args.seed}")
+
+        # Assume that the custom arg name already has the necessary information
+        if not self.hyperparams.env_id:
+            make_kwargs = self.env_hyperparams.get("make_kwargs", {})
+            if make_kwargs:
+                for k, v in make_kwargs.items():
+                    if type(v) == bool and v:
+                        parts.append(k)
+                    elif type(v) == int and v:
+                        parts.append(f"{k}{v}")
+                    else:
+                        parts.append(str(v))
+
+        return "-".join(parts)
+
+    def run_name(self, include_seed: bool = True) -> str:
+        parts = [self.model_name(include_seed=include_seed), self.run_id]
+        return "-".join(parts)
+
+    @property
+    def saved_models_dir(self) -> str:
+        return os.path.join(self.root_dir, "saved_models")
+
+    @property
+    def downloaded_models_dir(self) -> str:
+        return os.path.join(self.root_dir, "downloaded_models")
+
+    def model_dir_name(
+        self,
+        best: bool = False,
+        extension: str = "",
+    ) -> str:
+        return self.model_name() + ("-best" if best else "") + extension
+
+    def model_dir_path(self, best: bool = False, downloaded: bool = False) -> str:
+        return os.path.join(
+            self.saved_models_dir if not downloaded else self.downloaded_models_dir,
+            self.model_dir_name(best=best),
+        )
+
+    @property
+    def runs_dir(self) -> str:
+        return os.path.join(self.root_dir, "runs")
+
+    @property
+    def tensorboard_summary_path(self) -> str:
+        return os.path.join(self.runs_dir, self.run_name())
+
+    @property
+    def logs_path(self) -> str:
+        return os.path.join(self.runs_dir, f"log.yml")
+
+    @property
+    def videos_dir(self) -> str:
+        return os.path.join(self.root_dir, "videos")
+
+    @property
+    def video_prefix(self) -> str:
+        return os.path.join(self.videos_dir, self.model_name())
+
+    @property
+    def best_videos_dir(self) -> str:
+        return os.path.join(self.videos_dir, f"{self.model_name()}-best")

rl_algo_impls/runner/evaluate.py

@@ -0,0 +1,102 @@
+import os
+import shutil
+from dataclasses import dataclass
+from typing import NamedTuple, Optional
+
+from rl_algo_impls.runner.config import Config, EnvHyperparams, Hyperparams, RunArgs
+from rl_algo_impls.runner.running_utils import (
+    get_device,
+    load_hyperparams,
+    make_policy,
+    set_seeds,
+)
+from rl_algo_impls.shared.callbacks.eval_callback import evaluate
+from rl_algo_impls.shared.policy.policy import Policy
+from rl_algo_impls.shared.stats import EpisodesStats
+from rl_algo_impls.shared.vec_env import make_eval_env
+
+
+@dataclass
+class EvalArgs(RunArgs):
+    render: bool = True
+    best: bool = True
+    n_envs: Optional[int] = 1
+    n_episodes: int = 3
+    deterministic_eval: Optional[bool] = None
+    no_print_returns: bool = False
+    wandb_run_path: Optional[str] = None
+
+
+class Evaluation(NamedTuple):
+    policy: Policy
+    stats: EpisodesStats
+    config: Config
+
+
+def evaluate_model(args: EvalArgs, root_dir: str) -> Evaluation:
+    if args.wandb_run_path:
+        import wandb
+
+        api = wandb.Api()
+        run = api.run(args.wandb_run_path)
+        params = run.config
+
+        args.algo = params["algo"]
+        args.env = params["env"]
+        args.seed = params.get("seed", None)
+        args.use_deterministic_algorithms = params.get(
+            "use_deterministic_algorithms", True
+        )
+
+        config = Config(args, Hyperparams.from_dict_with_extra_fields(params), root_dir)
+        model_path = config.model_dir_path(best=args.best, downloaded=True)
+
+        model_archive_name = config.model_dir_name(best=args.best, extension=".zip")
+        run.file(model_archive_name).download()
+        if os.path.isdir(model_path):
+            shutil.rmtree(model_path)
+        shutil.unpack_archive(model_archive_name, model_path)
+        os.remove(model_archive_name)
+    else:
+        hyperparams = load_hyperparams(args.algo, args.env)
+
+        config = Config(args, hyperparams, root_dir)
+        model_path = config.model_dir_path(best=args.best)
+
+    print(args)
+
+    set_seeds(args.seed, args.use_deterministic_algorithms)
+
+    env = make_eval_env(
+        config,
+        EnvHyperparams(**config.env_hyperparams),
+        override_n_envs=args.n_envs,
+        render=args.render,
+        normalize_load_path=model_path,
+    )
+    device = get_device(config, env)
+    policy = make_policy(
+        args.algo,
+        env,
+        device,
+        load_path=model_path,
+        **config.policy_hyperparams,
+    ).eval()
+
+    deterministic = (
+        args.deterministic_eval
+        if args.deterministic_eval is not None
+        else config.eval_hyperparams.get("deterministic", True)
+    )
+    return Evaluation(
+        policy,
+        evaluate(
+            env,
+            policy,
+            args.n_episodes,
+            render=args.render,
+            deterministic=deterministic,
+            print_returns=not args.no_print_returns,
+        ),
+        config,
+    )

rl_algo_impls/runner/running_utils.py

@@ -0,0 +1,199 @@
+import argparse
+import json
+import os
+import random
+from dataclasses import asdict
+from pathlib import Path
+from typing import Dict, Optional, Type, Union
+
+import gym
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.backends.cudnn
+import yaml
+from gym.spaces import Box, Discrete
+from torch.utils.tensorboard.writer import SummaryWriter
+
+from rl_algo_impls.a2c.a2c import A2C
+from rl_algo_impls.dqn.dqn import DQN
+from rl_algo_impls.dqn.policy import DQNPolicy
+from rl_algo_impls.ppo.ppo import PPO
+from rl_algo_impls.runner.config import Config, Hyperparams
+from rl_algo_impls.shared.algorithm import Algorithm
+from rl_algo_impls.shared.callbacks.eval_callback import EvalCallback
+from rl_algo_impls.shared.policy.actor_critic import ActorCritic
+from rl_algo_impls.shared.policy.policy import Policy
+from rl_algo_impls.shared.vec_env.utils import import_for_env_id, is_microrts
+from rl_algo_impls.vpg.policy import VPGActorCritic
+from rl_algo_impls.vpg.vpg import VanillaPolicyGradient
+from rl_algo_impls.wrappers.vectorable_wrapper import VecEnv, single_observation_space
+
+ALGOS: Dict[str, Type[Algorithm]] = {
+    "dqn": DQN,
+    "vpg": VanillaPolicyGradient,
+    "ppo": PPO,
+    "a2c": A2C,
+}
+POLICIES: Dict[str, Type[Policy]] = {
+    "dqn": DQNPolicy,
+    "vpg": VPGActorCritic,
+    "ppo": ActorCritic,
+    "a2c": ActorCritic,
+}
+
+HYPERPARAMS_PATH = "hyperparams"
+
+
+def base_parser(multiple: bool = True) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--algo",
+        default=["dqn"],
+        type=str,
+        choices=list(ALGOS.keys()),
+        nargs="+" if multiple else 1,
+        help="Abbreviation(s) of algorithm(s)",
+    )
+    parser.add_argument(
+        "--env",
+        default=["CartPole-v1"],
+        type=str,
+        nargs="+" if multiple else 1,
+        help="Name of environment(s) in gym",
+    )
+    parser.add_argument(
+        "--seed",
+        default=[1],
+        type=int,
+        nargs="*" if multiple else "?",
+        help="Seeds to run experiment. Unset will do one run with no set seed",
+    )
+    return parser
+
+
+def load_hyperparams(algo: str, env_id: str) -> Hyperparams:
+    root_path = Path(__file__).parent.parent
+    hyperparams_path = os.path.join(root_path, HYPERPARAMS_PATH, f"{algo}.yml")
+    with open(hyperparams_path, "r") as f:
+        hyperparams_dict = yaml.safe_load(f)
+
+    if env_id in hyperparams_dict:
+        return Hyperparams(**hyperparams_dict[env_id])
+
+    import_for_env_id(env_id)
+    spec = gym.spec(env_id)
+    entry_point_name = str(spec.entry_point)  # type: ignore
+    if "AtariEnv" in entry_point_name and "_atari" in hyperparams_dict:
+        return Hyperparams(**hyperparams_dict["_atari"])
+    elif "gym_microrts" in entry_point_name and "_microrts" in hyperparams_dict:
+        return Hyperparams(**hyperparams_dict["_microrts"])
+    else:
+        raise ValueError(f"{env_id} not specified in {algo} hyperparameters file")
+
+
+def get_device(config: Config, env: VecEnv) -> torch.device:
+    device = config.device
+    # cuda by default
+    if device == "auto":
+        device = "cuda"
+        # Apple MPS is a second choice (sometimes)
+        if device == "cuda" and not torch.cuda.is_available():
+            device = "mps"
+        # If no MPS, fallback to cpu
+        if device == "mps" and not torch.backends.mps.is_available():
+            device = "cpu"
+        # Simple environments like Discreet and 1-D Boxes might also be better
+        # served with the CPU.
+        if device == "mps":
+            obs_space = single_observation_space(env)
+            if isinstance(obs_space, Discrete):
+                device = "cpu"
+            elif isinstance(obs_space, Box) and len(obs_space.shape) == 1:
+                device = "cpu"
+            if is_microrts(config):
+                device = "cpu"
+    print(f"Device: {device}")
+    return torch.device(device)
+
+
+def set_seeds(seed: Optional[int], use_deterministic_algorithms: bool) -> None:
+    if seed is None:
+        return
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.backends.cudnn.benchmark = False
+    torch.use_deterministic_algorithms(use_deterministic_algorithms)
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+    # Stop warning and it would introduce stochasticity if I was using TF
+    os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
+
+
+def make_policy(
+    algo: str,
+    env: VecEnv,
+    device: torch.device,
+    load_path: Optional[str] = None,
+    **kwargs,
+) -> Policy:
+    policy = POLICIES[algo](env, **kwargs).to(device)
+    if load_path:
+        policy.load(load_path)
+    return policy
+
+
+def plot_eval_callback(callback: EvalCallback, tb_writer: SummaryWriter, run_name: str):
+    figure = plt.figure()
+    cumulative_steps = [
+        (idx + 1) * callback.step_freq for idx in range(len(callback.stats))
+    ]
+    plt.plot(
+        cumulative_steps,
+        [s.score.mean for s in callback.stats],
+        "b-",
+        label="mean",
+    )
+    plt.plot(
+        cumulative_steps,
+        [s.score.mean - s.score.std for s in callback.stats],
+        "g--",
+        label="mean-std",
+    )
+    plt.fill_between(
+        cumulative_steps,
+        [s.score.min for s in callback.stats],  # type: ignore
+        [s.score.max for s in callback.stats],  # type: ignore
+        facecolor="cyan",
+        label="range",
+    )
+    plt.xlabel("Steps")
+    plt.ylabel("Score")
+    plt.legend()
+    plt.title(f"Eval {run_name}")
+    tb_writer.add_figure("eval", figure)
+
+
+Scalar = Union[bool, str, float, int, None]
+
+
+def hparam_dict(
+    hyperparams: Hyperparams, args: Dict[str, Union[Scalar, list]]
+) -> Dict[str, Scalar]:
+    flattened = args.copy()
+    for k, v in flattened.items():
+        if isinstance(v, list):
+            flattened[k] = json.dumps(v)
+    for k, v in asdict(hyperparams).items():
+        if isinstance(v, dict):
+            for sk, sv in v.items():
+                key = f"{k}/{sk}"
+                if isinstance(sv, dict) or isinstance(sv, list):
+                    flattened[key] = str(sv)
+                else:
+                    flattened[key] = sv
+        elif isinstance(v, list):
+            flattened[k] = json.dumps(v)
+        else:
+            flattened[k] = v  # type: ignore
+    return flattened  # type: ignore

rl_algo_impls/runner/train.py

@@ -0,0 +1,161 @@
+# Support for PyTorch mps mode (https://pytorch.org/docs/stable/notes/mps.html)
+import os
+
+from rl_algo_impls.shared.callbacks import Callback
+from rl_algo_impls.shared.callbacks.self_play_callback import SelfPlayCallback
+from rl_algo_impls.wrappers.self_play_wrapper import SelfPlayWrapper
+from rl_algo_impls.wrappers.vectorable_wrapper import find_wrapper
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+import dataclasses
+import shutil
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, List, Optional, Sequence
+
+import yaml
+from torch.utils.tensorboard.writer import SummaryWriter
+
+import wandb
+from rl_algo_impls.runner.config import Config, EnvHyperparams, RunArgs
+from rl_algo_impls.runner.running_utils import (
+    ALGOS,
+    get_device,
+    hparam_dict,
+    load_hyperparams,
+    make_policy,
+    plot_eval_callback,
+    set_seeds,
+)
+from rl_algo_impls.shared.callbacks.eval_callback import EvalCallback
+from rl_algo_impls.shared.callbacks.microrts_reward_decay_callback import (
+    MicrortsRewardDecayCallback,
+)
+from rl_algo_impls.shared.stats import EpisodesStats
+from rl_algo_impls.shared.vec_env import make_env, make_eval_env
+
+
+@dataclass
+class TrainArgs(RunArgs):
+    wandb_project_name: Optional[str] = None
+    wandb_entity: Optional[str] = None
+    wandb_tags: Sequence[str] = dataclasses.field(default_factory=list)
+    wandb_group: Optional[str] = None
+
+
+def train(args: TrainArgs):
+    print(args)
+    hyperparams = load_hyperparams(args.algo, args.env)
+    print(hyperparams)
+    config = Config(args, hyperparams, os.getcwd())
+
+    wandb_enabled = args.wandb_project_name
+    if wandb_enabled:
+        wandb.tensorboard.patch(
+            root_logdir=config.tensorboard_summary_path, pytorch=True
+        )
+        wandb.init(
+            project=args.wandb_project_name,
+            entity=args.wandb_entity,
+            config=asdict(hyperparams),
+            name=config.run_name(),
+            monitor_gym=True,
+            save_code=True,
+            tags=args.wandb_tags,
+            group=args.wandb_group,
+        )
+        wandb.config.update(args)
+
+    tb_writer = SummaryWriter(config.tensorboard_summary_path)
+
+    set_seeds(args.seed, args.use_deterministic_algorithms)
+
+    env = make_env(
+        config, EnvHyperparams(**config.env_hyperparams), tb_writer=tb_writer
+    )
+    device = get_device(config, env)
+    policy_factory = lambda: make_policy(
+        args.algo, env, device, **config.policy_hyperparams
+    )
+    policy = policy_factory()
+    algo = ALGOS[args.algo](policy, env, device, tb_writer, **config.algo_hyperparams)
+
+    num_parameters = policy.num_parameters()
+    num_trainable_parameters = policy.num_trainable_parameters()
+    if wandb_enabled:
+        wandb.run.summary["num_parameters"] = num_parameters  # type: ignore
+        wandb.run.summary["num_trainable_parameters"] = num_trainable_parameters  # type: ignore
+    else:
+        print(
+            f"num_parameters = {num_parameters} ; "
+            f"num_trainable_parameters = {num_trainable_parameters}"
+        )
+
+    eval_env = make_eval_env(config, EnvHyperparams(**config.env_hyperparams))
+    record_best_videos = config.eval_hyperparams.get("record_best_videos", True)
+    eval_callback = EvalCallback(
+        policy,
+        eval_env,
+        tb_writer,
+        best_model_path=config.model_dir_path(best=True),
+        **config.eval_callback_params(),
+        video_env=make_eval_env(
+            config, EnvHyperparams(**config.env_hyperparams), override_n_envs=1
+        )
+        if record_best_videos
+        else None,
+        best_video_dir=config.best_videos_dir,
+        additional_keys_to_log=config.additional_keys_to_log,
+    )
+    callbacks: List[Callback] = [eval_callback]
+    if config.hyperparams.microrts_reward_decay_callback:
+        callbacks.append(MicrortsRewardDecayCallback(config, env))
+    selfPlayWrapper = find_wrapper(env, SelfPlayWrapper)
+    if selfPlayWrapper:
+        callbacks.append(SelfPlayCallback(policy, policy_factory, selfPlayWrapper))
+    algo.learn(config.n_timesteps, callbacks=callbacks)
+
+    policy.save(config.model_dir_path(best=False))
+
+    eval_stats = eval_callback.evaluate(n_episodes=10, print_returns=True)
+
+    plot_eval_callback(eval_callback, tb_writer, config.run_name())
+
+    log_dict: Dict[str, Any] = {
+        "eval": eval_stats._asdict(),
+    }
+    if eval_callback.best:
+        log_dict["best_eval"] = eval_callback.best._asdict()
+    log_dict.update(asdict(hyperparams))
+    log_dict.update(vars(args))
+    with open(config.logs_path, "a") as f:
+        yaml.dump({config.run_name(): log_dict}, f)
+
+    best_eval_stats: EpisodesStats = eval_callback.best  # type: ignore
+    tb_writer.add_hparams(
+        hparam_dict(hyperparams, vars(args)),
+        {
+            "hparam/best_mean": best_eval_stats.score.mean,
+            "hparam/best_result": best_eval_stats.score.mean
+            - best_eval_stats.score.std,
+            "hparam/last_mean": eval_stats.score.mean,
+            "hparam/last_result": eval_stats.score.mean - eval_stats.score.std,
+        },
+        None,
+        config.run_name(),
+    )
+
+    tb_writer.close()
+
+    if wandb_enabled:
+        shutil.make_archive(
+            os.path.join(wandb.run.dir, config.model_dir_name()),
+            "zip",
+            config.model_dir_path(),
+        )
+        shutil.make_archive(
+            os.path.join(wandb.run.dir, config.model_dir_name(best=True)),
+            "zip",
+            config.model_dir_path(best=True),
+        )
+        wandb.finish()

rl_algo_impls/shared/actor/__init__.py

@@ -0,0 +1,2 @@
+from rl_algo_impls.shared.actor.actor import Actor, PiForward, pi_forward
+from rl_algo_impls.shared.actor.make_actor import actor_head

rl_algo_impls/shared/actor/actor.py

@@ -0,0 +1,43 @@
+from abc import ABC, abstractmethod
+from typing import NamedTuple, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.distributions import Distribution
+
+
+class PiForward(NamedTuple):
+    pi: Distribution
+    logp_a: Optional[torch.Tensor]
+    entropy: Optional[torch.Tensor]
+
+
+class Actor(nn.Module, ABC):
+    @abstractmethod
+    def forward(
+        self,
+        obs: torch.Tensor,
+        actions: Optional[torch.Tensor] = None,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> PiForward:
+        ...
+
+    def sample_weights(self, batch_size: int = 1) -> None:
+        pass
+
+    @property
+    @abstractmethod
+    def action_shape(self) -> Tuple[int, ...]:
+        ...
+
+
+def pi_forward(
+    distribution: Distribution, actions: Optional[torch.Tensor] = None
+) -> PiForward:
+    logp_a = None
+    entropy = None
+    if actions is not None:
+        logp_a = distribution.log_prob(actions)
+        entropy = distribution.entropy()
+    return PiForward(distribution, logp_a, entropy)

rl_algo_impls/shared/actor/categorical.py

@@ -0,0 +1,64 @@
+from typing import Optional, Tuple, Type
+
+import torch
+import torch.nn as nn
+from torch.distributions import Categorical
+
+from rl_algo_impls.shared.actor import Actor, PiForward, pi_forward
+from rl_algo_impls.shared.module.utils import mlp
+
+
+class MaskedCategorical(Categorical):
+    def __init__(
+        self,
+        probs=None,
+        logits=None,
+        validate_args=None,
+        mask: Optional[torch.Tensor] = None,
+    ):
+        if mask is not None:
+            assert logits is not None, "mask requires logits and not probs"
+            logits = torch.where(mask, logits, -1e8)
+        self.mask = mask
+        super().__init__(probs, logits, validate_args)
+
+    def entropy(self) -> torch.Tensor:
+        if self.mask is None:
+            return super().entropy()
+        # If mask set, then use approximation for entropy
+        p_log_p = self.logits * self.probs  # type: ignore
+        masked = torch.where(self.mask, p_log_p, 0)
+        return -masked.sum(-1)
+
+
+class CategoricalActorHead(Actor):
+    def __init__(
+        self,
+        act_dim: int,
+        in_dim: int,
+        hidden_sizes: Tuple[int, ...] = (32,),
+        activation: Type[nn.Module] = nn.Tanh,
+        init_layers_orthogonal: bool = True,
+    ) -> None:
+        super().__init__()
+        layer_sizes = (in_dim,) + hidden_sizes + (act_dim,)
+        self._fc = mlp(
+            layer_sizes,
+            activation,
+            init_layers_orthogonal=init_layers_orthogonal,
+            final_layer_gain=0.01,
+        )
+
+    def forward(
+        self,
+        obs: torch.Tensor,
+        actions: Optional[torch.Tensor] = None,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> PiForward:
+        logits = self._fc(obs)
+        pi = MaskedCategorical(logits=logits, mask=action_masks)
+        return pi_forward(pi, actions)
+
+    @property
+    def action_shape(self) -> Tuple[int, ...]:
+        return ()

rl_algo_impls/shared/actor/gaussian.py

@@ -0,0 +1,61 @@
+from typing import Optional, Tuple, Type
+
+import torch
+import torch.nn as nn
+from torch.distributions import Distribution, Normal
+
+from rl_algo_impls.shared.actor.actor import Actor, PiForward, pi_forward
+from rl_algo_impls.shared.module.utils import mlp
+
+
+class GaussianDistribution(Normal):
+    def log_prob(self, a: torch.Tensor) -> torch.Tensor:
+        return super().log_prob(a).sum(axis=-1)
+
+    def sample(self) -> torch.Tensor:
+        return self.rsample()
+
+
+class GaussianActorHead(Actor):
+    def __init__(
+        self,
+        act_dim: int,
+        in_dim: int,
+        hidden_sizes: Tuple[int, ...] = (32,),
+        activation: Type[nn.Module] = nn.Tanh,
+        init_layers_orthogonal: bool = True,
+        log_std_init: float = -0.5,
+    ) -> None:
+        super().__init__()
+        self.act_dim = act_dim
+        layer_sizes = (in_dim,) + hidden_sizes + (act_dim,)
+        self.mu_net = mlp(
+            layer_sizes,
+            activation,
+            init_layers_orthogonal=init_layers_orthogonal,
+            final_layer_gain=0.01,
+        )
+        self.log_std = nn.Parameter(
+            torch.ones(act_dim, dtype=torch.float32) * log_std_init
+        )
+
+    def _distribution(self, obs: torch.Tensor) -> Distribution:
+        mu = self.mu_net(obs)
+        std = torch.exp(self.log_std)
+        return GaussianDistribution(mu, std)
+
+    def forward(
+        self,
+        obs: torch.Tensor,
+        actions: Optional[torch.Tensor] = None,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> PiForward:
+        assert (
+            not action_masks
+        ), f"{self.__class__.__name__} does not support action_masks"
+        pi = self._distribution(obs)
+        return pi_forward(pi, actions)
+
+    @property
+    def action_shape(self) -> Tuple[int, ...]:
+        return (self.act_dim,)

rl_algo_impls/shared/actor/gridnet.py

@@ -0,0 +1,108 @@
+from typing import Dict, Optional, Tuple, Type
+
+import numpy as np
+import torch
+import torch.nn as nn
+from numpy.typing import NDArray
+from torch.distributions import Distribution, constraints
+
+from rl_algo_impls.shared.actor import Actor, PiForward, pi_forward
+from rl_algo_impls.shared.actor.categorical import MaskedCategorical
+from rl_algo_impls.shared.encoder import EncoderOutDim
+from rl_algo_impls.shared.module.utils import mlp
+
+
+class GridnetDistribution(Distribution):
+    def __init__(
+        self,
+        map_size: int,
+        action_vec: NDArray[np.int64],
+        logits: torch.Tensor,
+        masks: torch.Tensor,
+        validate_args: Optional[bool] = None,
+    ) -> None:
+        self.map_size = map_size
+        self.action_vec = action_vec
+
+        masks = masks.view(-1, masks.shape[-1])
+        split_masks = torch.split(masks, action_vec.tolist(), dim=1)
+
+        grid_logits = logits.reshape(-1, action_vec.sum())
+        split_logits = torch.split(grid_logits, action_vec.tolist(), dim=1)
+        self.categoricals = [
+            MaskedCategorical(logits=lg, validate_args=validate_args, mask=m)
+            for lg, m in zip(split_logits, split_masks)
+        ]
+
+        batch_shape = logits.size()[:-1] if logits.ndimension() > 1 else torch.Size()
+        super().__init__(batch_shape=batch_shape, validate_args=validate_args)
+
+    def log_prob(self, action: torch.Tensor) -> torch.Tensor:
+        prob_stack = torch.stack(
+            [
+                c.log_prob(a)
+                for a, c in zip(action.view(-1, action.shape[-1]).T, self.categoricals)
+            ],
+            dim=-1,
+        )
+        logprob = prob_stack.view(-1, self.map_size, len(self.action_vec))
+        return logprob.sum(dim=(1, 2))
+
+    def entropy(self) -> torch.Tensor:
+        ent = torch.stack([c.entropy() for c in self.categoricals], dim=-1)
+        ent = ent.view(-1, self.map_size, len(self.action_vec))
+        return ent.sum(dim=(1, 2))
+
+    def sample(self, sample_shape: torch.Size = torch.Size()) -> torch.Tensor:
+        s = torch.stack([c.sample(sample_shape) for c in self.categoricals], dim=-1)
+        return s.view(-1, self.map_size, len(self.action_vec))
+
+    @property
+    def mode(self) -> torch.Tensor:
+        m = torch.stack([c.mode for c in self.categoricals], dim=-1)
+        return m.view(-1, self.map_size, len(self.action_vec))
+
+    @property
+    def arg_constraints(self) -> Dict[str, constraints.Constraint]:
+        # Constraints handled by child distributions in dist
+        return {}
+
+
+class GridnetActorHead(Actor):
+    def __init__(
+        self,
+        map_size: int,
+        action_vec: NDArray[np.int64],
+        in_dim: EncoderOutDim,
+        hidden_sizes: Tuple[int, ...] = (32,),
+        activation: Type[nn.Module] = nn.ReLU,
+        init_layers_orthogonal: bool = True,
+    ) -> None:
+        super().__init__()
+        self.map_size = map_size
+        self.action_vec = action_vec
+        assert isinstance(in_dim, int)
+        layer_sizes = (in_dim,) + hidden_sizes + (map_size * action_vec.sum(),)
+        self._fc = mlp(
+            layer_sizes,
+            activation,
+            init_layers_orthogonal=init_layers_orthogonal,
+            final_layer_gain=0.01,
+        )
+
+    def forward(
+        self,
+        obs: torch.Tensor,
+        actions: Optional[torch.Tensor] = None,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> PiForward:
+        assert (
+            action_masks is not None
+        ), f"No mask case unhandled in {self.__class__.__name__}"
+        logits = self._fc(obs)
+        pi = GridnetDistribution(self.map_size, self.action_vec, logits, action_masks)
+        return pi_forward(pi, actions)
+
+    @property
+    def action_shape(self) -> Tuple[int, ...]:
+        return (self.map_size, len(self.action_vec))

rl_algo_impls/shared/actor/gridnet_decoder.py

@@ -0,0 +1,79 @@
+from typing import Optional, Tuple, Type
+
+import numpy as np
+import torch
+import torch.nn as nn
+from numpy.typing import NDArray
+
+from rl_algo_impls.shared.actor import Actor, PiForward, pi_forward
+from rl_algo_impls.shared.actor.gridnet import GridnetDistribution
+from rl_algo_impls.shared.encoder import EncoderOutDim
+from rl_algo_impls.shared.module.utils import layer_init
+
+
+class Transpose(nn.Module):
+    def __init__(self, permutation: Tuple[int, ...]) -> None:
+        super().__init__()
+        self.permutation = permutation
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.permute(self.permutation)
+
+
+class GridnetDecoder(Actor):
+    def __init__(
+        self,
+        map_size: int,
+        action_vec: NDArray[np.int64],
+        in_dim: EncoderOutDim,
+        activation: Type[nn.Module] = nn.ReLU,
+        init_layers_orthogonal: bool = True,
+    ) -> None:
+        super().__init__()
+        self.map_size = map_size
+        self.action_vec = action_vec
+        assert isinstance(in_dim, tuple)
+        self.deconv = nn.Sequential(
+            layer_init(
+                nn.ConvTranspose2d(
+                    in_dim[0], 128, 3, stride=2, padding=1, output_padding=1
+                ),
+                init_layers_orthogonal=init_layers_orthogonal,
+            ),
+            activation(),
+            layer_init(
+                nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1),
+                init_layers_orthogonal=init_layers_orthogonal,
+            ),
+            activation(),
+            layer_init(
+                nn.ConvTranspose2d(64, 32, 3, stride=2, padding=1, output_padding=1),
+                init_layers_orthogonal=init_layers_orthogonal,
+            ),
+            activation(),
+            layer_init(
+                nn.ConvTranspose2d(
+                    32, action_vec.sum(), 3, stride=2, padding=1, output_padding=1
+                ),
+                init_layers_orthogonal=init_layers_orthogonal,
+                std=0.01,
+            ),
+            Transpose((0, 2, 3, 1)),
+        )
+
+    def forward(
+        self,
+        obs: torch.Tensor,
+        actions: Optional[torch.Tensor] = None,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> PiForward:
+        assert (
+            action_masks is not None
+        ), f"No mask case unhandled in {self.__class__.__name__}"
+        logits = self.deconv(obs)
+        pi = GridnetDistribution(self.map_size, self.action_vec, logits, action_masks)
+        return pi_forward(pi, actions)
+
+    @property
+    def action_shape(self) -> Tuple[int, ...]:
+        return (self.map_size, len(self.action_vec))

rl_algo_impls/shared/actor/make_actor.py

@@ -0,0 +1,98 @@
+from typing import Optional, Tuple, Type
+
+import gym
+import torch.nn as nn
+from gym.spaces import Box, Discrete, MultiDiscrete
+
+from rl_algo_impls.shared.actor.actor import Actor
+from rl_algo_impls.shared.actor.categorical import CategoricalActorHead
+from rl_algo_impls.shared.actor.gaussian import GaussianActorHead
+from rl_algo_impls.shared.actor.gridnet import GridnetActorHead
+from rl_algo_impls.shared.actor.gridnet_decoder import GridnetDecoder
+from rl_algo_impls.shared.actor.multi_discrete import MultiDiscreteActorHead
+from rl_algo_impls.shared.actor.state_dependent_noise import (
+    StateDependentNoiseActorHead,
+)
+from rl_algo_impls.shared.encoder import EncoderOutDim
+
+
+def actor_head(
+    action_space: gym.Space,
+    in_dim: EncoderOutDim,
+    hidden_sizes: Tuple[int, ...],
+    init_layers_orthogonal: bool,
+    activation: Type[nn.Module],
+    log_std_init: float = -0.5,
+    use_sde: bool = False,
+    full_std: bool = True,
+    squash_output: bool = False,
+    actor_head_style: str = "single",
+    action_plane_space: Optional[bool] = None,
+) -> Actor:
+    assert not use_sde or isinstance(
+        action_space, Box
+    ), "use_sde only valid if Box action_space"
+    assert not squash_output or use_sde, "squash_output only valid if use_sde"
+    if isinstance(action_space, Discrete):
+        assert isinstance(in_dim, int)
+        return CategoricalActorHead(
+            action_space.n,  # type: ignore
+            in_dim=in_dim,
+            hidden_sizes=hidden_sizes,
+            activation=activation,
+            init_layers_orthogonal=init_layers_orthogonal,
+        )
+    elif isinstance(action_space, Box):
+        assert isinstance(in_dim, int)
+        if use_sde:
+            return StateDependentNoiseActorHead(
+                action_space.shape[0],  # type: ignore
+                in_dim=in_dim,
+                hidden_sizes=hidden_sizes,
+                activation=activation,
+                init_layers_orthogonal=init_layers_orthogonal,
+                log_std_init=log_std_init,
+                full_std=full_std,
+                squash_output=squash_output,
+            )
+        else:
+            return GaussianActorHead(
+                action_space.shape[0],  # type: ignore
+                in_dim=in_dim,
+                hidden_sizes=hidden_sizes,
+                activation=activation,
+                init_layers_orthogonal=init_layers_orthogonal,
+                log_std_init=log_std_init,
+            )
+    elif isinstance(action_space, MultiDiscrete):
+        if actor_head_style == "single":
+            return MultiDiscreteActorHead(
+                action_space.nvec,  # type: ignore
+                in_dim=in_dim,
+                hidden_sizes=hidden_sizes,
+                activation=activation,
+                init_layers_orthogonal=init_layers_orthogonal,
+            )
+        elif actor_head_style == "gridnet":
+            assert isinstance(action_plane_space, MultiDiscrete)
+            return GridnetActorHead(
+                len(action_space.nvec) // len(action_plane_space.nvec),  # type: ignore
+                action_plane_space.nvec,  # type: ignore
+                in_dim=in_dim,
+                hidden_sizes=hidden_sizes,
+                activation=activation,
+                init_layers_orthogonal=init_layers_orthogonal,
+            )
+        elif actor_head_style == "gridnet_decoder":
+            assert isinstance(action_plane_space, MultiDiscrete)
+            return GridnetDecoder(
+                len(action_space.nvec) // len(action_plane_space.nvec),  # type: ignore
+                action_plane_space.nvec,  # type: ignore
+                in_dim=in_dim,
+                activation=activation,
+                init_layers_orthogonal=init_layers_orthogonal,
+            )
+        else:
+            raise ValueError(f"Doesn't support actor_head_style {actor_head_style}")
+    else:
+        raise ValueError(f"Unsupported action space: {action_space}")

rl_algo_impls/shared/actor/multi_discrete.py

@@ -0,0 +1,101 @@
+from typing import Dict, Optional, Tuple, Type
+
+import numpy as np
+import torch
+import torch.nn as nn
+from numpy.typing import NDArray
+from torch.distributions import Distribution, constraints
+
+from rl_algo_impls.shared.actor.actor import Actor, PiForward, pi_forward
+from rl_algo_impls.shared.actor.categorical import MaskedCategorical
+from rl_algo_impls.shared.encoder import EncoderOutDim
+from rl_algo_impls.shared.module.utils import mlp
+
+
+class MultiCategorical(Distribution):
+    def __init__(
+        self,
+        nvec: NDArray[np.int64],
+        probs=None,
+        logits=None,
+        validate_args=None,
+        masks: Optional[torch.Tensor] = None,
+    ):
+        # Either probs or logits should be set
+        assert (probs is None) != (logits is None)
+        masks_split = (
+            torch.split(masks, nvec.tolist(), dim=1)
+            if masks is not None
+            else [None] * len(nvec)
+        )
+        if probs:
+            self.dists = [
+                MaskedCategorical(probs=p, validate_args=validate_args, mask=m)
+                for p, m in zip(torch.split(probs, nvec.tolist(), dim=1), masks_split)
+            ]
+            param = probs
+        else:
+            assert logits is not None
+            self.dists = [
+                MaskedCategorical(logits=lg, validate_args=validate_args, mask=m)
+                for lg, m in zip(torch.split(logits, nvec.tolist(), dim=1), masks_split)
+            ]
+            param = logits
+        batch_shape = param.size()[:-1] if param.ndimension() > 1 else torch.Size()
+        super().__init__(batch_shape=batch_shape, validate_args=validate_args)
+
+    def log_prob(self, action: torch.Tensor) -> torch.Tensor:
+        prob_stack = torch.stack(
+            [c.log_prob(a) for a, c in zip(action.T, self.dists)], dim=-1
+        )
+        return prob_stack.sum(dim=-1)
+
+    def entropy(self) -> torch.Tensor:
+        return torch.stack([c.entropy() for c in self.dists], dim=-1).sum(dim=-1)
+
+    def sample(self, sample_shape: torch.Size = torch.Size()) -> torch.Tensor:
+        return torch.stack([c.sample(sample_shape) for c in self.dists], dim=-1)
+
+    @property
+    def mode(self) -> torch.Tensor:
+        return torch.stack([c.mode for c in self.dists], dim=-1)
+
+    @property
+    def arg_constraints(self) -> Dict[str, constraints.Constraint]:
+        # Constraints handled by child distributions in dist
+        return {}
+
+
+class MultiDiscreteActorHead(Actor):
+    def __init__(
+        self,
+        nvec: NDArray[np.int64],
+        in_dim: EncoderOutDim,
+        hidden_sizes: Tuple[int, ...] = (32,),
+        activation: Type[nn.Module] = nn.ReLU,
+        init_layers_orthogonal: bool = True,
+    ) -> None:
+        super().__init__()
+        self.nvec = nvec
+        assert isinstance(in_dim, int)
+        layer_sizes = (in_dim,) + hidden_sizes + (nvec.sum(),)
+        self._fc = mlp(
+            layer_sizes,
+            activation,
+            init_layers_orthogonal=init_layers_orthogonal,
+            final_layer_gain=0.01,
+        )
+
+    def forward(
+        self,
+        obs: torch.Tensor,
+        actions: Optional[torch.Tensor] = None,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> PiForward:
+        logits = self._fc(obs)
+        pi = MultiCategorical(self.nvec, logits=logits, masks=action_masks)
+        return pi_forward(pi, actions)
+
+    @property
+    def action_shape(self) -> Tuple[int, ...]:
+        return (len(self.nvec),)

rl_algo_impls/shared/actor/state_dependent_noise.py

@@ -0,0 +1,199 @@
+from typing import Optional, Tuple, Type, TypeVar, Union
+
+import torch
+import torch.nn as nn
+from torch.distributions import Distribution, Normal
+
+from rl_algo_impls.shared.actor.actor import Actor, PiForward
+from rl_algo_impls.shared.module.utils import mlp
+
+
+class TanhBijector:
+    def __init__(self, epsilon: float = 1e-6) -> None:
+        self.epsilon = epsilon
+
+    @staticmethod
+    def forward(x: torch.Tensor) -> torch.Tensor:
+        return torch.tanh(x)
+
+    @staticmethod
+    def inverse(y: torch.Tensor) -> torch.Tensor:
+        eps = torch.finfo(y.dtype).eps
+        clamped_y = y.clamp(min=-1.0 + eps, max=1.0 - eps)
+        return torch.atanh(clamped_y)
+
+    def log_prob_correction(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.log(1.0 - torch.tanh(x) ** 2 + self.epsilon)
+
+
+def sum_independent_dims(tensor: torch.Tensor) -> torch.Tensor:
+    if len(tensor.shape) > 1:
+        return tensor.sum(dim=1)
+    return tensor.sum()
+
+
+class StateDependentNoiseDistribution(Normal):
+    def __init__(
+        self,
+        loc,
+        scale,
+        latent_sde: torch.Tensor,
+        exploration_mat: torch.Tensor,
+        exploration_matrices: torch.Tensor,
+        bijector: Optional[TanhBijector] = None,
+        validate_args=None,
+    ):
+        super().__init__(loc, scale, validate_args)
+        self.latent_sde = latent_sde
+        self.exploration_mat = exploration_mat
+        self.exploration_matrices = exploration_matrices
+        self.bijector = bijector
+
+    def log_prob(self, a: torch.Tensor) -> torch.Tensor:
+        gaussian_a = self.bijector.inverse(a) if self.bijector else a
+        log_prob = sum_independent_dims(super().log_prob(gaussian_a))
+        if self.bijector:
+            log_prob -= torch.sum(self.bijector.log_prob_correction(gaussian_a), dim=1)
+        return log_prob
+
+    def sample(self) -> torch.Tensor:
+        noise = self._get_noise()
+        actions = self.mean + noise
+        return self.bijector.forward(actions) if self.bijector else actions
+
+    def _get_noise(self) -> torch.Tensor:
+        if len(self.latent_sde) == 1 or len(self.latent_sde) != len(
+            self.exploration_matrices
+        ):
+            return torch.mm(self.latent_sde, self.exploration_mat)
+        # (batch_size, n_features) -> (batch_size, 1, n_features)
+        latent_sde = self.latent_sde.unsqueeze(dim=1)
+        # (batch_size, 1, n_actions)
+        noise = torch.bmm(latent_sde, self.exploration_matrices)
+        return noise.squeeze(dim=1)
+
+    @property
+    def mode(self) -> torch.Tensor:
+        mean = super().mode
+        return self.bijector.forward(mean) if self.bijector else mean
+
+
+StateDependentNoiseActorHeadSelf = TypeVar(
+    "StateDependentNoiseActorHeadSelf", bound="StateDependentNoiseActorHead"
+)
+
+
+class StateDependentNoiseActorHead(Actor):
+    def __init__(
+        self,
+        act_dim: int,
+        in_dim: int,
+        hidden_sizes: Tuple[int, ...] = (32,),
+        activation: Type[nn.Module] = nn.Tanh,
+        init_layers_orthogonal: bool = True,
+        log_std_init: float = -0.5,
+        full_std: bool = True,
+        squash_output: bool = False,
+        learn_std: bool = False,
+    ) -> None:
+        super().__init__()
+        self.act_dim = act_dim
+        layer_sizes = (in_dim,) + hidden_sizes + (act_dim,)
+        if len(layer_sizes) == 2:
+            self.latent_net = nn.Identity()
+        elif len(layer_sizes) > 2:
+            self.latent_net = mlp(
+                layer_sizes[:-1],
+                activation,
+                output_activation=activation,
+                init_layers_orthogonal=init_layers_orthogonal,
+            )
+        self.mu_net = mlp(
+            layer_sizes[-2:],
+            activation,
+            init_layers_orthogonal=init_layers_orthogonal,
+            final_layer_gain=0.01,
+        )
+        self.full_std = full_std
+        std_dim = (layer_sizes[-2], act_dim if self.full_std else 1)
+        self.log_std = nn.Parameter(
+            torch.ones(std_dim, dtype=torch.float32) * log_std_init
+        )
+        self.bijector = TanhBijector() if squash_output else None
+        self.learn_std = learn_std
+        self.device = None
+
+        self.exploration_mat = None
+        self.exploration_matrices = None
+        self.sample_weights()
+
+    def to(
+        self: StateDependentNoiseActorHeadSelf,
+        device: Optional[torch.device] = None,
+        dtype: Optional[Union[torch.dtype, str]] = None,
+        non_blocking: bool = False,
+    ) -> StateDependentNoiseActorHeadSelf:
+        super().to(device, dtype, non_blocking)
+        self.device = device
+        return self
+
+    def _distribution(self, obs: torch.Tensor) -> Distribution:
+        latent = self.latent_net(obs)
+        mu = self.mu_net(latent)
+        latent_sde = latent if self.learn_std else latent.detach()
+        variance = torch.mm(latent_sde**2, self._get_std() ** 2)
+        assert self.exploration_mat is not None
+        assert self.exploration_matrices is not None
+        return StateDependentNoiseDistribution(
+            mu,
+            torch.sqrt(variance + 1e-6),
+            latent_sde,
+            self.exploration_mat,
+            self.exploration_matrices,
+            self.bijector,
+        )
+
+    def _get_std(self) -> torch.Tensor:
+        std = torch.exp(self.log_std)
+        if self.full_std:
+            return std
+        ones = torch.ones(self.log_std.shape[0], self.act_dim)
+        if self.device:
+            ones = ones.to(self.device)
+        return ones * std
+
+    def forward(
+        self,
+        obs: torch.Tensor,
+        actions: Optional[torch.Tensor] = None,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> PiForward:
+        assert (
+            not action_masks
+        ), f"{self.__class__.__name__} does not support action_masks"
+        pi = self._distribution(obs)
+        return pi_forward(pi, actions)
+
+    def sample_weights(self, batch_size: int = 1) -> None:
+        std = self._get_std()
+        weights_dist = Normal(torch.zeros_like(std), std)
+        # Reparametrization trick to pass gradients
+        self.exploration_mat = weights_dist.rsample()
+        self.exploration_matrices = weights_dist.rsample(torch.Size((batch_size,)))
+
+    @property
+    def action_shape(self) -> Tuple[int, ...]:
+        return (self.act_dim,)
+
+
+def pi_forward(
+    distribution: Distribution, actions: Optional[torch.Tensor] = None
+) -> PiForward:
+    logp_a = None
+    entropy = None
+    if actions is not None:
+        logp_a = distribution.log_prob(actions)
+        entropy = (
+            -logp_a if self.bijector else sum_independent_dims(distribution.entropy())
+        )
+    return PiForward(distribution, logp_a, entropy)

rl_algo_impls/shared/algorithm.py

@@ -0,0 +1,39 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional, TypeVar
+
+import gym
+import torch
+from torch.utils.tensorboard.writer import SummaryWriter
+
+from rl_algo_impls.shared.callbacks import Callback
+from rl_algo_impls.shared.policy.policy import Policy
+from rl_algo_impls.wrappers.vectorable_wrapper import VecEnv
+
+AlgorithmSelf = TypeVar("AlgorithmSelf", bound="Algorithm")
+
+
+class Algorithm(ABC):
+    @abstractmethod
+    def __init__(
+        self,
+        policy: Policy,
+        env: VecEnv,
+        device: torch.device,
+        tb_writer: SummaryWriter,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+        self.policy = policy
+        self.env = env
+        self.device = device
+        self.tb_writer = tb_writer
+
+    @abstractmethod
+    def learn(
+        self: AlgorithmSelf,
+        train_timesteps: int,
+        callbacks: Optional[List[Callback]] = None,
+        total_timesteps: Optional[int] = None,
+        start_timesteps: int = 0,
+    ) -> AlgorithmSelf:
+        ...