Fix env checker to handle Sequence spaces in composite spaces (Dict, Tuple, OneOf) (#2174)

* Initial plan

* Fix env checker to handle Sequence spaces in composite spaces

Co-authored-by: araffin <[email protected]>

* Simplify tests and checks

* Add fix for mypy for gym<1.0

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: araffin <[email protected]>
Co-authored-by: Antonin RAFFIN <[email protected]>

Author: Copilot

docs/misc/changelog.rst

@@ -15,6 +15,8 @@ New Features:
 
 Bug Fixes:
 ^^^^^^^^^^
+- Fixed env checker to properly handle ``Sequence`` observation spaces when nested inside composite spaces (``Dict``, ``Tuple``, ``OneOf``) (@copilot)
+- Update env checker to warn users when using Graph space (@dhruvmalik007).
 
 `SB3-Contrib`_
 ^^^^^^^^^^^^^^
@@ -57,7 +59,6 @@ Bug Fixes:
 ^^^^^^^^^^
 - Fixed docker GPU image (PyTorch GPU was not installed)
 - Fixed segmentation faults caused by non-portable schedules during model loading (@akanto)
-- Update env checker to warn users when using Graph space (@dhruvmalik007).
 
 `SB3-Contrib`_
 ^^^^^^^^^^^^^^

stable_baselines3/common/env_checker.py

@@ -9,6 +9,18 @@
 from stable_baselines3.common.vec_env import DummyVecEnv, VecCheckNan
 
 
+def _is_oneof_space(space: spaces.Space) -> bool:
+    """
+    Return True if the provided space is a OneOf space,
+    False if not or if the current version of Gym doesn't support this space.
+    """
+    try:
+        return isinstance(space, spaces.OneOf)  # type: ignore[attr-defined]
+    except AttributeError:
+        # Gym < v1.0
+        return False
+
+
 def _is_numpy_array_space(space: spaces.Space) -> bool:
     """
     Returns False if provided space is not representable as a single numpy array
@@ -80,21 +92,23 @@ def _check_image_input(observation_space: spaces.Box, key: str = "") -> None:
         )
 
 
-def _check_unsupported_spaces(env: gym.Env, observation_space: spaces.Space, action_space: spaces.Space) -> bool:
+def _check_unsupported_spaces(env: gym.Env, observation_space: spaces.Space, action_space: spaces.Space) -> bool:  # noqa: C901
     """
     Emit warnings when the observation space or action space used is not supported by Stable-Baselines.
 
     :return: True if return value tests should be skipped.
     """
 
-    should_skip = graph_space = False
+    should_skip = graph_space = sequence_space = False
     if isinstance(observation_space, spaces.Dict):
         nested_dict = False
         for key, space in observation_space.spaces.items():
             if isinstance(space, spaces.Dict):
                 nested_dict = True
-            if isinstance(space, spaces.Graph):
+            elif isinstance(space, spaces.Graph):
                 graph_space = True
+            elif isinstance(space, spaces.Sequence):
+                sequence_space = True
             _check_non_zero_start(space, "observation", key)
 
         if nested_dict:
@@ -122,10 +136,24 @@ def _check_unsupported_spaces(env: gym.Env, observation_space: spaces.Space, act
             "(cf. https://gymnasium.farama.org/api/spaces/composite/#dict). "
             "which is supported by SB3."
         )
+        # Check for Sequence spaces inside Tuple
+        for space in observation_space.spaces:
+            if isinstance(space, spaces.Sequence):
+                sequence_space = True
+            elif isinstance(space, spaces.Graph):
+                graph_space = True
+
+    # Check for Sequence spaces inside OneOf
+    if _is_oneof_space(observation_space):
+        warnings.warn(
+            "OneOf observation space is not supported by Stable-Baselines3. "
+            "Note: The checks for returned values are skipped."
+        )
+        should_skip = True
 
     _check_non_zero_start(observation_space, "observation")
 
-    if isinstance(observation_space, spaces.Sequence):
+    if isinstance(observation_space, spaces.Sequence) or sequence_space:
         warnings.warn(
             "Sequence observation space is not supported by Stable-Baselines3. "
             "You can pad your observation to have a fixed size instead.\n"

tests/test_env_checker.py

@@ -53,13 +53,6 @@ def step(self, action):
         return self.observation_space.sample(), 1.0, False, False, {}
 
 
-def test_check_env_sequence_obs():
-    test_env = CustomEnv()
-
-    with pytest.warns(Warning, match="Sequence.*not supported"):
-        check_env(env=test_env, warn=True)
-
-
 @pytest.mark.parametrize(
     "obs_tuple",
     [
@@ -223,3 +216,58 @@ def test_check_env_graph_space():
 
     with pytest.warns(UserWarning, match="Graph.*not supported"):
         check_env(SimpleDictGraphEnv(), warn=True)
+
+
+class SequenceInDictEnv(CustomEnv):
+    """Test env with Sequence space inside Dict space."""
+
+    def __init__(self):
+        self.action_space = spaces.Discrete(2)
+        self.observation_space = spaces.Dict(
+            {"seq": spaces.Sequence(spaces.Box(low=-100, high=100, shape=(1,), dtype=np.float32))}
+        )
+
+
+class SequenceInTupleEnv(CustomEnv):
+    """Test env with Sequence space inside Tuple space."""
+
+    def __init__(self):
+        self.action_space = spaces.Discrete(2)
+        self.observation_space = spaces.Tuple((spaces.Sequence(spaces.Box(low=-100, high=100, shape=(1,), dtype=np.float32)),))
+
+
+class SequenceInOneOfEnv(CustomEnv):
+    """Test env with Sequence space inside OneOf space."""
+
+    def __init__(self):
+        self.action_space = spaces.Discrete(2)
+        self.observation_space = spaces.OneOf(
+            (
+                spaces.Sequence(spaces.Box(low=-100, high=100, shape=(1,), dtype=np.float32)),
+                spaces.Discrete(3),
+            )
+        )
+
+
+@pytest.mark.parametrize("env_class", [CustomEnv, SequenceInDictEnv])
+def test_check_env_sequence_obs(env_class):
+    with pytest.warns(Warning, match="Sequence.*not supported"):
+        check_env(env_class(), warn=True)
+
+
+def test_check_env_sequence_tuple():
+    with (
+        pytest.warns(Warning, match="Sequence.*not supported"),
+        pytest.warns(Warning, match="Tuple.*not supported"),
+    ):
+        check_env(SequenceInTupleEnv(), warn=True)
+
+
+def test_check_env_oneof():
+    try:
+        env = SequenceInOneOfEnv()
+    except AttributeError:
+        pytest.skip("OneOf not supported by current Gymnasium version")
+
+    with pytest.warns(Warning, match="OneOf.*not supported"):
+        check_env(env, warn=True)