saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-10 17:37:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix type hint for activation fn

Browse source
This commit is contained in:
Antonin RAFFIN 2020-03-24 10:10:37 +01:00
parent ba18258af6
commit 72a88a8d92
5 changed files with 27 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
docs/misc/changelog.rst
View file

@ -16,6 +16,7 @@ New Features:
Bug Fixes:
^^^^^^^^^^
- Fix type hint for activation functions
Deprecations:
^^^^^^^^^^^^^

14
torchy_baselines/common/policies.py
View file

@ -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__()

6
torchy_baselines/ppo/policies.py
View file

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

14
torchy_baselines/sac/policies.py
View file

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

18
torchy_baselines/td3/policies.py
View file

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