diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 5930131..09c8fde 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -16,6 +16,7 @@ New Features: Bug Fixes: ^^^^^^^^^^ +- Fix type hint for activation functions Deprecations: ^^^^^^^^^^^^^ diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py index 8711cfc..85f0b18 100644 --- a/torchy_baselines/common/policies.py +++ b/torchy_baselines/common/policies.py @@ -56,7 +56,7 @@ class BasePolicy(nn.Module): @staticmethod def init_weights(module: nn.Module, gain: float = 1): - if type(module) == nn.Linear: + if isinstance(module, nn.Linear): nn.init.orthogonal_(module.weight, gain=gain) module.bias.data.fill_(0.0) @@ -109,7 +109,7 @@ class BasePolicy(nn.Module): def create_mlp(input_dim: int, output_dim: int, net_arch: List[int], - activation_fn: nn.Module = nn.ReLU, + activation_fn: Type[nn.Module] = nn.ReLU, squash_output: bool = False) -> List[nn.Module]: """ Create a multi layer perceptron (MLP), which is @@ -120,7 +120,7 @@ def create_mlp(input_dim: int, :param net_arch: (List[int]) Architecture of the neural net It represents the number of units per layer. The length of this list is the number of layers. - :param activation_fn: (nn.Module) The activation function + :param activation_fn: (Type[nn.Module]) The activation function to use after each layer. :param squash_output: (bool) Whether to squash the output using a Tanh activation function @@ -145,14 +145,14 @@ def create_mlp(input_dim: int, def create_sde_features_extractor(features_dim: int, sde_net_arch: List[int], - activation_fn: nn.Module) -> Tuple[nn.Sequential, int]: + activation_fn: Type[nn.Module]) -> Tuple[nn.Sequential, int]: """ Create the neural network that will be used to extract features for the SDE exploration function. :param features_dim: (int) :param sde_net_arch: ([int]) - :param activation_fn: (nn.Module) + :param activation_fn: (Type[nn.Module]) :return: (nn.Sequential, int) """ # Special case: when using states as features (i.e. sde_net_arch is an empty list) @@ -233,12 +233,12 @@ class MlpExtractor(nn.Module): :param feature_dim: (int) Dimension of the feature vector (can be the output of a CNN) :param net_arch: ([int or dict]) The specification of the policy and value networks. See above for details on its formatting. - :param activation_fn: (nn.Module) The activation function to use for the networks. + :param activation_fn: (Type[nn.Module]) The activation function to use for the networks. :param device: (th.device) """ def __init__(self, feature_dim: int, net_arch: List[Union[int, Dict[str, List[int]]]], - activation_fn: nn.Module, + activation_fn: Type[nn.Module], device: Union[th.device, str] = 'cpu'): super(MlpExtractor, self).__init__() diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 750cf21..216790c 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -1,4 +1,4 @@ -from typing import Optional, List, Tuple, Callable, Union, Dict +from typing import Optional, List, Tuple, Callable, Union, Dict, Type from functools import partial import gym @@ -23,7 +23,7 @@ class PPOPolicy(BasePolicy): :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: (nn.Module) Activation function + :param activation_fn: (Type[nn.Module]) Activation function :param adam_epsilon: (float) Small values to avoid NaN in ADAM optimizer :param ortho_init: (bool) Whether to use or not orthogonal initialization :param use_sde: (bool) Whether to use State Dependent Exploration or not @@ -47,7 +47,7 @@ class PPOPolicy(BasePolicy): lr_schedule: Callable, net_arch: Optional[List[Union[int, Dict[str, List[int]]]]] = None, device: Union[th.device, str] = 'cpu', - activation_fn: nn.Module = nn.Tanh, + activation_fn: Type[nn.Module] = nn.Tanh, adam_epsilon: float = 1e-5, ortho_init: bool = True, use_sde: bool = False, diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index 5a90e62..db0f1a9 100644 --- a/torchy_baselines/sac/policies.py +++ b/torchy_baselines/sac/policies.py @@ -1,4 +1,4 @@ -from typing import Optional, List, Tuple, Callable, Union +from typing import Optional, List, Tuple, Callable, Union, Type import gym import torch as th @@ -24,7 +24,7 @@ class Actor(BasePolicy): :param features_extractor: (nn.Module) Network to extract features (a CNN when using images, a nn.Flatten() layer otherwise) :param features_dim: (int) Number of features - :param activation_fn: (nn.Module) Activation function + :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 :param full_std: (bool) Whether to use (n_features x n_actions) parameters @@ -44,7 +44,7 @@ class Actor(BasePolicy): net_arch: List[int], features_extractor: nn.Module, features_dim: int, - activation_fn: nn.Module = nn.ReLU, + activation_fn: Type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, full_std: bool = True, @@ -151,7 +151,7 @@ class Critic(BasePolicy): :param features_extractor: (nn.Module) Network to extract features (a CNN when using images, a nn.Flatten() layer otherwise) :param features_dim: (int) Number of features - :param activation_fn: (nn.Module) Activation function + :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) """ @@ -160,7 +160,7 @@ class Critic(BasePolicy): net_arch: List[int], features_extractor: nn.Module, features_dim: int, - activation_fn: nn.Module = nn.ReLU, + activation_fn: Type[nn.Module] = nn.ReLU, normalize_images: bool = True): super(Critic, self).__init__(observation_space, action_space, features_extractor=features_extractor, @@ -191,7 +191,7 @@ class SACPolicy(BasePolicy): :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: (nn.Module) Activation function + :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 :param sde_net_arch: ([int]) Network architecture for extracting features @@ -209,7 +209,7 @@ class SACPolicy(BasePolicy): lr_schedule: Callable, net_arch: Optional[List[int]] = None, device: Union[th.device, str] = 'cpu', - activation_fn: nn.Module = nn.ReLU, + activation_fn: Type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, sde_net_arch: Optional[List[int]] = None, diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 3c9698c..8b24832 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -1,4 +1,4 @@ -from typing import Optional, List, Tuple, Callable, Union +from typing import Optional, List, Tuple, Callable, Union, Type import gym import torch as th @@ -20,7 +20,7 @@ class Actor(BasePolicy): :param features_extractor: (nn.Module) Network to extract features (a CNN when using images, a nn.Flatten() layer otherwise) :param features_dim: (int) Number of features - :param activation_fn: (nn.Module) Activation function + :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 :param clip_noise: (float) Clip the magnitude of the noise @@ -42,7 +42,7 @@ class Actor(BasePolicy): net_arch: List[int], features_extractor: nn.Module, features_dim: int, - activation_fn: nn.Module = nn.ReLU, + activation_fn: Type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, clip_noise: Optional[float] = None, @@ -168,7 +168,7 @@ class Critic(BasePolicy): :param features_extractor: (nn.Module) Network to extract features (a CNN when using images, a nn.Flatten() layer otherwise) :param features_dim: (int) Number of features - :param activation_fn: (nn.Module) Activation function + :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) """ @@ -177,7 +177,7 @@ class Critic(BasePolicy): net_arch: List[int], features_extractor: nn.Module, features_dim: int, - activation_fn: nn.Module = nn.ReLU, + activation_fn: Type[nn.Module] = nn.ReLU, normalize_images: bool = True): super(Critic, self).__init__(observation_space, action_space, features_extractor=features_extractor, @@ -211,7 +211,7 @@ class ValueFunction(BasePolicy): (a CNN when using images, a nn.Flatten() layer otherwise) :param features_dim: (int) Number of features :param net_arch: (Optional[List[int]]) Network architecture - :param activation_fn: (nn.Module) Activation function + :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) """ @@ -220,7 +220,7 @@ class ValueFunction(BasePolicy): features_extractor: nn.Module, features_dim: int, net_arch: Optional[List[int]] = None, - activation_fn: nn.Module = nn.Tanh, + activation_fn: Type[nn.Module] = nn.Tanh, normalize_images: bool = True): super(ValueFunction, self).__init__(observation_space, action_space, features_extractor=features_extractor, @@ -246,7 +246,7 @@ class TD3Policy(BasePolicy): :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: (nn.Module) Activation function + :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 :param sde_net_arch: ([int]) Network architecture for extracting features @@ -263,7 +263,7 @@ class TD3Policy(BasePolicy): lr_schedule: Callable, net_arch: Optional[List[int]] = None, device: Union[th.device, str] = 'cpu', - activation_fn: nn.Module = nn.ReLU, + activation_fn: Type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, clip_noise: Optional[float] = None,