From c3187604bc2417001a13d2132ce160fb093de958 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Mon, 16 Mar 2020 14:01:32 +0100 Subject: [PATCH] Code cleanup: rename lr to lr_schedule + typing --- docs/conf.py | 4 +- docs/misc/changelog.rst | 2 + torchy_baselines/a2c/a2c.py | 2 +- torchy_baselines/cem_rl/cem_rl.py | 2 +- torchy_baselines/common/base_class.py | 48 ++++++++++++--------- torchy_baselines/common/noise.py | 41 ++++++++++-------- torchy_baselines/common/running_mean_std.py | 16 ++++--- torchy_baselines/common/utils.py | 13 +++--- torchy_baselines/ppo/policies.py | 10 ++--- torchy_baselines/sac/policies.py | 14 +++--- torchy_baselines/sac/sac.py | 5 +-- torchy_baselines/td3/policies.py | 14 +++--- torchy_baselines/td3/td3.py | 2 +- 13 files changed, 94 insertions(+), 79 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index ab7fbbc..e71d166 100644 --- a/docs/conf.py +++ b/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', diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index a0dac56..5c8f55c 100644 --- a/docs/misc/changelog.rst +++ b/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) diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py index e4e6293..634a941 100644 --- a/torchy_baselines/a2c/a2c.py +++ b/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: diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index 585db55..1e78055 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/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 diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 3692cbe..0462002 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -27,25 +27,25 @@ class BaseRLModel(ABC): """ The base RL model - :param policy: Policy object - :param env: The environment to learn from + :param policy: (Type[BasePolicy]) Policy object + :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_kwargs: Additional arguments to be passed to the policy on creation - :param verbose: The verbosity level: 0 none, 1 training information, 2 debug - :param device: Device on which the code should run. + :param policy_base: (Type[BasePolicy]) The base policy used by this method + :param policy_kwargs: (Dict[str, Any]) Additional arguments to be passed to the policy on creation + :param verbose: (int) The verbosity level: 0 none, 1 training information, 2 debug + :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 use_sde: Whether to use State Dependent Exploration (SDE) + :param seed: (Optional[int]) Seed for the pseudo random generators + :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.action_space = None + self.observation_space = None # type: Optional[gym.spaces.Space] + 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: diff --git a/torchy_baselines/common/noise.py b/torchy_baselines/common/noise.py index 8511010..a8ff63c 100644 --- a/torchy_baselines/common/noise.py +++ b/torchy_baselines/common/noise.py @@ -1,7 +1,6 @@ -""" -Taken from stable-baselines -""" +from typing import Optional 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): - pass + def __call__(self) -> np.ndarray: + raise NotImplementedError() class NormalActionNoise(ActionNoise): """ A Gaussian action noise - :param mean: (float) the mean value of the noise - :param sigma: (float) the scale of the noise (std here) + :param mean: (np.ndarray) the mean value of the noise + :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 sigma: (float) the scale of the noise + :param mean: (np.ndarray) the mean 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})' diff --git a/torchy_baselines/common/running_mean_std.py b/torchy_baselines/common/running_mean_std.py index d6a03d6..f503106 100644 --- a/torchy_baselines/common/running_mean_std.py +++ b/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.var = np.ones(shape, 'float64') + self.mean = np.zeros(shape, np.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 diff --git a/torchy_baselines/common/utils.py b/torchy_baselines/common/utils.py index 151333b..1da28c2 100644 --- a/torchy_baselines/common/utils.py +++ b/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(_): diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 2d4904a..8487f6d 100644 --- a/torchy_baselines/ppo/policies.py +++ b/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): diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index 46d8a2e..b88f613 100644 --- a/torchy_baselines/sac/policies.py +++ b/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 net_arch: ([int or dict]) The specification of the policy and value networks. + :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 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) diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index 5bdc999..907aa05 100644 --- a/torchy_baselines/sac/sac.py +++ b/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() diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index d7d3e81..4daa02d 100644 --- a/torchy_baselines/td3/policies.py +++ b/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 net_arch: ([int or dict]) The specification of the policy and value networks. + :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 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) diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index d3c8672..331a6e0 100644 --- a/torchy_baselines/td3/td3.py +++ b/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()