From c97dbb31730a80e76ae0528a89646cbe738c7664 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Mon, 6 Apr 2020 15:12:20 +0200 Subject: [PATCH] Fix device arg that was not pass to children networks --- torchy_baselines/sac/policies.py | 14 ++++++++++---- torchy_baselines/td3/policies.py | 19 +++++++++++++------ 2 files changed, 23 insertions(+), 10 deletions(-) diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index c7417fc..6ed115e 100644 --- a/torchy_baselines/sac/policies.py +++ b/torchy_baselines/sac/policies.py @@ -38,6 +38,7 @@ class Actor(BasePolicy): :param clip_mean: (float) Clip the mean output when using SDE 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__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, @@ -51,10 +52,12 @@ class Actor(BasePolicy): sde_net_arch: Optional[List[int]] = None, use_expln: bool = False, clip_mean: float = 2.0, - normalize_images: bool = True): + normalize_images: bool = True, + device: Union[th.device, str] = 'cpu'): super(Actor, self).__init__(observation_space, action_space, features_extractor=features_extractor, - normalize_images=normalize_images) + normalize_images=normalize_images, + device=device) action_dim = get_action_dim(self.action_space) @@ -159,6 +162,7 @@ class Critic(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__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, @@ -166,10 +170,12 @@ class Critic(BasePolicy): features_extractor: nn.Module, features_dim: int, activation_fn: Type[nn.Module] = nn.ReLU, - normalize_images: bool = True): + normalize_images: bool = True, + device: Union[th.device, str] = 'cpu'): super(Critic, self).__init__(observation_space, action_space, features_extractor=features_extractor, - normalize_images=normalize_images) + normalize_images=normalize_images, + device=device) action_dim = get_action_dim(self.action_space) diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 30705b3..994650c 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -35,6 +35,7 @@ class Actor(BasePolicy): above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough. :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__(self, observation_space: gym.spaces.Space, @@ -50,10 +51,12 @@ class Actor(BasePolicy): full_std: bool = False, sde_net_arch: Optional[List[int]] = None, use_expln: bool = False, - normalize_images: bool = True): + normalize_images: bool = True, + device: Union[th.device, str] = 'cpu'): super(Actor, self).__init__(observation_space, action_space, features_extractor=features_extractor, - normalize_images=normalize_images) + normalize_images=normalize_images, + device=device) self.latent_pi, self.log_std = None, None self.weights_dist, self.exploration_mat = None, None @@ -167,6 +170,7 @@ class Critic(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__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, @@ -174,10 +178,12 @@ class Critic(BasePolicy): features_extractor: nn.Module, features_dim: int, activation_fn: Type[nn.Module] = nn.ReLU, - normalize_images: bool = True): + normalize_images: bool = True, + device: Union[th.device, str] = 'cpu'): super(Critic, self).__init__(observation_space, action_space, features_extractor=features_extractor, - normalize_images=normalize_images) + normalize_images=normalize_images, + device=device) action_dim = get_action_dim(self.action_space) @@ -241,7 +247,7 @@ 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: (str or th.device) Device on which the code should run. + :param device: (Union[th.device, str]) 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 @@ -286,7 +292,8 @@ class TD3Policy(BasePolicy): 'features_dim': self.features_dim, 'net_arch': self.net_arch, 'activation_fn': self.activation_fn, - 'normalize_images': normalize_images + 'normalize_images': normalize_images, + 'device': device } self.actor_kwargs = self.net_args.copy() sde_kwargs = {