Add new parameter to pytest function

Farama-Foundation · Oct 22, 2024 · 613a329 · 613a329
1 parent 0bb14ed
commit 613a329
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/gymnasium/wrappers/vector/stateful_observation.py b/gymnasium/wrappers/vector/stateful_observation.py
@@ -102,7 +102,11 @@ def reset(
         seed: int | list[int] | None = None,
         options: dict[str, Any] | None = None,
     ) -> tuple[ObsType, dict[str, Any]]:
-        assert options is None or "reset_mask" not in options or np.all(options["reset_mask"])
+        assert (
+            options is None
+            or "reset_mask" not in options
+            or np.all(options["reset_mask"])
+        )
         return super().reset(seed=seed, options=options)
 
     def observations(self, observations: ObsType) -> ObsType:

diff --git a/tests/wrappers/vector/test_vector_wrappers.py b/tests/wrappers/vector/test_vector_wrappers.py
@@ -37,7 +37,9 @@ def custom_environments():
     del gym.registry["DictObsEnv-v0"]
 
 
-@pytest.mark.parametrize("autoreset_mode", [AutoresetMode.NEXT_STEP, AutoresetMode.SAME_STEP])
+@pytest.mark.parametrize(
+    "autoreset_mode", [AutoresetMode.NEXT_STEP, AutoresetMode.SAME_STEP]
+)
 @pytest.mark.parametrize("num_envs", (1, 3))
 @pytest.mark.parametrize(
     "env_id, wrapper_name, kwargs",
@@ -70,6 +72,7 @@ def custom_environments():
     ),
 )
 def test_vector_wrapper_equivalence(
+    autoreset_mode: AutoresetMode,
     env_id: str,
     wrapper_name: str,
     kwargs: dict[str, Any],