PPO playing MicrortsAttackShapedReward-v1 from https://github.com/sgoodfriend/rl-algo-impls/tree/fb34ab86707f5e2db85e821ff7dbdc624072d640

README.md

@@ -10,7 +10,7 @@ model-index:
   results:
   - metrics:
     - type: mean_reward
-      value: 8.11 +/- 0.27
+      value: 19.22 +/- 0.0
       name: mean_reward
     task:
       type: reinforcement-learning
@@ -23,17 +23,17 @@ model-index:
 
 This is a trained model of a **PPO** agent playing **MicrortsAttackShapedReward-v1** 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/hz9h6o30.
+All models trained at this commit can be found at https://api.wandb.ai/links/sgoodfriend/z3kioih3.
 
 ## Training Results
 
-This model was trained from 3 trainings of **PPO** agents using different initial seeds. These agents were trained by checking out [587a52b](https://github.com/sgoodfriend/rl-algo-impls/tree/587a52bc38901314c7c1b5c6892acf9315796cf3). 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 **PPO** agents using different initial seeds. These agents were trained by checking out [fb34ab8](https://github.com/sgoodfriend/rl-algo-impls/tree/fb34ab86707f5e2db85e821ff7dbdc624072d640). The best and last models were kept from each training. This submission has loaded the best models from each training, reevaluates them, and selects the best model from these latest evaluations (mean - std).
 
 | algo   | env                           |   seed |   reward_mean |   reward_std |   eval_episodes | best   | wandb_url                                                                    |
 |:-------|:------------------------------|-------:|--------------:|-------------:|----------------:|:-------|:-----------------------------------------------------------------------------|
-| ppo    | MicrortsAttackShapedReward-v1 |      1 |       7.826   |    0.0610015 |              16 |        | [wandb](https://wandb.ai/sgoodfriend/rl-algo-impls-benchmarks/runs/67g3x4i9) |
-| ppo    | MicrortsAttackShapedReward-v1 |      2 |       8.10527 |    0.266247  |              16 | *      | [wandb](https://wandb.ai/sgoodfriend/rl-algo-impls-benchmarks/runs/usawt7bt) |
-| ppo    | MicrortsAttackShapedReward-v1 |      3 |       7.7645  |    0.318334  |              16 |        | [wandb](https://wandb.ai/sgoodfriend/rl-algo-impls-benchmarks/runs/g3rmjrcf) |
+| ppo    | MicrortsAttackShapedReward-v1 |      1 |       19.2195 |            0 |              16 |        | [wandb](https://wandb.ai/sgoodfriend/rl-algo-impls-benchmarks/runs/7x622tan) |
+| ppo    | MicrortsAttackShapedReward-v1 |      2 |       19.2195 |            0 |              16 | *      | [wandb](https://wandb.ai/sgoodfriend/rl-algo-impls-benchmarks/runs/i8e9nqxz) |
+| ppo    | MicrortsAttackShapedReward-v1 |      3 |       19.2195 |            0 |              16 |        | [wandb](https://wandb.ai/sgoodfriend/rl-algo-impls-benchmarks/runs/285khfoz) |
 
 
 ### 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:
-[587a52b](https://github.com/sgoodfriend/rl-algo-impls/tree/587a52bc38901314c7c1b5c6892acf9315796cf3).
+[fb34ab8](https://github.com/sgoodfriend/rl-algo-impls/tree/fb34ab86707f5e2db85e821ff7dbdc624072d640).
 ```
 # Downloads the model, sets hyperparameters, and runs agent for 3 episodes
-python enjoy.py --wandb-run-path=sgoodfriend/rl-algo-impls-benchmarks/usawt7bt
+python enjoy.py --wandb-run-path=sgoodfriend/rl-algo-impls-benchmarks/i8e9nqxz
 ```
 
 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: [587a52b](https://github.com/sgoodfriend/rl-algo-impls/tree/587a52bc38901314c7c1b5c6892acf9315796cf3). While
+commit the agent was trained on: [fb34ab8](https://github.com/sgoodfriend/rl-algo-impls/tree/fb34ab86707f5e2db85e821ff7dbdc624072d640). 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/hz9h6o30 were generated by running a script on a Lambda
+This and other models from https://api.wandb.ai/links/sgoodfriend/z3kioih3 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
@@ -120,6 +120,7 @@ algo_hyperparams:
 device: auto
 env: MicrortsAttackShapedReward-v1
 env_hyperparams:
+  mask_actions: true
   n_envs: 8
   vec_env_class: sync
 env_id: null
@@ -136,7 +137,7 @@ wandb_entity: null
 wandb_group: null
 wandb_project_name: rl-algo-impls-benchmarks
 wandb_tags:
-- benchmark_587a52b
-- host_192-9-151-119
+- benchmark_fb34ab8
+- host_155-248-210-13
 
 ```

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rl_algo_impls"
-version = "0.0.6"
+version = "0.0.7"
 description = "Implementations of reinforcement learning algorithms"
 authors = [
     {name = "Scott Goodfriend", email = "[email protected]"},

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

rl_algo_impls/a2c/a2c.py

@@ -10,6 +10,7 @@ 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.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.stats import log_scalars
@@ -84,12 +85,12 @@ class A2C(Algorithm):
         obs = np.zeros(epoch_dim + obs_space.shape, dtype=obs_space.dtype)
         actions = np.zeros(epoch_dim + act_space.shape, dtype=act_space.dtype)
         rewards = np.zeros(epoch_dim, dtype=np.float32)
-        episode_starts = np.zeros(epoch_dim, dtype=np.byte)
+        episode_starts = np.zeros(epoch_dim, dtype=np.bool8)
         values = np.zeros(epoch_dim, dtype=np.float32)
         logprobs = np.zeros(epoch_dim, dtype=np.float32)
 
         next_obs = self.env.reset()
-        next_episode_starts = np.ones(step_dim, dtype=np.byte)
+        next_episode_starts = np.full(step_dim, True, dtype=np.bool8)
 
         timesteps_elapsed = start_timesteps
         while timesteps_elapsed < start_timesteps + train_timesteps:
@@ -126,23 +127,16 @@ class A2C(Algorithm):
                     clamped_action
                 )
 
-            advantages = np.zeros(epoch_dim, dtype=np.float32)
-            last_gae_lam = 0
-            for t in reversed(range(self.n_steps)):
-                if t == self.n_steps - 1:
-                    next_nonterminal = 1.0 - next_episode_starts
-                    next_value = self.policy.value(next_obs)
-                else:
-                    next_nonterminal = 1.0 - episode_starts[t + 1]
-                    next_value = values[t + 1]
-                delta = (
-                    rewards[t] + self.gamma * next_value * next_nonterminal - values[t]
-                )
-                last_gae_lam = (
-                    delta
-                    + self.gamma * self.gae_lambda * next_nonterminal * last_gae_lam
-                )
-                advantages[t] = last_gae_lam
+            advantages = compute_advantages(
+                rewards,
+                values,
+                episode_starts,
+                next_episode_starts,
+                next_obs,
+                self.policy,
+                self.gamma,
+                self.gae_lambda,
+            )
             returns = advantages + values
 
             b_obs = torch.tensor(obs.reshape((-1,) + obs_space.shape)).to(self.device)

rl_algo_impls/compare_runs.py

@@ -194,5 +194,6 @@ def compare_runs() -> None:
     df.loc["mean"] = df.mean(numeric_only=True)
     print(df.to_markdown())
 
+
 if __name__ == "__main__":
-    compare_runs()
+    compare_runs()

rl_algo_impls/huggingface_publish.py

@@ -162,6 +162,7 @@ def publish(
             path_in_repo="",
             commit_message=f"{algo.upper()} playing {env_id} from {github_url}/tree/{commit_hash}",
             token=huggingface_token,
+            delete_patterns="*",
         )
         print(f"Pushed model to the hub: {repo_url}")
 

rl_algo_impls/hyperparams/ppo.yml

@@ -218,6 +218,7 @@ _microrts: &microrts-defaults
   env_hyperparams: &microrts-env-defaults
     n_envs: 8
     vec_env_class: sync
+    mask_actions: true
   policy_hyperparams:
     <<: *atari-policy-defaults
     cnn_style: microrts
@@ -227,10 +228,23 @@ _microrts: &microrts-defaults
     clip_range_decay: none
     clip_range_vf: 0.1
 
-debug-MicrortsMining-v1:
+_no-mask-microrts: &no-mask-microrts-defaults
   <<: *microrts-defaults
+  env_hyperparams:
+    <<: *microrts-env-defaults
+    mask_actions: false
+
+MicrortsMining-v1-NoMask:
+  <<: *no-mask-microrts-defaults
   env_id: MicrortsMining-v1
-  device: cpu
+
+MicrortsAttackShapedReward-v1-NoMask:
+  <<: *no-mask-microrts-defaults
+  env_id: MicrortsAttackShapedReward-v1
+
+MicrortsRandomEnemyShapedReward3-v1-NoMask:
+  <<: *no-mask-microrts-defaults
+  env_id: MicrortsRandomEnemyShapedReward3-v1
 
 HalfCheetahBulletEnv-v0: &pybullet-defaults
   n_timesteps: !!float 2e6

rl_algo_impls/optimize.py

@@ -194,7 +194,7 @@ def simple_optimize(trial: optuna.Trial, args: RunArgs, study_args: StudyArgs) -
     env = make_env(
         config, EnvHyperparams(**config.env_hyperparams), tb_writer=tb_writer
     )
-    device = get_device(config.device, env)
+    device = get_device(config, env)
     policy = make_policy(args.algo, env, device, **config.policy_hyperparams)
     algo = ALGOS[args.algo](policy, env, device, tb_writer, **config.algo_hyperparams)
 
@@ -298,7 +298,7 @@ def stepwise_optimize(
                 normalize_load_path=config.model_dir_path() if i > 0 else None,
                 tb_writer=tb_writer,
             )
-            device = get_device(config.device, env)
+            device = get_device(config, env)
             policy = make_policy(arg.algo, env, device, **config.policy_hyperparams)
             if i > 0:
                 policy.load(config.model_dir_path())

rl_algo_impls/ppo/ppo.py

@@ -1,8 +1,9 @@
+import logging
 import numpy as np
 import torch
 import torch.nn as nn
 
-from dataclasses import asdict, dataclass, field
+from dataclasses import asdict, dataclass
 from time import perf_counter
 from torch.optim import Adam
 from torch.utils.tensorboard.writer import SummaryWriter
@@ -10,49 +11,22 @@ from typing import List, Optional, NamedTuple, TypeVar
 
 from rl_algo_impls.shared.algorithm import Algorithm
 from rl_algo_impls.shared.callbacks.callback import Callback
-from rl_algo_impls.shared.gae import compute_advantage, compute_rtg_and_advantage
+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.trajectory import Trajectory, TrajectoryAccumulator
-from rl_algo_impls.wrappers.vectorable_wrapper import VecEnv, VecEnvObs
-
-
-@dataclass
-class PPOTrajectory(Trajectory):
-    logp_a: List[float] = field(default_factory=list)
-
-    def add(
-        self,
-        obs: np.ndarray,
-        act: np.ndarray,
-        next_obs: np.ndarray,
-        rew: float,
-        terminated: bool,
-        v: float,
-        logp_a: float,
-    ):
-        super().add(obs, act, next_obs, rew, terminated, v)
-        self.logp_a.append(logp_a)
-
-
-class PPOTrajectoryAccumulator(TrajectoryAccumulator):
-    def __init__(self, num_envs: int) -> None:
-        super().__init__(num_envs, PPOTrajectory)
-
-    def step(
-        self,
-        obs: VecEnvObs,
-        action: np.ndarray,
-        next_obs: VecEnvObs,
-        reward: np.ndarray,
-        done: np.ndarray,
-        val: np.ndarray,
-        logp_a: np.ndarray,
-    ) -> None:
-        super().step(obs, action, next_obs, reward, done, val, logp_a)
+from rl_algo_impls.shared.stats import log_scalars
+from rl_algo_impls.wrappers.action_mask_wrapper import ActionMaskWrapper
+from rl_algo_impls.wrappers.vectorable_wrapper import (
+    VecEnv,
+    find_wrapper,
+    single_observation_space,
+    single_action_space,
+)
 
 
 class TrainStepStats(NamedTuple):
@@ -131,11 +105,11 @@ class PPO(Algorithm):
         vf_coef: float = 0.5,
         ppo2_vf_coef_halving: bool = True,
         max_grad_norm: float = 0.5,
-        update_rtg_between_epochs: bool = False,
         sde_sample_freq: int = -1,
     ) -> None:
         super().__init__(policy, env, device, tb_writer)
         self.policy = policy
+        self.action_masker = find_wrapper(env, ActionMaskWrapper)
 
         self.gamma = gamma
         self.gae_lambda = gae_lambda
@@ -146,7 +120,13 @@ class PPO(Algorithm):
         self.clip_range_vf_schedule = None
         if clip_range_vf:
             self.clip_range_vf_schedule = schedule(clip_range_vf_decay, clip_range_vf)
+
+        if normalize_advantage:
+            assert (
+                env.num_envs * n_steps > 1 and batch_size > 1
+            ), f"Each minibatch must be larger than 1 to support normalization"
         self.normalize_advantage = normalize_advantage
+
         self.ent_coef_schedule = schedule(ent_coef_decay, ent_coef)
         self.vf_coef = vf_coef
         self.ppo2_vf_coef_halving = ppo2_vf_coef_halving
@@ -156,181 +136,235 @@ class PPO(Algorithm):
         self.n_epochs = n_epochs
         self.sde_sample_freq = sde_sample_freq
 
-        self.update_rtg_between_epochs = update_rtg_between_epochs
-
     def learn(
         self: PPOSelf,
-        total_timesteps: int,
+        train_timesteps: int,
         callback: Optional[Callback] = None,
+        total_timesteps: Optional[int] = None,
+        start_timesteps: int = 0,
     ) -> PPOSelf:
-        obs = self.env.reset()
-        ts_elapsed = 0
-        while ts_elapsed < total_timesteps:
-            start_time = perf_counter()
-            accumulator = self._collect_trajectories(obs)
-            rollout_steps = self.n_steps * self.env.num_envs
-            ts_elapsed += rollout_steps
-            progress = ts_elapsed / total_timesteps
-            train_stats = self.train(accumulator.all_trajectories, progress, ts_elapsed)
-            train_stats.write_to_tensorboard(self.tb_writer, ts_elapsed)
-            end_time = perf_counter()
-            self.tb_writer.add_scalar(
-                "train/steps_per_second",
-                rollout_steps / (end_time - start_time),
-                ts_elapsed,
+        if total_timesteps is None:
+            total_timesteps = train_timesteps
+        assert start_timesteps + train_timesteps <= total_timesteps
+
+        epoch_dim = (self.n_steps, self.env.num_envs)
+        step_dim = (self.env.num_envs,)
+        obs_space = single_observation_space(self.env)
+        act_space = single_action_space(self.env)
+
+        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)
+
+        obs = np.zeros(epoch_dim + obs_space.shape, dtype=obs_space.dtype)  # type: ignore
+        actions = np.zeros(epoch_dim + act_space.shape, dtype=act_space.dtype)  # type: ignore
+        rewards = np.zeros(epoch_dim, dtype=np.float32)
+        episode_starts = np.zeros(epoch_dim, dtype=np.bool8)
+        values = np.zeros(epoch_dim, dtype=np.float32)
+        logprobs = np.zeros(epoch_dim, dtype=np.float32)
+        action_masks = (
+            np.zeros(
+                (self.n_steps,) + next_action_masks.shape, dtype=next_action_masks.dtype
             )
-            if callback:
-                callback.on_step(timesteps_elapsed=rollout_steps)
-
-        return self
-
-    def _collect_trajectories(self, obs: VecEnvObs) -> PPOTrajectoryAccumulator:
-        self.policy.eval()
-        accumulator = PPOTrajectoryAccumulator(self.env.num_envs)
-        self.policy.reset_noise()
-        for i in range(self.n_steps):
-            if self.sde_sample_freq > 0 and i > 0 and i % self.sde_sample_freq == 0:
-                self.policy.reset_noise()
-            action, value, logp_a, clamped_action = self.policy.step(obs)
-            next_obs, reward, done, _ = self.env.step(clamped_action)
-            accumulator.step(obs, action, next_obs, reward, done, value, logp_a)
-            obs = next_obs
-        return accumulator
-
-    def train(
-        self, trajectories: List[PPOTrajectory], progress: float, timesteps_elapsed: int
-    ) -> TrainStats:
-        self.policy.train()
-        learning_rate = self.lr_schedule(progress)
-        update_learning_rate(self.optimizer, learning_rate)
-        self.tb_writer.add_scalar(
-            "charts/learning_rate",
-            self.optimizer.param_groups[0]["lr"],
-            timesteps_elapsed,
+            if next_action_masks is not None
+            else None
         )
 
-        pi_clip = self.clip_range_schedule(progress)
-        self.tb_writer.add_scalar("charts/pi_clip", pi_clip, timesteps_elapsed)
-        if self.clip_range_vf_schedule:
-            v_clip = self.clip_range_vf_schedule(progress)
-            self.tb_writer.add_scalar("charts/v_clip", v_clip, timesteps_elapsed)
-        else:
-            v_clip = None
-        ent_coef = self.ent_coef_schedule(progress)
-        self.tb_writer.add_scalar("charts/ent_coef", ent_coef, timesteps_elapsed)
-
-        obs = torch.as_tensor(
-            np.concatenate([np.array(t.obs) for t in trajectories]), device=self.device
-        )
-        act = torch.as_tensor(
-            np.concatenate([np.array(t.act) for t in trajectories]), device=self.device
-        )
-        rtg, adv = compute_rtg_and_advantage(
-            trajectories, self.policy, self.gamma, self.gae_lambda, self.device
-        )
-        orig_v = torch.as_tensor(
-            np.concatenate([np.array(t.v) for t in trajectories]), device=self.device
-        )
-        orig_logp_a = torch.as_tensor(
-            np.concatenate([np.array(t.logp_a) for t in trajectories]),
-            device=self.device,
-        )
+        timesteps_elapsed = start_timesteps
+        while timesteps_elapsed < start_timesteps + train_timesteps:
+            start_time = perf_counter()
 
-        step_stats = []
-        for _ in range(self.n_epochs):
-            step_stats.clear()
-            if self.update_rtg_between_epochs:
-                rtg, adv = compute_rtg_and_advantage(
-                    trajectories, self.policy, self.gamma, self.gae_lambda, self.device
-                )
+            progress = timesteps_elapsed / total_timesteps
+            ent_coef = self.ent_coef_schedule(progress)
+            learning_rate = self.lr_schedule(progress)
+            update_learning_rate(self.optimizer, learning_rate)
+            pi_clip = self.clip_range_schedule(progress)
+            chart_scalars = {
+                "learning_rate": self.optimizer.param_groups[0]["lr"],
+                "ent_coef": ent_coef,
+                "pi_clip": pi_clip,
+            }
+            if self.clip_range_vf_schedule:
+                v_clip = self.clip_range_vf_schedule(progress)
+                chart_scalars["v_clip"] = v_clip
             else:
-                adv = compute_advantage(
-                    trajectories, self.policy, self.gamma, self.gae_lambda, self.device
+                v_clip = None
+            log_scalars(self.tb_writer, "charts", chart_scalars, timesteps_elapsed)
+
+            self.policy.eval()
+            self.policy.reset_noise()
+            for s in range(self.n_steps):
+                timesteps_elapsed += self.env.num_envs
+                if self.sde_sample_freq > 0 and s > 0 and s % self.sde_sample_freq == 0:
+                    self.policy.reset_noise()
+
+                obs[s] = next_obs
+                episode_starts[s] = next_episode_starts
+                if action_masks is not None:
+                    action_masks[s] = next_action_masks
+
+                (
+                    actions[s],
+                    values[s],
+                    logprobs[s],
+                    clamped_action,
+                ) = self.policy.step(next_obs, action_masks=next_action_masks)
+                next_obs, rewards[s], next_episode_starts, _ = self.env.step(
+                    clamped_action
                 )
-            idxs = torch.randperm(len(obs))
-            for i in range(0, len(obs), self.batch_size):
-                mb_idxs = idxs[i : i + self.batch_size]
-                mb_adv = adv[mb_idxs]
-                if self.normalize_advantage:
-                    mb_adv = (mb_adv - mb_adv.mean(-1)) / (mb_adv.std(-1) + 1e-8)
-                self.policy.reset_noise(self.batch_size)
-                step_stats.append(
-                    self._train_step(
-                        pi_clip,
-                        v_clip,
-                        ent_coef,
-                        obs[mb_idxs],
-                        act[mb_idxs],
-                        rtg[mb_idxs],
-                        mb_adv,
-                        orig_v[mb_idxs],
-                        orig_logp_a[mb_idxs],
-                    )
+                next_action_masks = (
+                    self.action_masker.action_masks() if self.action_masker else None
                 )
 
-        y_pred, y_true = orig_v.cpu().numpy(), rtg.cpu().numpy()
-        var_y = np.var(y_true).item()
-        explained_var = (
-            np.nan if var_y == 0 else 1 - np.var(y_true - y_pred).item() / var_y
-        )
+            self.policy.train()
+
+            advantages = compute_advantages(
+                rewards,
+                values,
+                episode_starts,
+                next_episode_starts,
+                next_obs,
+                self.policy,
+                self.gamma,
+                self.gae_lambda,
+            )
+            returns = advantages + values
+
+            b_obs = torch.tensor(obs.reshape((-1,) + obs_space.shape)).to(self.device)  # type: ignore
+            b_actions = torch.tensor(actions.reshape((-1,) + act_space.shape)).to(  # type: ignore
+                self.device
+            )
+            b_logprobs = torch.tensor(logprobs.reshape(-1)).to(self.device)
+            b_action_masks = (
+                torch.tensor(action_masks.reshape((-1,) + next_action_masks.shape[1:])).to(  # type: ignore
+                    self.device
+                )
+                if action_masks is not None
+                else None
+            )
+
+            b_advantages = torch.tensor(advantages.reshape(-1)).to(self.device)
+
+            y_pred = values.reshape(-1)
+            b_values = torch.tensor(y_pred).to(self.device)
+            y_true = returns.reshape(-1)
+            b_returns = torch.tensor(y_true).to(self.device)
+
+            step_stats = []
+            for _ in range(self.n_epochs):
+                b_idxs = torch.randperm(len(b_obs))
+                # Only record last epoch's stats
+                step_stats.clear()
+                for i in range(0, len(b_obs), self.batch_size):
+                    self.policy.reset_noise(self.batch_size)
+
+                    mb_idxs = b_idxs[i : i + self.batch_size]
+
+                    mb_obs = b_obs[mb_idxs]
+                    mb_actions = b_actions[mb_idxs]
+                    mb_values = b_values[mb_idxs]
+                    mb_logprobs = b_logprobs[mb_idxs]
+                    mb_action_masks = (
+                        b_action_masks[mb_idxs] if b_action_masks is not None else None
+                    )
 
-        return TrainStats(step_stats, explained_var)
+                    mb_adv = b_advantages[mb_idxs]
+                    if self.normalize_advantage:
+                        mb_adv = (mb_adv - mb_adv.mean()) / (mb_adv.std() + 1e-8)
+                    mb_returns = b_returns[mb_idxs]
 
-    def _train_step(
-        self,
-        pi_clip: float,
-        v_clip: Optional[float],
-        ent_coef: float,
-        obs: torch.Tensor,
-        act: torch.Tensor,
-        rtg: torch.Tensor,
-        adv: torch.Tensor,
-        orig_v: torch.Tensor,
-        orig_logp_a: torch.Tensor,
-    ) -> TrainStepStats:
-        logp_a, entropy, v = self.policy(obs, act)
-        logratio = logp_a - orig_logp_a
-        ratio = torch.exp(logratio)
-        clip_ratio = torch.clamp(ratio, min=1 - pi_clip, max=1 + pi_clip)
-        pi_loss = torch.maximum(-ratio * adv, -clip_ratio * adv).mean()
-
-        v_loss_unclipped = (v - rtg) ** 2
-        if v_clip:
-            v_loss_clipped = (
-                orig_v + torch.clamp(v - orig_v, -v_clip, v_clip) - rtg
-            ) ** 2
-            v_loss = torch.max(v_loss_unclipped, v_loss_clipped).mean()
-        else:
-            v_loss = v_loss_unclipped.mean()
-        if self.ppo2_vf_coef_halving:
-            v_loss *= 0.5
-
-        entropy_loss = -entropy.mean()
-
-        loss = pi_loss + ent_coef * entropy_loss + self.vf_coef * v_loss
-
-        self.optimizer.zero_grad()
-        loss.backward()
-        nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
-        self.optimizer.step()
-
-        with torch.no_grad():
-            approx_kl = ((ratio - 1) - logratio).mean().cpu().numpy().item()
-            clipped_frac = (
-                ((ratio - 1).abs() > pi_clip).float().mean().cpu().numpy().item()
+                    new_logprobs, entropy, new_values = self.policy(
+                        mb_obs, mb_actions, action_masks=mb_action_masks
+                    )
+
+                    logratio = new_logprobs - mb_logprobs
+                    ratio = torch.exp(logratio)
+                    clipped_ratio = torch.clamp(ratio, min=1 - pi_clip, max=1 + pi_clip)
+                    pi_loss = torch.max(
+                        -ratio * mb_adv, -clipped_ratio * mb_adv
+                    ).mean()
+
+                    v_loss_unclipped = (new_values - mb_returns) ** 2
+                    if v_clip:
+                        v_loss_clipped = (
+                            mb_values
+                            + torch.clamp(new_values - mb_values, -v_clip, v_clip)
+                            - mb_returns
+                        ) ** 2
+                        v_loss = torch.max(v_loss_unclipped, v_loss_clipped).mean()
+                    else:
+                        v_loss = v_loss_unclipped.mean()
+
+                    if self.ppo2_vf_coef_halving:
+                        v_loss *= 0.5
+
+                    entropy_loss = -entropy.mean()
+
+                    loss = pi_loss + ent_coef * entropy_loss + self.vf_coef * v_loss
+
+                    self.optimizer.zero_grad()
+                    loss.backward()
+                    nn.utils.clip_grad_norm_(
+                        self.policy.parameters(), self.max_grad_norm
+                    )
+                    self.optimizer.step()
+
+                    with torch.no_grad():
+                        approx_kl = ((ratio - 1) - logratio).mean().cpu().numpy().item()
+                        clipped_frac = (
+                            ((ratio - 1).abs() > pi_clip)
+                            .float()
+                            .mean()
+                            .cpu()
+                            .numpy()
+                            .item()
+                        )
+                        val_clipped_frac = (
+                            ((new_values - mb_values).abs() > v_clip)
+                            .float()
+                            .mean()
+                            .cpu()
+                            .numpy()
+                            .item()
+                            if v_clip
+                            else 0
+                        )
+
+                    step_stats.append(
+                        TrainStepStats(
+                            loss.item(),
+                            pi_loss.item(),
+                            v_loss.item(),
+                            entropy_loss.item(),
+                            approx_kl,
+                            clipped_frac,
+                            val_clipped_frac,
+                        )
+                    )
+
+            var_y = np.var(y_true).item()
+            explained_var = (
+                np.nan if var_y == 0 else 1 - np.var(y_true - y_pred).item() / var_y
             )
-            val_clipped_frac = (
-                (((v - orig_v).abs() > v_clip).float().mean().cpu().numpy().item())
-                if v_clip
-                else 0
+            TrainStats(step_stats, explained_var).write_to_tensorboard(
+                self.tb_writer, timesteps_elapsed
             )
 
-        return TrainStepStats(
-            loss.item(),
-            pi_loss.item(),
-            v_loss.item(),
-            entropy_loss.item(),
-            approx_kl,
-            clipped_frac,
-            val_clipped_frac,
-        )
+            end_time = perf_counter()
+            rollout_steps = self.n_steps * self.env.num_envs
+            self.tb_writer.add_scalar(
+                "train/steps_per_second",
+                rollout_steps / (end_time - start_time),
+                timesteps_elapsed,
+            )
+
+            if callback:
+                if not callback.on_step(timesteps_elapsed=rollout_steps):
+                    logging.info(
+                        f"Callback terminated training at {timesteps_elapsed} timesteps"
+                    )
+                    break
+
+        return self

rl_algo_impls/runner/config.py

@@ -36,7 +36,7 @@ class RunArgs:
 
 @dataclass
 class EnvHyperparams:
-    env_type: str = "sb3vec"
+    env_type: str = "gymvec"
     n_envs: int = 1
     frame_stack: int = 1
     make_kwargs: Optional[Dict[str, Any]] = None
@@ -50,7 +50,8 @@ class EnvHyperparams:
     video_step_interval: Union[int, float] = 1_000_000
     initial_steps_to_truncate: Optional[int] = None
     clip_atari_rewards: bool = True
-    normalize_type: Optional[str] = "gymlike"
+    normalize_type: Optional[str] = None
+    mask_actions: bool = False
 
 
 HyperparamsSelf = TypeVar("HyperparamsSelf", bound="Hyperparams")

rl_algo_impls/runner/env.py

@@ -20,6 +20,7 @@ from typing import Callable, Optional
 
 from rl_algo_impls.runner.config import Config, EnvHyperparams
 from rl_algo_impls.shared.policy.policy import VEC_NORMALIZE_FILENAME
+from rl_algo_impls.wrappers.action_mask_wrapper import ActionMaskWrapper
 from rl_algo_impls.wrappers.atari_wrappers import (
     EpisodicLifeEnv,
     FireOnLifeStarttEnv,
@@ -113,21 +114,20 @@ def _make_vec_env(
         initial_steps_to_truncate,
         clip_atari_rewards,
         normalize_type,
+        mask_actions,
     ) = astuple(hparams)
 
     import_for_env_id(config.env_id)
 
-    spec = gym.spec(config.env_id)
     seed = config.seed(training=training)
 
     make_kwargs = make_kwargs.copy() if make_kwargs is not None else {}
-    if "BulletEnv" in config.env_id and render:
+    if is_bullet_env(config) and render:
         make_kwargs["render"] = True
-    if "CarRacing" in config.env_id:
+    if is_car_racing(config):
         make_kwargs["verbose"] = 0
-    if "procgen" in config.env_id:
-        if not render:
-            make_kwargs["render_mode"] = "rgb_array"
+    if is_gym_procgen(config) and not render:
+        make_kwargs["render_mode"] = "rgb_array"
 
     def make(idx: int) -> Callable[[], gym.Env]:
         def _make() -> gym.Env:
@@ -145,7 +145,7 @@ def _make_vec_env(
                 env = InitialStepTruncateWrapper(
                     env, idx * initial_steps_to_truncate // n_envs
                 )
-            if "AtariEnv" in spec.entry_point:  # type: ignore
+            if is_atari(config):  # type: ignore
                 env = NoopResetEnv(env, noop_max=30)
                 env = MaxAndSkipEnv(env, skip=4)
                 env = EpisodicLifeEnv(env, training=training)
@@ -157,17 +157,17 @@ def _make_vec_env(
                 env = ResizeObservation(env, (84, 84))
                 env = GrayScaleObservation(env, keep_dim=False)
                 env = FrameStack(env, frame_stack)
-            elif "CarRacing" in config.env_id:
+            elif is_car_racing(config):
                 env = ResizeObservation(env, (64, 64))
                 env = GrayScaleObservation(env, keep_dim=False)
                 env = FrameStack(env, frame_stack)
-            elif "procgen" in config.env_id:
+            elif is_gym_procgen(config):
                 # env = GrayScaleObservation(env, keep_dim=False)
                 env = NoopEnvSeed(env)
                 env = HwcToChwObservation(env)
                 if frame_stack > 1:
                     env = FrameStack(env, frame_stack)
-            elif "Microrts" in config.env_id:
+            elif is_microrts(config):
                 env = HwcToChwObservation(env)
 
             if no_reward_timeout_steps:
@@ -195,6 +195,8 @@ def _make_vec_env(
         envs = SyncVectorEnvRenderCompat(envs)
     if env_type == "sb3vec":
         envs = IsVectorEnv(envs)
+    if mask_actions:
+        envs = ActionMaskWrapper(envs)
     if training:
         assert tb_writer
         envs = EpisodeStatsWriter(
@@ -262,6 +264,8 @@ def _make_procgen_env(
         _,  # video_step_interval
         _,  # initial_steps_to_truncate
         _,  # clip_atari_rewards
+        _,  # normalize_type
+        _,  # mask_actions
     ) = astuple(hparams)
 
     seed = config.seed(training=training)
@@ -307,3 +311,24 @@ def import_for_env_id(env_id: str) -> None:
         import pybullet_envs
     if "Microrts" in env_id:
         import gym_microrts
+
+
+def is_atari(config: Config) -> bool:
+    spec = gym.spec(config.env_id)
+    return "AtariEnv" in str(spec.entry_point)
+
+
+def is_bullet_env(config: Config) -> bool:
+    return "BulletEnv" in config.env_id
+
+
+def is_car_racing(config: Config) -> bool:
+    return "CarRacing" in config.env_id
+
+
+def is_gym_procgen(config: Config) -> bool:
+    return "procgen" in config.env_id
+
+
+def is_microrts(config: Config) -> bool:
+    return "Microrts" in config.env_id

rl_algo_impls/runner/evaluate.py

@@ -75,7 +75,7 @@ def evaluate_model(args: EvalArgs, root_dir: str) -> Evaluation:
         render=args.render,
         normalize_load_path=model_path,
     )
-    device = get_device(config.device, env)
+    device = get_device(config, env)
     policy = make_policy(
         args.algo,
         env,

rl_algo_impls/runner/running_utils.py

@@ -15,8 +15,8 @@ from pathlib import Path
 from torch.utils.tensorboard.writer import SummaryWriter
 from typing import Dict, Optional, Type, Union
 
-from rl_algo_impls.runner.config import Hyperparams
-from rl_algo_impls.runner.env import import_for_env_id
+from rl_algo_impls.runner.config import Config, Hyperparams
+from rl_algo_impls.runner.env import import_for_env_id, is_microrts
 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
@@ -93,7 +93,8 @@ def load_hyperparams(algo: str, env_id: str) -> Hyperparams:
         raise ValueError(f"{env_id} not specified in {algo} hyperparameters file")
 
 
-def get_device(device: str, env: VecEnv) -> torch.device:
+def get_device(config: Config, env: VecEnv) -> torch.device:
+    device = config.device
     # cuda by default
     if device == "auto":
         device = "cuda"
@@ -111,6 +112,8 @@ def get_device(device: str, env: VecEnv) -> torch.device:
             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/train.py

@@ -65,7 +65,7 @@ def train(args: TrainArgs):
     env = make_env(
         config, EnvHyperparams(**config.env_hyperparams), tb_writer=tb_writer
     )
-    device = get_device(config.device, env)
+    device = get_device(config, env)
     policy = make_policy(args.algo, env, device, **config.policy_hyperparams)
     algo = ALGOS[args.algo](policy, env, device, tb_writer, **config.algo_hyperparams)
 

rl_algo_impls/shared/callbacks/eval_callback.py

@@ -9,8 +9,9 @@ from typing import List, Optional, Union
 from rl_algo_impls.shared.callbacks.callback 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 ActionMaskWrapper
 from rl_algo_impls.wrappers.vec_episode_recorder import VecEpisodeRecorder
-from rl_algo_impls.wrappers.vectorable_wrapper import VecEnv
+from rl_algo_impls.wrappers.vectorable_wrapper import VecEnv, find_wrapper
 
 
 class EvaluateAccumulator(EpisodeAccumulator):
@@ -83,8 +84,13 @@ def evaluate(
     )
 
     obs = env.reset()
+    action_masker = find_wrapper(env, ActionMaskWrapper)
     while not episodes.is_done():
-        act = policy.act(obs, deterministic=deterministic)
+        act = policy.act(
+            obs,
+            deterministic=deterministic,
+            action_masks=action_masker.action_masks() if action_masker else None,
+        )
         obs, rew, done, _ = env.step(act)
         episodes.step(rew, done)
         if render:

rl_algo_impls/shared/gae.py

@@ -5,6 +5,7 @@ from typing import NamedTuple, Sequence
 
 from rl_algo_impls.shared.policy.on_policy import OnPolicy
 from rl_algo_impls.shared.trajectory import Trajectory
+from rl_algo_impls.wrappers.vectorable_wrapper import VecEnvObs
 
 
 class RtgAdvantage(NamedTuple):
@@ -19,7 +20,7 @@ def discounted_cumsum(x: np.ndarray, gamma: float) -> np.ndarray:
     return dc
 
 
-def compute_advantage(
+def compute_advantage_from_trajectories(
     trajectories: Sequence[Trajectory],
     policy: OnPolicy,
     gamma: float,
@@ -40,7 +41,7 @@ def compute_advantage(
     )
 
 
-def compute_rtg_and_advantage(
+def compute_rtg_and_advantage_from_trajectories(
     trajectories: Sequence[Trajectory],
     policy: OnPolicy,
     gamma: float,
@@ -65,3 +66,29 @@ def compute_rtg_and_advantage(
         ),
         torch.as_tensor(np.concatenate(advantages), dtype=torch.float32, device=device),
     )
+
+
+def compute_advantages(
+    rewards: np.ndarray,
+    values: np.ndarray,
+    episode_starts: np.ndarray,
+    next_episode_starts: np.ndarray,
+    next_obs: VecEnvObs,
+    policy: OnPolicy,
+    gamma: float,
+    gae_lambda: float,
+) -> np.ndarray:
+    advantages = np.zeros_like(rewards)
+    last_gae_lam = 0
+    n_steps = advantages.shape[0]
+    for t in reversed(range(n_steps)):
+        if t == n_steps - 1:
+            next_nonterminal = 1.0 - next_episode_starts
+            next_value = policy.value(next_obs)
+        else:
+            next_nonterminal = 1.0 - episode_starts[t + 1]
+            next_value = values[t + 1]
+        delta = rewards[t] + gamma * next_value * next_nonterminal - values[t]
+        last_gae_lam = delta + gamma * gae_lambda * next_nonterminal * last_gae_lam
+        advantages[t] = last_gae_lam
+    return advantages

rl_algo_impls/shared/policy/actor.py

@@ -6,8 +6,8 @@ import torch.nn as nn
 from abc import ABC, abstractmethod
 from gym.spaces import Box, Discrete, MultiDiscrete
 from numpy.typing import NDArray
-from torch.distributions import Categorical, Distribution, Normal
-from typing import NamedTuple, Optional, Sequence, Type, TypeVar, Union
+from torch.distributions import Categorical, Distribution, Normal, constraints
+from typing import Dict, NamedTuple, Optional, Sequence, Type, TypeVar, Union
 
 from rl_algo_impls.shared.module.module import mlp
 
@@ -20,7 +20,12 @@ class PiForward(NamedTuple):
 
 class Actor(nn.Module, ABC):
     @abstractmethod
-    def forward(self, obs: torch.Tensor, a: Optional[torch.Tensor] = None) -> PiForward:
+    def forward(
+        self,
+        obs: torch.Tensor,
+        actions: Optional[torch.Tensor] = None,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> PiForward:
         ...
 
 
@@ -41,34 +46,53 @@ class CategoricalActorHead(Actor):
             final_layer_gain=0.01,
         )
 
-    def forward(self, obs: torch.Tensor, a: Optional[torch.Tensor] = None) -> PiForward:
+    def forward(
+        self,
+        obs: torch.Tensor,
+        actions: Optional[torch.Tensor] = None,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> PiForward:
         logits = self._fc(obs)
-        pi = Categorical(logits=logits)
+        pi = MaskedCategorical(logits=logits, mask=action_masks)
         logp_a = None
         entropy = None
-        if a is not None:
-            logp_a = pi.log_prob(a)
+        if actions is not None:
+            logp_a = pi.log_prob(actions)
             entropy = pi.entropy()
         return PiForward(pi, logp_a, entropy)
 
 
-class MultiCategorical(Categorical):
+class MultiCategorical(Distribution):
     def __init__(
-        self, nvec: NDArray[np.int64], probs=None, logits=None, validate_args=None
+        self,
+        nvec: NDArray[np.int64],
+        probs=None,
+        logits=None,
+        validate_args=None,
+        masks: Optional[torch.Tensor] = None,
     ):
         # Either probs or logits should be set
-        assert (probs is not None) != (logits is not None)
+        assert (probs is None) != (logits is None)
+        masks_split = (
+            torch.split(masks, nvec.tolist(), dim=1)
+            if masks is not None
+            else [None] * len(nvec)
+        )
         if probs:
             self.dists = [
-                Categorical(probs=p, validate_args=validate_args)
-                for p in torch.split(probs, nvec.tolist(), dim=1)
+                MaskedCategorical(probs=p, validate_args=validate_args, mask=m)
+                for p, m in zip(torch.split(probs, nvec.tolist(), dim=1), masks_split)
             ]
+            param = probs
         else:
             assert logits is not None
             self.dists = [
-                Categorical(logits=lg, validate_args=validate_args)
-                for lg in torch.split(logits, nvec.tolist(), dim=1)
+                MaskedCategorical(logits=lg, validate_args=validate_args, mask=m)
+                for lg, m in zip(torch.split(logits, nvec.tolist(), dim=1), masks_split)
             ]
+            param = logits
+        batch_shape = param.size()[:-1] if param.ndimension() > 1 else torch.Size()
+        super().__init__(batch_shape=batch_shape, validate_args=validate_args)
 
     def log_prob(self, action: torch.Tensor) -> torch.Tensor:
         prob_stack = torch.stack(
@@ -82,6 +106,34 @@ class MultiCategorical(Categorical):
     def sample(self, sample_shape: torch.Size = torch.Size()):
         return torch.stack([c.sample(sample_shape) for c in self.dists], dim=-1)
 
+    @property
+    def arg_constraints(self) -> Dict[str, constraints.Constraint]:
+        # Constraints handled by child distributions in dist
+        return {}
+
+
+class MaskedCategorical(Categorical):
+    def __init__(
+        self,
+        probs=None,
+        logits=None,
+        validate_args=None,
+        mask: Optional[torch.Tensor] = None,
+    ):
+        if mask is not None:
+            assert logits is not None, "mask requires logits and not probs"
+            logits = torch.where(mask, logits, -1e8)
+        self.mask = mask
+        super().__init__(probs, logits, validate_args)
+
+    def entropy(self) -> torch.Tensor:
+        if self.mask is None:
+            return super().entropy()
+        # If mask set, then use approximation for entropy
+        p_log_p = self.logits * self.probs
+        masked = torch.where(self.mask, p_log_p, 0)
+        return -masked.sum(-1)
+
 
 class MultiDiscreteActorHead(Actor):
     def __init__(
@@ -101,13 +153,18 @@ class MultiDiscreteActorHead(Actor):
             final_layer_gain=0.01,
         )
 
-    def forward(self, obs: torch.Tensor, a: Optional[torch.Tensor] = None) -> PiForward:
+    def forward(
+        self,
+        obs: torch.Tensor,
+        actions: Optional[torch.Tensor] = None,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> PiForward:
         logits = self._fc(obs)
-        pi = MultiCategorical(self.nvec, logits=logits)
+        pi = MultiCategorical(self.nvec, logits=logits, masks=action_masks)
         logp_a = None
         entropy = None
-        if a is not None:
-            logp_a = pi.log_prob(a)
+        if actions is not None:
+            logp_a = pi.log_prob(actions)
             entropy = pi.entropy()
         return PiForward(pi, logp_a, entropy)
 
@@ -146,12 +203,20 @@ class GaussianActorHead(Actor):
         std = torch.exp(self.log_std)
         return GaussianDistribution(mu, std)
 
-    def forward(self, obs: torch.Tensor, a: Optional[torch.Tensor] = None) -> PiForward:
+    def forward(
+        self,
+        obs: torch.Tensor,
+        actions: Optional[torch.Tensor] = None,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> PiForward:
+        assert (
+            not action_masks
+        ), f"{self.__class__.__name__} does not support action_masks"
         pi = self._distribution(obs)
         logp_a = None
         entropy = None
-        if a is not None:
-            logp_a = pi.log_prob(a)
+        if actions is not None:
+            logp_a = pi.log_prob(actions)
             entropy = pi.entropy()
         return PiForward(pi, logp_a, entropy)
 
@@ -311,12 +376,20 @@ class StateDependentNoiseActorHead(Actor):
             ones = ones.to(self.device)
         return ones * std
 
-    def forward(self, obs: torch.Tensor, a: Optional[torch.Tensor] = None) -> PiForward:
+    def forward(
+        self,
+        obs: torch.Tensor,
+        actions: Optional[torch.Tensor] = None,
+        action_masks: Optional[torch.Tensor] = None,
+    ) -> PiForward:
+        assert (
+            not action_masks
+        ), f"{self.__class__.__name__} does not support action_masks"
         pi = self._distribution(obs)
         logp_a = None
         entropy = None
-        if a is not None:
-            logp_a = pi.log_prob(a)
+        if actions is not None:
+            logp_a = pi.log_prob(actions)
             entropy = -logp_a if self.bijector else sum_independent_dims(pi.entropy())
         return PiForward(pi, logp_a, entropy)
 

rl_algo_impls/shared/policy/on_policy.py

@@ -77,7 +77,7 @@ class OnPolicy(Policy):
         ...
 
     @abstractmethod
-    def step(self, obs: VecEnvObs) -> Step:
+    def step(self, obs: VecEnvObs, action_masks: Optional[np.ndarray] = None) -> Step:
         ...
 
 
@@ -162,10 +162,13 @@ class ActorCritic(OnPolicy):
         )
 
     def _pi_forward(
-        self, obs: torch.Tensor, action: Optional[torch.Tensor] = None
+        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, action)
+        pi_forward = self._pi(p_fe, actions=action, action_masks=action_masks)
 
         return pi_forward, p_fe
 
@@ -173,8 +176,13 @@ class ActorCritic(OnPolicy):
         v_fe = self._v_feature_extractor(obs) if self._v_feature_extractor else p_fc
         return self._v(v_fe)
 
-    def forward(self, obs: torch.Tensor, action: torch.Tensor) -> ACForward:
-        (_, logp_a, entropy), p_fc = self._pi_forward(obs, action)
+    def forward(
+        self,
+        obs: torch.Tensor,
+        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)
 
         assert logp_a is not None
@@ -192,10 +200,11 @@ class ActorCritic(OnPolicy):
             v = self._v(fe)
         return v.cpu().numpy()
 
-    def step(self, obs: VecEnvObs) -> Step:
+    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)
+            (pi, _, _), p_fc = self._pi_forward(o, action_masks=a_masks)
             a = pi.sample()
             logp_a = pi.log_prob(a)
 
@@ -205,13 +214,21 @@ class ActorCritic(OnPolicy):
         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)
 
-    def act(self, obs: np.ndarray, deterministic: bool = True) -> np.ndarray:
+    def act(
+        self,
+        obs: np.ndarray,
+        deterministic: bool = True,
+        action_masks: Optional[np.ndarray] = None,
+    ) -> np.ndarray:
         if not deterministic:
-            return self.step(obs).clamped_a
+            return self.step(obs, action_masks=action_masks).clamped_a
         else:
             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, _, _), _ = self._pi_forward(o)
+                (pi, _, _), _ = self._pi_forward(o, action_masks=a_masks)
                 a = pi.mode
             return clamp_actions(a.cpu().numpy(), self.action_space, self.squash_output)
 

rl_algo_impls/shared/policy/policy.py

@@ -46,7 +46,12 @@ class Policy(nn.Module, ABC):
         return self
 
     @abstractmethod
-    def act(self, obs: VecEnvObs, deterministic: bool = True) -> np.ndarray:
+    def act(
+        self,
+        obs: VecEnvObs,
+        deterministic: bool = True,
+        action_masks: Optional[np.ndarray] = None,
+    ) -> np.ndarray:
         ...
 
     def save(self, path: str) -> None:

rl_algo_impls/vpg/vpg.py

@@ -10,7 +10,7 @@ from typing import Optional, Sequence, TypeVar
 
 from rl_algo_impls.shared.algorithm import Algorithm
 from rl_algo_impls.shared.callbacks.callback import Callback
-from rl_algo_impls.shared.gae import compute_rtg_and_advantage, compute_advantage
+from rl_algo_impls.shared.gae import compute_rtg_and_advantage_from_trajectories
 from rl_algo_impls.shared.trajectory import Trajectory, TrajectoryAccumulator
 from rl_algo_impls.vpg.policy import VPGActorCritic
 from rl_algo_impls.wrappers.vectorable_wrapper import VecEnv
@@ -58,7 +58,6 @@ class VanillaPolicyGradient(Algorithm):
         max_grad_norm: float = 10.0,
         n_steps: int = 4_000,
         sde_sample_freq: int = -1,
-        update_rtg_between_v_iters: bool = False,
         ent_coef: float = 0.0,
     ) -> None:
         super().__init__(policy, env, device, tb_writer)
@@ -73,7 +72,6 @@ class VanillaPolicyGradient(Algorithm):
         self.n_steps = n_steps
         self.train_v_iters = train_v_iters
         self.sde_sample_freq = sde_sample_freq
-        self.update_rtg_between_v_iters = update_rtg_between_v_iters
 
         self.ent_coef = ent_coef
 
@@ -118,7 +116,7 @@ class VanillaPolicyGradient(Algorithm):
         act = torch.as_tensor(
             np.concatenate([np.array(t.act) for t in trajectories]), device=self.device
         )
-        rtg, adv = compute_rtg_and_advantage(
+        rtg, adv = compute_rtg_and_advantage_from_trajectories(
             trajectories, self.policy, self.gamma, self.gae_lambda, self.device
         )
 
@@ -135,10 +133,6 @@ class VanillaPolicyGradient(Algorithm):
 
         v_loss = 0
         for _ in range(self.train_v_iters):
-            if self.update_rtg_between_v_iters:
-                rtg = compute_advantage(
-                    trajectories, self.policy, self.gamma, self.gae_lambda, self.device
-                )
             v = self.policy.v(obs)
             v_loss = ((v - rtg) ** 2).mean()
 

rl_algo_impls/wrappers/action_mask_wrapper.py

@@ -0,0 +1,22 @@
+import numpy as np
+
+from gym.vector.vector_env import VectorEnv
+from stable_baselines3.common.vec_env import VecEnv as SBVecEnv
+from typing import Optional
+
+from rl_algo_impls.wrappers.vectorable_wrapper import VecotarableWrapper
+
+
+class IncompleteArrayError(Exception):
+    pass
+
+
+class ActionMaskWrapper(VecotarableWrapper):
+    def action_masks(self) -> Optional[np.ndarray]:
+        envs = getattr(self.env.unwrapped, "envs")
+        assert (
+            envs
+        ), f"{self.__class__.__name__} expects to wrap synchronous vectorized env"
+        masks = [getattr(e.unwrapped, "action_mask") for e in envs]
+        assert all(m is not None for m in masks)
+        return np.array(masks, dtype=np.bool8)

saved_models/ppo-MicrortsAttackShapedReward-v1-S2-best/model.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:114140423904ae116b52975d9de49633eb55ff1c790fe97dded5c8a64f63484b
+oid sha256:0eff8cfad9e5927093804e190813c112c9484c2ddbf8c97abc7713bf4a01d27c
 size 294255

scripts/benchmark.sh

@@ -8,6 +8,7 @@ do
         -e) envs=$2 ;;
         --procgen) procgen=t ;;
         --microrts) microrts=t ;;
+        --no-mask-microrts) no_mask_microrts=t ;;
     esac
     shift
 done
@@ -58,6 +59,13 @@ for algo in $(echo $algos); do
             "MicrortsRandomEnemyShapedReward3-v1"
         )
         algo_envs=${MICRORTS_ENVS[*]}
+    elif [ "$no_mask_microrts" = "t" ]; then
+        NO_MASK_MICRORTS_ENVS=(
+            "MicrortsMining-v1-NoMask"
+            "MicrortsAttackShapedReward-v1-NoMask"
+            "MicrortsRandomEnemyShapedReward3-v1-NoMask"
+        )
+        algo_envs=${NO_MASK_MICRORTS_ENVS[*]}
     elif [ -z "$envs" ]; then
         if [ "$algo" = "dqn" ]; then
             BENCHMARK_ENVS="${DISCRETE_ENVS[*]}"