diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 0d1620d..13c7ca5 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -17,6 +17,7 @@ New Features: Bug Fixes: ^^^^^^^^^^ - Fixed ``reset_num_timesteps`` behavior, so ``env.reset()`` is not called if ``reset_num_timesteps=True`` +- Fixed ``squashed_output`` that was not pass to policy constructor for ``SAC`` and ``TD3`` (would result in scaled actions for unscaled action spaces) Deprecations: ^^^^^^^^^^^^^ @@ -24,6 +25,7 @@ Deprecations: Others: ^^^^^^^ - Cleanup rollout return +- Added ``get_device`` util to manage PyTorch devices Documentation: ^^^^^^^^^^^^^^ diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index fe8f97e..0efd7d5 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -13,7 +13,7 @@ import numpy as np from torchy_baselines.common import logger from torchy_baselines.common.policies import BasePolicy, get_policy_from_name -from torchy_baselines.common.utils import set_random_seed, get_schedule_fn, update_learning_rate +from torchy_baselines.common.utils import set_random_seed, get_schedule_fn, update_learning_rate, get_device from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, VecNormalize from torchy_baselines.common.save_util import data_to_json, json_to_data, recursive_getattr, recursive_setattr from torchy_baselines.common.type_aliases import GymEnv, TensorDict, RolloutReturn, MaybeCallback @@ -71,10 +71,7 @@ class BaseRLModel(ABC): else: self.policy_class = policy - if device == 'auto': - device = 'cuda' if th.cuda.is_available() else 'cpu' - - self.device = th.device(device) + self.device = get_device(device) if verbose > 0: print(f"Using {self.device} device") @@ -387,7 +384,7 @@ class BaseRLModel(ABC): raise ValueError(f"Error: the file {load_path} could not be found") # set device to cpu if cuda is not available - device = th.device('cuda') if th.cuda.is_available() else th.device('cpu') + device = get_device() # Open the zip archive and load data try: diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py index e881139..98ac717 100644 --- a/torchy_baselines/common/policies.py +++ b/torchy_baselines/common/policies.py @@ -8,6 +8,7 @@ import torch.nn as nn import numpy as np from torchy_baselines.common.preprocessing import preprocess_obs +from torchy_baselines.common.utils import get_device, get_schedule_fn class BasePolicy(nn.Module): @@ -18,7 +19,7 @@ class BasePolicy(nn.Module): :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 squash_output: (bool) For continuous actions, whether the output is squashed - or not using a `tanh()` function. + or not using a ``tanh()`` function. :param features_extractor: (nn.Module) Network to extract features (a CNN when using images, a nn.Flatten() layer otherwise) :param normalize_images: (bool) Whether to normalize images or not, @@ -26,14 +27,14 @@ class BasePolicy(nn.Module): """ def __init__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, - device: Union[th.device, str] = 'cpu', + device: Union[th.device, str] = 'auto', squash_output: bool = False, features_extractor: Optional[nn.Module] = None, normalize_images: bool = True): super(BasePolicy, self).__init__() self.observation_space = observation_space self.action_space = action_space - self.device = device + self.device = get_device(device) self.features_extractor = features_extractor self.normalize_images = normalize_images self._squash_output = squash_output @@ -359,8 +360,9 @@ class MlpExtractor(nn.Module): def __init__(self, feature_dim: int, net_arch: List[Union[int, Dict[str, List[int]]]], activation_fn: Type[nn.Module], - device: Union[th.device, str] = 'cpu'): + device: Union[th.device, str] = 'auto'): super(MlpExtractor, self).__init__() + device = get_device(device) shared_net, policy_net, value_net = [], [], [] policy_only_layers = [] # Layer sizes of the network that only belongs to the policy network diff --git a/torchy_baselines/common/utils.py b/torchy_baselines/common/utils.py index 1da28c2..b6fbb64 100644 --- a/torchy_baselines/common/utils.py +++ b/torchy_baselines/common/utils.py @@ -84,3 +84,26 @@ def constant_fn(val: float) -> Callable: return val return func + + +def get_device(device: Union[th.device, str] = 'auto') -> th.device: + """ + Retrieve PyTorch device. + It checks that the requested device is available first. + For now, it supports only cpu and cuda. + By default, it tries to use the gpu. + + :param device: (Union[str, th.device]) One for 'auto', 'cuda', 'cpu' + :return: (th.device) + """ + # Cuda by default + if device == 'auto': + device = 'cuda' + # Force conversion to th.device + device = th.device(device) + + # Cuda not available + if device == th.device('cuda') and not th.cuda.is_available(): + return th.device('cpu') + + return device diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 0dddf51..694428f 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -49,7 +49,7 @@ class PPOPolicy(BasePolicy): action_space: gym.spaces.Space, lr_schedule: Callable, net_arch: Optional[List[Union[int, Dict[str, List[int]]]]] = None, - device: Union[th.device, str] = 'cpu', + device: Union[th.device, str] = 'auto', activation_fn: Type[nn.Module] = nn.Tanh, ortho_init: bool = True, use_sde: bool = False, diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index 8f4ab51..dd128bb 100644 --- a/torchy_baselines/sac/policies.py +++ b/torchy_baselines/sac/policies.py @@ -53,11 +53,12 @@ class Actor(BasePolicy): use_expln: bool = False, clip_mean: float = 2.0, normalize_images: bool = True, - device: Union[th.device, str] = 'cpu'): + device: Union[th.device, str] = 'auto'): super(Actor, self).__init__(observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images, - device=device) + device=device, + squash_output=True) action_dim = get_action_dim(self.action_space) @@ -171,7 +172,7 @@ class Critic(BasePolicy): features_dim: int, activation_fn: Type[nn.Module] = nn.ReLU, normalize_images: bool = True, - device: Union[th.device, str] = 'cpu'): + device: Union[th.device, str] = 'auto'): super(Critic, self).__init__(observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images, @@ -223,7 +224,7 @@ class SACPolicy(BasePolicy): action_space: gym.spaces.Space, lr_schedule: Callable, net_arch: Optional[List[int]] = None, - device: Union[th.device, str] = 'cpu', + device: Union[th.device, str] = 'auto', activation_fn: Type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 7cc1dca..8a8459c 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -52,11 +52,12 @@ class Actor(BasePolicy): sde_net_arch: Optional[List[int]] = None, use_expln: bool = False, normalize_images: bool = True, - device: Union[th.device, str] = 'cpu'): + device: Union[th.device, str] = 'auto'): super(Actor, self).__init__(observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images, - device=device) + device=device, + squash_output=not use_sde) self.latent_pi, self.log_std = None, None self.weights_dist, self.exploration_mat = None, None @@ -179,7 +180,7 @@ class Critic(BasePolicy): features_dim: int, activation_fn: Type[nn.Module] = nn.ReLU, normalize_images: bool = True, - device: Union[th.device, str] = 'cpu'): + device: Union[th.device, str] = 'auto'): super(Critic, self).__init__(observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images, @@ -268,7 +269,7 @@ class TD3Policy(BasePolicy): action_space: gym.spaces.Space, lr_schedule: Callable, net_arch: Optional[List[int]] = None, - device: Union[th.device, str] = 'cpu', + device: Union[th.device, str] = 'auto', activation_fn: Type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, diff --git a/torchy_baselines/version.txt b/torchy_baselines/version.txt index 515423e..d443359 100644 --- a/torchy_baselines/version.txt +++ b/torchy_baselines/version.txt @@ -1 +1 @@ -0.5.0a0 +0.5.0a1