Improve tests and add check for float32 (#686)

* Add additional checks

* Improve tests and error message

* Update changelog

* Bump version

* Update doc

* Add tests for action space

* Improve test

docs/guide/rl_tips.rst

@@ -8,6 +8,11 @@ The aim of this section is to help you doing reinforcement learning experiments.
 It covers general advice about RL (where to start, which algorithm to choose, how to evaluate an algorithm, ...),
 as well as tips and tricks when using a custom environment or implementing an RL algorithm.
 
+.. note::
+
+  We have a `video on YouTube <https://www.youtube.com/watch?v=Ikngt0_DXJg>`_ that covers
+  this section in more details. You can also find the `slides here <https://araffin.github.io/slides/rlvs-tips-tricks/>`_.
+
 
 General advice when using Reinforcement Learning
 ================================================

docs/misc/changelog.rst

@@ -4,7 +4,7 @@ Changelog
 ==========
 
 
-Release 1.3.1a5 (WIP)
+Release 1.3.1a6 (WIP)
 ---------------------------
 
 Breaking Changes:
@@ -27,12 +27,16 @@ Bug Fixes:
 - Fixed a bug where ``set_env()`` with ``VecNormalize`` would result in an error with off-policy algorithms (thanks @cleversonahum)
 - FPS calculation is now performed based on number of steps performed during last ``learn`` call, even when ``reset_num_timesteps`` is set to ``False`` (@kachayev)
 - Fixed evaluation script for recurrent policies (experimental feature in SB3 contrib)
+- Fixed a bug where the observation would be incorrectly detected as non-vectorized instead of throwing an error
+- The env checker now properly checks and warns about potential issues for continuous action spaces when the boundaries are too small or when the dtype is not float32
 
 Deprecations:
 ^^^^^^^^^^^^^
 
 Others:
 ^^^^^^^
+- Added a warning in the env checker when not using ``np.float32`` for continuous actions
+- Improved test coverage and error message when checking shape of observation
 
 Documentation:
 ^^^^^^^^^^^^^^
@@ -43,7 +47,7 @@ Documentation:
 - Update GAE computation docstring
 - Add documentation on exporting to TFLite/Coral
 - Added JMLR paper and updated citation
-
+- Added link to RL Tips and Tricks video
 
 Release 1.3.0 (2021-10-23)
 ---------------------------

stable_baselines3/common/env_checker.py

@@ -266,14 +266,19 @@ def check_env(env: gym.Env, warn: bool = True, skip_render_check: bool = True) -
         # Check for the action space, it may lead to hard-to-debug issues
         if isinstance(action_space, spaces.Box) and (
             np.any(np.abs(action_space.low) != np.abs(action_space.high))
-            or np.any(np.abs(action_space.low) > 1)
-            or np.any(np.abs(action_space.high) > 1)
+            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) 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 returned values ===============
     _check_returned_values(env, observation_space, action_space)
 

stable_baselines3/common/utils.py

@@ -262,7 +262,7 @@ def is_vectorized_discrete_observation(observation: Union[int, np.ndarray], obse
     else:
         raise ValueError(
             f"Error: Unexpected observation shape {observation.shape} for "
-            + "Discrete environment, please use (1,) or (n_env, 1) for the observation shape."
+            + "Discrete environment, please use () or (n_env,) for the observation shape."
         )
 
 
@@ -317,23 +317,38 @@ def is_vectorized_dict_observation(observation: np.ndarray, observation_space: g
     :param observation_space: the observation space
     :return: whether the given observation is vectorized or not
     """
+    # We first assume that all observations are not vectorized
+    all_non_vectorized = True
     for key, subspace in observation_space.spaces.items():
-        if observation[key].shape == subspace.shape:
-            return False
-
-    all_good = True
-
-    for key, subspace in observation_space.spaces.items():
-        if observation[key].shape[1:] != subspace.shape:
-            all_good = False
+        # This fails when the observation is not vectorized
+        # or when it has the wrong shape
+        if observation[key].shape != subspace.shape:
+            all_non_vectorized = False
             break
 
-    if all_good:
+    if all_non_vectorized:
+        return False
+
+    all_vectorized = True
+    # Now we check that all observation are vectorized and have the correct shape
+    for key, subspace in observation_space.spaces.items():
+        if observation[key].shape[1:] != subspace.shape:
+            all_vectorized = False
+            break
+
+    if all_vectorized:
         return True
     else:
+        # Retrieve error message
+        error_msg = ""
+        try:
+            is_vectorized_observation(observation[key], observation_space.spaces[key])
+        except ValueError as e:
+            error_msg = f"{e}"
         raise ValueError(
-            f"Error: Unexpected observation shape {observation[key].shape} for key {key}, "
-            + f"please use {observation_space.spaces[key]} "
+            f"There seems to be a mix of vectorized and non-vectorized observations. "
+            f"Unexpected observation shape {observation[key].shape} for key {key} "
+            f"of type {observation_space.spaces[key]}. {error_msg}"
         )
 
 

stable_baselines3/version.txt

@@ -1 +1 @@
-1.3.1a5
+1.3.1a6

tests/test_envs.py

@@ -125,6 +125,40 @@ def test_non_default_spaces(new_obs_space):
         check_env(env)
 
 
+@pytest.mark.parametrize(
+    "new_action_space",
+    [
+        # Not symmetric
+        spaces.Box(low=0, high=1, shape=(3,), dtype=np.float32),
+        # Wrong dtype
+        spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float64),
+        # Too big range
+        spaces.Box(low=-1000, high=1000, shape=(3,), dtype=np.float32),
+        # Too small range
+        spaces.Box(low=-0.1, high=0.1, shape=(2,), dtype=np.float32),
+        # Inverted boundaries
+        spaces.Box(low=1, high=-1, shape=(2,), dtype=np.float32),
+        # Same boundaries
+        spaces.Box(low=1, high=1, shape=(2,), dtype=np.float32),
+        # Almost good, except for one dim
+        spaces.Box(low=np.array([-1, -1, -1]), high=np.array([1, 1, 0.99]), dtype=np.float32),
+    ],
+)
+def test_non_default_action_spaces(new_action_space):
+    env = FakeImageEnv(discrete=False)
+    # Default, should pass the test
+    with pytest.warns(None) as record:
+        check_env(env)
+
+    # No warnings for custom envs
+    assert len(record) == 0
+    # Change the action space
+    env.action_space = new_action_space
+
+    with pytest.warns(UserWarning):
+        check_env(env)
+
+
 def check_reset_assert_error(env, new_reset_return):
     """
     Helper to check that the error is caught.

tests/test_utils.py

@@ -5,6 +5,7 @@ import gym
 import numpy as np
 import pytest
 import torch as th
+from gym import spaces
 
 import stable_baselines3 as sb3
 from stable_baselines3 import A2C, PPO
@@ -13,7 +14,7 @@ from stable_baselines3.common.env_util import is_wrapped, make_atari_env, make_v
 from stable_baselines3.common.evaluation import evaluate_policy
 from stable_baselines3.common.monitor import Monitor
 from stable_baselines3.common.noise import ActionNoise, OrnsteinUhlenbeckActionNoise, VectorizedActionNoise
-from stable_baselines3.common.utils import get_system_info, polyak_update, zip_strict
+from stable_baselines3.common.utils import get_system_info, is_vectorized_observation, polyak_update, zip_strict
 from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv
 
 
@@ -387,3 +388,79 @@ def test_get_system_info():
     assert "GPU Enabled" in info_str
     assert "Numpy" in info_str
     assert "Gym" in info_str
+
+
+def test_is_vectorized_observation():
+    # with pytest.raises("ValueError"):
+    #     pass
+    # All vectorized
+    box_space = spaces.Box(-1, 1, shape=(2,))
+    box_obs = np.ones((1,) + box_space.shape)
+    assert is_vectorized_observation(box_obs, box_space)
+
+    discrete_space = spaces.Discrete(2)
+    discrete_obs = np.ones((3,), dtype=np.int8)
+    assert is_vectorized_observation(discrete_obs, discrete_space)
+
+    multidiscrete_space = spaces.MultiDiscrete([2, 3])
+    multidiscrete_obs = np.ones((1, 2), dtype=np.int8)
+    assert is_vectorized_observation(multidiscrete_obs, multidiscrete_space)
+
+    multibinary_space = spaces.MultiBinary(3)
+    multibinary_obs = np.ones((1, 3), dtype=np.int8)
+    assert is_vectorized_observation(multibinary_obs, multibinary_space)
+
+    dict_space = spaces.Dict({"box": box_space, "discrete": discrete_space})
+    dict_obs = {"box": box_obs, "discrete": discrete_obs}
+    assert is_vectorized_observation(dict_obs, dict_space)
+
+    # All not vectorized
+    box_obs = np.ones(box_space.shape)
+    assert not is_vectorized_observation(box_obs, box_space)
+
+    discrete_obs = np.ones((), dtype=np.int8)
+    assert not is_vectorized_observation(discrete_obs, discrete_space)
+
+    multidiscrete_obs = np.ones((2,), dtype=np.int8)
+    assert not is_vectorized_observation(multidiscrete_obs, multidiscrete_space)
+
+    multibinary_obs = np.ones((3,), dtype=np.int8)
+    assert not is_vectorized_observation(multibinary_obs, multibinary_space)
+
+    dict_obs = {"box": box_obs, "discrete": discrete_obs}
+    assert not is_vectorized_observation(dict_obs, dict_space)
+
+    # A mix of vectorized and non-vectorized things
+    with pytest.raises(ValueError):
+        discrete_obs = np.ones((1,), dtype=np.int8)
+        dict_obs = {"box": box_obs, "discrete": discrete_obs}
+        is_vectorized_observation(dict_obs, dict_space)
+
+    # Vectorized with the wrong shape
+    with pytest.raises(ValueError):
+        discrete_obs = np.ones((1,), dtype=np.int8)
+        box_obs = np.ones((1, 2) + box_space.shape)
+        dict_obs = {"box": box_obs, "discrete": discrete_obs}
+        is_vectorized_observation(dict_obs, dict_space)
+
+    # Weird shape: error
+    with pytest.raises(ValueError):
+        discrete_obs = np.ones((1,) + box_space.shape, dtype=np.int8)
+        is_vectorized_observation(discrete_obs, discrete_space)
+
+    # wrong shape
+    with pytest.raises(ValueError):
+        multidiscrete_obs = np.ones((2, 1), dtype=np.int8)
+        is_vectorized_observation(multidiscrete_obs, multidiscrete_space)
+
+    # wrong shape
+    with pytest.raises(ValueError):
+        multibinary_obs = np.ones((2, 1), dtype=np.int8)
+        is_vectorized_observation(multidiscrete_obs, multibinary_space)
+
+    # Almost good shape: one dimension too much for Discrete obs
+    with pytest.raises(ValueError):
+        box_obs = np.ones((1,) + box_space.shape)
+        discrete_obs = np.ones((1, 1), dtype=np.int8)
+        dict_obs = {"box": box_obs, "discrete": discrete_obs}
+        is_vectorized_observation(dict_obs, dict_space)