Refactor VecEnv checker: move to separate files and add documentation

Co-authored-by: araffin <1973948+araffin@users.noreply.github.com>

Author: Copilot

docs/guide/vec_envs.rst

@@ -183,6 +183,59 @@ This callback can then be used to safely modify environment attributes during tr
 it calls the environment setter method.
 
 
+Checking VecEnv Implementation
+-----------------------------
+
+When implementing custom vectorized environments, it's easy to make mistakes that can lead to hard-to-debug issues.
+To help with this, Stable-Baselines3 provides a ``check_vecenv`` function that validates your VecEnv implementation
+and checks for common issues.
+
+The ``check_vecenv`` function verifies:
+
+* The VecEnv properly inherits from ``stable_baselines3.common.vec_env.VecEnv``
+* Required attributes (``num_envs``, ``observation_space``, ``action_space``) are present and valid
+* The ``reset()`` method returns observations with the correct vectorized shape (batch dimension first)
+* The ``step()`` method returns properly shaped observations, rewards, dones, and infos
+* All return values have the expected types and dimensions
+* Compatibility with Stable-Baselines3 algorithms
+
+**Usage:**
+
+.. code-block:: python
+
+    from stable_baselines3.common.vec_env import DummyVecEnv
+    from stable_baselines3.common.vec_env_checker import check_vecenv
+    import gymnasium as gym
+
+    def make_env():
+        return gym.make('CartPole-v1')
+
+    # Create your VecEnv
+    vec_env = DummyVecEnv([make_env for _ in range(4)])
+
+    # Check the VecEnv implementation
+    check_vecenv(vec_env, warn=True)
+
+    vec_env.close()
+
+**When to use:**
+
+* When implementing a custom VecEnv class
+* When debugging issues with vectorized environments
+* When contributing new VecEnv implementations to ensure they follow the API
+* As a sanity check before training to catch potential issues early
+
+**Note:** Similar to ``check_env`` for single environments, ``check_vecenv`` is particularly useful during development
+and debugging. It helps catch common vectorization mistakes like incorrect batch dimensions, wrong return types, or
+missing required methods.
+
+
+VecEnv Checker
+~~~~~~~~~~~~~~
+
+.. autofunction:: stable_baselines3.common.vec_env_checker.check_vecenv
+
+
 Vectorized Environments Wrappers
 --------------------------------
 

stable_baselines3/common/env_checker.py

@@ -7,7 +7,6 @@
 
 from stable_baselines3.common.preprocessing import check_for_nested_spaces, is_image_space_channels_first
 from stable_baselines3.common.vec_env import DummyVecEnv, VecCheckNan
-from stable_baselines3.common.vec_env.base_vec_env import VecEnv
 
 
 def _is_oneof_space(space: spaces.Space) -> bool:
@@ -538,257 +537,3 @@ def check_env(env: gym.Env, warn: bool = True, skip_render_check: bool = True) -
         _check_nan(env)
     except NotImplementedError:
         pass
-
-
-def _check_vecenv_spaces(vec_env: VecEnv) -> None:
-    """
-    Check that the VecEnv has valid observation and action spaces.
-    """
-    assert hasattr(vec_env, "observation_space"), "VecEnv must have an observation_space attribute"
-    assert hasattr(vec_env, "action_space"), "VecEnv must have an action_space attribute"
-    assert hasattr(vec_env, "num_envs"), "VecEnv must have a num_envs attribute"
-
-    assert isinstance(
-        vec_env.observation_space, spaces.Space
-    ), "The observation space must inherit from gymnasium.spaces"
-    assert isinstance(vec_env.action_space, spaces.Space), "The action space must inherit from gymnasium.spaces"
-    assert isinstance(vec_env.num_envs, int) and vec_env.num_envs > 0, "num_envs must be a positive integer"
-
-
-def _check_vecenv_reset(vec_env: VecEnv) -> Any:
-    """
-    Check that VecEnv reset method works correctly and returns properly shaped observations.
-    """
-    try:
-        obs = vec_env.reset()
-    except Exception as e:
-        raise RuntimeError(f"VecEnv reset() failed: {e}") from e
-
-    # Check observation shape matches expected vectorized shape
-    if isinstance(vec_env.observation_space, spaces.Box):
-        assert isinstance(obs, np.ndarray), f"For Box observation space, reset() must return np.ndarray, got {type(obs)}"
-        expected_shape = (vec_env.num_envs,) + vec_env.observation_space.shape
-        assert obs.shape == expected_shape, (
-            f"Expected observation shape {expected_shape}, got {obs.shape}. "
-            f"VecEnv observations should have batch dimension first."
-        )
-    elif isinstance(vec_env.observation_space, spaces.Dict):
-        assert isinstance(obs, dict), f"For Dict observation space, reset() must return dict, got {type(obs)}"
-        for key, space in vec_env.observation_space.spaces.items():
-            assert key in obs, f"Missing key '{key}' in observation dict"
-            if isinstance(space, spaces.Box):
-                expected_shape = (vec_env.num_envs,) + space.shape
-                assert obs[key].shape == expected_shape, (
-                    f"Expected observation['{key}'] shape {expected_shape}, got {obs[key].shape}"
-                )
-    elif isinstance(vec_env.observation_space, spaces.Discrete):
-        assert isinstance(obs, np.ndarray), f"For Discrete observation space, reset() must return np.ndarray, got {type(obs)}"
-        expected_shape = (vec_env.num_envs,)
-        assert obs.shape == expected_shape, f"Expected observation shape {expected_shape}, got {obs.shape}"
-
-    return obs
-
-
-def _check_vecenv_step(vec_env: VecEnv, obs: Any) -> None:
-    """
-    Check that VecEnv step method works correctly and returns properly shaped values.
-    """
-    # Generate valid actions
-    if isinstance(vec_env.action_space, spaces.Box):
-        actions = np.array([vec_env.action_space.sample() for _ in range(vec_env.num_envs)])
-    elif isinstance(vec_env.action_space, spaces.Discrete):
-        actions = np.array([vec_env.action_space.sample() for _ in range(vec_env.num_envs)])
-    elif isinstance(vec_env.action_space, spaces.MultiDiscrete):
-        actions = np.array([vec_env.action_space.sample() for _ in range(vec_env.num_envs)])
-    elif isinstance(vec_env.action_space, spaces.MultiBinary):
-        actions = np.array([vec_env.action_space.sample() for _ in range(vec_env.num_envs)])
-    else:
-        actions = np.array([vec_env.action_space.sample() for _ in range(vec_env.num_envs)])
-
-    try:
-        obs, rewards, dones, infos = vec_env.step(actions)
-    except Exception as e:
-        raise RuntimeError(f"VecEnv step() failed: {e}") from e
-
-    # Check rewards
-    assert isinstance(rewards, np.ndarray), f"step() must return rewards as np.ndarray, got {type(rewards)}"
-    assert rewards.shape == (vec_env.num_envs,), f"Expected rewards shape ({vec_env.num_envs},), got {rewards.shape}"
-
-    # Check dones
-    assert isinstance(dones, np.ndarray), f"step() must return dones as np.ndarray, got {type(dones)}"
-    assert dones.shape == (vec_env.num_envs,), f"Expected dones shape ({vec_env.num_envs},), got {dones.shape}"
-    assert dones.dtype == bool, f"dones must have dtype bool, got {dones.dtype}"
-
-    # Check infos
-    assert isinstance(infos, (list, tuple)), f"step() must return infos as list or tuple, got {type(infos)}"
-    assert len(infos) == vec_env.num_envs, f"Expected infos length {vec_env.num_envs}, got {len(infos)}"
-    for i, info in enumerate(infos):
-        assert isinstance(info, dict), f"infos[{i}] must be dict, got {type(info)}"
-
-    # Check observation shape consistency (similar to reset)
-    if isinstance(vec_env.observation_space, spaces.Box):
-        assert isinstance(obs, np.ndarray), f"For Box observation space, step() must return np.ndarray, got {type(obs)}"
-        expected_shape = (vec_env.num_envs,) + vec_env.observation_space.shape
-        assert obs.shape == expected_shape, (
-            f"Expected observation shape {expected_shape}, got {obs.shape}. "
-            f"VecEnv observations should have batch dimension first."
-        )
-    elif isinstance(vec_env.observation_space, spaces.Dict):
-        assert isinstance(obs, dict), f"For Dict observation space, step() must return dict, got {type(obs)}"
-        for key, space in vec_env.observation_space.spaces.items():
-            assert key in obs, f"Missing key '{key}' in observation dict"
-            if isinstance(space, spaces.Box):
-                expected_shape = (vec_env.num_envs,) + space.shape
-                assert obs[key].shape == expected_shape, (
-                    f"Expected observation['{key}'] shape {expected_shape}, got {obs[key].shape}"
-                )
-
-
-def _check_vecenv_unsupported_spaces(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
-    for VecEnv. This is a VecEnv-specific version of _check_unsupported_spaces.
-
-    :return: True if return value tests should be skipped.
-    """
-    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
-            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:
-            warnings.warn(
-                "Nested observation spaces are not supported by Stable Baselines3 "
-                "(Dict spaces inside Dict space). "
-                "You should flatten it to have only one level of keys."
-                "For example, `dict(space1=dict(space2=Box(), space3=Box()), spaces4=Discrete())` "
-                "is not supported but `dict(space2=Box(), spaces3=Box(), spaces4=Discrete())` is."
-            )
-
-    if isinstance(observation_space, spaces.MultiDiscrete) and len(observation_space.nvec.shape) > 1:
-        warnings.warn(
-            f"The MultiDiscrete observation space uses a multidimensional array {observation_space.nvec} "
-            "which is currently not supported by Stable-Baselines3. "
-            "Please convert it to a 1D array using a wrapper: "
-            "https://github.com/DLR-RM/stable-baselines3/issues/1836."
-        )
-
-    if isinstance(observation_space, spaces.Tuple):
-        warnings.warn(
-            "The observation space is a Tuple, "
-            "this is currently not supported by Stable Baselines3. "
-            "However, you can convert it to a Dict observation space "
-            "(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) 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"
-            "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")
-
-    if not _is_numpy_array_space(action_space):
-        warnings.warn(
-            "The action space is not based off a numpy array. Typically this means it's either a Dict or Tuple space. "
-            "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_vecenv(vec_env: VecEnv, warn: bool = True) -> None:
-    """
-    Check that a VecEnv follows the VecEnv API and is compatible with Stable-Baselines3.
-    
-    This checker verifies that:
-    - The VecEnv has proper observation_space, action_space, and num_envs attributes
-    - The reset() method returns observations with correct vectorized shape
-    - The step() method returns observations, rewards, dones, and infos with correct shapes
-    - All return values have the expected types and dimensions
-    
-    :param vec_env: The vectorized environment to check
-    :param warn: Whether to output additional warnings mainly related to 
-        the interaction with Stable Baselines
-    """
-    assert isinstance(vec_env, VecEnv), (
-        "Your environment must inherit from stable_baselines3.common.vec_env.VecEnv"
-    )
-
-    # ============= Check basic VecEnv attributes ================
-    _check_vecenv_spaces(vec_env)
-
-    # Define aliases for convenience  
-    observation_space = vec_env.observation_space
-    action_space = vec_env.action_space
-
-    # Warn the user if needed - reuse existing space checking logic
-    if warn:
-        should_skip = _check_vecenv_unsupported_spaces(observation_space, action_space)
-        if should_skip:
-            warnings.warn("VecEnv contains unsupported spaces, skipping some checks")
-            return
-
-        obs_spaces = observation_space.spaces if isinstance(observation_space, spaces.Dict) else {"": observation_space}
-        for key, space in obs_spaces.items():
-            if isinstance(space, spaces.Box):
-                _check_box_obs(space, key)
-
-        # Check for the action space
-        if isinstance(action_space, spaces.Box) and (
-            np.any(np.abs(action_space.low) != np.abs(action_space.high))
-            or np.any(action_space.low != -1)
-            or np.any(action_space.high != 1)
-        ):
-            warnings.warn(
-                "We recommend you to use a symmetric and normalized Box action space (range=[-1, 1]) "
-                "cf. https://stable-baselines3.readthedocs.io/en/master/guide/rl_tips.html"
-            )
-
-        if isinstance(action_space, spaces.Box):
-            assert np.all(
-                np.isfinite(np.array([action_space.low, action_space.high]))
-            ), "Continuous action space must have a finite lower and upper bound"
-
-        if isinstance(action_space, spaces.Box) and action_space.dtype != np.dtype(np.float32):
-            warnings.warn(
-                f"Your action space has dtype {action_space.dtype}, we recommend using np.float32 to avoid cast errors."
-            )
-
-    # ============ Check the VecEnv methods ===============
-    obs = _check_vecenv_reset(vec_env)
-    _check_vecenv_step(vec_env, obs)

stable_baselines3/common/vec_env/__init__.py

@@ -13,6 +13,9 @@
 from stable_baselines3.common.vec_env.vec_transpose import VecTransposeImage
 from stable_baselines3.common.vec_env.vec_video_recorder import VecVideoRecorder
 
+# Avoid circular import by importing the vec_env_checker here  
+from stable_baselines3.common.vec_env_checker import check_vecenv
+
 VecEnvWrapperT = TypeVar("VecEnvWrapperT", bound=VecEnvWrapper)
 
 
@@ -98,6 +101,7 @@ def sync_envs_normalization(env: VecEnv, eval_env: VecEnv) -> None:
     "VecNormalize",
     "VecTransposeImage",
     "VecVideoRecorder",
+    "check_vecenv",
     "is_vecenv_wrapped",
     "sync_envs_normalization",
     "unwrap_vec_normalize",