From 2b9fc1f923a764443068aaaaa9341478d9feb92c Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Sun, 6 Dec 2020 13:05:10 +0100 Subject: [PATCH] Add supported action spaces checks (#254) * Add supported action spaces checks * Address comment --- docs/misc/changelog.rst | 1 + stable_baselines3/a2c/a2c.py | 6 ++++++ stable_baselines3/common/base_class.py | 12 ++++++++++-- .../common/off_policy_algorithm.py | 3 +++ .../common/on_policy_algorithm.py | 3 +++ stable_baselines3/dqn/dqn.py | 2 ++ stable_baselines3/ppo/ppo.py | 6 ++++++ stable_baselines3/sac/sac.py | 2 ++ stable_baselines3/td3/td3.py | 2 ++ tests/test_spaces.py | 19 ++++++++++++++++++- 10 files changed, 53 insertions(+), 3 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 77b02cd..a690651 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -42,6 +42,7 @@ Others: - Add more issue templates - Add signatures to callable type annotations (@erniejunior) - Improve error message in ``NatureCNN`` +- Added checks for supported action spaces to improve clarity of error messages for the user Documentation: ^^^^^^^^^^^^^^ diff --git a/stable_baselines3/a2c/a2c.py b/stable_baselines3/a2c/a2c.py index d2dd7f1..b88b01f 100644 --- a/stable_baselines3/a2c/a2c.py +++ b/stable_baselines3/a2c/a2c.py @@ -96,6 +96,12 @@ class A2C(OnPolicyAlgorithm): create_eval_env=create_eval_env, seed=seed, _init_setup_model=False, + supported_action_spaces=( + spaces.Box, + spaces.Discrete, + spaces.MultiDiscrete, + spaces.MultiBinary, + ), ) self.normalize_advantage = normalize_advantage diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index 32c1ce5..59f1a3b 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -79,6 +79,7 @@ class BaseAlgorithm(ABC): instead of action noise exploration (default: False) :param sde_sample_freq: Sample a new noise matrix every n steps when using gSDE Default: -1 (only sample at the beginning of the rollout) + :param supported_action_spaces: The action spaces supported by the algorithm. """ def __init__( @@ -97,6 +98,7 @@ class BaseAlgorithm(ABC): seed: Optional[int] = None, use_sde: bool = False, sde_sample_freq: int = -1, + supported_action_spaces: Optional[Tuple[gym.spaces.Space, ...]] = None, ): if isinstance(policy, str) and policy_base is not None: @@ -158,13 +160,19 @@ class BaseAlgorithm(ABC): self.n_envs = env.num_envs self.env = env + if supported_action_spaces is not None: + assert isinstance(self.action_space, supported_action_spaces), ( + f"The algorithm only supports {supported_action_spaces} as action spaces " + f"but {self.action_space} was provided" + ) + if not support_multi_env and self.n_envs > 1: raise ValueError( "Error: the model does not support multiple envs; it requires " "a single vectorized environment." ) - if self.use_sde and not isinstance(self.action_space, gym.spaces.Box): - raise ValueError("generalized State-Dependent Exploration (gSDE) can only be used with continuous actions.") + if self.use_sde and not isinstance(self.action_space, gym.spaces.Box): + raise ValueError("generalized State-Dependent Exploration (gSDE) can only be used with continuous actions.") @staticmethod def _wrap_env(env: GymEnv, verbose: int = 0, monitor_wrapper: bool = True) -> VecEnv: diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 4545610..6cea454 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/stable_baselines3/common/off_policy_algorithm.py @@ -69,6 +69,7 @@ class OffPolicyAlgorithm(BaseAlgorithm): :param sde_support: Whether the model support gSDE or not :param remove_time_limit_termination: Remove terminations (dones) that are due to time limit. See https://github.com/hill-a/stable-baselines/issues/863 + :param supported_action_spaces: The action spaces supported by the algorithm. """ def __init__( @@ -100,6 +101,7 @@ class OffPolicyAlgorithm(BaseAlgorithm): use_sde_at_warmup: bool = False, sde_support: bool = True, remove_time_limit_termination: bool = False, + supported_action_spaces: Optional[Tuple[gym.spaces.Space, ...]] = None, ): super(OffPolicyAlgorithm, self).__init__( @@ -117,6 +119,7 @@ class OffPolicyAlgorithm(BaseAlgorithm): seed=seed, use_sde=use_sde, sde_sample_freq=sde_sample_freq, + supported_action_spaces=supported_action_spaces, ) self.buffer_size = buffer_size self.batch_size = batch_size diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index 9f7a665..0a930d9 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/stable_baselines3/common/on_policy_algorithm.py @@ -46,6 +46,7 @@ class OnPolicyAlgorithm(BaseAlgorithm): :param device: Device (cpu, cuda, ...) on which the code should be run. Setting it to auto, the code will be run on the GPU if possible. :param _init_setup_model: Whether or not to build the network at the creation of the instance + :param supported_action_spaces: The action spaces supported by the algorithm. """ def __init__( @@ -69,6 +70,7 @@ class OnPolicyAlgorithm(BaseAlgorithm): seed: Optional[int] = None, device: Union[th.device, str] = "auto", _init_setup_model: bool = True, + supported_action_spaces: Optional[Tuple[gym.spaces.Space, ...]] = None, ): super(OnPolicyAlgorithm, self).__init__( @@ -85,6 +87,7 @@ class OnPolicyAlgorithm(BaseAlgorithm): support_multi_env=True, seed=seed, tensorboard_log=tensorboard_log, + supported_action_spaces=supported_action_spaces, ) self.n_steps = n_steps diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index 180651a..045c377 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/stable_baselines3/dqn/dqn.py @@ -1,5 +1,6 @@ from typing import Any, Dict, List, Optional, Tuple, Type, Union +import gym import numpy as np import torch as th from torch.nn import functional as F @@ -105,6 +106,7 @@ class DQN(OffPolicyAlgorithm): seed=seed, sde_support=False, optimize_memory_usage=optimize_memory_usage, + supported_action_spaces=(gym.spaces.Discrete,), ) self.exploration_initial_eps = exploration_initial_eps diff --git a/stable_baselines3/ppo/ppo.py b/stable_baselines3/ppo/ppo.py index a2b6aea..52579b8 100644 --- a/stable_baselines3/ppo/ppo.py +++ b/stable_baselines3/ppo/ppo.py @@ -108,6 +108,12 @@ class PPO(OnPolicyAlgorithm): create_eval_env=create_eval_env, seed=seed, _init_setup_model=False, + supported_action_spaces=( + spaces.Box, + spaces.Discrete, + spaces.MultiDiscrete, + spaces.MultiBinary, + ), ) self.batch_size = batch_size diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index e94249f..a0c299a 100644 --- a/stable_baselines3/sac/sac.py +++ b/stable_baselines3/sac/sac.py @@ -1,5 +1,6 @@ from typing import Any, Dict, List, Optional, Tuple, Type, Union +import gym import numpy as np import torch as th from torch.nn import functional as F @@ -124,6 +125,7 @@ class SAC(OffPolicyAlgorithm): sde_sample_freq=sde_sample_freq, use_sde_at_warmup=use_sde_at_warmup, optimize_memory_usage=optimize_memory_usage, + supported_action_spaces=(gym.spaces.Box), ) self.target_entropy = target_entropy diff --git a/stable_baselines3/td3/td3.py b/stable_baselines3/td3/td3.py index 2c2d273..ed74830 100644 --- a/stable_baselines3/td3/td3.py +++ b/stable_baselines3/td3/td3.py @@ -1,5 +1,6 @@ from typing import Any, Dict, List, Optional, Tuple, Type, Union +import gym import numpy as np import torch as th from torch.nn import functional as F @@ -107,6 +108,7 @@ class TD3(OffPolicyAlgorithm): seed=seed, sde_support=False, optimize_memory_usage=optimize_memory_usage, + supported_action_spaces=(gym.spaces.Box), ) self.policy_delay = policy_delay diff --git a/tests/test_spaces.py b/tests/test_spaces.py index 8b1feb3..b98ec81 100644 --- a/tests/test_spaces.py +++ b/tests/test_spaces.py @@ -2,7 +2,7 @@ import gym import numpy as np import pytest -from stable_baselines3 import DQN, SAC, TD3 +from stable_baselines3 import A2C, DDPG, DQN, PPO, SAC, TD3 from stable_baselines3.common.evaluation import evaluate_policy @@ -49,3 +49,20 @@ def test_identity_spaces(model_class, env): model.learn(total_timesteps=500) evaluate_policy(model, env, n_eval_episodes=5, warn=False) + + +@pytest.mark.parametrize("model_class", [A2C, DDPG, DQN, PPO, SAC, TD3]) +@pytest.mark.parametrize("env", ["Pendulum-v0", "CartPole-v1"]) +def test_action_spaces(model_class, env): + if model_class in [SAC, DDPG, TD3]: + supported_action_space = env == "Pendulum-v0" + elif model_class == DQN: + supported_action_space = env == "CartPole-v1" + elif model_class in [A2C, PPO]: + supported_action_space = True + + if supported_action_space: + model_class("MlpPolicy", env) + else: + with pytest.raises(AssertionError): + model_class("MlpPolicy", env)