From 42ef6d4677511d5d35d4f911447ce37739660f90 Mon Sep 17 00:00:00 2001 From: Sam Toyer Date: Sun, 23 Aug 2020 04:27:52 -0700 Subject: [PATCH] Remove "device" argument from policies (#141) * Remove device arg from policies * Clean up for PR * Update test and doc * Fix codestyle Co-authored-by: Antonin RAFFIN --- docs/index.rst | 9 ++++++ docs/misc/changelog.rst | 3 +- .../common/off_policy_algorithm.py | 1 - .../common/on_policy_algorithm.py | 1 - stable_baselines3/common/policies.py | 31 +++++++------------ stable_baselines3/dqn/policies.py | 17 ++-------- stable_baselines3/sac/policies.py | 12 +------ stable_baselines3/td3/policies.py | 12 +------ stable_baselines3/version.txt | 2 +- tests/test_predict.py | 15 +++++++-- 10 files changed, 40 insertions(+), 63 deletions(-) diff --git a/docs/index.rst b/docs/index.rst index 5dcb89e..939655a 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -98,6 +98,15 @@ To cite this project in publications: howpublished = {\url{https://github.com/DLR-RM/stable-baselines3}}, } +Contributing +------------ + +To any interested in making the rl baselines better, there are still some improvements +that need to be done. +You can check issues in the `repo `_. + +If you want to contribute, please read `CONTRIBUTING.md `_ first. + Indices and tables ------------------- diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 8e78485..35b9ec9 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,11 +3,12 @@ Changelog ========== -Pre-Release 0.9.0a0 (WIP) +Pre-Release 0.9.0a1 (WIP) ------------------------------ Breaking Changes: ^^^^^^^^^^^^^^^^^ +- Removed ``device`` keyword argument of policies; use ``policy.to(device)`` instead. (@qxcv) New Features: ^^^^^^^^^^^^^ diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index b2814df..167f34d 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/stable_baselines3/common/off_policy_algorithm.py @@ -141,7 +141,6 @@ class OffPolicyAlgorithm(BaseAlgorithm): # Update policy keyword arguments if sde_support: self.policy_kwargs["use_sde"] = self.use_sde - self.policy_kwargs["device"] = self.device # For gSDE only self.use_sde_at_warmup = use_sde_at_warmup diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index 8c9cb8b..5189991 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/stable_baselines3/common/on_policy_algorithm.py @@ -116,7 +116,6 @@ class OnPolicyAlgorithm(BaseAlgorithm): self.action_space, self.lr_schedule, use_sde=self.use_sde, - device=self.device, **self.policy_kwargs # pytype:disable=not-instantiable ) self.policy = self.policy.to(self.device) diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index 01b788f..41280a8 100644 --- a/stable_baselines3/common/policies.py +++ b/stable_baselines3/common/policies.py @@ -34,7 +34,6 @@ class BaseModel(nn.Module, ABC): :param observation_space: (gym.spaces.Space) The observation space of the environment :param action_space: (gym.spaces.Space) The action space of the environment - :param device: (Union[th.device, str]) Device on which the code should run. :param features_extractor_class: (Type[BaseFeaturesExtractor]) Features extractor to use. :param features_extractor_kwargs: (Optional[Dict[str, Any]]) Keyword arguments to pass to the feature extractor. @@ -52,7 +51,6 @@ class BaseModel(nn.Module, ABC): self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, - device: Union[th.device, str] = "auto", features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor, features_extractor_kwargs: Optional[Dict[str, Any]] = None, features_extractor: Optional[nn.Module] = None, @@ -70,7 +68,6 @@ class BaseModel(nn.Module, ABC): self.observation_space = observation_space self.action_space = action_space - self.device = get_device(device) self.features_extractor = features_extractor self.normalize_images = normalize_images @@ -112,6 +109,16 @@ class BaseModel(nn.Module, ABC): normalize_images=self.normalize_images, ) + @property + def device(self) -> th.device: + """Infer which device this policy lives on by inspecting its parameters. + If it has no parameters, the 'auto' device is used as a fallback. + + :return: (th.device)""" + for param in self.parameters(): + return param.device + return get_device("auto") + def save(self, path: str) -> None: """ Save model to a given location. @@ -300,7 +307,6 @@ class ActorCriticPolicy(BasePolicy): :param action_space: (gym.spaces.Space) Action space :param lr_schedule: (Callable) Learning rate schedule (could be constant) :param net_arch: ([int or dict]) The specification of the policy and value networks. - :param device: (str or th.device) Device on which the code should run. :param activation_fn: (Type[nn.Module]) Activation function :param ortho_init: (bool) Whether to use or not orthogonal initialization :param use_sde: (bool) Whether to use State Dependent Exploration or not @@ -332,7 +338,6 @@ class ActorCriticPolicy(BasePolicy): action_space: gym.spaces.Space, lr_schedule: Callable[[float], float], net_arch: Optional[List[Union[int, Dict[str, List[int]]]]] = None, - device: Union[th.device, str] = "auto", activation_fn: Type[nn.Module] = nn.Tanh, ortho_init: bool = True, use_sde: bool = False, @@ -357,7 +362,6 @@ class ActorCriticPolicy(BasePolicy): super(ActorCriticPolicy, self).__init__( observation_space, action_space, - device, features_extractor_class, features_extractor_kwargs, optimizer_class=optimizer_class, @@ -445,9 +449,7 @@ class ActorCriticPolicy(BasePolicy): # Note: If net_arch is None and some features extractor is used, # net_arch here is an empty list and mlp_extractor does not # really contain any layers (acts like an identity module). - self.mlp_extractor = MlpExtractor( - self.features_dim, net_arch=self.net_arch, activation_fn=self.activation_fn, device=self.device - ) + self.mlp_extractor = MlpExtractor(self.features_dim, net_arch=self.net_arch, activation_fn=self.activation_fn) latent_dim_pi = self.mlp_extractor.latent_dim_pi @@ -594,7 +596,6 @@ class ActorCriticCnnPolicy(ActorCriticPolicy): :param action_space: (gym.spaces.Space) Action space :param lr_schedule: (Callable) Learning rate schedule (could be constant) :param net_arch: ([int or dict]) The specification of the policy and value networks. - :param device: (str or th.device) Device on which the code should run. :param activation_fn: (Type[nn.Module]) Activation function :param ortho_init: (bool) Whether to use or not orthogonal initialization :param use_sde: (bool) Whether to use State Dependent Exploration or not @@ -626,7 +627,6 @@ class ActorCriticCnnPolicy(ActorCriticPolicy): action_space: gym.spaces.Space, lr_schedule: Callable, net_arch: Optional[List[Union[int, Dict[str, List[int]]]]] = None, - device: Union[th.device, str] = "auto", activation_fn: Type[nn.Module] = nn.Tanh, ortho_init: bool = True, use_sde: bool = False, @@ -646,7 +646,6 @@ class ActorCriticCnnPolicy(ActorCriticPolicy): action_space, lr_schedule, net_arch, - device, activation_fn, ortho_init, use_sde, @@ -685,7 +684,6 @@ class ContinuousCritic(BaseModel): :param activation_fn: (Type[nn.Module]) Activation function :param normalize_images: (bool) Whether to normalize images or not, dividing by 255.0 (True by default) - :param device: (Union[th.device, str]) Device on which the code should run. :param n_critics: (int) Number of critic networks to create. """ @@ -698,15 +696,10 @@ class ContinuousCritic(BaseModel): features_dim: int, activation_fn: Type[nn.Module] = nn.ReLU, normalize_images: bool = True, - device: Union[th.device, str] = "auto", n_critics: int = 2, ): super().__init__( - observation_space, - action_space, - features_extractor=features_extractor, - normalize_images=normalize_images, - device=device, + observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images, ) action_dim = get_action_dim(self.action_space) diff --git a/stable_baselines3/dqn/policies.py b/stable_baselines3/dqn/policies.py index 495b61f..f5001c7 100644 --- a/stable_baselines3/dqn/policies.py +++ b/stable_baselines3/dqn/policies.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Dict, List, Optional, Type, Union +from typing import Any, Callable, Dict, List, Optional, Type import gym import torch as th @@ -15,7 +15,6 @@ class QNetwork(BasePolicy): :param observation_space: (gym.spaces.Space) Observation space :param action_space: (gym.spaces.Space) Action space :param net_arch: (Optional[List[int]]) The specification of the policy and value networks. - :param device: (str or th.device) Device on which the code should run. :param activation_fn: (Type[nn.Module]) Activation function :param normalize_images: (bool) Whether to normalize images or not, dividing by 255.0 (True by default) @@ -28,16 +27,11 @@ class QNetwork(BasePolicy): features_extractor: nn.Module, features_dim: int, net_arch: Optional[List[int]] = None, - device: Union[th.device, str] = "auto", activation_fn: Type[nn.Module] = nn.ReLU, normalize_images: bool = True, ): super(QNetwork, self).__init__( - observation_space, - action_space, - features_extractor=features_extractor, - normalize_images=normalize_images, - device=device, + observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images, ) if net_arch is None: @@ -90,7 +84,6 @@ class DQNPolicy(BasePolicy): :param action_space: (gym.spaces.Space) Action space :param lr_schedule: (callable) Learning rate schedule (could be constant) :param net_arch: (Optional[List[int]]) The specification of the policy and value networks. - :param device: (str or th.device) Device on which the code should run. :param activation_fn: (Type[nn.Module]) Activation function :param features_extractor_class: (Type[BaseFeaturesExtractor]) Features extractor to use. :param features_extractor_kwargs: (Optional[Dict[str, Any]]) Keyword arguments @@ -109,7 +102,6 @@ class DQNPolicy(BasePolicy): action_space: gym.spaces.Space, lr_schedule: Callable, net_arch: Optional[List[int]] = None, - device: Union[th.device, str] = "auto", activation_fn: Type[nn.Module] = nn.ReLU, features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor, features_extractor_kwargs: Optional[Dict[str, Any]] = None, @@ -120,7 +112,6 @@ class DQNPolicy(BasePolicy): super(DQNPolicy, self).__init__( observation_space, action_space, - device, features_extractor_class, features_extractor_kwargs, optimizer_class=optimizer_class, @@ -143,7 +134,6 @@ class DQNPolicy(BasePolicy): "net_arch": self.net_arch, "activation_fn": self.activation_fn, "normalize_images": normalize_images, - "device": device, } self.q_net, self.q_net_target = None, None @@ -204,7 +194,6 @@ class CnnPolicy(DQNPolicy): :param action_space: (gym.spaces.Space) Action space :param lr_schedule: (callable) Learning rate schedule (could be constant) :param net_arch: (Optional[List[int]]) The specification of the policy and value networks. - :param device: (str or th.device) Device on which the code should run. :param activation_fn: (Type[nn.Module]) Activation function :param features_extractor_class: (Type[BaseFeaturesExtractor]) Features extractor to use. :param normalize_images: (bool) Whether to normalize images or not, @@ -221,7 +210,6 @@ class CnnPolicy(DQNPolicy): action_space: gym.spaces.Space, lr_schedule: Callable, net_arch: Optional[List[int]] = None, - device: Union[th.device, str] = "auto", activation_fn: Type[nn.Module] = nn.ReLU, features_extractor_class: Type[BaseFeaturesExtractor] = NatureCNN, features_extractor_kwargs: Optional[Dict[str, Any]] = None, @@ -234,7 +222,6 @@ class CnnPolicy(DQNPolicy): action_space, lr_schedule, net_arch, - device, activation_fn, features_extractor_class, features_extractor_kwargs, diff --git a/stable_baselines3/sac/policies.py b/stable_baselines3/sac/policies.py index b5a6667..843ee06 100644 --- a/stable_baselines3/sac/policies.py +++ b/stable_baselines3/sac/policies.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, Type import gym import torch as th @@ -38,7 +38,6 @@ class Actor(BasePolicy): :param clip_mean: (float) Clip the mean output when using gSDE to avoid numerical instability. :param normalize_images: (bool) Whether to normalize images or not, dividing by 255.0 (True by default) - :param device: (Union[th.device, str]) Device on which the code should run. """ def __init__( @@ -56,14 +55,12 @@ class Actor(BasePolicy): use_expln: bool = False, clip_mean: float = 2.0, normalize_images: bool = True, - device: Union[th.device, str] = "auto", ): super(Actor, self).__init__( observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images, - device=device, squash_output=True, ) @@ -196,7 +193,6 @@ class SACPolicy(BasePolicy): :param action_space: (gym.spaces.Space) Action space :param lr_schedule: (callable) Learning rate schedule (could be constant) :param net_arch: (Optional[List[int]]) The specification of the policy and value networks. - :param device: (str or th.device) Device on which the code should run. :param activation_fn: (Type[nn.Module]) Activation function :param use_sde: (bool) Whether to use State Dependent Exploration or not :param log_std_init: (float) Initial value for the log standard deviation @@ -225,7 +221,6 @@ class SACPolicy(BasePolicy): action_space: gym.spaces.Space, lr_schedule: Callable, net_arch: Optional[List[int]] = None, - device: Union[th.device, str] = "auto", activation_fn: Type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, @@ -242,7 +237,6 @@ class SACPolicy(BasePolicy): super(SACPolicy, self).__init__( observation_space, action_space, - device, features_extractor_class, features_extractor_kwargs, optimizer_class=optimizer_class, @@ -270,7 +264,6 @@ class SACPolicy(BasePolicy): "net_arch": self.net_arch, "activation_fn": self.activation_fn, "normalize_images": normalize_images, - "device": device, } self.actor_kwargs = self.net_args.copy() sde_kwargs = { @@ -356,7 +349,6 @@ class CnnPolicy(SACPolicy): :param action_space: (gym.spaces.Space) Action space :param lr_schedule: (callable) Learning rate schedule (could be constant) :param net_arch: (Optional[List[int]]) The specification of the policy and value networks. - :param device: (str or th.device) Device on which the code should run. :param activation_fn: (Type[nn.Module]) Activation function :param use_sde: (bool) Whether to use State Dependent Exploration or not :param log_std_init: (float) Initial value for the log standard deviation @@ -383,7 +375,6 @@ class CnnPolicy(SACPolicy): action_space: gym.spaces.Space, lr_schedule: Callable, net_arch: Optional[List[int]] = None, - device: Union[th.device, str] = "auto", activation_fn: Type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, @@ -402,7 +393,6 @@ class CnnPolicy(SACPolicy): action_space, lr_schedule, net_arch, - device, activation_fn, use_sde, log_std_init, diff --git a/stable_baselines3/td3/policies.py b/stable_baselines3/td3/policies.py index bcc64a1..2a60043 100644 --- a/stable_baselines3/td3/policies.py +++ b/stable_baselines3/td3/policies.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Dict, List, Optional, Type, Union +from typing import Any, Callable, Dict, List, Optional, Type import gym import torch as th @@ -22,7 +22,6 @@ class Actor(BasePolicy): :param activation_fn: (Type[nn.Module]) Activation function :param normalize_images: (bool) Whether to normalize images or not, dividing by 255.0 (True by default) - :param device: (Union[th.device, str]) Device on which the code should run. """ def __init__( @@ -34,14 +33,12 @@ class Actor(BasePolicy): features_dim: int, activation_fn: Type[nn.Module] = nn.ReLU, normalize_images: bool = True, - device: Union[th.device, str] = "auto", ): super(Actor, self).__init__( observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images, - device=device, squash_output=True, ) @@ -86,7 +83,6 @@ class TD3Policy(BasePolicy): :param action_space: (gym.spaces.Space) Action space :param lr_schedule: (Callable) Learning rate schedule (could be constant) :param net_arch: (Optional[List[int]]) The specification of the policy and value networks. - :param device: (Union[th.device, str]) Device on which the code should run. :param activation_fn: (Type[nn.Module]) Activation function :param features_extractor_class: (Type[BaseFeaturesExtractor]) Features extractor to use. :param features_extractor_kwargs: (Optional[Dict[str, Any]]) Keyword arguments @@ -106,7 +102,6 @@ class TD3Policy(BasePolicy): action_space: gym.spaces.Space, lr_schedule: Callable, net_arch: Optional[List[int]] = None, - device: Union[th.device, str] = "auto", activation_fn: Type[nn.Module] = nn.ReLU, features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor, features_extractor_kwargs: Optional[Dict[str, Any]] = None, @@ -118,7 +113,6 @@ class TD3Policy(BasePolicy): super(TD3Policy, self).__init__( observation_space, action_space, - device, features_extractor_class, features_extractor_kwargs, optimizer_class=optimizer_class, @@ -146,7 +140,6 @@ class TD3Policy(BasePolicy): "net_arch": self.net_arch, "activation_fn": self.activation_fn, "normalize_images": normalize_images, - "device": device, } self.critic_kwargs = self.net_args.copy() self.critic_kwargs.update({"n_critics": n_critics}) @@ -206,7 +199,6 @@ class CnnPolicy(TD3Policy): :param action_space: (gym.spaces.Space) Action space :param lr_schedule: (Callable) Learning rate schedule (could be constant) :param net_arch: (Optional[List[int]]) The specification of the policy and value networks. - :param device: (Union[th.device, str]) Device on which the code should run. :param activation_fn: (Type[nn.Module]) Activation function :param features_extractor_class: (Type[BaseFeaturesExtractor]) Features extractor to use. :param features_extractor_kwargs: (Optional[Dict[str, Any]]) Keyword arguments @@ -226,7 +218,6 @@ class CnnPolicy(TD3Policy): action_space: gym.spaces.Space, lr_schedule: Callable, net_arch: Optional[List[int]] = None, - device: Union[th.device, str] = "auto", activation_fn: Type[nn.Module] = nn.ReLU, features_extractor_class: Type[BaseFeaturesExtractor] = NatureCNN, features_extractor_kwargs: Optional[Dict[str, Any]] = None, @@ -240,7 +231,6 @@ class CnnPolicy(TD3Policy): action_space, lr_schedule, net_arch, - device, activation_fn, features_extractor_class, features_extractor_kwargs, diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 657e7c0..8d8f786 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -0.9.0a0 +0.9.0a1 diff --git a/tests/test_predict.py b/tests/test_predict.py index 4f60f1b..288ab0d 100644 --- a/tests/test_predict.py +++ b/tests/test_predict.py @@ -1,7 +1,9 @@ import gym import pytest +import torch as th from stable_baselines3 import A2C, DQN, PPO, SAC, TD3 +from stable_baselines3.common.utils import get_device from stable_baselines3.common.vec_env import DummyVecEnv MODEL_LIST = [ @@ -30,15 +32,22 @@ def test_auto_wrap(model_class): @pytest.mark.parametrize("model_class", MODEL_LIST) @pytest.mark.parametrize("env_id", ["Pendulum-v0", "CartPole-v1"]) -def test_predict(model_class, env_id): +@pytest.mark.parametrize("device", ["cpu", "cuda", "auto"]) +def test_predict(model_class, env_id, device): + if device == "cuda" and not th.cuda.is_available(): + pytest.skip("CUDA not available") + if env_id == "CartPole-v1": if model_class in [SAC, TD3]: return elif model_class in [DQN]: return - # test detection of different shapes by the predict method - model = model_class("MlpPolicy", env_id) + # Test detection of different shapes by the predict method + model = model_class("MlpPolicy", env_id, device=device) + # Check that the policy is on the right device + assert get_device(device) == model.policy.device + env = gym.make(env_id) vec_env = DummyVecEnv([lambda: gym.make(env_id), lambda: gym.make(env_id)])