A2C playing BipedalWalker-v3 from https://github.com/sgoodfriend/rl-algo-impls/tree/983cb75e43e51cf4ef57f177194ab9a4a1a8808b

README.md

@@ -23,17 +23,17 @@ model-index:
 
 This is a trained model of a **A2C** agent playing **BipedalWalker-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/7lx79bf0.
+All models trained at this commit can be found at https://api.wandb.ai/links/sgoodfriend/ysd5gj7p.
 
 ## Training Results
 
-This model was trained from 3 trainings of **A2C** agents using different initial seeds. These agents were trained by checking out [0511de3](https://github.com/sgoodfriend/rl-algo-impls/tree/0511de345b17175b7cf1ea706c3e05981f11761c). 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).
+This model was trained from 3 trainings of **A2C** agents using different initial seeds. These agents were trained by checking out [983cb75](https://github.com/sgoodfriend/rl-algo-impls/tree/983cb75e43e51cf4ef57f177194ab9a4a1a8808b). 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                                                                    |
 |:-------|:-----------------|-------:|--------------:|-------------:|----------------:|:-------|:-----------------------------------------------------------------------------|
-| a2c    | BipedalWalker-v3 |      1 |       273.388 |    73.7341   |              16 |        | [wandb](https://wandb.ai/sgoodfriend/rl-algo-impls-benchmarks/runs/7tqykx8n) |
-| a2c    | BipedalWalker-v3 |      2 |       285.59  |    50.0351   |              16 |        | [wandb](https://wandb.ai/sgoodfriend/rl-algo-impls-benchmarks/runs/pp7btarg) |
-| a2c    | BipedalWalker-v3 |      3 |       297.937 |     0.587948 |              16 | *      | [wandb](https://wandb.ai/sgoodfriend/rl-algo-impls-benchmarks/runs/vxh9e1hm) |
+| a2c    | BipedalWalker-v3 |      1 |       273.388 |    73.7341   |              16 |        | [wandb](https://wandb.ai/sgoodfriend/rl-algo-impls-benchmarks/runs/o8o1rqmd) |
+| a2c    | BipedalWalker-v3 |      2 |       285.59  |    50.0351   |              16 |        | [wandb](https://wandb.ai/sgoodfriend/rl-algo-impls-benchmarks/runs/o6e9e32z) |
+| a2c    | BipedalWalker-v3 |      3 |       297.937 |     0.587948 |              16 | *      | [wandb](https://wandb.ai/sgoodfriend/rl-algo-impls-benchmarks/runs/phljdsth) |
 
 
 ### Prerequisites: Weights & Biases (WandB)
@@ -53,10 +53,10 @@ login`.
 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:
-[0511de3](https://github.com/sgoodfriend/rl-algo-impls/tree/0511de345b17175b7cf1ea706c3e05981f11761c).
+[983cb75](https://github.com/sgoodfriend/rl-algo-impls/tree/983cb75e43e51cf4ef57f177194ab9a4a1a8808b).
 ```
 # Downloads the model, sets hyperparameters, and runs agent for 3 episodes
-python enjoy.py --wandb-run-path=sgoodfriend/rl-algo-impls-benchmarks/vxh9e1hm
+python enjoy.py --wandb-run-path=sgoodfriend/rl-algo-impls-benchmarks/phljdsth
 ```
 
 Setup hasn't been completely worked out yet, so you might be best served by using Google
@@ -68,7 +68,7 @@ notebook.
 
 ## Training
 If you want the highest chance to reproduce these results, you'll want to checkout the
-commit the agent was trained on: [0511de3](https://github.com/sgoodfriend/rl-algo-impls/tree/0511de345b17175b7cf1ea706c3e05981f11761c). While
+commit the agent was trained on: [983cb75](https://github.com/sgoodfriend/rl-algo-impls/tree/983cb75e43e51cf4ef57f177194ab9a4a1a8808b). While
 training is deterministic, different hardware will give different results.
 
 ```
@@ -83,7 +83,7 @@ notebook.
 
 
 ## Benchmarking (with Lambda Labs instance)
-This and other models from https://api.wandb.ai/links/sgoodfriend/7lx79bf0 were generated by running a script on a Lambda
+This and other models from https://api.wandb.ai/links/sgoodfriend/ysd5gj7p were generated by running a script on a Lambda
 Labs instance. In a Lambda Labs instance terminal:
 ```
 git clone [email protected]:sgoodfriend/rl-algo-impls.git
@@ -122,7 +122,8 @@ env_hyperparams:
   n_envs: 16
   normalize: true
 env_id: null
-eval_params: {}
+eval_hyperparams: {}
+microrts_reward_decay_callback: false
 n_timesteps: 5000000
 policy_hyperparams:
   init_layers_orthogonal: false
@@ -134,9 +135,9 @@ wandb_entity: null
 wandb_group: null
 wandb_project_name: rl-algo-impls-benchmarks
 wandb_tags:
-- benchmark_0511de3
-- host_152-67-249-42
+- benchmark_983cb75
+- host_129-159-43-75
 - branch_main
-- v0.0.8
+- v0.0.9
 
 ```

environment.yml

@@ -4,7 +4,7 @@ channels:
   - conda-forge
   - nodefaults
 dependencies:
-  - python>=3.8, <3.11
+  - python>=3.8, <3.10
   - mamba
   - pip
   - pytorch

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rl_algo_impls"
-version = "0.0.8"
+version = "0.0.9"
 description = "Implementations of reinforcement learning algorithms"
 authors = [
     {name = "Scott Goodfriend", email = "[email protected]"},
@@ -56,14 +56,17 @@ procgen = [
     "glfw >= 1.12.0, < 1.13",
     "procgen; platform_machine=='x86_64'",
 ]
-microrts-old = [
+microrts-ppo = [
     "numpy < 1.24.0", # Support for gym-microrts < 0.6.0
     "gym-microrts == 0.2.0", # Match ppo-implementation-details
 ]
-microrts = [
+microrts-paper = [
     "numpy < 1.24.0", # Support for gym-microrts < 0.6.0
     "gym-microrts == 0.3.2",
 ]
+microrts = [
+    "gym-microrts",
+]
 jupyter = [
     "jupyter",
     "notebook"

replay.meta.json

@@ -1 +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", "600x400", "-pix_fmt", "rgb24", "-framerate", "50", "-i", "-", "-vf", "scale=trunc(iw/2)*2:trunc(ih/2)*2", "-vcodec", "libx264", "-pix_fmt", "yuv420p", "-r", "50", "/tmp/tmpobbokqe8/a2c-BipedalWalker-v3/replay.mp4"]}, "episode": {"r": 298.23712158203125, "l": 895, "t": 9.969387}}
+{"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", "600x400", "-pix_fmt", "rgb24", "-framerate", "50", "-i", "-", "-vf", "scale=trunc(iw/2)*2:trunc(ih/2)*2", "-vcodec", "libx264", "-pix_fmt", "yuv420p", "-r", "50", "/tmp/tmpl5jgowib/a2c-BipedalWalker-v3/replay.mp4"]}, "episodes": [{"r": 298.23712158203125, "l": 895, "t": 10.336128}]}

rl_algo_impls/a2c/a2c.py

@@ -1,23 +1,23 @@
 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 time import perf_counter
 from torch.utils.tensorboard.writer import SummaryWriter
-from typing import Optional, TypeVar
 
 from rl_algo_impls.shared.algorithm import Algorithm
-from rl_algo_impls.shared.callbacks.callback import Callback
+from rl_algo_impls.shared.callbacks import Callback
 from rl_algo_impls.shared.gae import compute_advantages
-from rl_algo_impls.shared.policy.on_policy import ActorCritic
+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_observation_space,
     single_action_space,
+    single_observation_space,
 )
 
 A2CSelf = TypeVar("A2CSelf", bound="A2C")
@@ -70,7 +70,7 @@ class A2C(Algorithm):
     def learn(
         self: A2CSelf,
         train_timesteps: int,
-        callback: Optional[Callback] = None,
+        callbacks: Optional[List[Callback]] = None,
         total_timesteps: Optional[int] = None,
         start_timesteps: int = 0,
     ) -> A2CSelf:
@@ -193,8 +193,10 @@ class A2C(Algorithm):
                 timesteps_elapsed,
             )
 
-            if callback:
-                if not callback.on_step(timesteps_elapsed=rollout_steps):
+            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"
                     )

rl_algo_impls/a2c/optimize.py

@@ -1,10 +1,10 @@
-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
+import optuna
+
+from rl_algo_impls.runner.config import Config, EnvHyperparams, Hyperparams
 from rl_algo_impls.shared.policy.optimize_on_policy import sample_on_policy_hyperparams
+from rl_algo_impls.shared.vec_env import make_eval_env
 from rl_algo_impls.tuning.optimize_env import sample_env_hyperparams
 
 
@@ -16,7 +16,11 @@ def sample_params(
     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 = make_eval_env(
+        base_config,
+        base_env_hyperparams,
+        override_hparams={"n_envs": 1},
+    )
 
     # env_hyperparams
     env_hyperparams = sample_env_hyperparams(trial, hyperparams.env_hyperparams, env)

rl_algo_impls/dqn/dqn.py

@@ -1,18 +1,19 @@
 import copy
-import numpy as np
+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 collections import deque
 from torch.optim import Adam
 from torch.utils.tensorboard.writer import SummaryWriter
-from typing import NamedTuple, Optional, TypeVar
 
 from rl_algo_impls.dqn.policy import DQNPolicy
 from rl_algo_impls.shared.algorithm import Algorithm
-from rl_algo_impls.shared.callbacks.callback import Callback
+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
 
@@ -118,7 +119,7 @@ class DQN(Algorithm):
         self.max_grad_norm = max_grad_norm
 
     def learn(
-        self: DQNSelf, total_timesteps: int, callback: Optional[Callback] = None
+        self: DQNSelf, total_timesteps: int, callbacks: Optional[List[Callback]] = None
     ) -> DQNSelf:
         self.policy.train(True)
         obs = self.env.reset()
@@ -140,8 +141,14 @@ class DQN(Algorithm):
             if steps_since_target_update >= self.target_update_interval:
                 self._update_target()
                 steps_since_target_update = 0
-            if callback:
-                callback.on_step(timesteps_elapsed=rollout_steps)
+            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:

rl_algo_impls/dqn/q_net.py

@@ -6,7 +6,7 @@ import torch.nn as nn
 from gym.spaces import Discrete
 
 from rl_algo_impls.shared.encoder import Encoder
-from rl_algo_impls.shared.module.module import mlp
+from rl_algo_impls.shared.module.utils import mlp
 
 
 class QNetwork(nn.Module):

rl_algo_impls/huggingface_publish.py

@@ -3,24 +3,23 @@ import os
 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
 
 import argparse
-import requests
 import shutil
 import subprocess
 import tempfile
-import wandb
-import wandb.apis.public
-
 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.vec_env import make_eval_env
 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
 
 
@@ -134,7 +133,7 @@ def publish(
             make_eval_env(
                 config,
                 EnvHyperparams(**config.env_hyperparams),
-                override_n_envs=1,
+                override_hparams={"n_envs": 1},
                 normalize_load_path=model_path,
             ),
             os.path.join(repo_dir_path, "replay"),
@@ -144,7 +143,7 @@ def publish(
             video_env,
             policy,
             1,
-            deterministic=config.eval_params.get("deterministic", True),
+            deterministic=config.eval_hyperparams.get("deterministic", True),
         )
 
         api = HfApi()

rl_algo_impls/hyperparams/a2c.yml

@@ -101,31 +101,32 @@ HopperBulletEnv-v0:
 CarRacing-v0:
   n_timesteps: !!float 4e6
   env_hyperparams:
-    n_envs: 16
+    n_envs: 4
     frame_stack: 4
     normalize: true
     normalize_kwargs:
       norm_obs: false
       norm_reward: true
   policy_hyperparams:
-    use_sde: false
-    log_std_init: -1.3502584927786276
+    use_sde: true
+    log_std_init: -4.839609092563
     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
+    n_steps: 64
+    learning_rate: 0.000018971962220405576
+    gamma: 0.9942776405534832
+    gae_lambda: 0.9549244758833236
+    ent_coef: 0.0000015666550584860516
+    ent_coef_decay: linear
+    vf_coef: 0.12164696385898476
+    max_grad_norm: 2.2574480552177127
+    normalize_advantage: false
     use_rms_prop: false
+    sde_sample_freq: 16
 
 _atari: &atari-defaults
   n_timesteps: !!float 1e7

rl_algo_impls/hyperparams/dqn.yml

@@ -15,7 +15,7 @@ CartPole-v1: &cartpole-defaults
     gradient_steps: 128
     exploration_fraction: 0.16
     exploration_final_eps: 0.04
-  eval_params:
+  eval_hyperparams:
     step_freq: !!float 1e4
 
 CartPole-v0:
@@ -76,7 +76,7 @@ LunarLander-v2:
     exploration_fraction: 0.12
     exploration_final_eps: 0.1
     max_grad_norm: 0.5
-  eval_params:
+  eval_hyperparams:
     step_freq: 25_000
 
 _atari: &atari-defaults
@@ -97,7 +97,7 @@ _atari: &atari-defaults
     gradient_steps: 2
     exploration_fraction: 0.1
     exploration_final_eps: 0.01
-  eval_params:
+  eval_hyperparams:
     deterministic: false
 
 PongNoFrameskip-v4:

rl_algo_impls/hyperparams/ppo.yml

@@ -13,7 +13,7 @@ CartPole-v1: &cartpole-defaults
     learning_rate_decay: linear
     clip_range: 0.2
     clip_range_decay: linear
-  eval_params:
+  eval_hyperparams:
     step_freq: !!float 2.5e4
 
 CartPole-v0:
@@ -52,7 +52,7 @@ MountainCarContinuous-v0:
     gae_lambda: 0.9
     max_grad_norm: 5
     vf_coef: 0.19
-  eval_params:
+  eval_hyperparams:
     step_freq: 5000
 
 Acrobot-v1:
@@ -162,7 +162,7 @@ _atari: &atari-defaults
     clip_range_decay: linear
     vf_coef: 0.5
     ent_coef: 0.01
-  eval_params:
+  eval_hyperparams:
     deterministic: false
 
 _norm-rewards-atari: &norm-rewards-atari-default
@@ -228,7 +228,7 @@ _microrts: &microrts-defaults
     clip_range_decay: none
     clip_range_vf: 0.1
     ppo2_vf_coef_halving: true
-  eval_params:
+  eval_hyperparams: &microrts-eval-defaults
     deterministic: false # Good idea because MultiCategorical mode isn't great
 
 _no-mask-microrts: &no-mask-microrts-defaults
@@ -252,15 +252,15 @@ MicrortsRandomEnemyShapedReward3-v1-NoMask:
 _microrts_ai: &microrts-ai-defaults
   <<: *microrts-defaults
   n_timesteps: !!float 100e6
-  additional_keys_to_log: ["microrts_stats"]
+  additional_keys_to_log: ["microrts_stats", "microrts_results"]
   env_hyperparams: &microrts-ai-env-defaults
     n_envs: 24
     env_type: microrts
-    make_kwargs:
+    make_kwargs: &microrts-ai-env-make-kwargs-defaults
       num_selfplay_envs: 0
-      max_steps: 2000
+      max_steps: 4000
       render_theme: 2
-      map_path: maps/16x16/basesWorkers16x16.xml
+      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
@@ -278,6 +278,15 @@ _microrts_ai: &microrts-ai-defaults
     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
@@ -305,6 +314,18 @@ enc-dec-MicrortsDefeatRandomEnemySparseReward-v3:
     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
@@ -313,6 +334,27 @@ MicrortsDefeatCoacAIShaped-v3: &microrts-coacai-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
@@ -325,6 +367,7 @@ MicrortsDefeatCoacAIShaped-v3-diverseBots: &microrts-diverse-defaults
       workerRushAI: 2
 
 enc-dec-MicrortsDefeatCoacAIShaped-v3-diverseBots:
+  &microrts-env-dec-diverse-defaults
   <<: *microrts-diverse-defaults
   policy_hyperparams:
     <<: *microrts-ai-policy-defaults
@@ -332,6 +375,76 @@ enc-dec-MicrortsDefeatCoacAIShaped-v3-diverseBots:
     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: 300000
+      swap_steps: 6000
+      swap_window_size: 4
+      window: 33
+  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
@@ -418,7 +531,7 @@ _procgen: &procgen-defaults
     learning_rate: !!float 5e-4
     # learning_rate_decay: linear
     vf_coef: 0.5
-  eval_params: &procgen-eval-defaults
+  eval_hyperparams: &procgen-eval-defaults
     ignore_first_episode: true
     # deterministic: false
     step_freq: !!float 1e5
@@ -466,7 +579,7 @@ _procgen-hard: &procgen-hard-defaults
     batch_size: 8192
     clip_range_decay: linear
     learning_rate_decay: linear
-  eval_params:
+  eval_hyperparams:
     <<: *procgen-eval-defaults
     step_freq: !!float 5e5
 

rl_algo_impls/hyperparams/vpg.yml

@@ -7,7 +7,7 @@ CartPole-v1: &cartpole-defaults
     gae_lambda: 1
     val_lr: 0.01
     train_v_iters: 80
-  eval_params:
+  eval_hyperparams:
     step_freq: !!float 2.5e4
 
 CartPole-v0:
@@ -52,7 +52,7 @@ MountainCarContinuous-v0:
     val_lr: !!float 1e-3
     train_v_iters: 80
     max_grad_norm: 5
-  eval_params:
+  eval_hyperparams:
     step_freq: 5000
 
 Acrobot-v1:
@@ -78,7 +78,7 @@ LunarLander-v2:
     val_lr: 0.0001
     train_v_iters: 80
     max_grad_norm: 0.5
-  eval_params:
+  eval_hyperparams:
     deterministic: false
 
 BipedalWalker-v3:
@@ -96,7 +96,7 @@ BipedalWalker-v3:
     val_lr: !!float 1e-4
     train_v_iters: 80
     max_grad_norm: 0.5
-  eval_params:
+  eval_hyperparams:
     deterministic: false
 
 CarRacing-v0:
@@ -169,7 +169,7 @@ FrozenLake-v1:
     val_lr: 0.01
     train_v_iters: 80
     max_grad_norm: 0.5
-  eval_params:
+  eval_hyperparams:
     step_freq: !!float 5e4
     n_episodes: 10
     save_best: true
@@ -193,5 +193,5 @@ _atari: &atari-defaults
     train_v_iters: 80
     max_grad_norm: 0.5
     ent_coef: 0.01
-  eval_params:
+  eval_hyperparams:
     deterministic: false

rl_algo_impls/optimize.py

@@ -2,37 +2,44 @@ 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 os
 import torch
-import wandb
-
-from dataclasses import asdict, dataclass
 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
-from typing import Callable, List, NamedTuple, Optional, Sequence, Union
 
+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.shared.vec_env import make_env, make_eval_env
 from rl_algo_impls.runner.running_utils import (
+    ALGOS,
     base_parser,
-    load_hyperparams,
-    set_seeds,
     get_device,
-    make_policy,
-    ALGOS,
     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
@@ -195,29 +202,38 @@ def simple_optimize(trial: optuna.Trial, args: RunArgs, study_args: StudyArgs) -
         config, EnvHyperparams(**config.env_hyperparams), tb_writer=tb_writer
     )
     device = get_device(config, env)
-    policy = make_policy(args.algo, env, device, **config.policy_hyperparams)
+    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,
+        override_hparams={"n_envs": study_args.n_eval_envs},
     )
-    callback = OptimizeCallback(
+    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_params.get("deterministic", True),
+        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, callback=callback)
+        algo.learn(config.n_timesteps, callbacks=callbacks)
 
-        if not callback.is_pruned:
-            callback.evaluate()
-            if not callback.is_pruned:
+        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
@@ -230,8 +246,8 @@ def simple_optimize(trial: optuna.Trial, args: RunArgs, study_args: StudyArgs) -
                 "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": callback.last_score,
-                "hparam/is_pruned": callback.is_pruned,
+                "hparam/score": optimize_callback.last_score,
+                "hparam/is_pruned": optimize_callback.is_pruned,
             },
             None,
             config.run_name(),
@@ -239,13 +255,15 @@ def simple_optimize(trial: optuna.Trial, args: RunArgs, study_args: StudyArgs) -
         tb_writer.close()
 
         if wandb_enabled:
-            wandb.run.summary["state"] = "Pruned" if callback.is_pruned else "Complete"
+            wandb.run.summary["state"] = (  # type: ignore
+                "Pruned" if optimize_callback.is_pruned else "Complete"
+            )
             wandb.finish(quiet=True)
 
-        if callback.is_pruned:
+        if optimize_callback.is_pruned:
             raise optuna.exceptions.TrialPruned()
 
-        return callback.last_score
+        return optimize_callback.last_score
     except AssertionError as e:
         logging.warning(e)
         return np.nan
@@ -299,7 +317,10 @@ def stepwise_optimize(
                 tb_writer=tb_writer,
             )
             device = get_device(config, env)
-            policy = make_policy(arg.algo, env, device, **config.policy_hyperparams)
+            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](
@@ -310,7 +331,7 @@ def stepwise_optimize(
                 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,
+                override_hparams={"n_envs": study_args.n_eval_envs},
             )
 
             start_timesteps = int(i * config.n_timesteps / study_args.n_evaluations)
@@ -319,10 +340,22 @@ def stepwise_optimize(
                 - 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,
-                    callback=None,
+                    callbacks=callbacks,
                     total_timesteps=config.n_timesteps,
                     start_timesteps=start_timesteps,
                 )
@@ -333,7 +366,7 @@ def stepwise_optimize(
                         eval_env,
                         tb_writer,
                         study_args.n_eval_episodes,
-                        config.eval_params.get("deterministic", True),
+                        config.eval_hyperparams.get("deterministic", True),
                         start_timesteps + train_timesteps,
                     )
                 )
@@ -379,7 +412,7 @@ def stepwise_optimize(
 
 
 def wandb_finish(state: str) -> None:
-    wandb.run.summary["state"] = state
+    wandb.run.summary["state"] = state  # type: ignore
     wandb.finish(quiet=True)
 
 

rl_algo_impls/ppo/ppo.py

@@ -10,12 +10,16 @@ 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.callback import Callback
+from rl_algo_impls.shared.callbacks import Callback
 from rl_algo_impls.shared.gae import compute_advantages
-from rl_algo_impls.shared.policy.on_policy import ActorCritic
-from rl_algo_impls.shared.schedule import schedule, update_learning_rate
+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.action_mask_wrapper import find_action_masker
 from rl_algo_impls.wrappers.vectorable_wrapper import (
     VecEnv,
     single_action_space,
@@ -102,12 +106,17 @@ class PPO(Algorithm):
         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.action_masker = find_action_masker(env)
+        self.get_action_mask = getattr(env, "get_action_mask", None)
 
-        self.gamma = gamma
+        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)
@@ -138,7 +147,7 @@ class PPO(Algorithm):
     def learn(
         self: PPOSelf,
         train_timesteps: int,
-        callback: Optional[Callback] = None,
+        callbacks: Optional[List[Callback]] = None,
         total_timesteps: Optional[int] = None,
         start_timesteps: int = 0,
     ) -> PPOSelf:
@@ -153,15 +162,13 @@ class PPO(Algorithm):
         act_shape = self.policy.action_shape
 
         next_obs = self.env.reset()
-        next_action_masks = (
-            self.action_masker.action_masks() if self.action_masker else None
-        )
-        next_episode_starts = np.full(step_dim, True, dtype=np.bool8)
+        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.bool8)
+        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 = (
@@ -181,10 +188,12 @@ class PPO(Algorithm):
             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)
@@ -215,7 +224,7 @@ class PPO(Algorithm):
                     clamped_action
                 )
                 next_action_masks = (
-                    self.action_masker.action_masks() if self.action_masker else None
+                    self.get_action_mask() if self.get_action_mask else None
                 )
 
             self.policy.train()
@@ -251,7 +260,7 @@ class PPO(Algorithm):
                         next_episode_starts,
                         next_obs,
                         self.policy,
-                        self.gamma,
+                        gamma,
                         self.gae_lambda,
                     )
                     b_advantages = torch.tensor(advantages.reshape(-1)).to(self.device)
@@ -364,8 +373,10 @@ class PPO(Algorithm):
                 timesteps_elapsed,
             )
 
-            if callback:
-                if not callback.on_step(timesteps_elapsed=rollout_steps):
+            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"
                     )

rl_algo_impls/runner/config.py

@@ -51,6 +51,8 @@ class EnvHyperparams:
     normalize_type: Optional[str] = None
     mask_actions: bool = False
     bots: Optional[Dict[str, int]] = None
+    self_play_kwargs: Optional[Dict[str, Any]] = None
+    selfplay_bots: Optional[Dict[str, int]] = None
 
 
 HyperparamsSelf = TypeVar("HyperparamsSelf", bound="Hyperparams")
@@ -63,9 +65,10 @@ class Hyperparams:
     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_params: 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(
@@ -110,8 +113,14 @@ class Config:
         return self.hyperparams.algo_hyperparams
 
     @property
-    def eval_params(self) -> Dict[str, Any]:
-        return self.hyperparams.eval_params
+    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:

rl_algo_impls/runner/evaluate.py

@@ -1,20 +1,19 @@
 import os
 import shutil
-
 from dataclasses import dataclass
 from typing import NamedTuple, Optional
 
-from rl_algo_impls.shared.vec_env import make_eval_env
 from rl_algo_impls.runner.config import Config, EnvHyperparams, Hyperparams, RunArgs
 from rl_algo_impls.runner.running_utils import (
-    load_hyperparams,
-    set_seeds,
     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
@@ -71,7 +70,7 @@ def evaluate_model(args: EvalArgs, root_dir: str) -> Evaluation:
     env = make_eval_env(
         config,
         EnvHyperparams(**config.env_hyperparams),
-        override_n_envs=args.n_envs,
+        override_hparams={"n_envs": args.n_envs} if args.n_envs else None,
         render=args.render,
         normalize_load_path=model_path,
     )
@@ -87,7 +86,7 @@ def evaluate_model(args: EvalArgs, root_dir: str) -> Evaluation:
     deterministic = (
         args.deterministic_eval
         if args.deterministic_eval is not None
-        else config.eval_params.get("deterministic", True)
+        else config.eval_hyperparams.get("deterministic", True)
     )
     return Evaluation(
         policy,

rl_algo_impls/runner/running_utils.py

@@ -22,7 +22,7 @@ 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.on_policy import ActorCritic
+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
@@ -97,29 +97,21 @@ def get_device(config: Config, env: VecEnv) -> torch.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):
+        # 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"
-        elif isinstance(obs_space, Box) and len(obs_space.shape) == 1:
-            device = "cpu"
-        if is_microrts(config):
-            try:
-                from gym_microrts.envs.vec_env import MicroRTSGridModeVecEnv
-
-                # Models that move more than one unit at a time should use mps
-                if not isinstance(env.unwrapped, MicroRTSGridModeVecEnv):
-                    device = "cpu"
-            except ModuleNotFoundError:
-                # Likely on gym_microrts v0.0.2 to match ppo-implementation-details
+        # 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)

rl_algo_impls/runner/selfplay_evaluate.py

@@ -0,0 +1,142 @@
+import copy
+import dataclasses
+import os
+import shutil
+from dataclasses import dataclass
+from typing import List, NamedTuple, Optional
+
+import numpy as np
+
+import wandb
+from rl_algo_impls.runner.config import Config, EnvHyperparams, Hyperparams, RunArgs
+from rl_algo_impls.runner.evaluate import Evaluation
+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.vec_env import make_eval_env
+from rl_algo_impls.wrappers.vec_episode_recorder import VecEpisodeRecorder
+
+
+@dataclass
+class SelfplayEvalArgs(RunArgs):
+    # Either wandb_run_paths or model_file_paths must have 2 elements in it.
+    wandb_run_paths: List[str] = dataclasses.field(default_factory=list)
+    model_file_paths: List[str] = dataclasses.field(default_factory=list)
+    render: bool = False
+    best: bool = True
+    n_envs: int = 1
+    n_episodes: int = 1
+    deterministic_eval: Optional[bool] = None
+    no_print_returns: bool = False
+    video_path: Optional[str] = None
+
+
+def selfplay_evaluate(args: SelfplayEvalArgs, root_dir: str) -> Evaluation:
+    if args.wandb_run_paths:
+        api = wandb.Api()
+        args, config, player_1_model_path = load_player(
+            api, args.wandb_run_paths[0], args, root_dir
+        )
+        _, _, player_2_model_path = load_player(
+            api, args.wandb_run_paths[1], args, root_dir
+        )
+    elif args.model_file_paths:
+        hyperparams = load_hyperparams(args.algo, args.env)
+
+        config = Config(args, hyperparams, root_dir)
+        player_1_model_path, player_2_model_path = args.model_file_paths
+    else:
+        raise ValueError("Must specify 2 wandb_run_paths or 2 model_file_paths")
+
+    print(args)
+
+    set_seeds(args.seed, args.use_deterministic_algorithms)
+
+    env_make_kwargs = (
+        config.eval_hyperparams.get("env_overrides", {}).get("make_kwargs", {}).copy()
+    )
+    env_make_kwargs["num_selfplay_envs"] = args.n_envs * 2
+    env = make_eval_env(
+        config,
+        EnvHyperparams(**config.env_hyperparams),
+        override_hparams={
+            "n_envs": args.n_envs,
+            "selfplay_bots": {
+                player_2_model_path: args.n_envs,
+            },
+            "self_play_kwargs": {
+                "num_old_policies": 0,
+                "save_steps": np.inf,
+                "swap_steps": np.inf,
+                "bot_always_player_2": True,
+            },
+            "bots": None,
+            "make_kwargs": env_make_kwargs,
+        },
+        render=args.render,
+        normalize_load_path=player_1_model_path,
+    )
+    if args.video_path:
+        env = VecEpisodeRecorder(
+            env, args.video_path, max_video_length=18000, num_episodes=args.n_episodes
+        )
+    device = get_device(config, env)
+    policy = make_policy(
+        args.algo,
+        env,
+        device,
+        load_path=player_1_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,
+    )
+
+
+class PlayerData(NamedTuple):
+    args: SelfplayEvalArgs
+    config: Config
+    model_path: str
+
+
+def load_player(
+    api: wandb.Api, run_path: str, args: SelfplayEvalArgs, root_dir: str
+) -> PlayerData:
+    args = copy.copy(args)
+
+    run = api.run(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)
+
+    return PlayerData(args, config, model_path)

rl_algo_impls/runner/train.py

@@ -1,12 +1,17 @@
 # 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, Optional, Sequence
+from typing import Any, Dict, List, Optional, Sequence
 
 import yaml
 from torch.utils.tensorboard.writer import SummaryWriter
@@ -23,6 +28,9 @@ from rl_algo_impls.runner.running_utils import (
     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
 
@@ -41,7 +49,7 @@ def train(args: TrainArgs):
     print(hyperparams)
     config = Config(args, hyperparams, os.getcwd())
 
-    wandb_enabled = args.wandb_project_name
+    wandb_enabled = bool(args.wandb_project_name)
     if wandb_enabled:
         wandb.tensorboard.patch(
             root_logdir=config.tensorboard_summary_path, pytorch=True
@@ -66,14 +74,17 @@ def train(args: TrainArgs):
         config, EnvHyperparams(**config.env_hyperparams), tb_writer=tb_writer
     )
     device = get_device(config, env)
-    policy = make_policy(args.algo, env, device, **config.policy_hyperparams)
+    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
-        wandb.run.summary["num_trainable_parameters"] = num_trainable_parameters
+        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} ; "
@@ -81,40 +92,49 @@ def train(args: TrainArgs):
         )
 
     eval_env = make_eval_env(config, EnvHyperparams(**config.env_hyperparams))
-    record_best_videos = config.eval_params.get("record_best_videos", True)
-    callback = EvalCallback(
+    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_params,
+        **config.eval_callback_params(),
         video_env=make_eval_env(
-            config, EnvHyperparams(**config.env_hyperparams), override_n_envs=1
+            config,
+            EnvHyperparams(**config.env_hyperparams),
+            override_hparams={"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,
+        wandb_enabled=wandb_enabled,
     )
-    algo.learn(config.n_timesteps, callback=callback)
+    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 = callback.evaluate(n_episodes=10, print_returns=True)
+    eval_stats = eval_callback.evaluate(n_episodes=10, print_returns=True)
 
-    plot_eval_callback(callback, tb_writer, config.run_name())
+    plot_eval_callback(eval_callback, tb_writer, config.run_name())
 
     log_dict: Dict[str, Any] = {
         "eval": eval_stats._asdict(),
     }
-    if callback.best:
-        log_dict["best_eval"] = callback.best._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 = callback.best  # type: ignore
+    best_eval_stats: EpisodesStats = eval_callback.best  # type: ignore
     tb_writer.add_hparams(
         hparam_dict(hyperparams, vars(args)),
         {
@@ -132,13 +152,8 @@ def train(args: TrainArgs):
 
     if wandb_enabled:
         shutil.make_archive(
-            os.path.join(wandb.run.dir, config.model_dir_name()),
+            os.path.join(wandb.run.dir, config.model_dir_name()),  # type: ignore
             "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/selfplay_enjoy.py

@@ -0,0 +1,53 @@
+# 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.running_utils import base_parser
+from rl_algo_impls.runner.selfplay_evaluate import SelfplayEvalArgs, selfplay_evaluate
+
+
+def selfplay_enjoy() -> None:
+    parser = base_parser(multiple=False)
+    parser.add_argument(
+        "--wandb-run-paths",
+        type=str,
+        nargs="*",
+        help="WandB run paths to load players from. Must be 0 or 2",
+    )
+    parser.add_argument(
+        "--model-file-paths",
+        type=str,
+        help="File paths to load players from. Must be 0 or 2",
+    )
+    parser.add_argument("--render", action="store_true")
+    parser.add_argument("--n-envs", default=1, type=int)
+    parser.add_argument("--n-episodes", default=1, type=int)
+    parser.add_argument("--deterministic-eval", default=None, type=bool)
+    parser.add_argument(
+        "--no-print-returns", action="store_true", help="Limit printing"
+    )
+    parser.add_argument(
+        "--video-path", type=str, help="Path to save video of all plays"
+    )
+    # parser.set_defaults(
+    #     algo=["ppo"],
+    #     env=["Microrts-selfplay-unet-decay"],
+    #     n_episodes=10,
+    #     model_file_paths=[
+    #         "downloaded_models/ppo-Microrts-selfplay-unet-decay-S3-best",
+    #         "downloaded_models/ppo-Microrts-selfplay-unet-decay-S2-best",
+    #     ],
+    #     video_path="/Users/sgoodfriend/Desktop/decay3-vs-decay2",
+    # )
+    args = parser.parse_args()
+    args.algo = args.algo[0]
+    args.env = args.env[0]
+    args.seed = args.seed[0]
+    args = SelfplayEvalArgs(**vars(args))
+
+    selfplay_evaluate(args, os.getcwd())
+
+
+if __name__ == "__main__":
+    selfplay_enjoy()

rl_algo_impls/shared/actor/__init__.py

@@ -1,2 +1,2 @@
-from rl_algo_impls.shared.actor.actor import Actor, PiForward
+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

@@ -31,12 +31,13 @@ class Actor(nn.Module, ABC):
     def action_shape(self) -> Tuple[int, ...]:
         ...
 
-    def pi_forward(
-        self, 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)
+
+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

@@ -4,8 +4,8 @@ import torch
 import torch.nn as nn
 from torch.distributions import Categorical
 
-from rl_algo_impls.shared.actor import Actor, PiForward
-from rl_algo_impls.shared.module.module import mlp
+from rl_algo_impls.shared.actor import Actor, PiForward, pi_forward
+from rl_algo_impls.shared.module.utils import mlp
 
 
 class MaskedCategorical(Categorical):
@@ -57,7 +57,7 @@ class CategoricalActorHead(Actor):
     ) -> PiForward:
         logits = self._fc(obs)
         pi = MaskedCategorical(logits=logits, mask=action_masks)
-        return self.pi_forward(pi, actions)
+        return pi_forward(pi, actions)
 
     @property
     def action_shape(self) -> Tuple[int, ...]:

rl_algo_impls/shared/actor/gaussian.py

@@ -4,8 +4,8 @@ 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.module import mlp
+from rl_algo_impls.shared.actor.actor import Actor, PiForward, pi_forward
+from rl_algo_impls.shared.module.utils import mlp
 
 
 class GaussianDistribution(Normal):
@@ -54,7 +54,7 @@ class GaussianActorHead(Actor):
             not action_masks
         ), f"{self.__class__.__name__} does not support action_masks"
         pi = self._distribution(obs)
-        return self.pi_forward(pi, actions)
+        return pi_forward(pi, actions)
 
     @property
     def action_shape(self) -> Tuple[int, ...]:

rl_algo_impls/shared/actor/gridnet.py

@@ -6,10 +6,10 @@ 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
+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.module import mlp
+from rl_algo_impls.shared.module.utils import mlp
 
 
 class GridnetDistribution(Distribution):
@@ -25,7 +25,7 @@ class GridnetDistribution(Distribution):
         self.action_vec = action_vec
 
         masks = masks.view(-1, masks.shape[-1])
-        split_masks = torch.split(masks[:, 1:], action_vec.tolist(), dim=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)
@@ -101,7 +101,7 @@ class GridnetActorHead(Actor):
         ), 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 self.pi_forward(pi, actions)
+        return pi_forward(pi, actions)
 
     @property
     def action_shape(self) -> Tuple[int, ...]:

rl_algo_impls/shared/actor/gridnet_decoder.py

@@ -5,11 +5,10 @@ import torch
 import torch.nn as nn
 from numpy.typing import NDArray
 
-from rl_algo_impls.shared.actor import Actor, PiForward
-from rl_algo_impls.shared.actor.categorical import MaskedCategorical
+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.module import layer_init
+from rl_algo_impls.shared.module.utils import layer_init
 
 
 class Transpose(nn.Module):
@@ -73,7 +72,7 @@ class GridnetDecoder(Actor):
         ), 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 self.pi_forward(pi, actions)
+        return pi_forward(pi, actions)
 
     @property
     def action_shape(self) -> Tuple[int, ...]:

rl_algo_impls/shared/actor/make_actor.py

@@ -1,4 +1,4 @@
-from typing import Tuple, Type
+from typing import Optional, Tuple, Type
 
 import gym
 import torch.nn as nn
@@ -27,6 +27,7 @@ def actor_head(
     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
@@ -73,18 +74,20 @@ def actor_head(
                 init_layers_orthogonal=init_layers_orthogonal,
             )
         elif actor_head_style == "gridnet":
+            assert isinstance(action_plane_space, MultiDiscrete)
             return GridnetActorHead(
-                action_space.nvec[0],  # type: ignore
-                action_space.nvec[1:],  # type: ignore
+                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(
-                action_space.nvec[0],  # type: ignore
-                action_space.nvec[1:],  # type: ignore
+                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,

rl_algo_impls/shared/actor/multi_discrete.py

@@ -6,10 +6,10 @@ 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
+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.module import mlp
+from rl_algo_impls.shared.module.utils import mlp
 
 
 class MultiCategorical(Distribution):
@@ -94,7 +94,7 @@ class MultiDiscreteActorHead(Actor):
     ) -> PiForward:
         logits = self._fc(obs)
         pi = MultiCategorical(self.nvec, logits=logits, masks=action_masks)
-        return self.pi_forward(pi, actions)
+        return pi_forward(pi, actions)
 
     @property
     def action_shape(self) -> Tuple[int, ...]:

rl_algo_impls/shared/actor/state_dependent_noise.py

@@ -5,7 +5,7 @@ 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.module import mlp
+from rl_algo_impls.shared.module.utils import mlp
 
 
 class TanhBijector:
@@ -172,7 +172,7 @@ class StateDependentNoiseActorHead(Actor):
             not action_masks
         ), f"{self.__class__.__name__} does not support action_masks"
         pi = self._distribution(obs)
-        return self.pi_forward(pi, actions)
+        return pi_forward(pi, actions, self.bijector)
 
     def sample_weights(self, batch_size: int = 1) -> None:
         std = self._get_std()
@@ -185,16 +185,15 @@ class StateDependentNoiseActorHead(Actor):
     def action_shape(self) -> Tuple[int, ...]:
         return (self.act_dim,)
 
-    def pi_forward(
-        self, 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)
+
+def pi_forward(
+    distribution: Distribution,
+    actions: Optional[torch.Tensor] = None,
+    bijector: Optional[TanhBijector] = None,
+) -> PiForward:
+    logp_a = None
+    entropy = None
+    if actions is not None:
+        logp_a = distribution.log_prob(actions)
+        entropy = -logp_a if bijector else sum_independent_dims(distribution.entropy())
+    return PiForward(distribution, logp_a, entropy)

rl_algo_impls/shared/algorithm.py

@@ -1,11 +1,11 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional, TypeVar
+
 import gym
 import torch
-
-from abc import ABC, abstractmethod
 from torch.utils.tensorboard.writer import SummaryWriter
-from typing import Optional, TypeVar
 
-from rl_algo_impls.shared.callbacks.callback import Callback
+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
 
@@ -32,7 +32,7 @@ class Algorithm(ABC):
     def learn(
         self: AlgorithmSelf,
         train_timesteps: int,
-        callback: Optional[Callback] = None,
+        callbacks: Optional[List[Callback]] = None,
         total_timesteps: Optional[int] = None,
         start_timesteps: int = 0,
     ) -> AlgorithmSelf:

rl_algo_impls/shared/callbacks/__init__.py

@@ -0,0 +1 @@
+from rl_algo_impls.shared.callbacks.callback import Callback

rl_algo_impls/shared/callbacks/eval_callback.py

@@ -1,12 +1,13 @@
 import itertools
 import os
+import shutil
 from time import perf_counter
 from typing import Dict, List, Optional, Union
 
 import numpy as np
 from torch.utils.tensorboard.writer import SummaryWriter
 
-from rl_algo_impls.shared.callbacks.callback import Callback
+from rl_algo_impls.shared.callbacks import Callback
 from rl_algo_impls.shared.policy.policy import Policy
 from rl_algo_impls.shared.stats import Episode, EpisodeAccumulator, EpisodesStats
 from rl_algo_impls.wrappers.action_mask_wrapper import find_action_masker
@@ -80,6 +81,7 @@ def evaluate(
     print_returns: bool = True,
     ignore_first_episode: bool = False,
     additional_keys_to_log: Optional[List[str]] = None,
+    score_function: str = "mean-std",
 ) -> EpisodesStats:
     policy.sync_normalization(env)
     policy.eval()
@@ -93,18 +95,21 @@ def evaluate(
     )
 
     obs = env.reset()
-    action_masker = find_action_masker(env)
+    get_action_mask = getattr(env, "get_action_mask", None)
     while not episodes.is_done():
         act = policy.act(
             obs,
             deterministic=deterministic,
-            action_masks=action_masker.action_masks() if action_masker else None,
+            action_masks=get_action_mask() if get_action_mask else None,
         )
         obs, rew, done, info = env.step(act)
         episodes.step(rew, done, info)
         if render:
             env.render()
-    stats = EpisodesStats(episodes.episodes)
+    stats = EpisodesStats(
+        episodes.episodes,
+        score_function=score_function,
+    )
     if print_returns:
         print(stats)
     return stats
@@ -127,6 +132,8 @@ class EvalCallback(Callback):
         max_video_length: int = 3600,
         ignore_first_episode: bool = False,
         additional_keys_to_log: Optional[List[str]] = None,
+        score_function: str = "mean-std",
+        wandb_enabled: bool = False,
     ) -> None:
         super().__init__()
         self.policy = policy
@@ -151,6 +158,8 @@ class EvalCallback(Callback):
         self.best_video_base_path = None
         self.ignore_first_episode = ignore_first_episode
         self.additional_keys_to_log = additional_keys_to_log
+        self.score_function = score_function
+        self.wandb_enabled = wandb_enabled
 
     def on_step(self, timesteps_elapsed: int = 1) -> bool:
         super().on_step(timesteps_elapsed)
@@ -170,6 +179,7 @@ class EvalCallback(Callback):
             print_returns=print_returns or False,
             ignore_first_episode=self.ignore_first_episode,
             additional_keys_to_log=self.additional_keys_to_log,
+            score_function=self.score_function,
         )
         end_time = perf_counter()
         self.tb_writer.add_scalar(
@@ -189,6 +199,15 @@ class EvalCallback(Callback):
                 assert self.best_model_path
                 self.policy.save(self.best_model_path)
                 print("Saved best model")
+                if self.wandb_enabled:
+                    import wandb
+
+                    best_model_name = os.path.split(self.best_model_path)[-1]
+                    shutil.make_archive(
+                        os.path.join(wandb.run.dir, best_model_name),  # type: ignore
+                        "zip",
+                        self.best_model_path,
+                    )
             self.best.write_to_tensorboard(
                 self.tb_writer, "best_eval", self.timesteps_elapsed
             )
@@ -208,6 +227,7 @@ class EvalCallback(Callback):
                     1,
                     deterministic=self.deterministic,
                     print_returns=False,
+                    score_function=self.score_function,
                 )
                 print(f"Saved best video: {video_stats}")
 

rl_algo_impls/shared/callbacks/microrts_reward_decay_callback.py

@@ -0,0 +1,36 @@
+import numpy as np
+
+from rl_algo_impls.runner.config import Config
+from rl_algo_impls.shared.callbacks import Callback
+from rl_algo_impls.wrappers.vectorable_wrapper import VecEnv
+
+
+class MicrortsRewardDecayCallback(Callback):
+    def __init__(
+        self,
+        config: Config,
+        env: VecEnv,
+        start_timesteps: int = 0,
+    ) -> None:
+        super().__init__()
+        from gym_microrts.envs.vec_env import MicroRTSGridModeVecEnv
+
+        unwrapped = env.unwrapped
+        assert isinstance(unwrapped, MicroRTSGridModeVecEnv)
+        self.microrts_env = unwrapped
+        self.base_reward_weights = self.microrts_env.reward_weight
+
+        self.total_train_timesteps = config.n_timesteps
+        self.timesteps_elapsed = start_timesteps
+
+    def on_step(self, timesteps_elapsed: int = 1) -> bool:
+        super().on_step(timesteps_elapsed)
+
+        progress = self.timesteps_elapsed / self.total_train_timesteps
+        # Decay all rewards except WinLoss
+        reward_weights = self.base_reward_weights * np.array(
+            [1] + [1 - progress] * (len(self.base_reward_weights) - 1)
+        )
+        self.microrts_env.reward_weight = reward_weights
+
+        return True

rl_algo_impls/shared/callbacks/optimize_callback.py

@@ -5,7 +5,7 @@ from time import perf_counter
 from torch.utils.tensorboard.writer import SummaryWriter
 from typing import NamedTuple, Union
 
-from rl_algo_impls.shared.callbacks.callback import Callback
+from rl_algo_impls.shared.callbacks import Callback
 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

rl_algo_impls/shared/callbacks/self_play_callback.py

@@ -0,0 +1,34 @@
+from typing import Callable
+
+from rl_algo_impls.shared.callbacks import Callback
+from rl_algo_impls.shared.policy.policy import Policy
+from rl_algo_impls.wrappers.self_play_wrapper import SelfPlayWrapper
+
+
+class SelfPlayCallback(Callback):
+    def __init__(
+        self,
+        policy: Policy,
+        policy_factory: Callable[[], Policy],
+        selfPlayWrapper: SelfPlayWrapper,
+    ) -> None:
+        super().__init__()
+        self.policy = policy
+        self.policy_factory = policy_factory
+        self.selfPlayWrapper = selfPlayWrapper
+        self.checkpoint_policy()
+
+    def on_step(self, timesteps_elapsed: int = 1) -> bool:
+        super().on_step(timesteps_elapsed)
+        if (
+            self.timesteps_elapsed
+            >= self.last_checkpoint_step + self.selfPlayWrapper.save_steps
+        ):
+            self.checkpoint_policy()
+        return True
+
+    def checkpoint_policy(self):
+        self.selfPlayWrapper.checkpoint_policy(
+            self.policy_factory().load_from(self.policy)
+        )
+        self.last_checkpoint_step = self.timesteps_elapsed

rl_algo_impls/shared/encoder/cnn.py

@@ -6,7 +6,7 @@ import numpy as np
 import torch
 import torch.nn as nn
 
-from rl_algo_impls.shared.module.module import layer_init
+from rl_algo_impls.shared.module.utils import layer_init
 
 EncoderOutDim = Union[int, Tuple[int, ...]]
 

rl_algo_impls/shared/encoder/encoder.py

@@ -12,7 +12,7 @@ from rl_algo_impls.shared.encoder.gridnet_encoder import GridnetEncoder
 from rl_algo_impls.shared.encoder.impala_cnn import ImpalaCnn
 from rl_algo_impls.shared.encoder.microrts_cnn import MicrortsCnn
 from rl_algo_impls.shared.encoder.nature_cnn import NatureCnn
-from rl_algo_impls.shared.module.module import layer_init
+from rl_algo_impls.shared.module.utils import layer_init
 
 CNN_EXTRACTORS_BY_STYLE: Dict[str, Type[CnnEncoder]] = {
     "nature": NatureCnn,

rl_algo_impls/shared/encoder/gridnet_encoder.py

@@ -5,7 +5,7 @@ import torch
 import torch.nn as nn
 
 from rl_algo_impls.shared.encoder.cnn import CnnEncoder, EncoderOutDim
-from rl_algo_impls.shared.module.module import layer_init
+from rl_algo_impls.shared.module.utils import layer_init
 
 
 class GridnetEncoder(CnnEncoder):

rl_algo_impls/shared/encoder/impala_cnn.py

@@ -5,7 +5,7 @@ import torch
 import torch.nn as nn
 
 from rl_algo_impls.shared.encoder.cnn import FlattenedCnnEncoder
-from rl_algo_impls.shared.module.module import layer_init
+from rl_algo_impls.shared.module.utils import layer_init
 
 
 class ResidualBlock(nn.Module):

rl_algo_impls/shared/encoder/microrts_cnn.py

@@ -5,7 +5,7 @@ import torch
 import torch.nn as nn
 
 from rl_algo_impls.shared.encoder.cnn import FlattenedCnnEncoder
-from rl_algo_impls.shared.module.module import layer_init
+from rl_algo_impls.shared.module.utils import layer_init
 
 
 class MicrortsCnn(FlattenedCnnEncoder):

rl_algo_impls/shared/encoder/nature_cnn.py

@@ -4,7 +4,7 @@ import gym
 import torch.nn as nn
 
 from rl_algo_impls.shared.encoder.cnn import FlattenedCnnEncoder
-from rl_algo_impls.shared.module.module import layer_init
+from rl_algo_impls.shared.module.utils import layer_init
 
 
 class NatureCnn(FlattenedCnnEncoder):

rl_algo_impls/shared/gae.py

@@ -3,7 +3,7 @@ import torch
 
 from typing import NamedTuple, Sequence
 
-from rl_algo_impls.shared.policy.on_policy import OnPolicy
+from rl_algo_impls.shared.policy.actor_critic import OnPolicy
 from rl_algo_impls.shared.trajectory import Trajectory
 from rl_algo_impls.wrappers.vectorable_wrapper import VecEnvObs
 

rl_algo_impls/shared/policy/{on_policy.py → actor_critic.py}

@@ -4,12 +4,14 @@ from typing import NamedTuple, Optional, Sequence, Tuple, TypeVar
 import gym
 import numpy as np
 import torch
-from gym.spaces import Box, Discrete, Space
+from gym.spaces import Box, Space
 
-from rl_algo_impls.shared.actor import PiForward, actor_head
-from rl_algo_impls.shared.encoder import Encoder
-from rl_algo_impls.shared.policy.critic import CriticHead
-from rl_algo_impls.shared.policy.policy import ACTIVATION, Policy
+from rl_algo_impls.shared.policy.actor_critic_network import (
+    ConnectedTrioActorCriticNetwork,
+    SeparateActorCriticNetwork,
+    UNetActorCriticNetwork,
+)
+from rl_algo_impls.shared.policy.policy import Policy
 from rl_algo_impls.wrappers.vectorable_wrapper import (
     VecEnv,
     VecEnvObs,
@@ -52,21 +54,6 @@ def clamp_actions(
     return actions
 
 
-def default_hidden_sizes(obs_space: Space) -> Sequence[int]:
-    if isinstance(obs_space, Box):
-        if len(obs_space.shape) == 3:
-            # By default feature extractor to output has no hidden layers
-            return []
-        elif len(obs_space.shape) == 1:
-            return [64, 64]
-        else:
-            raise ValueError(f"Unsupported observation space: {obs_space}")
-    elif isinstance(obs_space, Discrete):
-        return [64]
-    else:
-        raise ValueError(f"Unsupported observation space: {obs_space}")
-
-
 class OnPolicy(Policy):
     @abstractmethod
     def value(self, obs: VecEnvObs) -> np.ndarray:
@@ -106,78 +93,59 @@ class ActorCritic(OnPolicy):
 
         observation_space = single_observation_space(env)
         action_space = single_action_space(env)
+        action_plane_space = getattr(env, "action_plane_space", None)
 
-        pi_hidden_sizes = (
-            pi_hidden_sizes
-            if pi_hidden_sizes is not None
-            else default_hidden_sizes(observation_space)
-        )
-        v_hidden_sizes = (
-            v_hidden_sizes
-            if v_hidden_sizes is not None
-            else default_hidden_sizes(observation_space)
-        )
-
-        activation = ACTIVATION[activation_fn]
         self.action_space = action_space
         self.squash_output = squash_output
-        self.share_features_extractor = share_features_extractor
-        self._feature_extractor = Encoder(
-            observation_space,
-            activation,
-            init_layers_orthogonal=init_layers_orthogonal,
-            cnn_flatten_dim=cnn_flatten_dim,
-            cnn_style=cnn_style,
-            cnn_layers_init_orthogonal=cnn_layers_init_orthogonal,
-            impala_channels=impala_channels,
-        )
-        self._pi = actor_head(
-            self.action_space,
-            self._feature_extractor.out_dim,
-            tuple(pi_hidden_sizes),
-            init_layers_orthogonal,
-            activation,
-            log_std_init=log_std_init,
-            use_sde=use_sde,
-            full_std=full_std,
-            squash_output=squash_output,
-            actor_head_style=actor_head_style,
-        )
 
-        if not share_features_extractor:
-            self._v_feature_extractor = Encoder(
+        if actor_head_style == "unet":
+            self.network = UNetActorCriticNetwork(
                 observation_space,
-                activation,
+                action_space,
+                action_plane_space,
+                v_hidden_sizes=v_hidden_sizes,
                 init_layers_orthogonal=init_layers_orthogonal,
+                activation_fn=activation_fn,
+                cnn_layers_init_orthogonal=cnn_layers_init_orthogonal,
+            )
+        elif share_features_extractor:
+            self.network = ConnectedTrioActorCriticNetwork(
+                observation_space,
+                action_space,
+                pi_hidden_sizes=pi_hidden_sizes,
+                v_hidden_sizes=v_hidden_sizes,
+                init_layers_orthogonal=init_layers_orthogonal,
+                activation_fn=activation_fn,
+                log_std_init=log_std_init,
+                use_sde=use_sde,
+                full_std=full_std,
+                squash_output=squash_output,
                 cnn_flatten_dim=cnn_flatten_dim,
                 cnn_style=cnn_style,
                 cnn_layers_init_orthogonal=cnn_layers_init_orthogonal,
+                impala_channels=impala_channels,
+                actor_head_style=actor_head_style,
+                action_plane_space=action_plane_space,
             )
-            critic_in_dim = self._v_feature_extractor.out_dim
         else:
-            self._v_feature_extractor = None
-            critic_in_dim = self._feature_extractor.out_dim
-        self._v = CriticHead(
-            in_dim=critic_in_dim,
-            hidden_sizes=v_hidden_sizes,
-            activation=activation,
-            init_layers_orthogonal=init_layers_orthogonal,
-        )
-
-    def _pi_forward(
-        self,
-        obs: torch.Tensor,
-        action_masks: Optional[torch.Tensor],
-        action: Optional[torch.Tensor] = None,
-    ) -> Tuple[PiForward, torch.Tensor]:
-        p_fe = self._feature_extractor(obs)
-        pi_forward = self._pi(p_fe, actions=action, action_masks=action_masks)
-
-        return pi_forward, p_fe
-
-    def _v_forward(self, obs: torch.Tensor, p_fc: torch.Tensor) -> torch.Tensor:
-        v_fe = self._v_feature_extractor(obs) if self._v_feature_extractor else p_fc
-        return self._v(v_fe)
+            self.network = SeparateActorCriticNetwork(
+                observation_space,
+                action_space,
+                pi_hidden_sizes=pi_hidden_sizes,
+                v_hidden_sizes=v_hidden_sizes,
+                init_layers_orthogonal=init_layers_orthogonal,
+                activation_fn=activation_fn,
+                log_std_init=log_std_init,
+                use_sde=use_sde,
+                full_std=full_std,
+                squash_output=squash_output,
+                cnn_flatten_dim=cnn_flatten_dim,
+                cnn_style=cnn_style,
+                cnn_layers_init_orthogonal=cnn_layers_init_orthogonal,
+                impala_channels=impala_channels,
+                actor_head_style=actor_head_style,
+                action_plane_space=action_plane_space,
+            )
 
     def forward(
         self,
@@ -185,8 +153,7 @@ class ActorCritic(OnPolicy):
         action: torch.Tensor,
         action_masks: Optional[torch.Tensor] = None,
     ) -> ACForward:
-        (_, logp_a, entropy), p_fc = self._pi_forward(obs, action_masks, action=action)
-        v = self._v_forward(obs, p_fc)
+        (_, logp_a, entropy), v = self.network(obs, action, action_masks=action_masks)
 
         assert logp_a is not None
         assert entropy is not None
@@ -195,24 +162,17 @@ class ActorCritic(OnPolicy):
     def value(self, obs: VecEnvObs) -> np.ndarray:
         o = self._as_tensor(obs)
         with torch.no_grad():
-            fe = (
-                self._v_feature_extractor(o)
-                if self._v_feature_extractor
-                else self._feature_extractor(o)
-            )
-            v = self._v(fe)
+            v = self.network.value(o)
         return v.cpu().numpy()
 
     def step(self, obs: VecEnvObs, action_masks: Optional[np.ndarray] = None) -> Step:
         o = self._as_tensor(obs)
         a_masks = self._as_tensor(action_masks) if action_masks is not None else None
         with torch.no_grad():
-            (pi, _, _), p_fc = self._pi_forward(o, action_masks=a_masks)
+            (pi, _, _), v = self.network.distribution_and_value(o, action_masks=a_masks)
             a = pi.sample()
             logp_a = pi.log_prob(a)
 
-            v = self._v_forward(o, p_fc)
-
         a_np = a.cpu().numpy()
         clamped_a_np = clamp_actions(a_np, self.action_space, self.squash_output)
         return Step(a_np, v.cpu().numpy(), logp_a.cpu().numpy(), clamped_a_np)
@@ -231,7 +191,9 @@ class ActorCritic(OnPolicy):
                 self._as_tensor(action_masks) if action_masks is not None else None
             )
             with torch.no_grad():
-                (pi, _, _), _ = self._pi_forward(o, action_masks=a_masks)
+                (pi, _, _), _ = self.network.distribution_and_value(
+                    o, action_masks=a_masks
+                )
                 a = pi.mode
             return clamp_actions(a.cpu().numpy(), self.action_space, self.squash_output)
 
@@ -239,11 +201,16 @@ class ActorCritic(OnPolicy):
         super().load(path)
         self.reset_noise()
 
+    def load_from(self: ActorCriticSelf, policy: ActorCriticSelf) -> ActorCriticSelf:
+        super().load_from(policy)
+        self.reset_noise()
+        return self
+
     def reset_noise(self, batch_size: Optional[int] = None) -> None:
-        self._pi.sample_weights(
+        self.network.reset_noise(
             batch_size=batch_size if batch_size else self.env.num_envs
         )
 
     @property
     def action_shape(self) -> Tuple[int, ...]:
-        return self._pi.action_shape
+        return self.network.action_shape

rl_algo_impls/shared/policy/actor_critic_network/__init__.py

@@ -0,0 +1,11 @@
+from rl_algo_impls.shared.policy.actor_critic_network.connected_trio import (
+    ConnectedTrioActorCriticNetwork,
+)
+from rl_algo_impls.shared.policy.actor_critic_network.network import (
+    ActorCriticNetwork,
+    default_hidden_sizes,
+)
+from rl_algo_impls.shared.policy.actor_critic_network.separate_actor_critic import (
+    SeparateActorCriticNetwork,
+)
+from rl_algo_impls.shared.policy.actor_critic_network.unet import UNetActorCriticNetwork

rl_algo_impls/shared/policy/actor_critic_network/connected_trio.py

@@ -0,0 +1,118 @@
+from typing import Optional, Sequence, Tuple
+
+import torch
+from gym.spaces import Space
+
+from rl_algo_impls.shared.actor import actor_head
+from rl_algo_impls.shared.encoder import Encoder
+from rl_algo_impls.shared.policy.actor_critic_network.network import (
+    ACNForward,
+    ActorCriticNetwork,
+    default_hidden_sizes,
+)
+from rl_algo_impls.shared.policy.critic import CriticHead
+from rl_algo_impls.shared.policy.policy import ACTIVATION
+
+
+class ConnectedTrioActorCriticNetwork(ActorCriticNetwork):
+    """Encode (feature extractor), decoder (actor head), critic head networks"""
+
+    def __init__(
+        self,
+        observation_space: Space,
+        action_space: Space,
+        pi_hidden_sizes: Optional[Sequence[int]] = None,
+        v_hidden_sizes: Optional[Sequence[int]] = None,
+        init_layers_orthogonal: bool = True,
+        activation_fn: str = "tanh",
+        log_std_init: float = -0.5,
+        use_sde: bool = False,
+        full_std: bool = True,
+        squash_output: bool = False,
+        cnn_flatten_dim: int = 512,
+        cnn_style: str = "nature",
+        cnn_layers_init_orthogonal: Optional[bool] = None,
+        impala_channels: Sequence[int] = (16, 32, 32),
+        actor_head_style: str = "single",
+        action_plane_space: Optional[Space] = None,
+    ) -> None:
+        super().__init__()
+
+        pi_hidden_sizes = (
+            pi_hidden_sizes
+            if pi_hidden_sizes is not None
+            else default_hidden_sizes(observation_space)
+        )
+        v_hidden_sizes = (
+            v_hidden_sizes
+            if v_hidden_sizes is not None
+            else default_hidden_sizes(observation_space)
+        )
+
+        activation = ACTIVATION[activation_fn]
+        self._feature_extractor = Encoder(
+            observation_space,
+            activation,
+            init_layers_orthogonal=init_layers_orthogonal,
+            cnn_flatten_dim=cnn_flatten_dim,
+            cnn_style=cnn_style,
+            cnn_layers_init_orthogonal=cnn_layers_init_orthogonal,
+            impala_channels=impala_channels,
+        )
+        self._pi = actor_head(
+            action_space,
+            self._feature_extractor.out_dim,
+            tuple(pi_hidden_sizes),
+            init_layers_orthogonal,
+            activation,
+            log_std_init=log_std_init,
+            use_sde=use_sde,
+            full_std=full_std,
+            squash_output=squash_output,
+            actor_head_style=actor_head_style,
+            action_plane_space=action_plane_space,
+        )
+
+        self._v = CriticHead(
+            in_dim=self._feature_extractor.out_dim,
+            hidden_sizes=v_hidden_sizes,
+            activation=activation,
+            init_layers_orthogonal=init_layers_orthogonal,
+        )
+
+    def forward(
+        self,
+        obs: torch.Tensor,
+        action: torch.Tensor,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> ACNForward:
+        return self._distribution_and_value(
+            obs, action=action, action_masks=action_masks
+        )
+
+    def distribution_and_value(
+        self, obs: torch.Tensor, action_masks: Optional[torch.Tensor] = None
+    ) -> ACNForward:
+        return self._distribution_and_value(obs, action_masks=action_masks)
+
+    def _distribution_and_value(
+        self,
+        obs: torch.Tensor,
+        action: Optional[torch.Tensor] = None,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> ACNForward:
+        encoded = self._feature_extractor(obs)
+        pi_forward = self._pi(encoded, actions=action, action_masks=action_masks)
+        v = self._v(encoded)
+        return ACNForward(pi_forward, v)
+
+    def value(self, obs: torch.Tensor) -> torch.Tensor:
+        encoded = self._feature_extractor(obs)
+        return self._v(encoded)
+
+    def reset_noise(self, batch_size: int) -> None:
+        self._pi.sample_weights(batch_size=batch_size)
+
+    @property
+    def action_shape(self) -> Tuple[int, ...]:
+        return self._pi.action_shape

rl_algo_impls/shared/policy/actor_critic_network/network.py

@@ -0,0 +1,57 @@
+from abc import ABC, abstractmethod
+from typing import NamedTuple, Optional, Sequence, Tuple
+
+import torch
+import torch.nn as nn
+from gym.spaces import Box, Discrete, Space
+
+from rl_algo_impls.shared.actor import PiForward
+
+
+class ACNForward(NamedTuple):
+    pi_forward: PiForward
+    v: torch.Tensor
+
+
+class ActorCriticNetwork(nn.Module, ABC):
+    @abstractmethod
+    def forward(
+        self,
+        obs: torch.Tensor,
+        action: torch.Tensor,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> ACNForward:
+        ...
+
+    @abstractmethod
+    def distribution_and_value(
+        self, obs: torch.Tensor, action_masks: Optional[torch.Tensor] = None
+    ) -> ACNForward:
+        ...
+
+    @abstractmethod
+    def value(self, obs: torch.Tensor) -> torch.Tensor:
+        ...
+
+    @abstractmethod
+    def reset_noise(self, batch_size: Optional[int] = None) -> None:
+        ...
+
+    @property
+    def action_shape(self) -> Tuple[int, ...]:
+        ...
+
+
+def default_hidden_sizes(obs_space: Space) -> Sequence[int]:
+    if isinstance(obs_space, Box):
+        if len(obs_space.shape) == 3:  # type: ignore
+            # By default feature extractor to output has no hidden layers
+            return []
+        elif len(obs_space.shape) == 1:  # type: ignore
+            return [64, 64]
+        else:
+            raise ValueError(f"Unsupported observation space: {obs_space}")
+    elif isinstance(obs_space, Discrete):
+        return [64]
+    else:
+        raise ValueError(f"Unsupported observation space: {obs_space}")

rl_algo_impls/shared/policy/actor_critic_network/separate_actor_critic.py

@@ -0,0 +1,128 @@
+from typing import Optional, Sequence, Tuple
+
+import torch
+import torch.nn as nn
+from gym.spaces import Space
+
+from rl_algo_impls.shared.actor import actor_head
+from rl_algo_impls.shared.encoder import Encoder
+from rl_algo_impls.shared.policy.actor_critic_network.network import (
+    ACNForward,
+    ActorCriticNetwork,
+    default_hidden_sizes,
+)
+from rl_algo_impls.shared.policy.critic import CriticHead
+from rl_algo_impls.shared.policy.policy import ACTIVATION
+
+
+class SeparateActorCriticNetwork(ActorCriticNetwork):
+    def __init__(
+        self,
+        observation_space: Space,
+        action_space: Space,
+        pi_hidden_sizes: Optional[Sequence[int]] = None,
+        v_hidden_sizes: Optional[Sequence[int]] = None,
+        init_layers_orthogonal: bool = True,
+        activation_fn: str = "tanh",
+        log_std_init: float = -0.5,
+        use_sde: bool = False,
+        full_std: bool = True,
+        squash_output: bool = False,
+        cnn_flatten_dim: int = 512,
+        cnn_style: str = "nature",
+        cnn_layers_init_orthogonal: Optional[bool] = None,
+        impala_channels: Sequence[int] = (16, 32, 32),
+        actor_head_style: str = "single",
+        action_plane_space: Optional[Space] = None,
+    ) -> None:
+        super().__init__()
+
+        pi_hidden_sizes = (
+            pi_hidden_sizes
+            if pi_hidden_sizes is not None
+            else default_hidden_sizes(observation_space)
+        )
+        v_hidden_sizes = (
+            v_hidden_sizes
+            if v_hidden_sizes is not None
+            else default_hidden_sizes(observation_space)
+        )
+
+        activation = ACTIVATION[activation_fn]
+        self._feature_extractor = Encoder(
+            observation_space,
+            activation,
+            init_layers_orthogonal=init_layers_orthogonal,
+            cnn_flatten_dim=cnn_flatten_dim,
+            cnn_style=cnn_style,
+            cnn_layers_init_orthogonal=cnn_layers_init_orthogonal,
+            impala_channels=impala_channels,
+        )
+        self._pi = actor_head(
+            action_space,
+            self._feature_extractor.out_dim,
+            tuple(pi_hidden_sizes),
+            init_layers_orthogonal,
+            activation,
+            log_std_init=log_std_init,
+            use_sde=use_sde,
+            full_std=full_std,
+            squash_output=squash_output,
+            actor_head_style=actor_head_style,
+            action_plane_space=action_plane_space,
+        )
+
+        v_encoder = Encoder(
+            observation_space,
+            activation,
+            init_layers_orthogonal=init_layers_orthogonal,
+            cnn_flatten_dim=cnn_flatten_dim,
+            cnn_style=cnn_style,
+            cnn_layers_init_orthogonal=cnn_layers_init_orthogonal,
+        )
+        self._v = nn.Sequential(
+            v_encoder,
+            CriticHead(
+                in_dim=v_encoder.out_dim,
+                hidden_sizes=v_hidden_sizes,
+                activation=activation,
+                init_layers_orthogonal=init_layers_orthogonal,
+            ),
+        )
+
+    def forward(
+        self,
+        obs: torch.Tensor,
+        action: torch.Tensor,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> ACNForward:
+        return self._distribution_and_value(
+            obs, action=action, action_masks=action_masks
+        )
+
+    def distribution_and_value(
+        self, obs: torch.Tensor, action_masks: Optional[torch.Tensor] = None
+    ) -> ACNForward:
+        return self._distribution_and_value(obs, action_masks=action_masks)
+
+    def _distribution_and_value(
+        self,
+        obs: torch.Tensor,
+        action: Optional[torch.Tensor] = None,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> ACNForward:
+        pi_forward = self._pi(
+            self._feature_extractor(obs), actions=action, action_masks=action_masks
+        )
+        v = self._v(obs)
+        return ACNForward(pi_forward, v)
+
+    def value(self, obs: torch.Tensor) -> torch.Tensor:
+        return self._v(obs)
+
+    def reset_noise(self, batch_size: int) -> None:
+        self._pi.sample_weights(batch_size=batch_size)
+
+    @property
+    def action_shape(self) -> Tuple[int, ...]:
+        return self._pi.action_shape

rl_algo_impls/shared/policy/actor_critic_network/unet.py

@@ -0,0 +1,196 @@
+from typing import Optional, Sequence, Tuple, Type
+
+import numpy as np
+import torch
+import torch.nn as nn
+from gym.spaces import MultiDiscrete, Space
+
+from rl_algo_impls.shared.actor import pi_forward
+from rl_algo_impls.shared.actor.gridnet import GridnetDistribution
+from rl_algo_impls.shared.actor.gridnet_decoder import Transpose
+from rl_algo_impls.shared.module.utils import layer_init
+from rl_algo_impls.shared.policy.actor_critic_network.network import (
+    ACNForward,
+    ActorCriticNetwork,
+    default_hidden_sizes,
+)
+from rl_algo_impls.shared.policy.critic import CriticHead
+from rl_algo_impls.shared.policy.policy import ACTIVATION
+
+
+class UNetActorCriticNetwork(ActorCriticNetwork):
+    def __init__(
+        self,
+        observation_space: Space,
+        action_space: Space,
+        action_plane_space: Space,
+        v_hidden_sizes: Optional[Sequence[int]] = None,
+        init_layers_orthogonal: bool = True,
+        activation_fn: str = "tanh",
+        cnn_layers_init_orthogonal: Optional[bool] = None,
+    ) -> None:
+        if cnn_layers_init_orthogonal is None:
+            cnn_layers_init_orthogonal = True
+        super().__init__()
+        assert isinstance(action_space, MultiDiscrete)
+        assert isinstance(action_plane_space, MultiDiscrete)
+        self.range_size = np.max(observation_space.high) - np.min(observation_space.low)  # type: ignore
+        self.map_size = len(action_space.nvec) // len(action_plane_space.nvec)  # type: ignore
+        self.action_vec = action_plane_space.nvec  # type: ignore
+
+        activation = ACTIVATION[activation_fn]
+
+        def conv_relu(
+            in_channels: int, out_channels: int, kernel_size: int = 3, padding: int = 1
+        ) -> nn.Module:
+            return nn.Sequential(
+                layer_init(
+                    nn.Conv2d(
+                        in_channels,
+                        out_channels,
+                        kernel_size=kernel_size,
+                        padding=padding,
+                    ),
+                    cnn_layers_init_orthogonal,
+                ),
+                activation(),
+            )
+
+        def up_conv_relu(in_channels: int, out_channels: int) -> nn.Module:
+            return nn.Sequential(
+                layer_init(
+                    nn.ConvTranspose2d(
+                        in_channels,
+                        out_channels,
+                        kernel_size=3,
+                        stride=2,
+                        padding=1,
+                        output_padding=1,
+                    ),
+                    cnn_layers_init_orthogonal,
+                ),
+                activation(),
+            )
+
+        in_channels = observation_space.shape[0]  # type: ignore
+        self.enc1 = conv_relu(in_channels, 32)
+        self.enc2 = nn.Sequential(max_pool(), conv_relu(32, 64))
+        self.enc3 = nn.Sequential(max_pool(), conv_relu(64, 128))
+        self.enc4 = nn.Sequential(max_pool(), conv_relu(128, 256))
+        self.enc5 = nn.Sequential(
+            max_pool(), conv_relu(256, 512, kernel_size=1, padding=0)
+        )
+
+        self.dec4 = up_conv_relu(512, 256)
+        self.dec3 = nn.Sequential(conv_relu(512, 256), up_conv_relu(256, 128))
+        self.dec2 = nn.Sequential(conv_relu(256, 128), up_conv_relu(128, 64))
+        self.dec1 = nn.Sequential(conv_relu(128, 64), up_conv_relu(64, 32))
+        self.out = nn.Sequential(
+            conv_relu(64, 32),
+            layer_init(
+                nn.Conv2d(32, self.action_vec.sum(), kernel_size=1, padding=0),
+                cnn_layers_init_orthogonal,
+                std=0.01,
+            ),
+            Transpose((0, 2, 3, 1)),
+        )
+
+        with torch.no_grad():
+            cnn_out = torch.flatten(
+                self.enc5(
+                    self.enc4(
+                        self.enc3(
+                            self.enc2(
+                                self.enc1(
+                                    self._preprocess(
+                                        torch.as_tensor(observation_space.sample())
+                                    )
+                                )
+                            )
+                        )
+                    )
+                ),
+                start_dim=1,
+            )
+
+        v_hidden_sizes = (
+            v_hidden_sizes
+            if v_hidden_sizes is not None
+            else default_hidden_sizes(observation_space)
+        )
+        self.critic_head = CriticHead(
+            in_dim=cnn_out.shape[1:],
+            hidden_sizes=v_hidden_sizes,
+            activation=activation,
+            init_layers_orthogonal=init_layers_orthogonal,
+        )
+
+    def _preprocess(self, obs: torch.Tensor) -> torch.Tensor:
+        if len(obs.shape) == 3:
+            obs = obs.unsqueeze(0)
+        return obs.float() / self.range_size
+
+    def forward(
+        self,
+        obs: torch.Tensor,
+        action: torch.Tensor,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> ACNForward:
+        return self._distribution_and_value(
+            obs, action=action, action_masks=action_masks
+        )
+
+    def distribution_and_value(
+        self, obs: torch.Tensor, action_masks: Optional[torch.Tensor] = None
+    ) -> ACNForward:
+        return self._distribution_and_value(obs, action_masks=action_masks)
+
+    def _distribution_and_value(
+        self,
+        obs: torch.Tensor,
+        action: Optional[torch.Tensor] = None,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> ACNForward:
+        assert (
+            action_masks is not None
+        ), f"No mask case unhandled in {self.__class__.__name__}"
+
+        obs = self._preprocess(obs)
+        e1 = self.enc1(obs)
+        e2 = self.enc2(e1)
+        e3 = self.enc3(e2)
+        e4 = self.enc4(e3)
+        e5 = self.enc5(e4)
+
+        v = self.critic_head(e5)
+
+        d4 = self.dec4(e5)
+        d3 = self.dec3(torch.cat((d4, e4), dim=1))
+        d2 = self.dec2(torch.cat((d3, e3), dim=1))
+        d1 = self.dec1(torch.cat((d2, e2), dim=1))
+        logits = self.out(torch.cat((d1, e1), dim=1))
+
+        pi = GridnetDistribution(self.map_size, self.action_vec, logits, action_masks)
+
+        return ACNForward(pi_forward(pi, action), v)
+
+    def value(self, obs: torch.Tensor) -> torch.Tensor:
+        obs = self._preprocess(obs)
+        e1 = self.enc1(obs)
+        e2 = self.enc2(e1)
+        e3 = self.enc3(e2)
+        e4 = self.enc4(e3)
+        e5 = self.enc5(e4)
+
+        return self.critic_head(e5)
+
+    def reset_noise(self, batch_size: Optional[int] = None) -> None:
+        pass
+
+    @property
+    def action_shape(self) -> Tuple[int, ...]:
+        return (self.map_size, len(self.action_vec))
+
+
+def max_pool() -> nn.MaxPool2d:
+    return nn.MaxPool2d(3, stride=2, padding=1)