Add supported action spaces checks (#254)

* Add supported action spaces checks

* Address comment

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:
 ^^^^^^^^^^^^^^

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

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,6 +160,12 @@ 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."

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

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

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

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

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

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

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)