Code cleanup: rename lr to lr_schedule + typing

docs/conf.py

@@ -50,7 +50,7 @@ import torchy_baselines
 # -- Project information -----------------------------------------------------
 
 project = 'Torchy Baselines'
-copyright = '2019, Torchy Baselines'
+copyright = '2020, Torchy Baselines'
 author = 'Torchy Baselines Contributors'
 
 # The short X.Y version
@@ -70,7 +70,7 @@ release = torchy_baselines.__version__
 # ones.
 extensions = [
     'sphinx.ext.autodoc',
-    'sphinx_autodoc_typehints',
+    # 'sphinx_autodoc_typehints',
     'sphinx.ext.autosummary',
     'sphinx.ext.mathjax',
     'sphinx.ext.ifconfig',

docs/misc/changelog.rst

@@ -31,9 +31,11 @@ Others:
 - Buffers now return ``NamedTuple``
 - More typing
 - Add test for ``expln``
+- Renamed ``learning_rate`` to ``lr_schedule``
 
 Documentation:
 ^^^^^^^^^^^^^^
+- Deactivated ``sphinx_autodoc_typehints`` extension
 
 
 Pre-Release 0.2.0 (2020-02-14)

torchy_baselines/a2c/a2c.py

@@ -91,7 +91,7 @@ class A2C(PPO):
         super(A2C, self)._setup_model()
         if self.use_rms_prop:
             self.policy.optimizer = th.optim.RMSprop(self.policy.parameters(),
-                                                     lr=self.learning_rate(1), alpha=0.99,
+                                                     lr=self.lr_schedule(1), alpha=0.99,
                                                      eps=self.rms_prop_eps, weight_decay=0)
 
     def train(self, gradient_steps: int, batch_size: Optional[int] = None) -> None:

torchy_baselines/cem_rl/cem_rl.py

@@ -149,7 +149,7 @@ class CEMRL(TD3):
                     self.actor.load_from_vector(self.es_params[i])
                     self.actor_target.load_from_vector(self.es_params[i])
                     self.actor.optimizer = th.optim.Adam(self.actor.parameters(),
-                                                         lr=self.learning_rate(self._current_progress))
+                                                         lr=self.lr_schedule(self._current_progress))
 
                     # In the paper: 2 * actor_steps // self.n_grad
                     # In the original implementation: actor_steps // self.n_grad

torchy_baselines/common/base_class.py

@@ -27,25 +27,25 @@ class BaseRLModel(ABC):
     """
     The base RL model
 
-    :param policy: Policy object
+    :param policy: (Type[BasePolicy]) Policy object
-    :param env: The environment to learn from
+    :param env: (Union[GymEnv, str]) The environment to learn from
                 (if registered in Gym, can be str. Can be None for loading trained models)
-    :param policy_base: The base policy used by this method
+    :param policy_base: (Type[BasePolicy]) The base policy used by this method
-    :param policy_kwargs: Additional arguments to be passed to the policy on creation
+    :param policy_kwargs: (Dict[str, Any]) Additional arguments to be passed to the policy on creation
-    :param verbose: The verbosity level: 0 none, 1 training information, 2 debug
+    :param verbose: (int) The verbosity level: 0 none, 1 training information, 2 debug
-    :param device: Device on which the code should run.
+    :param device: (Union[th.device, str]) Device on which the code should run.
         By default, it will try to use a Cuda compatible device and fallback to cpu
         if it is not possible.
-    :param support_multi_env: Whether the algorithm supports training
+    :param support_multi_env: (bool) Whether the algorithm supports training
         with multiple environments (as in A2C)
-    :param create_eval_env: Whether to create a second environment that will be
+    :param create_eval_env: (bool) Whether to create a second environment that will be
         used for evaluating the agent periodically. (Only available when passing string for the environment)
-    :param monitor_wrapper: When creating an environment, whether to wrap it
+    :param monitor_wrapper: (bool) When creating an environment, whether to wrap it
         or not in a Monitor wrapper.
-    :param seed: Seed for the pseudo random generators
+    :param seed: (Optional[int]) Seed for the pseudo random generators
-    :param use_sde: Whether to use State Dependent Exploration (SDE)
+    :param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
         instead of action noise exploration (default: False)
-    :param sde_sample_freq: Sample a new noise matrix every n steps when using SDE
+    :param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE
         Default: -1 (only sample at the beginning of the rollout)
     """
 
@@ -80,8 +80,8 @@ class BaseRLModel(ABC):
         self._vec_normalize_env = unwrap_vec_normalize(env)
         self.verbose = verbose
         self.policy_kwargs = {} if policy_kwargs is None else policy_kwargs
-        self.observation_space = None
+        self.observation_space = None  # type: Optional[gym.spaces.Space]
-        self.action_space = None
+        self.action_space = None  # type: Optional[gym.spaces.Space]
         self.n_envs = None
         self.num_timesteps = 0
         self.eval_env = None
@@ -89,7 +89,8 @@ class BaseRLModel(ABC):
         self.action_noise = None  # type: Optional[ActionNoise]
         self.start_time = None
         self.policy = None
-        self.learning_rate = None
+        self.learning_rate = None  # type: Optional[float]
+        self.lr_schedule = None   # type: Optional[Callable]
         # Used for SDE only
         self.use_sde = use_sde
         self.sde_sample_freq = sde_sample_freq
@@ -134,13 +135,16 @@ class BaseRLModel(ABC):
     @abstractmethod
     def _setup_model(self) -> None:
         """
-        Setup model so state_dict can be loaded
+        Create networks and optimizers
         """
         raise NotImplementedError()
 
     def _get_eval_env(self, eval_env: Optional[GymEnv]) -> Optional[GymEnv]:
         """
         Return the environment that will be used for evaluation.
+
+        :param eval_env: (Optional[GymEnv]))
+        :return: (Optional[GymEnv])
         """
         if eval_env is None:
             eval_env = self.eval_env
@@ -156,7 +160,8 @@ class BaseRLModel(ABC):
         Rescale the action from [low, high] to [-1, 1]
         (no need for symmetric action space)
 
-        :param action: Action to scale
+        :param action: (np.ndarray) Action to scale
+        :return: (np.ndarray) Scaled action
         """
         low, high = self.action_space.low, self.action_space.high
         return 2.0 * ((action - low) / (high - low)) - 1.0
@@ -173,7 +178,7 @@ class BaseRLModel(ABC):
 
     def _setup_learning_rate(self) -> None:
         """Transform to callable if needed."""
-        self.learning_rate = get_schedule_fn(self.learning_rate)
+        self.lr_schedule = get_schedule_fn(self.learning_rate)
 
     def _update_current_progress(self, num_timesteps: int, total_timesteps: int) -> None:
         """
@@ -189,15 +194,16 @@ class BaseRLModel(ABC):
         Update the optimizers learning rate using the current learning rate schedule
         and the current progress (from 1 to 0).
 
-        :param optimizers: An optimizer or a list of optimizer.
+        :param optimizers: (Union[List[th.optim.Optimizer], th.optim.Optimizer])
+            An optimizer or a list of optimizers.
         """
         # Log the current learning rate
-        logger.logkv("learning_rate", self.learning_rate(self._current_progress))
+        logger.logkv("learning_rate", self.lr_schedule(self._current_progress))
 
         if not isinstance(optimizers, list):
             optimizers = [optimizers]
         for optimizer in optimizers:
-            update_learning_rate(optimizer, self.learning_rate(self._current_progress))
+            update_learning_rate(optimizer, self.lr_schedule(self._current_progress))
 
     @staticmethod
     def safe_mean(arr: Union[np.ndarray, list, deque]) -> np.ndarray:

torchy_baselines/common/noise.py

@@ -1,7 +1,6 @@
-"""
+from typing import Optional
-Taken from stable-baselines
-"""
 from abc import ABC, abstractmethod
+
 import numpy as np
 
 
@@ -13,34 +12,34 @@ class ActionNoise(ABC):
     def __init__(self):
         super(ActionNoise, self).__init__()
 
-    def reset(self):
+    def reset(self) -> None:
         """
         call end of episode reset for the noise
         """
         pass
 
     @abstractmethod
-    def __call__(self):
+    def __call__(self) -> np.ndarray:
-        pass
+        raise NotImplementedError()
 
 
 class NormalActionNoise(ActionNoise):
     """
     A Gaussian action noise
 
-    :param mean: (float) the mean value of the noise
+    :param mean: (np.ndarray) the mean value of the noise
-    :param sigma: (float) the scale of the noise (std here)
+    :param sigma: (np.ndarray) the scale of the noise (std here)
     """
 
-    def __init__(self, mean, sigma):
+    def __init__(self, mean: np.ndarray, sigma: np.ndarray):
         self._mu = mean
         self._sigma = sigma
         super(NormalActionNoise, self).__init__()
 
-    def __call__(self):
+    def __call__(self) -> np.ndarray:
         return np.random.normal(self._mu, self._sigma)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f'NormalActionNoise(mu={self._mu}, sigma={self._sigma})'
 
 
@@ -50,34 +49,38 @@ class OrnsteinUhlenbeckActionNoise(ActionNoise):
 
     Based on http://math.stackexchange.com/questions/1287634/implementing-ornstein-uhlenbeck-in-matlab
 
-    :param mean: (float) the mean of the noise
+    :param mean: (np.ndarray) the mean of the noise
-    :param sigma: (float) the scale of the noise
+    :param sigma: (np.ndarray) the scale of the noise
     :param theta: (float) the rate of mean reversion
     :param dt: (float) the timestep for the noise
-    :param initial_noise: ([float]) the initial value for the noise output, (if None: 0)
+    :param initial_noise: (Optional[np.ndarray]) the initial value for the noise output, (if None: 0)
     """
 
-    def __init__(self, mean, sigma, theta=.15, dt=1e-2, initial_noise=None):
+    def __init__(self, mean: np.ndarray,
+                 sigma: np.ndarray,
+                 theta: float = .15,
+                 dt: float = 1e-2,
+                 initial_noise: Optional[np.ndarray] = None):
         self._theta = theta
         self._mu = mean
         self._sigma = sigma
         self._dt = dt
         self.initial_noise = initial_noise
-        self.noise_prev = None
+        self.noise_prev = np.zeros_like(self._mu)
         self.reset()
         super(OrnsteinUhlenbeckActionNoise, self).__init__()
 
-    def __call__(self):
+    def __call__(self) -> np.ndarray:
         noise = self.noise_prev + self._theta * (self._mu - self.noise_prev) * self._dt + \
                 self._sigma * np.sqrt(self._dt) * np.random.normal(size=self._mu.shape)
         self.noise_prev = noise
         return noise
 
-    def reset(self):
+    def reset(self) -> None:
         """
         reset the Ornstein Uhlenbeck noise, to the initial position
         """
         self.noise_prev = self.initial_noise if self.initial_noise is not None else np.zeros_like(self._mu)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f'OrnsteinUhlenbeckActionNoise(mu={self._mu}, sigma={self._sigma})'

torchy_baselines/common/running_mean_std.py

@@ -1,26 +1,30 @@
+from typing import Tuple
+
 import numpy as np
 
 
 class RunningMeanStd(object):
-    def __init__(self, epsilon=1e-4, shape=()):
+    def __init__(self, epsilon: float = 1e-4, shape: Tuple[int, ...] = ()):
         """
-        calulates the running mean and std of a data stream
+        Calulates the running mean and std of a data stream
         https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm
 
         :param epsilon: (float) helps with arithmetic issues
         :param shape: (tuple) the shape of the data stream's output
         """
-        self.mean = np.zeros(shape, 'float64')
+        self.mean = np.zeros(shape, np.float64)
-        self.var = np.ones(shape, 'float64')
+        self.var = np.ones(shape, np.float64)
         self.count = epsilon
 
-    def update(self, arr):
+    def update(self, arr: np.ndarray) -> None:
         batch_mean = np.mean(arr, axis=0)
         batch_var = np.var(arr, axis=0)
         batch_count = arr.shape[0]
         self.update_from_moments(batch_mean, batch_var, batch_count)
 
-    def update_from_moments(self, batch_mean, batch_var, batch_count):
+    def update_from_moments(self, batch_mean: np.ndarray,
+                            batch_var: np.ndarray,
+                            batch_count: int) -> None:
         delta = batch_mean - self.mean
         tot_count = self.count + batch_count
 

torchy_baselines/common/utils.py

@@ -1,10 +1,11 @@
+from typing import Callable, Union
 import random
 
 import numpy as np
 import torch as th
 
 
-def set_random_seed(seed, using_cuda=False):
+def set_random_seed(seed: int, using_cuda: bool = False) -> None:
     """
     Seed the different random generators
     :param seed: (int)
@@ -21,7 +22,7 @@ def set_random_seed(seed, using_cuda=False):
 
 
 # From stable baselines
-def explained_variance(y_pred, y_true):
+def explained_variance(y_pred: np.ndarray, y_true: np.ndarray) -> np.ndarray:
     """
     Computes fraction of variance that ypred explains about y.
     Returns 1 - Var[y-ypred] / Var[y]
@@ -40,7 +41,7 @@ def explained_variance(y_pred, y_true):
     return np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y
 
 
-def update_learning_rate(optimizer, learning_rate):
+def update_learning_rate(optimizer: th.optim.Optimizer, learning_rate: float) -> None:
     """
     Update the learning rate for a given optimizer.
     Useful when doing linear schedule.
@@ -52,7 +53,7 @@ def update_learning_rate(optimizer, learning_rate):
         param_group['lr'] = learning_rate
 
 
-def get_schedule_fn(value_schedule):
+def get_schedule_fn(value_schedule: Union[Callable, float]) -> Callable:
     """
     Transform (if needed) learning rate and clip range (for PPO)
     to callable.
@@ -70,13 +71,13 @@ def get_schedule_fn(value_schedule):
     return value_schedule
 
 
-def constant_fn(val):
+def constant_fn(val: float) -> Callable:
     """
     Create a function that returns a constant
     It is useful for learning rate schedule (to avoid code duplication)
 
     :param val: (float)
-    :return: (function)
+    :return: (Callable)
     """
 
     def func(_):

torchy_baselines/ppo/policies.py

@@ -19,7 +19,7 @@ class PPOPolicy(BasePolicy):
 
     :param observation_space: (gym.spaces.Space) Observation space
     :param action_space: (gym.spaces.Space) Action space
-    :param learning_rate: (callable) Learning rate schedule (could be constant)
+    :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
@@ -41,7 +41,7 @@ class PPOPolicy(BasePolicy):
     def __init__(self,
                  observation_space: gym.spaces.Space,
                  action_space: gym.spaces.Space,
-                 learning_rate: Callable,
+                 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,
@@ -93,7 +93,7 @@ class PPOPolicy(BasePolicy):
         # Action distribution
         self.action_dist = make_proba_distribution(action_space, use_sde=use_sde, dist_kwargs=dist_kwargs)
 
-        self._build(learning_rate)
+        self._build(lr_schedule)
 
     def reset_noise(self, n_envs: int = 1) -> None:
         """
@@ -104,7 +104,7 @@ class PPOPolicy(BasePolicy):
         assert isinstance(self.action_dist, StateDependentNoiseDistribution), 'reset_noise() is only available when using SDE'
         self.action_dist.sample_weights(self.log_std, batch_size=n_envs)
 
-    def _build(self, learning_rate: Callable) -> None:
+    def _build(self, lr_schedule: Callable) -> None:
         self.mlp_extractor = MlpExtractor(self.features_dim, net_arch=self.net_arch,
                                           activation_fn=self.activation_fn, device=self.device)
 
@@ -139,7 +139,7 @@ class PPOPolicy(BasePolicy):
                     self.value_net: 1
                 }[module]
                 module.apply(partial(self.init_weights, gain=gain))
-        self.optimizer = th.optim.Adam(self.parameters(), lr=learning_rate(1), eps=self.adam_epsilon)
+        self.optimizer = th.optim.Adam(self.parameters(), lr=lr_schedule(1), eps=self.adam_epsilon)
 
     def forward(self, obs: th.Tensor, deterministic: bool = False) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:
         if not isinstance(obs, th.Tensor):

torchy_baselines/sac/policies.py

@@ -161,8 +161,8 @@ class SACPolicy(BasePolicy):
 
     :param observation_space: (gym.spaces.Space) Observation space
     :param action_space: (gym.spaces.Space) Action space
-    :param learning_rate: (callable) Learning rate schedule (could be constant)
+    :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 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 use_sde: (bool) Whether to use State Dependent Exploration or not
@@ -177,7 +177,7 @@ class SACPolicy(BasePolicy):
     """
     def __init__(self, observation_space: gym.spaces.Space,
                  action_space: gym.spaces.Space,
-                 learning_rate: Callable,
+                 lr_schedule: Callable,
                  net_arch: Optional[List[int]] = None,
                  device: Union[th.device, str] = 'cpu',
                  activation_fn: nn.Module = nn.ReLU,
@@ -213,16 +213,16 @@ class SACPolicy(BasePolicy):
         self.actor, self.actor_target = None, None
         self.critic, self.critic_target = None, None
 
-        self._build(learning_rate)
+        self._build(lr_schedule)
 
-    def _build(self, learning_rate: Callable) -> None:
+    def _build(self, lr_schedule: Callable) -> None:
         self.actor = self.make_actor()
-        self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=learning_rate(1))
+        self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=lr_schedule(1))
 
         self.critic = self.make_critic()
         self.critic_target = self.make_critic()
         self.critic_target.load_state_dict(self.critic.state_dict())
-        self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=learning_rate(1))
+        self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=lr_schedule(1))
 
     def make_actor(self) -> Actor:
         return Actor(**self.actor_kwargs).to(self.device)

torchy_baselines/sac/sac.py

@@ -94,7 +94,6 @@ class SAC(OffPolicyRLModel):
                                   use_sde=use_sde, sde_sample_freq=sde_sample_freq,
                                   use_sde_at_warmup=use_sde_at_warmup)
 
-        self.learning_rate = learning_rate
         self.target_entropy = target_entropy
         self.log_ent_coef = None  # type: Optional[th.Tensor]
         self.target_update_interval = target_update_interval
@@ -146,7 +145,7 @@ class SAC(OffPolicyRLModel):
             # Note: we optimize the log of the entropy coeff which is slightly different from the paper
             # as discussed in https://github.com/rail-berkeley/softlearning/issues/37
             self.log_ent_coef = th.log(th.ones(1, device=self.device) * init_value).requires_grad_(True)
-            self.ent_coef_optimizer = th.optim.Adam([self.log_ent_coef], lr=self.learning_rate(1))
+            self.ent_coef_optimizer = th.optim.Adam([self.log_ent_coef], lr=self.lr_schedule(1))
         else:
             # Force conversion to float
             # this will throw an error if a malformed string (different from 'auto')
@@ -155,7 +154,7 @@ class SAC(OffPolicyRLModel):
 
         self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
         self.policy = self.policy_class(self.observation_space, self.action_space,
-                                        self.learning_rate, use_sde=self.use_sde,
+                                        self.lr_schedule, use_sde=self.use_sde,
                                         device=self.device, **self.policy_kwargs)
         self.policy = self.policy.to(self.device)
         self._create_aliases()

torchy_baselines/td3/policies.py

@@ -199,8 +199,8 @@ class TD3Policy(BasePolicy):
 
     :param observation_space: (gym.spaces.Space) Observation space
     :param action_space: (gym.spaces.Space) Action space
-    :param learning_rate: (callable) Learning rate schedule (could be constant)
+    :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 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 use_sde: (bool) Whether to use State Dependent Exploration or not
@@ -214,7 +214,7 @@ class TD3Policy(BasePolicy):
     """
     def __init__(self, observation_space: gym.spaces.Space,
                  action_space: gym.spaces.Space,
-                 learning_rate: Callable,
+                 lr_schedule: Callable,
                  net_arch: Optional[List[int]] = None,
                  device: Union[th.device, str] = 'cpu',
                  activation_fn: nn.Module = nn.ReLU,
@@ -257,18 +257,18 @@ class TD3Policy(BasePolicy):
         self.use_sde = use_sde
         self.vf_net = None
         self.log_std_init = log_std_init
-        self._build(learning_rate)
+        self._build(lr_schedule)
 
-    def _build(self, learning_rate: Callable) -> None:
+    def _build(self, lr_schedule: Callable) -> None:
         self.actor = self.make_actor()
         self.actor_target = self.make_actor()
         self.actor_target.load_state_dict(self.actor.state_dict())
-        self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=learning_rate(1))
+        self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=lr_schedule(1))
 
         self.critic = self.make_critic()
         self.critic_target = self.make_critic()
         self.critic_target.load_state_dict(self.critic.state_dict())
-        self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=learning_rate(1))
+        self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=lr_schedule(1))
 
         if self.use_sde:
             self.vf_net = ValueFunction(self.obs_dim)

torchy_baselines/td3/td3.py

@@ -123,7 +123,7 @@ class TD3(OffPolicyRLModel):
         self.set_random_seed(self.seed)
         self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
         self.policy = self.policy_class(self.observation_space, self.action_space,
-                                        self.learning_rate, use_sde=self.use_sde,
+                                        self.lr_schedule, use_sde=self.use_sde,
                                         device=self.device, **self.policy_kwargs)
         self.policy = self.policy.to(self.device)
         self._create_aliases()