Fixes in is_vectorized_observation (#587)

* Fix is vectorized bug in DQN

* Fix sub-classed obs

docs/misc/changelog.rst

@@ -23,6 +23,8 @@ Bug Fixes:
 - Fixed ``dtype`` of observations for ``SimpleMultiObsEnv``
 - Allow `VecNormalize` to wrap discrete-observation environments to normalize reward
   when observation normalization is disabled.
+- Fixed a bug where ``DQN`` would throw an error when using ``Discrete`` observation and stochastic actions
+- Fixed a bug where sub-classed observation spaces could not be used
 
 Deprecations:
 ^^^^^^^^^^^^^

stable_baselines3/common/utils.py

@@ -243,7 +243,7 @@ def is_vectorized_box_observation(observation: np.ndarray, observation_space: gy
         )
 
 
-def is_vectorized_discrete_observation(observation: np.ndarray, observation_space: gym.spaces.Discrete) -> bool:
+def is_vectorized_discrete_observation(observation: Union[int, np.ndarray], observation_space: gym.spaces.Discrete) -> bool:
     """
     For discrete observation type, detects and validates the shape,
     then returns whether or not the observation is vectorized.
@@ -252,7 +252,7 @@ def is_vectorized_discrete_observation(observation: np.ndarray, observation_spac
     :param observation_space: the observation space
     :return: whether the given observation is vectorized or not
     """
-    if observation.shape == ():  # A numpy array of a number, has shape empty tuple '()'
+    if isinstance(observation, int) or observation.shape == ():  # A numpy array of a number, has shape empty tuple '()'
         return False
     elif len(observation.shape) == 1:
         return True
@@ -334,7 +334,7 @@ def is_vectorized_dict_observation(observation: np.ndarray, observation_space: g
         )
 
 
-def is_vectorized_observation(observation: np.ndarray, observation_space: gym.spaces.Space) -> bool:
+def is_vectorized_observation(observation: Union[int, np.ndarray], observation_space: gym.spaces.Space) -> bool:
     """
     For every observation type, detects and validates the shape,
     then returns whether or not the observation is vectorized.
@@ -352,13 +352,12 @@ def is_vectorized_observation(observation: np.ndarray, observation_space: gym.sp
         gym.spaces.Dict: is_vectorized_dict_observation,
     }
 
-    try:
-        is_vec_obs_func = is_vec_obs_func_dict[type(observation_space)]
-        return is_vec_obs_func(observation, observation_space)
-    except KeyError:
-        raise ValueError(
-            "Error: Cannot determine if the observation is vectorized " + f" with the space type {observation_space}."
-        )
+    for space_type, is_vec_obs_func in is_vec_obs_func_dict.items():
+        if isinstance(observation_space, space_type):
+            return is_vec_obs_func(observation, observation_space)
+    else:
+        # for-else happens if no break is called
+        raise ValueError(f"Error: Cannot determine if the observation is vectorized with the space type {observation_space}.")
 
 
 def safe_mean(arr: Union[np.ndarray, list, deque]) -> np.ndarray:

tests/test_predict.py

@@ -1,8 +1,10 @@
 import gym
+import numpy as np
 import pytest
 import torch as th
 
 from stable_baselines3 import A2C, DQN, PPO, SAC, TD3
+from stable_baselines3.common.envs import IdentityEnv
 from stable_baselines3.common.utils import get_device
 from stable_baselines3.common.vec_env import DummyVecEnv
 
@@ -15,6 +17,24 @@ MODEL_LIST = [
 ]
 
 
+class SubClassedBox(gym.spaces.Box):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+
+class CustomSubClassedSpaceEnv(gym.Env):
+    def __init__(self):
+        super().__init__()
+        self.observation_space = SubClassedBox(-1, 1, shape=(2,), dtype=np.float32)
+        self.action_space = SubClassedBox(-1, 1, shape=(2,), dtype=np.float32)
+
+    def reset(self):
+        return self.observation_space.sample()
+
+    def step(self, action):
+        return self.observation_space.sample(), 0.0, np.random.rand() > 0.5, {}
+
+
 @pytest.mark.parametrize("model_class", MODEL_LIST)
 def test_auto_wrap(model_class):
     # test auto wrapping of env into a VecEnv
@@ -69,3 +89,22 @@ def test_predict(model_class, env_id, device):
 
         action, _ = model.predict(vec_env_obs, deterministic=False)
         assert action.shape[0] == vec_env_obs.shape[0]
+
+
+def test_dqn_epsilon_greedy():
+    env = IdentityEnv(2)
+    model = DQN("MlpPolicy", env)
+    model.exploration_rate = 1.0
+    obs = env.reset()
+    # is vectorized should not crash with discrete obs
+    action, _ = model.predict(obs, deterministic=False)
+    assert env.action_space.contains(action)
+
+
+@pytest.mark.parametrize("model_class", [A2C, SAC, PPO, TD3])
+def test_subclassed_space_env(model_class):
+    env = CustomSubClassedSpaceEnv()
+    model = model_class("MlpPolicy", env, policy_kwargs=dict(net_arch=[32]))
+    model.learn(300)
+    obs = env.reset()
+    env.step(model.predict(obs))