Add full action masking support, expand and refactor testing (#227)

Farama-Foundation · Jul 20, 2023 · 7021b7f · 7021b7f
2 parents cbeeec2 + 22ae565
commit 7021b7f
Show file tree

Hide file tree

Showing 3 changed files with 342 additions and 84 deletions.
diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
@@ -18,4 +18,4 @@ jobs:
       - run: python -m pip install pre-commit
       - run: python -m pre_commit --version
       - run: python -m pre_commit install
-      - run: python -m pre_commit run --all-files
+      - run: python -m pre_commit run --all-files --show-diff-on-failure
diff --git a/supersuit/utils/obs_delay.py b/supersuit/utils/obs_delay.py
@@ -14,4 +14,18 @@ def add(self, in_obs):
         if len(self.obs_queue) > self.delay:
             return self.obs_queue.popleft()
         else:
-            return np.zeros_like(in_obs)
+            if isinstance(in_obs, np.ndarray):
+                return np.zeros_like(in_obs)
+            elif (
+                isinstance(in_obs, dict)
+                and "observation" in in_obs.keys()
+                and "action_mask" in in_obs.keys()
+            ):
+                return {
+                    "observation": np.zeros_like(in_obs["observation"]),
+                    "action_mask": np.ones_like(in_obs["action_mask"]),
+                }
+            else:
+                raise TypeError(
+                    "Observation must be of type np.ndarray or dictionary with keys 'observation' and 'action_mask'"
+                )