check_env warning for FrameStacked observation in stable_baselines3.common.env_checker

🐛 Bug

When using the "check_env "function of "stable_baselines3.common.env_checker" with an environment wrapped in a "FrameStack" wrapper from "gymnasium.wrappers", I get an error on the type of observation returned.

Indeed, the observation has a <class 'gymnasium.wrappers.frame_stack.LazyFrames'> instead of a np.array type (which is normal because it is indicated in the gymnasium documentation).

But then I can still use this environment with FrameStacked observation to train an agent. I know it is not critical but is this normal that the check_env warns about a problem while the environment can still be used to train agents with stable_baselines3 algorithms?

To Reproduce

import gymnasium as gym
from gymnasium.wrappers import FrameStack
from stable_baselines3.common.env_checker import check_env

env = gym.make("CartPole-v1")
env = FrameStack(env, n_stack)

# check the obs type
obs, info = env.reset()
print(type(obs))

check_env(env, warn = True)

Relevant log output / Error message

<class 'gymnasium.wrappers.frame_stack.LazyFrames'>

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[8], line 7
      4 obs, info = env.reset()
      5 print(type(obs))
----> 7 check_env(env, warn = True)

File ~/.local/lib/python3.10/site-packages/stable_baselines3/common/env_checker.py:433, in check_env(env, warn, skip_render_check)
    428         warnings.warn(
    429             f"Your action space has dtype {action_space.dtype}, we recommend using np.float32 to avoid cast errors."
    430         )
    432 # ============ Check the returned values ===============
--> 433 _check_returned_values(env, observation_space, action_space)
    435 # ==== Check the render method and the declared render modes ====
    436 if not skip_render_check:

File ~/.local/lib/python3.10/site-packages/stable_baselines3/common/env_checker.py:274, in _check_returned_values(env, observation_space, action_space)
    272             raise AssertionError(f"Error while checking key={key}: " + str(e)) from e
    273 else:
--> 274     _check_obs(obs, observation_space, "reset")
    276 # Sample a random action
    277 action = action_space.sample()

File ~/.local/lib/python3.10/site-packages/stable_baselines3/common/env_checker.py:189, in _check_obs(obs, observation_space, method_name)
    187     assert np.issubdtype(type(obs), np.integer), f"The observation returned by `{method_name}()` method must be an int"
    188 elif _is_numpy_array_space(observation_space):
--> 189     assert isinstance(obs, np.ndarray), f"The observation returned by `{method_name}()` method must be a numpy array"
    191 # Additional checks for numpy arrays, so the error message is clearer (see GH#1399)
    192 if isinstance(obs, np.ndarray):
    193     # check obs dimensions, dtype and bounds

AssertionError: The observation returned by `reset()` method must be a numpy array

System Info

I am using these versions :

• Python: 3.10.6
• stable_baselines : 2.0.0a5
• gymnasium : 0.28.1

Checklist

• I have checked that there is no similar issue in the repo
• I have read the documentation
• I have provided a minimal working example to reproduce the bug
• I've used the markdown code blocks for both code and stack traces.

Start with stable_baselines3/common/env_checker.py, especially _check_obs and _check_returned_values, and reproduce the FrameStack example from the issue. Compare the checker’s observation-type assumption with Gymnasium’s LazyFrames behavior; done means the reported case has an intentional, tested outcome without breaking existing observation checks.