From 306e49fda6a3f47027be96a78464e96e1cbf5e26 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Tue, 28 Sep 2021 21:57:49 +0200 Subject: [PATCH] Fixes in `is_vectorized_observation` (#587) * Fix is vectorized bug in DQN * Fix sub-classed obs --- docs/misc/changelog.rst | 2 ++ stable_baselines3/common/utils.py | 19 +++++++-------- tests/test_predict.py | 39 +++++++++++++++++++++++++++++++ 3 files changed, 50 insertions(+), 10 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 91d110e..d7f4775 100644 --- a/docs/misc/changelog.rst +++ b/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: ^^^^^^^^^^^^^ diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index 8711591..be5d488 100644 --- a/stable_baselines3/common/utils.py +++ b/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: diff --git a/tests/test_predict.py b/tests/test_predict.py index 2927796..436547b 100644 --- a/tests/test_predict.py +++ b/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))