diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 3f6e087..e86c8e2 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -20,6 +20,8 @@ Bug Fixes: - Fix ignoring the exclude parameter when recording logs using json, csv or log as logging format (@SwamyDev) - Make ``make_vec_env`` support the ``env_kwargs`` argument when using an env ID str (@ManifoldFR) - Fix model creation initializing CUDA even when `device="cpu"` is provided +- Fix ``check_env`` not checking if the env has a Dict actionspace before calling ``_check_nan`` (@wmmc88) +- Update the check for spaces unsupported by Stable Baselines 3 to include checks on the action space (@wmmc88) Deprecations: ^^^^^^^^^^^^^ @@ -457,4 +459,4 @@ And all the contributors: @MarvineGothic @jdossgollin @SyllogismRXS @rusu24edward @jbulow @Antymon @seheevic @justinkterry @edbeeching @flodorner @KuKuXia @NeoExtended @PartiallyTyped @mmcenta @richardwu @kinalmehta @rolandgvc @tkelestemur @mloo3 @tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag @xicocaio -@diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 @SwamyDev +@diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 @SwamyDev @wmmc88 diff --git a/stable_baselines3/common/env_checker.py b/stable_baselines3/common/env_checker.py index 4a4d075..262722f 100644 --- a/stable_baselines3/common/env_checker.py +++ b/stable_baselines3/common/env_checker.py @@ -8,12 +8,12 @@ from gym import spaces from stable_baselines3.common.vec_env import DummyVecEnv, VecCheckNan -def _enforce_array_obs(observation_space: spaces.Space) -> bool: +def _is_numpy_array_space(space: spaces.Space) -> bool: """ - Whether to check that the returned observation is a numpy array - it is not mandatory for `Dict` and `Tuple` spaces. + Returns False if provided space is not representable as a single numpy array + (e.g. Dict and Tuple spaces return False) """ - return not isinstance(observation_space, (spaces.Dict, spaces.Tuple)) + return not isinstance(space, (spaces.Dict, spaces.Tuple)) def _check_image_input(observation_space: spaces.Box) -> None: @@ -45,8 +45,8 @@ def _check_image_input(observation_space: spaces.Box) -> None: ) -def _check_unsupported_obs_spaces(env: gym.Env, observation_space: spaces.Space) -> None: - """Emit warnings when the observation space used is not supported by Stable-Baselines.""" +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.""" if isinstance(observation_space, spaces.Dict) and not isinstance(env, gym.GoalEnv): warnings.warn( @@ -65,6 +65,13 @@ def _check_unsupported_obs_spaces(env: gym.Env, observation_space: spaces.Space) "you will need to flatten the observation and maybe use a custom policy. " ) + 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." + ) + def _check_nan(env: gym.Env) -> None: """Check for Inf and NaN using the VecWrapper.""" @@ -87,7 +94,7 @@ def _check_obs(obs: Union[tuple, dict, np.ndarray, int], observation_space: spac # The check for a GoalEnv is done by the base class if isinstance(observation_space, spaces.Discrete): assert isinstance(obs, int), "The observation returned by `{}()` method must be an int".format(method_name) - elif _enforce_array_obs(observation_space): + elif _is_numpy_array_space(observation_space): assert isinstance(obs, np.ndarray), "The observation returned by `{}()` method must be a numpy array".format( method_name ) @@ -201,7 +208,7 @@ def check_env(env: gym.Env, warn: bool = True, skip_render_check: bool = True) - # Warn the user if needed. # A warning means that the environment may run but not work properly with Stable Baselines algorithms if warn: - _check_unsupported_obs_spaces(env, observation_space) + _check_unsupported_spaces(env, observation_space, action_space) # If image, check the low and high values, the type and the number of channels # and the shape (minimal value) @@ -234,5 +241,5 @@ def check_env(env: gym.Env, warn: bool = True, skip_render_check: bool = True) - _check_render(env, warn=warn) # The check only works with numpy arrays - if _enforce_array_obs(observation_space): + if _is_numpy_array_space(observation_space) and _is_numpy_array_space(action_space): _check_nan(env) diff --git a/tests/test_env_checker.py b/tests/test_env_checker.py new file mode 100644 index 0000000..6364bd4 --- /dev/null +++ b/tests/test_env_checker.py @@ -0,0 +1,31 @@ +import gym +import numpy as np +import pytest +from gym.spaces import Box, Dict, Discrete + +from stable_baselines3.common.env_checker import check_env + + +class ActionDictTestEnv(gym.Env): + action_space = Dict({"position": Discrete(1), "velocity": Discrete(1)}) + observation_space = Box(low=-1.0, high=2.0, shape=(3,), dtype=np.float32) + + def step(self, action): + observation = np.array([1.0, 1.5, 0.5]) + reward = 1 + done = True + info = {} + return observation, reward, done, info + + def reset(self): + return np.array([1.0, 1.5, 0.5]) + + def render(self, mode="human"): + pass + + +def test_check_env_dict_action(): + test_env = ActionDictTestEnv() + + with pytest.warns(Warning): + check_env(env=test_env, warn=True)