saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-09-15 22:10:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add get_device util and fix squash_output

Browse source
This commit is contained in:
Antonin RAFFIN 2020-04-20 15:43:11 +02:00
parent aa1026ee87
commit 17f9246257
8 changed files with 46 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
docs/misc/changelog.rst
View file

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

9
torchy_baselines/common/base_class.py
View file

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

10
torchy_baselines/common/policies.py
View file

@ -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

23
torchy_baselines/common/utils.py
View file

@ -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

2
torchy_baselines/ppo/policies.py
View file

@ -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,

9
torchy_baselines/sac/policies.py
View file

@ -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,

9
torchy_baselines/td3/policies.py
View file

@ -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,

2
torchy_baselines/version.txt
View file

@ -1 +1 @@
0.5.0a0
0.5.0a1