Fixes and simplify checks

Author: araffin

docs/misc/changelog.rst

@@ -3,7 +3,7 @@
 Changelog
 ==========
 
-Release 2.7.1a1 (WIP)
+Release 2.7.1a2 (WIP)
 --------------------------
 
 Breaking Changes:
@@ -57,7 +57,7 @@ Bug Fixes:
 ^^^^^^^^^^
 - Fixed docker GPU image (PyTorch GPU was not installed)
 - Fixed segmentation faults caused by non-portable schedules during model loading (@akanto)
-- Fixed NoneType error when importing Gymnasium.graph  for sb3 (@dhruvmalik007).
+- Update env checker to warn users when using Graph space (@dhruvmalik007).
 
 `SB3-Contrib`_
 ^^^^^^^^^^^^^^

stable_baselines3/common/env_checker.py

@@ -12,7 +12,7 @@
 def _is_numpy_array_space(space: spaces.Space) -> bool:
     """
     Returns False if provided space is not representable as a single numpy array
-    (e.g. Dict, Tuple spaces return False)
+    (e.g. Dict and Tuple spaces return False)
     """
     return not isinstance(space, (spaces.Dict, spaces.Tuple))
 
@@ -80,14 +80,21 @@ 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) -> None:
-    """Emit warnings when the observation space or action space used is not supported by Stable-Baselines."""
+def _check_unsupported_spaces(env: gym.Env, observation_space: spaces.Space, action_space: spaces.Space) -> bool:
+    """
+    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
     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):
+                graph_space = True
             _check_non_zero_start(space, "observation", key)
 
         if nested_dict:
@@ -124,6 +131,14 @@ def _check_unsupported_spaces(env: gym.Env, observation_space: spaces.Space, act
             "You can pad your observation to have a fixed size instead.\n"
             "Note: The checks for returned values are skipped."
         )
+        should_skip = True
+
+    if isinstance(observation_space, spaces.Graph) or graph_space:
+        warnings.warn(
+            "Graph observation space is not supported by Stable-Baselines3. "
+            "Note: The checks for returned values are skipped."
+        )
+        should_skip = True
 
     _check_non_zero_start(action_space, "action")
 
@@ -133,6 +148,7 @@ def _check_unsupported_spaces(env: gym.Env, observation_space: spaces.Space, act
             "This type of action space is currently not supported by Stable Baselines 3. You should try to flatten the "
             "action using a wrapper."
         )
+    return should_skip
 
 
 def _check_nan(env: gym.Env) -> None:
@@ -204,26 +220,10 @@ def _check_obs(obs: Union[tuple, dict, np.ndarray, int], observation_space: spac
     Check that the observation returned by the environment
     correspond to the declared one.
     """
-    # If the observation space is a Graph, return early with a warning.
-    if _is_graph_space(observation_space):
-        warnings.warn(
-            f"The observation space for `{method_name}()` is a Graph space, which is not supported the env checker.",
-            "Skipping further observation checks.",
-            UserWarning,
-        )
-        return
-
     if not isinstance(observation_space, spaces.Tuple):
         assert not isinstance(
             obs, tuple
         ), f"The observation returned by the `{method_name}()` method should be a single value, not a tuple"
-        # Graph spaces are not fully supported by the env checker
-        if isinstance(observation_space, spaces.Graph):
-            warnings.warn(
-                "Graph observation spaces are not fully supported by the env checker. "
-                "Skipping further observation checks."
-            )
-            return
     # The check for a GoalEnv is done by the base class
     if isinstance(observation_space, spaces.Discrete):
         # Since https://github.com/Farama-Foundation/Gymnasium/pull/141,
@@ -371,7 +371,6 @@ def _check_returned_values(env: gym.Env, observation_space: spaces.Space, action
         assert reward == env.compute_reward(obs["achieved_goal"], obs["desired_goal"], info)
 
 
-
 def _check_spaces(env: gym.Env) -> None:
     """
     Check that the observation and action spaces are defined and inherit from spaces.Space. For
@@ -400,18 +399,6 @@ def _check_spaces(env: gym.Env) -> None:
         )
 
 
-def _is_graph_space(space: Any) -> bool:
-    """
-    Check if the space is a gymnasium Graph space.
-    Handles the case where gymnasium.spaces.graph may not be available.
-    """
-    try:
-        from gymnasium.spaces.graph import Graph
-    except ImportError:
-        return False
-    return isinstance(space, Graph)
-
-
 # Check render cannot be covered by CI
 def _check_render(env: gym.Env, warn: bool = False) -> None:  # pragma: no cover
     """
@@ -470,8 +457,10 @@ def check_env(env: gym.Env, warn: bool = True, skip_render_check: bool = True) -
         raise TypeError("The reset() method must accept a `seed` parameter") from e
 
     # Warn the user if needed.
+    # A warning means that the environment may run but not work properly with Stable Baselines algorithms
+    should_skip = False
     if warn:
-        _check_unsupported_spaces(env, observation_space, action_space)
+        should_skip = _check_unsupported_spaces(env, observation_space, action_space)
 
         obs_spaces = observation_space.spaces if isinstance(observation_space, spaces.Dict) else {"": observation_space}
         for key, space in obs_spaces.items():
@@ -499,8 +488,8 @@ def check_env(env: gym.Env, warn: bool = True, skip_render_check: bool = True) -
                 f"Your action space has dtype {action_space.dtype}, we recommend using np.float32 to avoid cast errors."
             )
 
-    # If Sequence observation space, do not check the observation any further
-    if isinstance(observation_space, spaces.Sequence):
+    # If Sequence or Graph observation space, do not check the observation any further
+    if should_skip:
         return
 
     # ============ Check the returned values ===============
@@ -512,12 +501,8 @@ def check_env(env: gym.Env, warn: bool = True, skip_render_check: bool = True) -
 
     try:
         check_for_nested_spaces(env.observation_space)
-        # if there is None defined , then skip the test
-        if _is_graph_space(observation_space) or (
-            isinstance(observation_space, spaces.Dict) and any(_is_graph_space(s) for s in observation_space.spaces.values())
-        ):
-            pass  # skip _check_nan for Graph spaces
-        else:
-            _check_nan(env)
+        # The check doesn't support nested observations/dict actions
+        # A warning about it has already been emitted
+        _check_nan(env)
     except NotImplementedError:
         pass

stable_baselines3/common/utils.py

@@ -574,8 +574,8 @@ def obs_as_tensor(obs: Union[np.ndarray, dict[str, np.ndarray]], device: th.devi
     elif isinstance(obs, dict):
         return {key: th.as_tensor(_obs, device=device) for (key, _obs) in obs.items()}
     else:
-        raise Exception(f"Unrecognized type of observation {type(obs)}")
-        return obs
+        raise TypeError(f"Unrecognized type of observation {type(obs)}")
+
 
 def should_collect_more_steps(
     train_freq: TrainFreq,
@@ -604,6 +604,7 @@ def should_collect_more_steps(
             f"or TrainFrequencyUnit.EPISODE not '{train_freq.unit}'!"
         )
 
+
 def get_system_info(print_info: bool = True) -> tuple[dict[str, str], str]:
     """
     Retrieve system and python env info for the current system.
@@ -636,4 +637,4 @@ def get_system_info(print_info: bool = True) -> tuple[dict[str, str], str]:
         env_info_str += f"- {key}: {value}\n"
     if print_info:
         print(env_info_str)
-    return env_info, env_info_str
+    return env_info, env_info_str

stable_baselines3/version.txt

@@ -1 +1 @@
-2.7.1a0
+2.7.1a2

tests/test_env_checker.py

@@ -23,7 +23,7 @@ def step(self, action):
         info = {}
         return observation, reward, terminated, truncated, info
 
-    def reset(self, seed=None, options=None):
+    def reset(self, *, seed=None, options=None):
         return np.array([1.0, 1.5, 0.5], dtype=self.observation_space.dtype), {}
 
     def render(self):
@@ -37,14 +37,15 @@ def test_check_env_dict_action():
         check_env(env=test_env, warn=True)
 
 
-class SequenceObservationEnv(gym.Env):
+class CustomEnv(gym.Env):
     metadata = {"render_modes": [], "render_fps": 2}
 
     def __init__(self, render_mode=None):
+        # Test Sequence obs
         self.observation_space = spaces.Sequence(spaces.Discrete(8))
         self.action_space = spaces.Discrete(4)
 
-    def reset(self, seed=None, options=None):
+    def reset(self, *, seed=None, options=None):
         super().reset(seed=seed)
         return self.observation_space.sample(), {}
 
@@ -53,7 +54,7 @@ def step(self, action):
 
 
 def test_check_env_sequence_obs():
-    test_env = SequenceObservationEnv()
+    test_env = CustomEnv()
 
     with pytest.warns(Warning, match="Sequence.*not supported"):
         check_env(env=test_env, warn=True)
@@ -193,46 +194,32 @@ def test_check_env_single_step_env():
     check_env(env=test_env, warn=True)
 
 
-
-class SimpleGraphEnv(gym.Env):
+class SimpleGraphEnv(CustomEnv):
     def __init__(self):
         self.action_space = spaces.Discrete(2)
-        # Define a simple Graph observation space
-        node_shape = (2,)
-        edge_shape = (3,)
         self.observation_space = spaces.Graph(
-            node_space=spaces.Box(low=0, high=1, shape=node_shape),
-            edge_space=spaces.Box(low=0, high=1, shape=edge_shape),
+            node_space=spaces.Box(low=0, high=1, shape=(2,)),
+            edge_space=spaces.Box(low=0, high=1, shape=(3,)),
         )
 
-    def reset(self, seed=None, options=None):
-        # Just sample from the obs space
-        return self.observation_space.sample(), {}
 
-    def step(self, action):
-        return self.observation_space.sample(), 1.0, False, False, {}
+class SimpleDictGraphEnv(CustomEnv):
+    def __init__(self):
+        self.action_space = spaces.Discrete(2)
+        self.observation_space = spaces.Dict(
+            {
+                "test": spaces.Graph(
+                    node_space=spaces.Box(low=0, high=1, shape=(2,)),
+                    edge_space=spaces.Box(low=0, high=1, shape=(3,)),
+                )
+            }
+        )
 
 
-def test_check_env_simple_graph_space():
-    env = SimpleGraphEnv()
+def test_check_env_graph_space():
     # Should emit a warning about Graph space, but not fail
-    with pytest.warns(UserWarning, match="Graph space, which is not supported by the env checker"):
-        check_env(env, warn=True)
-
-
-def test_check_wrong_type_graph_space():
-    # Create a Graph space with a wrong type
-    node_shape = (2,)
-    edge_shape = (3,)
-    graph_space = spaces.Graph(
-        node_space=spaces.Box(low=0, high=1, shape=node_shape),
-        edge_space=spaces.Box(low=0, high=1, shape=edge_shape),
-    )
-    # Create an env with the wrong type
-    env = SimpleGraphEnv()
-    env.observation_space = graph_space
-
-    # Check that the env checker raises an error
-    with pytest.raises(AssertionError, match="incompatible w/ graph-obs"):
-        check_env(env, warn=True)
+    with pytest.warns(UserWarning, match="Graph.*not supported"):
+        check_env(SimpleGraphEnv(), warn=True)
 
+    with pytest.warns(UserWarning, match="Graph.*not supported"):
+        check_env(SimpleDictGraphEnv(), warn=True)