saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-09-15 22:10:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Code cleanup: rename lr to lr_schedule + typing

Browse source
This commit is contained in:
Antonin RAFFIN 2020-03-16 14:01:32 +01:00
parent a67bb75438
commit c3187604bc
13 changed files with 94 additions and 79 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
docs/conf.py
View file

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

2
docs/misc/changelog.rst
View file

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

2
torchy_baselines/a2c/a2c.py
View file

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

2
torchy_baselines/cem_rl/cem_rl.py
View file

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

48
torchy_baselines/common/base_class.py
View file

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

41
torchy_baselines/common/noise.py
View file

@ -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})'

16
torchy_baselines/common/running_mean_std.py
View file

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

13
torchy_baselines/common/utils.py
View file

@ -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(_):

10
torchy_baselines/ppo/policies.py
View file

@ -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):

14
torchy_baselines/sac/policies.py
View file

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

5
torchy_baselines/sac/sac.py
View file

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

14
torchy_baselines/td3/policies.py
View file

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

2
torchy_baselines/td3/td3.py
View file

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