Invalid masking support (#16)

* Support for using action_masks in ppo

Updated default microrts envs to use action_mask. Added non-masked
variants for comparison.

* Bump to v0.0.7

Author: sgoodfriend

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 = "goodfriend.scott@gmail.com"},

rl_algo_impls/a2c/a2c.py

@@ -10,6 +10,7 @@
 
 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,7 +85,7 @@ def learn(
         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)
 
@@ -126,23 +127,16 @@ def learn(
                     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/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