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

Cleanup + update doc

Browse source
This commit is contained in:
Antonin Raffin 2019-11-22 13:33:12 +01:00
parent b84e5e9e27
commit 604a19fbc3
7 changed files with 112 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
README.md
View file

@ -24,6 +24,9 @@ TODO:
- SDE: reduce the number of parameters (only n_features instead of n_features x n_actions) for A2C
(done for TD3)
- SDE: learn the feature extractor?
- Refactor: buffer with numpy array instead of pytorch
- Refactor: remove duplicated code for evaluation
- plotting? -> zoo
Later:
- get_parameters / set_parameters

1
torchy_baselines/a2c/a2c.py
View file

@ -25,6 +25,7 @@ class A2C(PPO):
(i.e. batch size is n_steps * n_env where n_env is number of environment copies running in parallel)
:param gamma: (float) Discount factor
:param gae_lambda: (float) Factor for trade-off of bias vs variance for Generalized Advantage Estimator
Equivalent to classic advantage when set to 1.
:param ent_coef: (float) Entropy coefficient for the loss calculation
:param vf_coef: (float) Value function coefficient for the loss calculation
:param max_grad_norm: (float) The maximum value for the gradient clipping

4
torchy_baselines/cem_rl/cem_rl.py
View file

@ -186,10 +186,6 @@ class CEMRL(TD3):
actor_steps += episode_timesteps
self.fitnesses.append(episode_reward)
if self.verbose > 1:
print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format(
self.num_timesteps, episode_num, episode_timesteps, episode_reward))
self._update_current_progress(self.num_timesteps, total_timesteps)
self.es.tell(self.es_params, self.fitnesses)
timesteps_since_eval += actor_steps

53
torchy_baselines/common/base_class.py
View file

@ -20,13 +20,19 @@ class BaseRLModel(object):
:param policy: (BasePolicy) Policy object
:param env: (Gym environment) The environment to learn from
(if registered in Gym, can be str. Can be None for loading trained models)
:param verbose: (int) the verbosity level: 0 none, 1 training information, 2 debug
:param policy_base: (BasePolicy) the base policy used by this method
:param policy_kwargs: (dict) 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: (str or th.device) Device on which the code should.
By default, it will try to use a Cuda compatible device and fallback to cpu
if it is not possible.
:param support_multi_env: (bool) Whether the algorithm supports training
with multiple environments (as in A2C)
: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: (bool) When creating an environment, whether to wrap it
or not in a Monitor wrapper.
:param seed: (int) Seed for the pseudo random generators
"""
__metaclass__ = ABCMeta
@ -65,6 +71,7 @@ class BaseRLModel(object):
# this is used to update the learning rate
self._current_progress = 1
# Create and wrap the env if needed
if env is not None:
if isinstance(env, str):
if create_eval_env:
@ -94,6 +101,12 @@ class BaseRLModel(object):
" environment.")
def _get_eval_env(self, eval_env):
"""
Return the environment that will be used for evaluation.
:param eval_env: (gym.Env or VecEnv)
:return: (VecEnv)
"""
if eval_env is None:
eval_env = self.eval_env
@ -133,7 +146,7 @@ class BaseRLModel(object):
"""
Compute current progress (from 1 to 0)
:param num_timesteps: (int)
:param num_timesteps: (int) current number of timesteps
:param total_timesteps: (int)
"""
self._current_progress = 1.0 - float(num_timesteps) / float(total_timesteps)
@ -169,7 +182,7 @@ class BaseRLModel(object):
"""
returns the current environment (can be None if not defined)
:return: (Gym Environment) The current environment
:return: (gym.Env) The current environment
"""
return self.env
@ -177,7 +190,7 @@ class BaseRLModel(object):
"""
Checks the validity of the environment, and if it is coherent, set it as the current environment.
:param env: (Gym Environment) The environment for learning a policy
:param env: (gym.Env) The environment for learning a policy
"""
pass
@ -311,21 +324,33 @@ class BaseRLModel(object):
self.eval_env.seed(seed)
def _setup_learn(self, eval_env):
"""
Initialize different variables needed for training.
:param eval_env: (gym.Env or VecEnv)
:return: (int, int, [float], np.ndarray, VecEnv)
"""
self.start_time = time.time()
self.ep_info_buffer = deque(maxlen=100)
if self.action_noise is not None:
self.action_noise.reset()
timesteps_since_eval, episode_num = 0, 0
evaluations = []
if eval_env is not None and self.seed is not None:
eval_env.seed(self.seed)
eval_env = self._get_eval_env(eval_env)
obs = self.env.reset()
return timesteps_since_eval, episode_num, evaluations, obs, eval_env
def _update_info_buffer(self, infos):
"""
Retrieve reward and episode length if using Monitor wrapper.
Retrieve reward and episode length and update the buffer
if using Monitor wrapper.
:param infos: ([dict])
"""
for info in infos:
@ -338,7 +363,23 @@ class BaseRLModel(object):
learning_starts=0, num_timesteps=0,
replay_buffer=None, obs=None,
episode_num=0, log_interval=None):
"""
Collect rollout using the current policy (and possibly fill the replay buffer)
TODO: move this method to off-policy base class.
:param env: (VecEnv)
:param n_episodes: (int)
:param n_steps: (int)
:param action_noise: (ActionNoise)
:param deterministic: (bool)
:param callback: (callable)
:param learning_starts: (int)
:param num_timesteps: (int)
:param replay_buffer: (ReplayBuffer)
:param obs: (np.ndarray)
:param episode_num: (int)
:param log_interval: (int)
"""
episode_rewards = []
total_timesteps = []
total_steps, total_episodes = 0, 0
@ -432,12 +473,10 @@ class BaseRLModel(object):
episode_num + total_episodes) % log_interval == 0:
fps = int(num_timesteps / (time.time() - self.start_time))
logger.logkv("episodes", episode_num + total_episodes)
# logger.logkv("mean 100 episode reward", mean_reward)
if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0:
logger.logkv('ep_rew_mean', self.safe_mean([ep_info['r'] for ep_info in self.ep_info_buffer]))
logger.logkv('ep_len_mean', self.safe_mean([ep_info['l'] for ep_info in self.ep_info_buffer]))
# logger.logkv("n_updates", n_updates)
# logger.logkv("current_lr", current_lr)
logger.logkv("fps", fps)
logger.logkv('time_elapsed', int(time.time() - self.start_time))
logger.logkv("total timesteps", num_timesteps)

68
torchy_baselines/common/buffers.py
View file

@ -5,6 +5,15 @@ from torchy_baselines.common.vec_env import unwrap_vec_normalize
class BaseBuffer(object):
"""
Base class that represent a buffer (rollout or replay)
:param buffer_size: (int) Max number of element in the buffer
:param obs_dim: (int) Dimension of the observation
:param action_dim: (int) Dimension of the action space
:param device: (th.device)
:param n_envs: (int) Number of parallel environments
"""
def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', n_envs=1):
super(BaseBuffer, self).__init__()
self.buffer_size = buffer_size
@ -31,27 +40,43 @@ class BaseBuffer(object):
return tensor.transpose(0, 1).reshape(shape[0] * shape[1], *shape[2:])
def size(self):
"""
:return: (int) The current size of the buffer
"""
if self.full:
return self.buffer_size
return self.pos
def get_pos(self):
return self.pos
def add(self, *args, **kwargs):
"""
Add elements to the buffer.
"""
raise NotImplementedError()
def reset(self):
"""
Reset the buffer.
"""
self.pos = 0
self.full = False
def sample(self, batch_size, env=None):
"""
:param batch_size: (int) Number of element to sample
:param env: (VecNormalize) [Optional] associated gym VecEnv
to normalize the observations/rewards when sampling
"""
upper_bound = self.buffer_size if self.full else self.pos
batch_inds = th.LongTensor(
np.random.randint(0, upper_bound, size=batch_size))
return self._get_samples(batch_inds, env=env)
def _get_samples(self, batch_inds, env=None):
"""
:param batch_inds: (th.Tensor)
:param env: (gym.Env)
:return: ([th.Tensor])
"""
raise NotImplementedError()
def _normalize_obs(self, obs, env=None):
@ -68,9 +93,15 @@ class BaseBuffer(object):
class ReplayBuffer(BaseBuffer):
"""
Taken from https://github.com/apourchot/CEM-RL
"""
Replay buffer used in off-policy algorithms like SAC/TD3.
Adapted from from https://github.com/apourchot/CEM-RL
:param buffer_size: (int) Max number of element in the buffer
:param obs_dim: (int) Dimension of the observation
:param action_dim: (int) Dimension of the action space
:param device: (th.device)
:param n_envs: (int) Number of parallel environments
"""
def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', n_envs=1):
super(ReplayBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs)
@ -103,6 +134,18 @@ class ReplayBuffer(BaseBuffer):
class RolloutBuffer(BaseBuffer):
"""
Rollout buffer used in on-policy algorithms like A2C/PPO.
:param buffer_size: (int) Max number of element in the buffer
:param obs_dim: (int) Dimension of the observation
:param action_dim: (int) Dimension of the action space
:param device: (th.device)
:param gae_lambda: (float) Factor for trade-off of bias vs variance for Generalized Advantage Estimator
Equivalent to classic advantage when set to 1.
:param gamma: (float) Discount factor
:param n_envs: (int) Number of parallel environments
"""
def __init__(self, buffer_size, obs_dim, action_dim, device='cpu',
gae_lambda=1, gamma=0.99, n_envs=1):
super(RolloutBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs)
@ -128,7 +171,10 @@ class RolloutBuffer(BaseBuffer):
def compute_returns_and_advantage(self, last_value, dones=False, use_gae=True):
"""
From Stable-Baselines PPO2
Post-processing step: compute the returns (sum of discounted rewards)
and advantage (A(s) = R - V(S)).
Adapted from Stable-Baselines PPO2.
:param last_value: (th.Tensor)
:param dones: ([bool])
:param use_gae: (bool) Whether to use Generalized Advantage Estimation
@ -164,6 +210,16 @@ class RolloutBuffer(BaseBuffer):
self.advantages = self.returns - self.values
def add(self, obs, action, reward, done, value, log_prob):
"""
:param obs: (np.ndarray) Observation
:param action: (np.ndarray) Action
:param reward: (np.ndarray)
:param done: (np.ndarray) End of episode signal.
:param value: (th.Tensor) estimated value of the current state
following the current policy.
:param log_prob: (th.Tensor) log probability of the action
following the current policy.
"""
if len(log_prob.shape) == 0:
# Reshape 0-d tensor to avoid error
log_prob = log_prob.reshape(-1, 1)

3
torchy_baselines/sac/sac.py
View file

@ -257,9 +257,6 @@ class SAC(BaseRLModel):
self._update_current_progress(self.num_timesteps, total_timesteps)
if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts:
if self.verbose > 1:
print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format(
self.num_timesteps, episode_num, episode_timesteps, episode_reward))
gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps
self.train(gradient_steps, batch_size=self.batch_size)

3
torchy_baselines/td3/td3.py
View file

@ -268,9 +268,6 @@ class TD3(BaseRLModel):
self._update_current_progress(self.num_timesteps, total_timesteps)
if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts:
if self.verbose > 1:
print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format(
self.num_timesteps, episode_num, episode_timesteps, episode_reward))
if self.use_sde:
if self.sde_log_std_scheduler is not None: