Add get_device util and fix squash_output

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

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':
+        self.device = get_device(device)
-            device = 'cuda' if th.cuda.is_available() else 'cpu'
-
-        self.device = th.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:

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

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

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,

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,

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,

torchy_baselines/version.txt

@@ -1 +1 @@
-0.5.0a0
+0.5.0a1