From db87e0d36a9dbcdc9fb2350319b9561df5435cd3 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Thu, 7 Nov 2019 17:31:52 +0100 Subject: [PATCH 01/24] Quick and dirty SDE version for TD3 --- tests/test_sde.py | 12 +++++- torchy_baselines/common/base_class.py | 5 +++ torchy_baselines/td3/policies.py | 56 ++++++++++++++++++++++----- torchy_baselines/td3/td3.py | 14 ++++--- 4 files changed, 70 insertions(+), 17 deletions(-) diff --git a/tests/test_sde.py b/tests/test_sde.py index f397558..955f9ab 100644 --- a/tests/test_sde.py +++ b/tests/test_sde.py @@ -3,7 +3,7 @@ import pytest import torch as th from torch.distributions import Normal -from torchy_baselines import A2C +from torchy_baselines import A2C, TD3 def test_state_dependent_exploration(): @@ -54,4 +54,12 @@ def test_state_dependent_noise(model_class): model = model_class('MlpPolicy', env, n_steps=200, max_grad_norm=1, use_rms_prop=False, use_sde=True, ent_coef=0.00, verbose=1, create_eval_env=True, learning_rate=3e-4, policy_kwargs=dict(log_std_init=0.0, ortho_init=False, net_arch=[256, dict(pi=[256], vf=[256])]), seed=None) - model.learn(total_timesteps=int(20000), log_interval=5, eval_freq=10000, eval_env=eval_env) + # model.learn(total_timesteps=int(20000), log_interval=5, eval_freq=10000, eval_env=eval_env) + model.learn(total_timesteps=int(1000), log_interval=5, eval_freq=500, eval_env=eval_env) + + +@pytest.mark.parametrize("model_class", [TD3]) +def test_state_dependent_offpolicy_noise(model_class): + model = model_class('MlpPolicy', 'Pendulum-v0', use_sde=True, seed=None, create_eval_env=True, + verbose=1, policy_kwargs=dict(log_std_init=-2)) + model.learn(total_timesteps=int(20000), eval_freq=1000) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index a6b9a41..d96950d 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -327,6 +327,9 @@ class BaseRLModel(object): assert isinstance(env, VecEnv) assert env.num_envs == 1 + if hasattr(self, 'use_sde') and self.use_sde: + self.policy.reset_noise() + while total_steps < n_steps or total_episodes < n_episodes: done = False # Reset environment: not needed for VecEnv @@ -338,6 +341,8 @@ class BaseRLModel(object): if num_timesteps < learning_starts: action = np.array([self.action_space.sample()]) else: + if hasattr(self, 'use_sde'): + deterministic = not self.use_sde action = self.predict(obs, deterministic=deterministic) # Rescale the action from [low, high] to [-1, 1] diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 4cf32f3..adbbda1 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -1,19 +1,46 @@ import torch as th import torch.nn as nn +from torch.distributions import Normal from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork class Actor(BaseNetwork): - def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU): + def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU, + use_sde=False, log_std_init=0.0, clip_noise=0.1): super(Actor, self).__init__() - # TODO: orthogonal initialization? - actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_out=True) - self.actor_net = nn.Sequential(*actor_net) + self.latent_pi, self.log_std = None, None + self.weights_dist, self.exploration_mat = None, None + self.use_sde = use_sde - def forward(self, obs): - return self.actor_net(obs) + if use_sde: + latent_dim = net_arch[-1] + latent_pi = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_out=False) + self.latent_pi = nn.Sequential(*latent_pi) + self.log_std = nn.Parameter(th.ones(latent_dim, action_dim) * log_std_init) + self.actor_net = nn.Sequential(nn.Linear(net_arch[-1], action_dim), nn.Tanh()) + self.clip_noise = clip_noise + self.reset_noise() + else: + actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_out=True) + self.actor_net = nn.Sequential(*actor_net) + + def reset_noise(self): + self.weights_dist = Normal(th.zeros_like(self.log_std), th.exp(self.log_std)) + self.exploration_mat = self.weights_dist.rsample() + + def forward(self, obs, deterministic=True): + if self.use_sde: + latent_pi = self.latent_pi(obs) + if deterministic: + return self.actor_net(latent_pi) + noise = th.mm(latent_pi.detach(), self.exploration_mat) + # noise = th.clamp(noise, -self.clip_noise, self.clip_noise) + # TODO: fix clipping + return th.clamp(self.actor_net(latent_pi) + noise, -1, 1) + else: + return self.actor_net(obs) class Critic(BaseNetwork): @@ -40,7 +67,7 @@ class Critic(BaseNetwork): class TD3Policy(BasePolicy): def __init__(self, observation_space, action_space, learning_rate, net_arch=None, device='cpu', - activation_fn=nn.ReLU): + activation_fn=nn.ReLU, use_sde=False, log_std_init=0.0): super(TD3Policy, self).__init__(observation_space, action_space, device) if net_arch is None: @@ -56,8 +83,14 @@ class TD3Policy(BasePolicy): 'net_arch': self.net_arch, 'activation_fn': self.activation_fn } + self.actor_kwargs = self.net_args.copy() + self.actor_kwargs['use_sde'] = use_sde + self.actor_kwargs['log_std_init'] = log_std_init + self.actor, self.actor_target = None, None self.critic, self.critic_target = None, None + self.use_sde = use_sde + self.log_std_init = log_std_init self._build(learning_rate) def _build(self, learning_rate): @@ -71,14 +104,17 @@ class TD3Policy(BasePolicy): self.critic_target.load_state_dict(self.critic.state_dict()) self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=learning_rate(1)) + def reset_noise(self): + return self.actor.reset_noise() + def make_actor(self): - return Actor(**self.net_args).to(self.device) + return Actor(**self.actor_kwargs).to(self.device) def make_critic(self): return Critic(**self.net_args).to(self.device) - def forward(self, obs): - return self.actor(obs) + def forward(self, obs, deterministic=True): + return self.actor(obs, deterministic=deterministic) MlpPolicy = TD3Policy diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 49cf16e..46943a9 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -37,6 +37,8 @@ class TD3(BaseRLModel): :param target_policy_noise: (float) Standard deviation of gaussian noise added to target policy (smoothing noise) :param target_noise_clip: (float) Limit for absolute value of target policy smoothing noise. + :param use_sde: (bool) Whether to use State Dependent Exploration (SDE) + instead of action noise exploration (default: False) :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 policy_kwargs: (dict) additional arguments to be passed to the policy on creation @@ -50,6 +52,7 @@ class TD3(BaseRLModel): policy_delay=2, learning_starts=100, gamma=0.99, batch_size=100, train_freq=-1, gradient_steps=-1, n_episodes_rollout=1, tau=0.005, action_noise=None, target_policy_noise=0.2, target_noise_clip=0.5, + use_sde=False, tensorboard_log=None, create_eval_env=False, policy_kwargs=None, verbose=0, seed=0, device='auto', _init_setup_model=True): @@ -70,6 +73,7 @@ class TD3(BaseRLModel): self.policy_delay = policy_delay self.target_noise_clip = target_noise_clip self.target_policy_noise = target_policy_noise + self.use_sde = use_sde if _init_setup_model: self._setup_model() @@ -79,8 +83,8 @@ class TD3(BaseRLModel): obs_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0] self.set_random_seed(self.seed) self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device) - self.policy = self.policy(self.observation_space, self.action_space, - self.learning_rate, device=self.device, **self.policy_kwargs) + self.policy = self.policy(self.observation_space, self.action_space, self.learning_rate, + use_sde=self.use_sde, device=self.device, **self.policy_kwargs) self.policy = self.policy.to(self.device) self._create_aliases() @@ -90,12 +94,12 @@ class TD3(BaseRLModel): self.critic = self.policy.critic self.critic_target = self.policy.critic_target - def select_action(self, observation): + def select_action(self, observation, deterministic=True): # Normally not needed observation = np.array(observation) with th.no_grad(): observation = th.FloatTensor(observation.reshape(1, -1)).to(self.device) - return self.actor(observation).cpu().numpy() + return self.actor(observation, deterministic=deterministic).cpu().numpy() def predict(self, observation, state=None, mask=None, deterministic=True): """ @@ -107,7 +111,7 @@ class TD3(BaseRLModel): :param deterministic: (bool) Whether or not to return deterministic actions. :return: (np.ndarray, np.ndarray) the model's action and the next state (used in recurrent policies) """ - return self.unscale_action(self.select_action(observation)) + return self.unscale_action(self.select_action(observation, deterministic=deterministic)) def train_critic(self, gradient_steps=1, batch_size=100, replay_data=None, tau=0.0): # Update optimizer learning rate From f4546837c37ad90a95b6f5ba0ef1a363ce533e48 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Thu, 7 Nov 2019 17:41:28 +0100 Subject: [PATCH 02/24] Add std to logger --- torchy_baselines/common/base_class.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index d96950d..2f8ef21 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -328,7 +328,7 @@ class BaseRLModel(object): assert env.num_envs == 1 if hasattr(self, 'use_sde') and self.use_sde: - self.policy.reset_noise() + self.actor.reset_noise() while total_steps < n_steps or total_episodes < n_episodes: done = False @@ -394,6 +394,8 @@ class BaseRLModel(object): logger.logkv("fps", fps) logger.logkv('time_elapsed', int(time.time() - self.start_time)) logger.logkv("total timesteps", num_timesteps) + if hasattr(self, 'use_sde') and self.use_sde: + logger.logkv("std", th.exp(self.actor.log_std).mean().item()) logger.dumpkvs() mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0 From 715865a0fe6acf605a9bfc03e62e6e6d25263e3a Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 8 Nov 2019 13:17:38 +0100 Subject: [PATCH 03/24] Add noise clipping --- torchy_baselines/common/base_class.py | 1 + torchy_baselines/td3/policies.py | 9 +++++---- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 2f8ef21..6396dcf 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -351,6 +351,7 @@ class BaseRLModel(object): # Add noise to the action (improve exploration) if action_noise is not None: # NOTE: in the original implementation of TD3, the noise was applied to the unscaled action + # Update(October 2019): Not anymore action = np.clip(action + action_noise(), -1, 1) # Rescale and perform action diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index adbbda1..2a39773 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -7,7 +7,7 @@ from torchy_baselines.common.policies import BasePolicy, register_policy, create class Actor(BaseNetwork): def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU, - use_sde=False, log_std_init=0.0, clip_noise=0.1): + use_sde=False, log_std_init=-2, clip_noise=0.1): super(Actor, self).__init__() self.latent_pi, self.log_std = None, None @@ -36,8 +36,8 @@ class Actor(BaseNetwork): if deterministic: return self.actor_net(latent_pi) noise = th.mm(latent_pi.detach(), self.exploration_mat) - # noise = th.clamp(noise, -self.clip_noise, self.clip_noise) - # TODO: fix clipping + noise = th.clamp(noise, -self.clip_noise, self.clip_noise) + # TODO: fix clipping with squashing ? return th.clamp(self.actor_net(latent_pi) + noise, -1, 1) else: return self.actor_net(obs) @@ -67,7 +67,7 @@ class Critic(BaseNetwork): class TD3Policy(BasePolicy): def __init__(self, observation_space, action_space, learning_rate, net_arch=None, device='cpu', - activation_fn=nn.ReLU, use_sde=False, log_std_init=0.0): + activation_fn=nn.ReLU, use_sde=False, log_std_init=-2, clip_noise=0.1): super(TD3Policy, self).__init__(observation_space, action_space, device) if net_arch is None: @@ -86,6 +86,7 @@ class TD3Policy(BasePolicy): self.actor_kwargs = self.net_args.copy() self.actor_kwargs['use_sde'] = use_sde self.actor_kwargs['log_std_init'] = log_std_init + self.actor_kwargs['clip_noise'] = clip_noise self.actor, self.actor_target = None, None self.critic, self.critic_target = None, None From f2a61949ae16b591de378d274fc694afaa84b693 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 8 Nov 2019 14:59:42 +0100 Subject: [PATCH 04/24] Bump version and change default noise clipping --- setup.py | 2 +- torchy_baselines/__init__.py | 2 +- torchy_baselines/td3/policies.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/setup.py b/setup.py index 89a3b57..640112e 100644 --- a/setup.py +++ b/setup.py @@ -34,7 +34,7 @@ setup(name='torchy_baselines', license="MIT", long_description="", long_description_content_type='text/markdown', - version="0.0.4", + version="0.0.6a", ) # python setup.py sdist diff --git a/torchy_baselines/__init__.py b/torchy_baselines/__init__.py index a5896e6..b8383db 100644 --- a/torchy_baselines/__init__.py +++ b/torchy_baselines/__init__.py @@ -4,4 +4,4 @@ from torchy_baselines.ppo import PPO from torchy_baselines.sac import SAC from torchy_baselines.td3 import TD3 -__version__ = "0.0.5a" +__version__ = "0.0.6a" diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 2a39773..19cff0f 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -7,7 +7,7 @@ from torchy_baselines.common.policies import BasePolicy, register_policy, create class Actor(BaseNetwork): def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU, - use_sde=False, log_std_init=-2, clip_noise=0.1): + use_sde=False, log_std_init=-2, clip_noise=0.5): super(Actor, self).__init__() self.latent_pi, self.log_std = None, None @@ -67,7 +67,7 @@ class Critic(BaseNetwork): class TD3Policy(BasePolicy): def __init__(self, observation_space, action_space, learning_rate, net_arch=None, device='cpu', - activation_fn=nn.ReLU, use_sde=False, log_std_init=-2, clip_noise=0.1): + activation_fn=nn.ReLU, use_sde=False, log_std_init=-2, clip_noise=0.5): super(TD3Policy, self).__init__(observation_space, action_space, device) if net_arch is None: From a08382faab0d350ef1e9a627bdeb95e66164bea3 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Tue, 12 Nov 2019 18:37:13 +0100 Subject: [PATCH 05/24] Add sde update for TD3 --- torchy_baselines/common/base_class.py | 40 ++++++++++++++++++++++++--- torchy_baselines/td3/policies.py | 23 ++++++++++++--- torchy_baselines/td3/td3.py | 38 ++++++++++++++++++++++++- 3 files changed, 92 insertions(+), 9 deletions(-) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 6396dcf..bf71232 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -57,6 +57,8 @@ class BaseRLModel(object): self.replay_buffer = None self.seed = seed self.action_noise = None + # Used for SDE only + self.rollout_data = None # Track the training progress (from 1 to 0) # this is used to update the learning rate self._current_progress = 1 @@ -113,7 +115,7 @@ class BaseRLModel(object): (no need for symmetric action space) """ low, high = self.action_space.low, self.action_space.high - return low + (0.5 * (scaled_action + 1.0) * (high - low)) + return low + (0.5 * (scaled_action + 1.0) * (high - low)) def _setup_learning_rate(self): """Transform to callable if needed.""" @@ -214,12 +216,14 @@ class BaseRLModel(object): Return a trained model. :param total_timesteps: (int) The total number of samples to train on - :param seed: (int) The initial seed for training, if None: keep current seed :param callback: (function (dict, dict)) -> boolean function called at every steps with state of the algorithm. It takes the local and global variables. If it returns False, training is aborted. :param log_interval: (int) The number of timesteps before logging. :param tb_log_name: (str) the name of the run for tensorboard log :param reset_num_timesteps: (bool) whether or not to reset the current timestep number (used in logging) + :param eval_env: (gym.Env) + :param eval_freq: (int) + :param n_eval_episodes: (int) :return: (BaseRLModel) the trained model """ pass @@ -327,8 +331,12 @@ class BaseRLModel(object): assert isinstance(env, VecEnv) assert env.num_envs == 1 + self.rollout_data = None if hasattr(self, 'use_sde') and self.use_sde: self.actor.reset_noise() + # Reset rollout data + self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones']} + # self.rollout_data = {'observations': [], 'actions': [], 'rewards': [], 'returns': [], 'dones': []} while total_steps < n_steps or total_episodes < n_episodes: done = False @@ -367,12 +375,19 @@ class BaseRLModel(object): if replay_buffer is not None: replay_buffer.add(obs, new_obs, action, reward, done_bool) + if self.rollout_data is not None: + # Assume only one env + self.rollout_data['observations'].append(obs[0].copy()) + self.rollout_data['actions'].append(action[0].copy()) + self.rollout_data['rewards'].append(reward[0].copy()) + self.rollout_data['dones'].append(np.array(done_bool[0]).copy()) + obs = new_obs num_timesteps += 1 episode_timesteps += 1 total_steps += 1 - if n_steps > 0 and total_steps >= n_steps: + if 0 < n_steps <= total_steps: break if done: @@ -383,7 +398,8 @@ class BaseRLModel(object): action_noise.reset() # Display training infos - if self.verbose >= 1 and log_interval is not None and (episode_num + total_episodes) % log_interval == 0: + if self.verbose >= 1 and log_interval is not None and ( + 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) @@ -401,4 +417,20 @@ class BaseRLModel(object): mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0 + # Post processing + if self.rollout_data is not None: + for key in ['observations', 'actions', 'rewards', 'dones']: + self.rollout_data[key] = th.FloatTensor(np.array(self.rollout_data[key])).to(self.device) + + self.rollout_data['returns'] = self.rollout_data['rewards'].clone() + # Compute return + last_return = 0.0 + for step in reversed(range(len(self.rollout_data['rewards']))): + if step == len(self.rollout_data['rewards']) - 1: + last_return = self.rollout_data['rewards'][step] + else: + next_non_terminal = 1.0 - self.rollout_data['dones'][step + 1] + last_return = self.rollout_data['rewards'][step] + self.gamma * last_return * next_non_terminal + self.rollout_data['returns'][step] = last_return + return mean_reward, total_steps, total_episodes, obs diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 19cff0f..26a7d8b 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -7,12 +7,12 @@ from torchy_baselines.common.policies import BasePolicy, register_policy, create class Actor(BaseNetwork): def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU, - use_sde=False, log_std_init=-2, clip_noise=0.5): + use_sde=False, log_std_init=-2, clip_noise=None, lr_sde=3e-4): super(Actor, self).__init__() self.latent_pi, self.log_std = None, None self.weights_dist, self.exploration_mat = None, None - self.use_sde = use_sde + self.use_sde, self.sde_optimizer = use_sde, None if use_sde: latent_dim = net_arch[-1] @@ -21,11 +21,25 @@ class Actor(BaseNetwork): self.log_std = nn.Parameter(th.ones(latent_dim, action_dim) * log_std_init) self.actor_net = nn.Sequential(nn.Linear(net_arch[-1], action_dim), nn.Tanh()) self.clip_noise = clip_noise + self.sde_optimizer = th.optim.Adam([self.log_std], lr=lr_sde) self.reset_noise() else: actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_out=True) self.actor_net = nn.Sequential(*actor_net) + def get_distribution_stats(self, obs, action): + with th.no_grad(): + latent_pi = self.latent_pi(obs) + mean_actions = self.actor_net(latent_pi) + variance = th.mm(latent_pi ** 2, th.exp(self.log_std) ** 2) + distribution = Normal(mean_actions, th.sqrt(variance)) + log_prob = distribution.log_prob(action) + if len(log_prob.shape) > 1: + log_prob = log_prob.sum(axis=1) + else: + log_prob = log_prob.sum() + return log_prob, distribution.entropy() + def reset_noise(self): self.weights_dist = Normal(th.zeros_like(self.log_std), th.exp(self.log_std)) self.exploration_mat = self.weights_dist.rsample() @@ -36,7 +50,8 @@ class Actor(BaseNetwork): if deterministic: return self.actor_net(latent_pi) noise = th.mm(latent_pi.detach(), self.exploration_mat) - noise = th.clamp(noise, -self.clip_noise, self.clip_noise) + if self.clip_noise is not None: + noise = th.clamp(noise, -self.clip_noise, self.clip_noise) # TODO: fix clipping with squashing ? return th.clamp(self.actor_net(latent_pi) + noise, -1, 1) else: @@ -67,7 +82,7 @@ class Critic(BaseNetwork): class TD3Policy(BasePolicy): def __init__(self, observation_space, action_space, learning_rate, net_arch=None, device='cpu', - activation_fn=nn.ReLU, use_sde=False, log_std_init=-2, clip_noise=0.5): + activation_fn=nn.ReLU, use_sde=False, log_std_init=-2, clip_noise=None): super(TD3Policy, self).__init__(observation_space, action_space, device) if net_arch is None: diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 46943a9..90ab127 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -52,7 +52,7 @@ class TD3(BaseRLModel): policy_delay=2, learning_starts=100, gamma=0.99, batch_size=100, train_freq=-1, gradient_steps=-1, n_episodes_rollout=1, tau=0.005, action_noise=None, target_policy_noise=0.2, target_noise_clip=0.5, - use_sde=False, + use_sde=False, sde_max_grad_norm=1, sde_ent_coef=0.0, tensorboard_log=None, create_eval_env=False, policy_kwargs=None, verbose=0, seed=0, device='auto', _init_setup_model=True): @@ -73,7 +73,10 @@ class TD3(BaseRLModel): self.policy_delay = policy_delay self.target_noise_clip = target_noise_clip self.target_policy_noise = target_policy_noise + self.use_sde = use_sde + self.sde_max_grad_norm = sde_max_grad_norm + self.sde_ent_coef = sde_ent_coef if _init_setup_model: self._setup_model() @@ -191,6 +194,37 @@ class TD3(BaseRLModel): if gradient_step % policy_delay == 0: self.train_actor(replay_data=replay_data, tau_actor=self.tau, tau_critic=self.tau) + def train_sde(self): + # Update optimizer learning rate + # self._update_learning_rate(self.policy.optimizer) + + # Unpack + obs, action, returns = self.rollout_data['observations'], self.rollout_data['actions'], self.rollout_data['returns'] + + # TODO: avoid second computation of everything because of the gradient + log_prob, entropy = self.actor.get_distribution_stats(obs, action) + + # Normalize returns + # returns = (returns - returns.mean()) / (returns.std() + 1e-8) + + policy_loss = -(returns * log_prob).mean() + + # Entropy loss favor exploration + entropy_loss = -th.mean(entropy) + + loss = policy_loss + self.sde_ent_coef * entropy_loss + + # Optimization step + self.actor.sde_optimizer.zero_grad() + loss.backward() + # print(self.actor.log_std.grad.mean().item(), self.actor.log_std.grad.max().item(), self.actor.log_std.grad.min().item()) + # print(self.actor.log_std.mean().item(), self.actor.log_std.max().item(), self.actor.log_std.min().item()) + # Clip grad norm + th.nn.utils.clip_grad_norm_([self.actor.log_std], self.sde_max_grad_norm) + self.actor.sde_optimizer.step() + + del self.rollout_data + def learn(self, total_timesteps, callback=None, log_interval=4, eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="TD3", reset_num_timesteps=True): @@ -226,6 +260,8 @@ class TD3(BaseRLModel): gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay) + if self.use_sde: + self.train_sde() # Evaluate episode if 0 < eval_freq <= timesteps_since_eval and eval_env is not None: From 623e3cf4f9076f6e2b5523efc8c4071b264b3e62 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Tue, 12 Nov 2019 18:38:42 +0100 Subject: [PATCH 06/24] Add SDE lr --- torchy_baselines/td3/policies.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 26a7d8b..9c93f10 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -82,7 +82,8 @@ class Critic(BaseNetwork): class TD3Policy(BasePolicy): def __init__(self, observation_space, action_space, learning_rate, net_arch=None, device='cpu', - activation_fn=nn.ReLU, use_sde=False, log_std_init=-2, clip_noise=None): + activation_fn=nn.ReLU, use_sde=False, log_std_init=-2, + clip_noise=None, lr_sde=3e-4): super(TD3Policy, self).__init__(observation_space, action_space, device) if net_arch is None: @@ -102,6 +103,7 @@ class TD3Policy(BasePolicy): self.actor_kwargs['use_sde'] = use_sde self.actor_kwargs['log_std_init'] = log_std_init self.actor_kwargs['clip_noise'] = clip_noise + self.actor_kwargs['lr_sde'] = lr_sde self.actor, self.actor_target = None, None self.critic, self.critic_target = None, None From d725d011866f6fa39cbf9701e02bd8a279cf1c64 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Tue, 12 Nov 2019 22:32:21 +0100 Subject: [PATCH 07/24] Normalize returns --- torchy_baselines/td3/td3.py | 1 + 1 file changed, 1 insertion(+) diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 90ab127..2dbdd5f 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -206,6 +206,7 @@ class TD3(BaseRLModel): # Normalize returns # returns = (returns - returns.mean()) / (returns.std() + 1e-8) + returns = (returns - returns.mean()) policy_loss = -(returns * log_prob).mean() From da325a0ba720c744b5fc61fd29df78b72e01f10a Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Wed, 13 Nov 2019 13:02:37 +0100 Subject: [PATCH 08/24] Solve NaN issue and reduce number of parameters --- README.md | 3 +++ torchy_baselines/td3/policies.py | 26 +++++++++++++++++++++----- torchy_baselines/td3/td3.py | 17 ++++++++++++++--- 3 files changed, 38 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index b5624ec..74b55c6 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,9 @@ TODO: - save/load - better predict - complete logger +- 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? Later: - get_parameters / set_parameters diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 9c93f10..89dc019 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -7,18 +7,27 @@ from torchy_baselines.common.policies import BasePolicy, register_policy, create class Actor(BaseNetwork): def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU, - use_sde=False, log_std_init=-2, clip_noise=None, lr_sde=3e-4): + use_sde=False, log_std_init=-2, clip_noise=None, + lr_sde=3e-4, full_std=False): super(Actor, self).__init__() self.latent_pi, self.log_std = None, None self.weights_dist, self.exploration_mat = None, None self.use_sde, self.sde_optimizer = use_sde, None + self.action_dim = action_dim + self.full_std = full_std if use_sde: latent_dim = net_arch[-1] latent_pi = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_out=False) self.latent_pi = nn.Sequential(*latent_pi) - self.log_std = nn.Parameter(th.ones(latent_dim, action_dim) * log_std_init) + if full_std: + self.log_std = nn.Parameter(th.ones(latent_dim, action_dim) * log_std_init) + else: + # Reduce the number of parameters: + self.log_std = nn.Parameter(th.ones(latent_dim, 1) * log_std_init) + + self.latent_dim = latent_dim self.actor_net = nn.Sequential(nn.Linear(net_arch[-1], action_dim), nn.Tanh()) self.clip_noise = clip_noise self.sde_optimizer = th.optim.Adam([self.log_std], lr=lr_sde) @@ -27,12 +36,19 @@ class Actor(BaseNetwork): actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_out=True) self.actor_net = nn.Sequential(*actor_net) + def get_log_std(self): + if self.full_std: + return self.log_std + # Reduce the number of parameters: + return th.ones((self.latent_dim, self.action_dim)) * self.log_std + def get_distribution_stats(self, obs, action): with th.no_grad(): latent_pi = self.latent_pi(obs) mean_actions = self.actor_net(latent_pi) - variance = th.mm(latent_pi ** 2, th.exp(self.log_std) ** 2) - distribution = Normal(mean_actions, th.sqrt(variance)) + + variance = th.mm(latent_pi ** 2, th.exp(self.get_log_std()) ** 2) + distribution = Normal(mean_actions, th.sqrt(variance + 1e-5)) log_prob = distribution.log_prob(action) if len(log_prob.shape) > 1: log_prob = log_prob.sum(axis=1) @@ -41,7 +57,7 @@ class Actor(BaseNetwork): return log_prob, distribution.entropy() def reset_noise(self): - self.weights_dist = Normal(th.zeros_like(self.log_std), th.exp(self.log_std)) + self.weights_dist = Normal(th.zeros_like(self.get_log_std()), th.exp(self.get_log_std())) self.exploration_mat = self.weights_dist.rsample() def forward(self, obs, deterministic=True): diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 2dbdd5f..ce6b9a4 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -206,7 +206,11 @@ class TD3(BaseRLModel): # Normalize returns # returns = (returns - returns.mean()) / (returns.std() + 1e-8) - returns = (returns - returns.mean()) + # returns = (returns - returns.mean()) + with th.no_grad(): + current_q1, current_q2 = self.critic(obs, action) + # Alternatively use the q value + returns = (returns - th.min(current_q1, current_q2)) policy_loss = -(returns * log_prob).mean() @@ -218,6 +222,11 @@ class TD3(BaseRLModel): # Optimization step self.actor.sde_optimizer.zero_grad() loss.backward() + + assert not th.isnan(log_prob).any(), log_prob + assert not th.isnan(entropy).any() + assert not th.isnan(self.actor.log_std.grad).any() + assert not th.isnan(self.actor.log_std).any() # print(self.actor.log_std.grad.mean().item(), self.actor.log_std.grad.max().item(), self.actor.log_std.grad.min().item()) # print(self.actor.log_std.mean().item(), self.actor.log_std.max().item(), self.actor.log_std.min().item()) # Clip grad norm @@ -259,11 +268,13 @@ class TD3(BaseRLModel): 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, policy_delay=self.policy_delay) if self.use_sde: self.train_sde() + gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps + self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay) + + # Evaluate episode if 0 < eval_freq <= timesteps_since_eval and eval_env is not None: timesteps_since_eval %= eval_freq From 8aac10f3fa24cf3051a0a67a39cbdd55270feb57 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Wed, 13 Nov 2019 14:38:18 +0100 Subject: [PATCH 09/24] Fix device --- torchy_baselines/td3/policies.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 89dc019..96fae45 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -40,7 +40,7 @@ class Actor(BaseNetwork): if self.full_std: return self.log_std # Reduce the number of parameters: - return th.ones((self.latent_dim, self.action_dim)) * self.log_std + return th.ones((self.latent_dim, self.action_dim)).to(self.log_std.device) * self.log_std def get_distribution_stats(self, obs, action): with th.no_grad(): From 5278a6f3f814ee352d769028ddf1bf11d6077691 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Thu, 14 Nov 2019 14:35:00 +0100 Subject: [PATCH 10/24] Testing off policy normalization --- README.md | 1 + tests/test_vec_normalize.py | 14 +++++++ torchy_baselines/cem_rl/cem_rl.py | 4 +- torchy_baselines/common/base_class.py | 11 +++++- torchy_baselines/common/buffers.py | 29 ++++++++++---- torchy_baselines/common/vec_env/__init__.py | 31 +++++++++++++++ .../common/vec_env/vec_normalize.py | 39 ++++++++++++++++--- torchy_baselines/ppo/ppo.py | 13 ++----- torchy_baselines/sac/sac.py | 4 +- torchy_baselines/td3/td3.py | 8 ++-- 10 files changed, 125 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index 74b55c6..c13574d 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,7 @@ 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? +- DEBUG normalization with replay buffer (apparently pb with observation normalization) Later: - get_parameters / set_parameters diff --git a/tests/test_vec_normalize.py b/tests/test_vec_normalize.py index 78437a3..7353de0 100644 --- a/tests/test_vec_normalize.py +++ b/tests/test_vec_normalize.py @@ -1,9 +1,11 @@ import gym +import pytest import numpy as np from torchy_baselines.common.running_mean_std import RunningMeanStd from torchy_baselines.common.vec_env.dummy_vec_env import DummyVecEnv from torchy_baselines.common.vec_env.vec_normalize import VecNormalize +from torchy_baselines import CEMRL, SAC, TD3 ENV_ID = 'Pendulum-v0' @@ -39,3 +41,15 @@ def test_vec_env(): actions = [env.action_space.sample()] obs, _, done, _ = env.step(actions) assert np.max(obs) <= 10 + + +@pytest.mark.parametrize("model_class", [SAC, TD3, CEMRL]) +def test_offpolicy_normalization(model_class): + env = DummyVecEnv([lambda: gym.make(ENV_ID)]) + env = VecNormalize(env, norm_obs=False, norm_reward=False, clip_obs=10., clip_reward=10.) + + eval_env = DummyVecEnv([lambda: gym.make(ENV_ID)]) + eval_env = VecNormalize(eval_env, norm_obs=True, norm_reward=False, clip_obs=10., clip_reward=10.) + + model = model_class('MlpPolicy', env, verbose=1) + model.learn(total_timesteps=10000, eval_env=eval_env, eval_freq=1000) diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index ad798ba..ff7100d 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -5,6 +5,7 @@ import torch as th from torchy_baselines.cem_rl.cem import CEM from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.td3.td3 import TD3 +from torchy_baselines.common.vec_env import sync_envs_normalization class CEMRL(TD3): @@ -102,7 +103,7 @@ class CEMRL(TD3): n_training_steps = 2 * (actor_steps // self.n_grad) for it in range(n_training_steps): # Sample replay buffer - replay_data = self.replay_buffer.sample(self.batch_size) + replay_data = self.replay_buffer.sample(self.batch_size, env=self._vec_normalize_env) self.train_critic(replay_data=replay_data) # Delayed policy updates @@ -117,6 +118,7 @@ class CEMRL(TD3): timesteps_since_eval %= eval_freq self.actor.load_from_vector(self.es.mu) + sync_envs_normalization(self.env, eval_env) mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes) evaluations.append(mean_reward) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index bf71232..fd82905 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -8,7 +8,7 @@ import numpy as np from torchy_baselines.common.policies import get_policy_from_name from torchy_baselines.common.utils import set_random_seed, get_schedule_fn, update_learning_rate -from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv +from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize from torchy_baselines.common.monitor import Monitor from torchy_baselines.common import logger @@ -46,6 +46,8 @@ class BaseRLModel(object): print("Using {} device".format(self.device)) self.env = env + # get VecNormalize object if needed + 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 @@ -373,6 +375,13 @@ class BaseRLModel(object): # Store data in replay buffer if replay_buffer is not None: + # Store only the unnormalized version + if self._vec_normalize_env is not None: + # TODO: save it instead of unnormalizing + obs = self._vec_normalize_env.unnormalize_obs(obs) + new_obs = self._vec_normalize_env.get_original_obs() + reward = self._vec_normalize_env.get_original_reward() + replay_buffer.add(obs, new_obs, action, reward, done_bool) if self.rollout_data is not None: diff --git a/torchy_baselines/common/buffers.py b/torchy_baselines/common/buffers.py index 7ca61b8..2c3d1fb 100644 --- a/torchy_baselines/common/buffers.py +++ b/torchy_baselines/common/buffers.py @@ -1,6 +1,8 @@ import numpy as np import torch as th +from torchy_baselines.common.vec_env import unwrap_vec_normalize + class BaseBuffer(object): def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', n_envs=1): @@ -43,15 +45,26 @@ class BaseBuffer(object): self.pos = 0 self.full = False - def sample(self, batch_size): + def sample(self, batch_size, env=None): 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) + return self._get_samples(batch_inds, env=env) - def _get_samples(self, batch_inds): + def _get_samples(self, batch_inds, env=None): raise NotImplementedError() + def _normalize_obs(self, obs, env=None): + if env is not None: + # TODO: get rid of pytorch - numpy conversion + return th.FloatTensor(env.normalize_obs(obs.numpy())) + return obs + + def _normalize_reward(self, reward, env=None): + if env is not None: + return th.FloatTensor(env.normalize_reward(reward.numpy())) + return reward + class ReplayBuffer(BaseBuffer): """ @@ -81,12 +94,12 @@ class ReplayBuffer(BaseBuffer): self.full = True self.pos = 0 - def _get_samples(self, batch_inds): - return (self.observations[batch_inds, 0, :].to(self.device), + def _get_samples(self, batch_inds, env=None): + return (self._normalize_obs(self.observations[batch_inds, 0, :], env).to(self.device), self.actions[batch_inds, 0, :].to(self.device), - self.next_observations[batch_inds, 0, :].to(self.device), + self._normalize_obs(self.next_observations[batch_inds, 0, :], env).to(self.device), self.dones[batch_inds].to(self.device), - self.rewards[batch_inds].to(self.device)) + self._normalize_reward(self.rewards[batch_inds], env).to(self.device)) class RolloutBuffer(BaseBuffer): @@ -184,7 +197,7 @@ class RolloutBuffer(BaseBuffer): yield self._get_samples(indices[start_idx:start_idx + batch_size]) start_idx += batch_size - def _get_samples(self, batch_inds): + def _get_samples(self, batch_inds, env=None): return (self.observations[batch_inds].to(self.device), self.actions[batch_inds].to(self.device), self.values[batch_inds].flatten().to(self.device), diff --git a/torchy_baselines/common/vec_env/__init__.py b/torchy_baselines/common/vec_env/__init__.py index 97f6022..2c542e5 100644 --- a/torchy_baselines/common/vec_env/__init__.py +++ b/torchy_baselines/common/vec_env/__init__.py @@ -1,7 +1,38 @@ # flake8: noqa F401 +from copy import deepcopy + from torchy_baselines.common.vec_env.base_vec_env import AlreadySteppingError, NotSteppingError,\ VecEnv, VecEnvWrapper, CloudpickleWrapper from torchy_baselines.common.vec_env.dummy_vec_env import DummyVecEnv from torchy_baselines.common.vec_env.subproc_vec_env import SubprocVecEnv from torchy_baselines.common.vec_env.vec_frame_stack import VecFrameStack from torchy_baselines.common.vec_env.vec_normalize import VecNormalize + + +def unwrap_vec_normalize(env): + """ + :param env: (gym.Env) + :return: (VecNormalize) + """ + env_tmp = env + while isinstance(env_tmp, VecEnvWrapper): + if isinstance(env_tmp, VecNormalize): + return env_tmp + env_tmp = env_tmp.venv + return None + + +# Define here to avoid circular import +def sync_envs_normalization(env, eval_env): + """ + Sync eval env and train env when using VecNormalize + + :param env: (gym.Env) + :param eval_env: (gym.Env) + """ + env_tmp, eval_env_tmp = env, eval_env + while isinstance(env_tmp, VecEnvWrapper): + if isinstance(env_tmp, VecNormalize): + eval_env_tmp.obs_rms = deepcopy(env_tmp.obs_rms) + env_tmp = env_tmp.venv + eval_env_tmp.venv diff --git a/torchy_baselines/common/vec_env/vec_normalize.py b/torchy_baselines/common/vec_env/vec_normalize.py index 6f2b2b3..0b9797f 100644 --- a/torchy_baselines/common/vec_env/vec_normalize.py +++ b/torchy_baselines/common/vec_env/vec_normalize.py @@ -36,6 +36,7 @@ class VecNormalize(VecEnvWrapper): self.norm_obs = norm_obs self.norm_reward = norm_reward self.old_obs = np.array([]) + self.old_reward = np.array([]) def step_wait(self): """ @@ -46,12 +47,13 @@ class VecNormalize(VecEnvWrapper): """ obs, rews, news, infos = self.venv.step_wait() self.ret = self.ret * self.gamma + rews - self.old_obs = obs + self.old_obs = obs.copy() + self.old_reward = rews.copy() obs = self._normalize_observation(obs) if self.norm_reward: if self.training: self.ret_rms.update(self.ret) - rews = np.clip(rews / np.sqrt(self.ret_rms.var + self.epsilon), -self.clip_reward, self.clip_reward) + rews = self.normalize_reward(rews) self.ret[news] = 0 return obs, rews, news, infos @@ -62,12 +64,31 @@ class VecNormalize(VecEnvWrapper): if self.norm_obs: if self.training: self.obs_rms.update(obs) - obs = np.clip((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.epsilon), -self.clip_obs, - self.clip_obs) - return obs + return self.normalize_obs(obs) else: return obs + def normalize_obs(self, obs): + if self.norm_obs: + return np.clip((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.epsilon), -self.clip_obs, + self.clip_obs) + return obs + + def normalize_reward(self, reward): + if self.norm_reward: + return np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon), -self.clip_reward, self.clip_reward) + return reward + + def unnormalize_obs(self, obs): + if self.norm_obs: + return (obs * np.sqrt(self.obs_rms.var + self.epsilon)) + self.obs_rms.mean + return obs + + def unnormalize_reward(self, reward): + if self.norm_reward: + return reward * np.sqrt(self.ret_rms.var + self.epsilon) + return reward + def get_original_obs(self): """ returns the unnormalized observation @@ -76,6 +97,14 @@ class VecNormalize(VecEnvWrapper): """ return self.old_obs + def get_original_reward(self): + """ + returns the unnormalized observation + + :return: (numpy float) + """ + return self.old_reward + def reset(self): """ Reset all environments diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 2b2a5e6..f3a738a 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -1,6 +1,5 @@ import os import time -from copy import deepcopy import gym from gym import spaces @@ -17,7 +16,7 @@ from torchy_baselines.common.base_class import BaseRLModel from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.common.buffers import RolloutBuffer from torchy_baselines.common.utils import explained_variance, get_schedule_fn -from torchy_baselines.common.vec_env import VecNormalize, VecEnvWrapper +from torchy_baselines.common.vec_env import sync_envs_normalization from torchy_baselines.common import logger from torchy_baselines.ppo.policies import PPOPolicy @@ -295,14 +294,8 @@ class PPO(BaseRLModel): # Evaluate agent if 0 < eval_freq <= timesteps_since_eval and eval_env is not None: timesteps_since_eval %= eval_freq - # TODO: move that to the base class - # Sync eval env and train env when using VecNormalize - env_tmp, eval_env_tmp = self.env, eval_env - while isinstance(env_tmp, VecEnvWrapper): - if isinstance(env_tmp, VecNormalize): - eval_env_tmp.obs_rms = deepcopy(env_tmp.obs_rms) - env_tmp = env_tmp.venv - eval_env_tmp.venv + sync_envs_normalization(self.env, eval_env) + mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes) if self.tb_writer is not None: self.tb_writer.add_scalar('Eval/reward', mean_reward, self.num_timesteps) diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index bad470a..f4c5c34 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -8,6 +8,7 @@ from torchy_baselines.common.base_class import BaseRLModel from torchy_baselines.common.buffers import ReplayBuffer from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.sac.policies import SACPolicy +from torchy_baselines.common.vec_env import sync_envs_normalization class SAC(BaseRLModel): @@ -162,7 +163,7 @@ class SAC(BaseRLModel): for gradient_step in range(gradient_steps): # Sample replay buffer - replay_data = self.replay_buffer.sample(batch_size) + replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env) obs, action_batch, next_obs, done, reward = replay_data @@ -266,6 +267,7 @@ class SAC(BaseRLModel): # Evaluate episode if 0 < eval_freq <= timesteps_since_eval and eval_env is not None: timesteps_since_eval %= eval_freq + sync_envs_normalization(self.env, eval_env) mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes) evaluations.append(mean_reward) if self.verbose > 0: diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index ce6b9a4..6c1dbcf 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -8,6 +8,7 @@ from torchy_baselines.common.base_class import BaseRLModel from torchy_baselines.common.buffers import ReplayBuffer from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.td3.policies import TD3Policy +from torchy_baselines.common.vec_env import sync_envs_normalization class TD3(BaseRLModel): @@ -123,7 +124,7 @@ class TD3(BaseRLModel): for gradient_step in range(gradient_steps): # Sample replay buffer if replay_data is None: - obs, action, next_obs, done, reward = self.replay_buffer.sample(batch_size) + obs, action, next_obs, done, reward = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env) else: obs, action, next_obs, done, reward = replay_data @@ -162,7 +163,7 @@ class TD3(BaseRLModel): for gradient_step in range(gradient_steps): # Sample replay buffer if replay_data is None: - obs, _, next_obs, done, reward = self.replay_buffer.sample(batch_size) + obs, _, next_obs, done, reward = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env) else: obs, _, next_obs, done, reward = replay_data @@ -187,7 +188,7 @@ class TD3(BaseRLModel): for gradient_step in range(gradient_steps): # Sample replay buffer - replay_data = self.replay_buffer.sample(batch_size) + replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env) self.train_critic(replay_data=replay_data) # Delayed policy updates @@ -278,6 +279,7 @@ class TD3(BaseRLModel): # Evaluate episode if 0 < eval_freq <= timesteps_since_eval and eval_env is not None: timesteps_since_eval %= eval_freq + sync_envs_normalization(self.env, eval_env) mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes) evaluations.append(mean_reward) if self.verbose > 0: From f7195443864255541bae3b92ff169e6dda7d8ef3 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Thu, 14 Nov 2019 14:35:47 +0100 Subject: [PATCH 11/24] Change default batch size for SAC --- torchy_baselines/sac/sac.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index f4c5c34..5def28f 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -54,7 +54,7 @@ class SAC(BaseRLModel): :param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance """ def __init__(self, policy, env, learning_rate=3e-4, buffer_size=int(1e6), - learning_starts=100, batch_size=64, + learning_starts=100, batch_size=256, tau=0.005, ent_coef='auto', target_update_interval=1, train_freq=1, gradient_steps=1, n_episodes_rollout=-1, target_entropy='auto', action_noise=None, From cdb62a93fea037b1cd046f7dd9ded4f7fd09ac5f Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 15 Nov 2019 11:00:31 +0100 Subject: [PATCH 12/24] Bug fix for off-policy normalization Now working properly --- README.md | 1 - tests/test_vec_normalize.py | 4 ++-- torchy_baselines/common/base_class.py | 21 ++++++++++++++++----- 3 files changed, 18 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index c13574d..74b55c6 100644 --- a/README.md +++ b/README.md @@ -24,7 +24,6 @@ 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? -- DEBUG normalization with replay buffer (apparently pb with observation normalization) Later: - get_parameters / set_parameters diff --git a/tests/test_vec_normalize.py b/tests/test_vec_normalize.py index 7353de0..181c42d 100644 --- a/tests/test_vec_normalize.py +++ b/tests/test_vec_normalize.py @@ -46,10 +46,10 @@ def test_vec_env(): @pytest.mark.parametrize("model_class", [SAC, TD3, CEMRL]) def test_offpolicy_normalization(model_class): env = DummyVecEnv([lambda: gym.make(ENV_ID)]) - env = VecNormalize(env, norm_obs=False, norm_reward=False, clip_obs=10., clip_reward=10.) + env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_obs=10., clip_reward=10.) eval_env = DummyVecEnv([lambda: gym.make(ENV_ID)]) eval_env = VecNormalize(eval_env, norm_obs=True, norm_reward=False, clip_obs=10., clip_reward=10.) model = model_class('MlpPolicy', env, verbose=1) - model.learn(total_timesteps=10000, eval_env=eval_env, eval_freq=1000) + model.learn(total_timesteps=1000, eval_env=eval_env, eval_freq=500) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index fd82905..a31c9e0 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -333,6 +333,11 @@ class BaseRLModel(object): assert isinstance(env, VecEnv) assert env.num_envs == 1 + # Retrieve unnormalized observation for saving into the buffer + if self._vec_normalize_env is not None: + obs_ = self._vec_normalize_env.get_original_obs() + + self.rollout_data = None if hasattr(self, 'use_sde') and self.use_sde: self.actor.reset_noise() @@ -377,12 +382,13 @@ class BaseRLModel(object): if replay_buffer is not None: # Store only the unnormalized version if self._vec_normalize_env is not None: - # TODO: save it instead of unnormalizing - obs = self._vec_normalize_env.unnormalize_obs(obs) - new_obs = self._vec_normalize_env.get_original_obs() - reward = self._vec_normalize_env.get_original_reward() + new_obs_ = self._vec_normalize_env.get_original_obs() + reward_ = self._vec_normalize_env.get_original_reward() + else: + # Avoid changing the original ones + obs_, new_obs_, reward_ = obs, new_obs, reward - replay_buffer.add(obs, new_obs, action, reward, done_bool) + replay_buffer.add(obs_, new_obs_, action, reward_, done_bool) if self.rollout_data is not None: # Assume only one env @@ -392,6 +398,11 @@ class BaseRLModel(object): self.rollout_data['dones'].append(np.array(done_bool[0]).copy()) obs = new_obs + # Save the true unnormalized observation + # otherwise obs_ = self._vec_normalize_env.unnormalize_obs(obs) + # is a good approximation + if self._vec_normalize_env is not None: + obs_ = new_obs_ num_timesteps += 1 episode_timesteps += 1 From fb6407285924af4d58e081e644fb9bba3d663947 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 15 Nov 2019 11:07:49 +0100 Subject: [PATCH 13/24] Update sde test --- tests/test_sde.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_sde.py b/tests/test_sde.py index 955f9ab..048469b 100644 --- a/tests/test_sde.py +++ b/tests/test_sde.py @@ -53,8 +53,8 @@ def test_state_dependent_noise(model_class): eval_env = VecNormalize(DummyVecEnv([lambda: Monitor(gym.make(env_id))]), training=False, norm_reward=False) model = model_class('MlpPolicy', env, n_steps=200, max_grad_norm=1, use_rms_prop=False, use_sde=True, ent_coef=0.00, verbose=1, create_eval_env=True, learning_rate=3e-4, - policy_kwargs=dict(log_std_init=0.0, ortho_init=False, net_arch=[256, dict(pi=[256], vf=[256])]), seed=None) - # model.learn(total_timesteps=int(20000), log_interval=5, eval_freq=10000, eval_env=eval_env) + policy_kwargs=dict(log_std_init=0.0, ortho_init=False, net_arch=[256, dict(pi=[256], vf=[256])]), + seed=None) model.learn(total_timesteps=int(1000), log_interval=5, eval_freq=500, eval_env=eval_env) @@ -62,4 +62,4 @@ def test_state_dependent_noise(model_class): def test_state_dependent_offpolicy_noise(model_class): model = model_class('MlpPolicy', 'Pendulum-v0', use_sde=True, seed=None, create_eval_env=True, verbose=1, policy_kwargs=dict(log_std_init=-2)) - model.learn(total_timesteps=int(20000), eval_freq=1000) + model.learn(total_timesteps=int(1000), eval_freq=500) From ad32aa60f3659b8bbb72fa44cdc4addaa73e85ad Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Mon, 18 Nov 2019 16:03:08 +0100 Subject: [PATCH 14/24] Add sde scheduler --- tests/test_sde.py | 10 ++++++++++ torchy_baselines/td3/td3.py | 17 +++++++++++++---- 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/tests/test_sde.py b/tests/test_sde.py index 7c6f5ff..09b48e8 100644 --- a/tests/test_sde.py +++ b/tests/test_sde.py @@ -65,3 +65,13 @@ def test_state_dependent_offpolicy_noise(model_class): model = model_class('MlpPolicy', 'Pendulum-v0', use_sde=True, seed=None, create_eval_env=True, verbose=1, policy_kwargs=dict(log_std_init=-2)) model.learn(total_timesteps=int(1000), eval_freq=500) + + +def test_scheduler(): + def scheduler(progress): + return -2.0 * progress + 1 + + model = TD3('MlpPolicy', 'Pendulum-v0', use_sde=True, seed=None, create_eval_env=True, + verbose=1, sde_log_std_scheduler=scheduler) + model.learn(total_timesteps=int(1000), eval_freq=500) + assert th.isclose(model.actor.log_std, th.ones_like(model.actor.log_std)).all() diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 6c1dbcf..ddd7e9d 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -40,6 +40,9 @@ class TD3(BaseRLModel): :param target_noise_clip: (float) Limit for absolute value of target policy smoothing noise. :param use_sde: (bool) Whether to use State Dependent Exploration (SDE) instead of action noise exploration (default: False) + :param sde_max_grad_norm: (float) + :param sde_ent_coef: (float) + :param sde_log_std_scheduler: (callable) :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 policy_kwargs: (dict) additional arguments to be passed to the policy on creation @@ -53,7 +56,7 @@ class TD3(BaseRLModel): policy_delay=2, learning_starts=100, gamma=0.99, batch_size=100, train_freq=-1, gradient_steps=-1, n_episodes_rollout=1, tau=0.005, action_noise=None, target_policy_noise=0.2, target_noise_clip=0.5, - use_sde=False, sde_max_grad_norm=1, sde_ent_coef=0.0, + use_sde=False, sde_max_grad_norm=1, sde_ent_coef=0.0, sde_log_std_scheduler=None, tensorboard_log=None, create_eval_env=False, policy_kwargs=None, verbose=0, seed=0, device='auto', _init_setup_model=True): @@ -78,6 +81,7 @@ class TD3(BaseRLModel): self.use_sde = use_sde self.sde_max_grad_norm = sde_max_grad_norm self.sde_ent_coef = sde_ent_coef + self.sde_log_std_scheduler = sde_log_std_scheduler if _init_setup_model: self._setup_model() @@ -228,8 +232,7 @@ class TD3(BaseRLModel): assert not th.isnan(entropy).any() assert not th.isnan(self.actor.log_std.grad).any() assert not th.isnan(self.actor.log_std).any() - # print(self.actor.log_std.grad.mean().item(), self.actor.log_std.grad.max().item(), self.actor.log_std.grad.min().item()) - # print(self.actor.log_std.mean().item(), self.actor.log_std.max().item(), self.actor.log_std.min().item()) + # Clip grad norm th.nn.utils.clip_grad_norm_([self.actor.log_std], self.sde_max_grad_norm) self.actor.sde_optimizer.step() @@ -270,7 +273,13 @@ class TD3(BaseRLModel): self.num_timesteps, episode_num, episode_timesteps, episode_reward)) if self.use_sde: - self.train_sde() + if self.sde_log_std_scheduler is not None: + # Call the scheduler + value = self.sde_log_std_scheduler(self._current_progress) + self.actor.log_std.data = th.ones_like(self.actor.log_std) * value + else: + # On-policy gradient + self.train_sde() gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay) From 99ea0b3a54cc892597bbd2f432212267afdd7284 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 22 Nov 2019 11:42:58 +0100 Subject: [PATCH 15/24] Cleanup --- torchy_baselines/a2c/a2c.py | 2 +- torchy_baselines/common/base_class.py | 8 ++++++-- torchy_baselines/ppo/policies.py | 6 +++--- torchy_baselines/ppo/ppo.py | 4 ++-- 4 files changed, 12 insertions(+), 8 deletions(-) diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py index 41ca60c..36ab9aa 100644 --- a/torchy_baselines/a2c/a2c.py +++ b/torchy_baselines/a2c/a2c.py @@ -92,7 +92,7 @@ class A2C(PPO): action = action.long().flatten() # TODO: avoid second computation of everything because of the gradient - values, log_prob, entropy = self.policy.get_policy_stats(obs, action) + values, log_prob, entropy = self.policy.evaluate_actions(obs, action) values = values.flatten() # Normalize advantage (not present in the original implementation) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index a8d56ae..36cd0c5 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -107,6 +107,9 @@ class BaseRLModel(object): """ Rescale the action from [low, high] to [-1, 1] (no need for symmetric action space) + + :param action: (np.ndarray) + :return: (np.ndarray) """ low, high = self.action_space.low, self.action_space.high return 2.0 * ((action - low) / (high - low)) - 1.0 @@ -115,6 +118,9 @@ class BaseRLModel(object): """ Rescale the action from [-1, 1] to [low, high] (no need for symmetric action space) + + :param scaled_action: (np.ndarray) + :return: (np.ndarray) """ low, high = self.action_space.low, self.action_space.high return low + (0.5 * (scaled_action + 1.0) * (high - low)) @@ -343,13 +349,11 @@ class BaseRLModel(object): if self._vec_normalize_env is not None: obs_ = self._vec_normalize_env.get_original_obs() - self.rollout_data = None if hasattr(self, 'use_sde') and self.use_sde: self.actor.reset_noise() # Reset rollout data self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones']} - # self.rollout_data = {'observations': [], 'actions': [], 'rewards': [], 'returns': [], 'dones': []} while total_steps < n_steps or total_episodes < n_episodes: done = False diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 1e3b25c..1cb7ca4 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -5,7 +5,7 @@ import torch as th import torch.nn as nn import numpy as np -from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp +from torchy_baselines.common.policies import BasePolicy, register_policy from torchy_baselines.common.distributions import make_proba_distribution,\ DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution @@ -30,7 +30,7 @@ class MlpExtractor(nn.Module): Adapted from Stable Baselines. - :param flat_observations: (th.Tensor) The observations to base policy and value function on. + :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. @@ -185,7 +185,7 @@ class PPOPolicy(BasePolicy): action, _ = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic) return action.detach().cpu().numpy() - def get_policy_stats(self, obs, action, deterministic=False): + def evaluate_actions(self, obs, action, deterministic=False): latent_pi, latent_vf = self._get_latent(obs) _, action_distribution = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic) log_prob = action_distribution.log_prob(action) diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index f3a738a..771b5d1 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -204,7 +204,7 @@ class PPO(BaseRLModel): # Convert discrete action for float to long action = action.long().flatten() - values, log_prob, entropy = self.policy.get_policy_stats(obs, action) + values, log_prob, entropy = self.policy.evaluate_actions(obs, action) values = values.flatten() # Normalize advantage advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8) @@ -241,7 +241,7 @@ class PPO(BaseRLModel): approx_kl_divs.append(th.mean(old_log_prob - log_prob).detach().cpu().numpy()) if self.target_kl is not None and np.mean(approx_kl_divs) > 1.5 * self.target_kl: - print("Early stopping at step {} due to reaching max kl: {:.2f}".format(it, np.mean(approx_kl_divs))) + print("Early stopping at step {} due to reaching max kl: {:.2f}".format(gradient_step, np.mean(approx_kl_divs))) break explained_var = explained_variance(self.rollout_buffer.returns.flatten().cpu().numpy(), From ea3902cd32fd10af16c4a806c96a3bc7be97241f Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 22 Nov 2019 13:03:57 +0100 Subject: [PATCH 16/24] Add doc for CEM-RL --- torchy_baselines/cem_rl/cem.py | 49 +++++++++++++++++----------- torchy_baselines/cem_rl/cem_rl.py | 53 +++++++++++++++++++++++++------ torchy_baselines/ppo/ppo.py | 6 ++-- 3 files changed, 78 insertions(+), 30 deletions(-) diff --git a/torchy_baselines/cem_rl/cem.py b/torchy_baselines/cem_rl/cem.py index 57513c5..bea75fc 100644 --- a/torchy_baselines/cem_rl/cem.py +++ b/torchy_baselines/cem_rl/cem.py @@ -1,24 +1,27 @@ import numpy as np - # TODO: add more from https://github.com/hardmaru/estool/blob/master/es.py # or https://github.com/facebookresearch/nevergrad class CEM(object): - - """ - Cross-entropy method with diagonal covariance (separable CEM) """ + Cross-entropy method with diagonal covariance (separable CEM). - def __init__(self, num_params, - mu_init=None, - sigma_init=1e-3, - pop_size=256, - damp=1e-3, - damp_limit=1e-5, - parents=None, - elitism=False, - antithetic=False): + :param num_params: (int) + :param mu_init: (np.ndarray) Initial mean of the population distribution + Taken to be zero if None is passed. + :param sigma_init: (float) Initial standard deviation of the population distribution + :param pop_size: (int) Number of individuals in the population + :param damp: (float) Damping for preventing from early convergence. + :param damp_limit: (float) Final value of damping + :param parents: (int) + :param elitism: (bool) + :param antithetic: (bool) Use a finite difference like method for sampling + (mu + epsilon, mu - epsilon) + """ + def __init__(self, num_params, mu_init=None, sigma_init=1e-3, + pop_size=256, damp=1e-3, damp_limit=1e-5, + parents=None, elitism=False, antithetic=False): super(CEM, self).__init__() # misc self.num_params = num_params @@ -31,6 +34,7 @@ class CEM(object): self.sigma = sigma_init self.damp = damp self.damp_limit = damp_limit + # Exponential moving average decay for damping self.tau = 0.95 self.cov = self.sigma * np.ones(self.num_params) @@ -56,6 +60,9 @@ class CEM(object): def ask(self, pop_size): """ Returns a list of candidates parameters + + :param pop_size: (int) + :return: ([np.ndarray]) """ if self.antithetic and not pop_size % 2: epsilon_half = np.random.randn(pop_size // 2, self.num_params) @@ -64,16 +71,20 @@ class CEM(object): else: epsilon = np.random.randn(pop_size, self.num_params) - inds = self.mu + epsilon * np.sqrt(self.cov) + individuals = self.mu + epsilon * np.sqrt(self.cov) if self.elitism: - inds[-1] = self.elite + individuals[-1] = self.elite - return inds + return individuals def tell(self, solutions, scores): """ Updates the distribution + + :param solutions: ([np.ndarray]) + :param scores: ([float]) episode reward. """ + # Convert rewards (we want to maximize) to cost (we want to minimize) scores = np.array(scores) scores *= -1 idx_sorted = np.argsort(scores) @@ -92,7 +103,9 @@ class CEM(object): def get_distrib_params(self): """ - Returns the parameters of the distrubtion: - the mean and sigma + Returns the parameters of the distribution: + the mean and standard deviation. + + :return: (np.ndarray, np.ndarray) """ return np.copy(self.mu), np.copy(self.cov) diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index ff7100d..024ac25 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -10,17 +10,48 @@ from torchy_baselines.common.vec_env import sync_envs_normalization class CEMRL(TD3): """ - Implementation of CEM-RL + Implementation of CEM-RL, in fact CEM combined with TD3. Paper: https://arxiv.org/abs/1810.01222 Code: https://github.com/apourchot/CEM-RL - """ + :param policy: (TD3Policy or str) The policy model to use (MlpPolicy, CnnPolicy, ...) + :param env: (Gym environment or str) The environment to learn from (if registered in Gym, can be str) + :param sigma_init: (float) Initial standard deviation of the population distribution + :param pop_size: (int) Number of individuals in the population + :param damp: (float) Damping for preventing from early convergence. + :param damp_limit: (float) Final value of damping + :param elitism: (bool) + :param n_grad: (int) Number of individuals that will receive a gradient update. + Half of the population size in the paper. + :param buffer_size: (int) size of the replay buffer + :param learning_rate: (float or callable) learning rate for adam optimizer, + the same learning rate will be used for all networks (Q-Values and Actor networks) + it can be a function of the current progress (from 1 to 0) + :param policy_delay: (int) Policy and target networks will only be updated once every policy_delay steps + per training steps. The Q values will be updated policy_delay more often (update every training step). + :param learning_starts: (int) how many steps of the model to collect transitions for before learning starts + :param gamma: (float) the discount factor + :param batch_size: (int) Minibatch size for each gradient update + :param tau: (float) the soft update coefficient ("polyak update" of the target networks, between 0 and 1) + :param action_noise: (ActionNoise) the action noise type. Cf common.noise for the different action noise type. + :param target_policy_noise: (float) Standard deviation of gaussian noise added to target policy + (smoothing noise) + :param target_noise_clip: (float) Limit for absolute value of target policy smoothing noise. + :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 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 tensorflow debug + :param seed: (int) Seed for the pseudo random generators + :param device: (str or th.device) Device (cpu, cuda, ...) on which the code should be run. + Setting it to auto, the code will be run on the GPU if possible. + :param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance + """ def __init__(self, policy, env, sigma_init=1e-3, pop_size=10, damp=1e-3, damp_limit=1e-5, elitism=False, n_grad=5, - policy_delay=2, batch_size=100, - buffer_size=int(1e6), learning_rate=1e-3, - action_noise=None, learning_starts=100, tau=0.005, + buffer_size=int(1e6), learning_rate=1e-3, policy_delay=2, + learning_starts=100, gamma=0.99, batch_size=100, tau=0.005, + action_noise=None, target_policy_noise=0.2, target_noise_clip=0.5, n_episodes_rollout=1, update_style='original', tensorboard_log=None, create_eval_env=False, policy_kwargs=None, verbose=0, seed=0, device='auto', @@ -28,13 +59,16 @@ class CEMRL(TD3): super(CEMRL, self).__init__(policy, env, buffer_size=buffer_size, learning_rate=learning_rate, seed=seed, device=device, - action_noise=action_noise, learning_starts=learning_starts, - n_episodes_rollout=n_episodes_rollout, tau=tau, + action_noise=action_noise, target_policy_noise=target_policy_noise, + target_noise_clip=target_noise_clip, learning_starts=learning_starts, + n_episodes_rollout=n_episodes_rollout, tau=tau, gamma=gamma, policy_kwargs=policy_kwargs, verbose=verbose, policy_delay=policy_delay, batch_size=batch_size, - create_eval_env=create_eval_env, + create_eval_env=create_eval_env, tensorboard_log=tensorboard_log, _init_setup_model=False) + # Evolution strategy method that follows cma-es interface (ask-tell) + # for now, only CEM is implemented self.es = None self.sigma_init = sigma_init self.pop_size = pop_size @@ -79,7 +113,8 @@ class CEMRL(TD3): # set params 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)) + self.actor.optimizer = th.optim.Adam(self.actor.parameters(), + lr=self.learning_rate(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/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 771b5d1..4257c1e 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -5,6 +5,7 @@ import gym from gym import spaces import torch as th import torch.nn.functional as F + # Check if tensorboard is available for pytorch try: from torch.utils.tensorboard import SummaryWriter @@ -192,7 +193,6 @@ class PPO(BaseRLModel): clip_range_vf = self.clip_range_vf(self._current_progress) logger.logkv("clip_range_vf", clip_range_vf) - for gradient_step in range(gradient_steps): approx_kl_divs = [] # Sample replay buffer @@ -226,7 +226,6 @@ class PPO(BaseRLModel): # Value loss using the TD(gae_lambda) target value_loss = F.mse_loss(return_batch, values_pred) - # Entropy loss favor exploration entropy_loss = -th.mean(entropy) @@ -241,7 +240,8 @@ class PPO(BaseRLModel): approx_kl_divs.append(th.mean(old_log_prob - log_prob).detach().cpu().numpy()) if self.target_kl is not None and np.mean(approx_kl_divs) > 1.5 * self.target_kl: - print("Early stopping at step {} due to reaching max kl: {:.2f}".format(gradient_step, np.mean(approx_kl_divs))) + print("Early stopping at step {} due to reaching max kl: {:.2f}".format(gradient_step, + np.mean(approx_kl_divs))) break explained_var = explained_variance(self.rollout_buffer.returns.flatten().cpu().numpy(), From b84e5e9e27c35fe0ad8180dd79f6f94e51211dbd Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 22 Nov 2019 13:06:41 +0100 Subject: [PATCH 17/24] Move flexible mlp to common --- torchy_baselines/cem_rl/cem.py | 2 + torchy_baselines/common/policies.py | 91 ++++++++++++++++++++++++++++ torchy_baselines/ppo/policies.py | 92 +---------------------------- 3 files changed, 94 insertions(+), 91 deletions(-) diff --git a/torchy_baselines/cem_rl/cem.py b/torchy_baselines/cem_rl/cem.py index bea75fc..bf6b6fa 100644 --- a/torchy_baselines/cem_rl/cem.py +++ b/torchy_baselines/cem_rl/cem.py @@ -1,5 +1,6 @@ import numpy as np + # TODO: add more from https://github.com/hardmaru/estool/blob/master/es.py # or https://github.com/facebookresearch/nevergrad @@ -19,6 +20,7 @@ class CEM(object): :param antithetic: (bool) Use a finite difference like method for sampling (mu + epsilon, mu - epsilon) """ + def __init__(self, num_params, mu_init=None, sigma_init=1e-3, pop_size=256, damp=1e-3, damp_limit=1e-5, parents=None, elitism=False, antithetic=False): diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py index 55e356c..4ac7e07 100644 --- a/torchy_baselines/common/policies.py +++ b/torchy_baselines/common/policies.py @@ -1,3 +1,5 @@ +from itertools import zip_longest + import torch as th import torch.nn as nn @@ -141,3 +143,92 @@ def register_policy(name, policy): raise ValueError("Error: the name {} is alreay registered for a different policy, will not override." .format(name)) _policy_registry[sub_class][name] = policy + + +class MlpExtractor(nn.Module): + """ + Constructs an MLP that receives observations as an input and outputs a latent representation for the policy and + a value network. The ``net_arch`` parameter allows to specify the amount and size of the hidden layers and how many + of them are shared between the policy network and the value network. It is assumed to be a list with the following + structure: + + 1. An arbitrary length (zero allowed) number of integers each specifying the number of units in a shared layer. + If the number of ints is zero, there will be no shared layers. + 2. An optional dict, to specify the following non-shared layers for the value network and the policy network. + It is formatted like ``dict(vf=[], pi=[])``. + If it is missing any of the keys (pi or vf), no non-shared layers (empty list) is assumed. + + For example to construct a network with one shared layer of size 55 followed by two non-shared layers for the value + network of size 255 and a single non-shared layer of size 128 for the policy network, the following layers_spec + would be used: ``[55, dict(vf=[255, 255], pi=[128])]``. A simple shared network topology with two layers of size 128 + would be specified as [128, 128]. + + Adapted from Stable Baselines. + + :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 device: (th.device) + """ + def __init__(self, feature_dim, net_arch, activation_fn, device='cpu'): + super(MlpExtractor, self).__init__() + + shared_net, policy_net, value_net = [], [], [] + policy_only_layers = [] # Layer sizes of the network that only belongs to the policy network + value_only_layers = [] # Layer sizes of the network that only belongs to the value network + last_layer_dim_shared = feature_dim + + # Iterate through the shared layers and build the shared parts of the network + for idx, layer in enumerate(net_arch): + if isinstance(layer, int): # Check that this is a shared layer + layer_size = layer + # TODO: give layer a meaningful name + shared_net.append(nn.Linear(last_layer_dim_shared, layer_size)) + shared_net.append(activation_fn()) + last_layer_dim_shared = layer_size + else: + assert isinstance(layer, dict), "Error: the net_arch list can only contain ints and dicts" + if 'pi' in layer: + assert isinstance(layer['pi'], list), "Error: net_arch[-1]['pi'] must contain a list of integers." + policy_only_layers = layer['pi'] + + if 'vf' in layer: + assert isinstance(layer['vf'], list), "Error: net_arch[-1]['vf'] must contain a list of integers." + value_only_layers = layer['vf'] + break # From here on the network splits up in policy and value network + + last_layer_dim_pi = last_layer_dim_shared + last_layer_dim_vf = last_layer_dim_shared + + # Build the non-shared part of the network + for idx, (pi_layer_size, vf_layer_size) in enumerate(zip_longest(policy_only_layers, value_only_layers)): + if pi_layer_size is not None: + assert isinstance(pi_layer_size, int), "Error: net_arch[-1]['pi'] must only contain integers." + policy_net.append(nn.Linear(last_layer_dim_pi, pi_layer_size)) + policy_net.append(activation_fn()) + last_layer_dim_pi = pi_layer_size + + if vf_layer_size is not None: + assert isinstance(vf_layer_size, int), "Error: net_arch[-1]['vf'] must only contain integers." + value_net.append(nn.Linear(last_layer_dim_vf, vf_layer_size)) + value_net.append(activation_fn()) + last_layer_dim_vf = vf_layer_size + + # Save dim, used to create the distributions + self.latent_dim_pi = last_layer_dim_pi + self.latent_dim_vf = last_layer_dim_vf + + # Create networks + # If the list of layers is empty, the network will just act as an Identity module + self.shared_net = nn.Sequential(*shared_net).to(device) + self.policy_net = nn.Sequential(*policy_net).to(device) + self.value_net = nn.Sequential(*value_net).to(device) + + def forward(self, features): + """ + :return: (th.Tensor, th.Tensor) latent_policy, latent_value of the specified network. + If all layers are shared, then ``latent_policy == latent_value`` + """ + shared_latent = self.shared_net(features) + return self.policy_net(shared_latent), self.value_net(shared_latent) diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 1cb7ca4..e5ac7af 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -1,104 +1,14 @@ from functools import partial -from itertools import zip_longest import torch as th import torch.nn as nn import numpy as np -from torchy_baselines.common.policies import BasePolicy, register_policy +from torchy_baselines.common.policies import BasePolicy, register_policy, MlpExtractor from torchy_baselines.common.distributions import make_proba_distribution,\ DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution -class MlpExtractor(nn.Module): - """ - Constructs an MLP that receives observations as an input and outputs a latent representation for the policy and - a value network. The ``net_arch`` parameter allows to specify the amount and size of the hidden layers and how many - of them are shared between the policy network and the value network. It is assumed to be a list with the following - structure: - - 1. An arbitrary length (zero allowed) number of integers each specifying the number of units in a shared layer. - If the number of ints is zero, there will be no shared layers. - 2. An optional dict, to specify the following non-shared layers for the value network and the policy network. - It is formatted like ``dict(vf=[], pi=[])``. - If it is missing any of the keys (pi or vf), no non-shared layers (empty list) is assumed. - - For example to construct a network with one shared layer of size 55 followed by two non-shared layers for the value - network of size 255 and a single non-shared layer of size 128 for the policy network, the following layers_spec - would be used: ``[55, dict(vf=[255, 255], pi=[128])]``. A simple shared network topology with two layers of size 128 - would be specified as [128, 128]. - - Adapted from Stable Baselines. - - :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 device: (th.device) - """ - def __init__(self, feature_dim, net_arch, activation_fn, device='cpu'): - super(MlpExtractor, self).__init__() - - shared_net, policy_net, value_net = [], [], [] - policy_only_layers = [] # Layer sizes of the network that only belongs to the policy network - value_only_layers = [] # Layer sizes of the network that only belongs to the value network - last_layer_dim_shared = feature_dim - - # Iterate through the shared layers and build the shared parts of the network - for idx, layer in enumerate(net_arch): - if isinstance(layer, int): # Check that this is a shared layer - layer_size = layer - # TODO: give layer a meaningful name - shared_net.append(nn.Linear(last_layer_dim_shared, layer_size)) - shared_net.append(activation_fn()) - last_layer_dim_shared = layer_size - else: - assert isinstance(layer, dict), "Error: the net_arch list can only contain ints and dicts" - if 'pi' in layer: - assert isinstance(layer['pi'], list), "Error: net_arch[-1]['pi'] must contain a list of integers." - policy_only_layers = layer['pi'] - - if 'vf' in layer: - assert isinstance(layer['vf'], list), "Error: net_arch[-1]['vf'] must contain a list of integers." - value_only_layers = layer['vf'] - break # From here on the network splits up in policy and value network - - last_layer_dim_pi = last_layer_dim_shared - last_layer_dim_vf = last_layer_dim_shared - - # Build the non-shared part of the network - for idx, (pi_layer_size, vf_layer_size) in enumerate(zip_longest(policy_only_layers, value_only_layers)): - if pi_layer_size is not None: - assert isinstance(pi_layer_size, int), "Error: net_arch[-1]['pi'] must only contain integers." - policy_net.append(nn.Linear(last_layer_dim_pi, pi_layer_size)) - policy_net.append(activation_fn()) - last_layer_dim_pi = pi_layer_size - - if vf_layer_size is not None: - assert isinstance(vf_layer_size, int), "Error: net_arch[-1]['vf'] must only contain integers." - value_net.append(nn.Linear(last_layer_dim_vf, vf_layer_size)) - value_net.append(activation_fn()) - last_layer_dim_vf = vf_layer_size - - # Save dim, used to create the distributions - self.latent_dim_pi = last_layer_dim_pi - self.latent_dim_vf = last_layer_dim_vf - - # Create networks - # If the list of layers is empty, the network will just act as an Identity module - self.shared_net = nn.Sequential(*shared_net).to(device) - self.policy_net = nn.Sequential(*policy_net).to(device) - self.value_net = nn.Sequential(*value_net).to(device) - - def forward(self, features): - """ - :return: (th.Tensor, th.Tensor) latent_policy, latent_value of the specified network. - If all layers are shared, then ``latent_policy == latent_value`` - """ - shared_latent = self.shared_net(features) - return self.policy_net(shared_latent), self.value_net(shared_latent) - - class PPOPolicy(BasePolicy): def __init__(self, observation_space, action_space, learning_rate, net_arch=None, device='cpu', From 604a19fbc3108d8f9a90d6cda85dd5a1c5d2641f Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 22 Nov 2019 13:33:12 +0100 Subject: [PATCH 18/24] Cleanup + update doc --- README.md | 3 ++ torchy_baselines/a2c/a2c.py | 1 + torchy_baselines/cem_rl/cem_rl.py | 4 -- torchy_baselines/common/base_class.py | 53 ++++++++++++++++++--- torchy_baselines/common/buffers.py | 68 ++++++++++++++++++++++++--- torchy_baselines/sac/sac.py | 3 -- torchy_baselines/td3/td3.py | 3 -- 7 files changed, 112 insertions(+), 23 deletions(-) diff --git a/README.md b/README.md index 74b55c6..cc9f015 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py index 36ab9aa..68a1c44 100644 --- a/torchy_baselines/a2c/a2c.py +++ b/torchy_baselines/a2c/a2c.py @@ -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 diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index 024ac25..0ada89c 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -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 diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 36cd0c5..34da520 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -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) diff --git a/torchy_baselines/common/buffers.py b/torchy_baselines/common/buffers.py index 2c3d1fb..ae31a1f 100644 --- a/torchy_baselines/common/buffers.py +++ b/torchy_baselines/common/buffers.py @@ -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) diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index 5def28f..c14623c 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -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) diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index ddd7e9d..2e312ac 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -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: From 03d2ab10f86412fd91c7686228b4f46ff90856d5 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 22 Nov 2019 15:04:34 +0100 Subject: [PATCH 19/24] Fix clipped action when adapting noise with TD3 --- torchy_baselines/common/base_class.py | 29 +++++++++++++++++---------- torchy_baselines/td3/policies.py | 12 ++++++----- torchy_baselines/td3/td3.py | 8 ++++---- 3 files changed, 29 insertions(+), 20 deletions(-) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 34da520..9841172 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -67,6 +67,7 @@ class BaseRLModel(object): self.action_noise = None # Used for SDE only self.rollout_data = None + self.use_sde = False # Track the training progress (from 1 to 0) # this is used to update the learning rate self._current_progress = 1 @@ -391,7 +392,7 @@ class BaseRLModel(object): obs_ = self._vec_normalize_env.get_original_obs() self.rollout_data = None - if hasattr(self, 'use_sde') and self.use_sde: + if self.use_sde: self.actor.reset_noise() # Reset rollout data self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones']} @@ -405,23 +406,29 @@ class BaseRLModel(object): while not done: # Select action randomly or according to policy if num_timesteps < learning_starts: - action = np.array([self.action_space.sample()]) + # Warmup phase + unscaled_action = np.array([self.action_space.sample()]) else: - if hasattr(self, 'use_sde'): - deterministic = not self.use_sde - action = self.predict(obs, deterministic=deterministic) + unscaled_action = self.predict(obs, deterministic=not self.use_sde) # Rescale the action from [low, high] to [-1, 1] - action = self.scale_action(action) + scaled_action = self.scale_action(unscaled_action) + + if self.use_sde: + # When using SDE, the action can be out of bounds + # TODO: fix with squashing and account for that in the proba distribution + clipped_action = np.clip(scaled_action, -1, 1) + else: + clipped_action = scaled_action # Add noise to the action (improve exploration) if action_noise is not None: # NOTE: in the original implementation of TD3, the noise was applied to the unscaled action # Update(October 2019): Not anymore - action = np.clip(action + action_noise(), -1, 1) + clipped_action = np.clip(clipped_action + action_noise(), -1, 1) # Rescale and perform action - new_obs, reward, done, infos = env.step(self.unscale_action(action)) + new_obs, reward, done, infos = env.step(self.unscale_action(clipped_action)) done_bool = [float(done[0])] episode_reward += reward @@ -439,12 +446,12 @@ class BaseRLModel(object): # Avoid changing the original ones obs_, new_obs_, reward_ = obs, new_obs, reward - replay_buffer.add(obs_, new_obs_, action, reward_, done_bool) + replay_buffer.add(obs_, new_obs_, clipped_action, reward_, done_bool) if self.rollout_data is not None: # Assume only one env self.rollout_data['observations'].append(obs[0].copy()) - self.rollout_data['actions'].append(action[0].copy()) + self.rollout_data['actions'].append(scaled_action[0].copy()) self.rollout_data['rewards'].append(reward[0].copy()) self.rollout_data['dones'].append(np.array(done_bool[0]).copy()) @@ -480,7 +487,7 @@ class BaseRLModel(object): logger.logkv("fps", fps) logger.logkv('time_elapsed', int(time.time() - self.start_time)) logger.logkv("total timesteps", num_timesteps) - if hasattr(self, 'use_sde') and self.use_sde: + if self.use_sde: logger.logkv("std", th.exp(self.actor.log_std).mean().item()) logger.dumpkvs() diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 96fae45..7ef7fd7 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -22,10 +22,10 @@ class Actor(BaseNetwork): latent_pi = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_out=False) self.latent_pi = nn.Sequential(*latent_pi) if full_std: - self.log_std = nn.Parameter(th.ones(latent_dim, action_dim) * log_std_init) + self.log_std = nn.Parameter(th.ones(latent_dim, action_dim) * log_std_init, requires_grad=True) else: # Reduce the number of parameters: - self.log_std = nn.Parameter(th.ones(latent_dim, 1) * log_std_init) + self.log_std = nn.Parameter(th.ones(latent_dim, 1) * log_std_init, requires_grad=True) self.latent_dim = latent_dim self.actor_net = nn.Sequential(nn.Linear(net_arch[-1], action_dim), nn.Tanh()) @@ -42,7 +42,7 @@ class Actor(BaseNetwork): # Reduce the number of parameters: return th.ones((self.latent_dim, self.action_dim)).to(self.log_std.device) * self.log_std - def get_distribution_stats(self, obs, action): + def evaluate_actions(self, obs, action): with th.no_grad(): latent_pi = self.latent_pi(obs) mean_actions = self.actor_net(latent_pi) @@ -68,8 +68,10 @@ class Actor(BaseNetwork): noise = th.mm(latent_pi.detach(), self.exploration_mat) if self.clip_noise is not None: noise = th.clamp(noise, -self.clip_noise, self.clip_noise) - # TODO: fix clipping with squashing ? - return th.clamp(self.actor_net(latent_pi) + noise, -1, 1) + # TODO: Replace with squashing -> need to account for that in the sde update + # return th.clamp(self.actor_net(latent_pi) + noise, -1, 1) + # NOTE: the clipping is done in the rollout for now + return self.actor_net(latent_pi) + noise else: return self.actor_net(obs) diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 2e312ac..5ecb38e 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -52,6 +52,7 @@ class TD3(BaseRLModel): Setting it to auto, the code will be run on the GPU if possible. :param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance """ + def __init__(self, policy, env, buffer_size=int(1e6), learning_rate=1e-3, policy_delay=2, learning_starts=100, gamma=0.99, batch_size=100, train_freq=-1, gradient_steps=-1, n_episodes_rollout=1, @@ -64,7 +65,6 @@ class TD3(BaseRLModel): create_eval_env=create_eval_env, seed=seed) self.buffer_size = buffer_size - # TODO: accept callables self.learning_rate = learning_rate self.learning_starts = learning_starts self.train_freq = train_freq @@ -78,6 +78,7 @@ class TD3(BaseRLModel): self.target_noise_clip = target_noise_clip self.target_policy_noise = target_policy_noise + # State Dependent Exploration self.use_sde = use_sde self.sde_max_grad_norm = sde_max_grad_norm self.sde_ent_coef = sde_ent_coef @@ -204,10 +205,10 @@ class TD3(BaseRLModel): # self._update_learning_rate(self.policy.optimizer) # Unpack - obs, action, returns = self.rollout_data['observations'], self.rollout_data['actions'], self.rollout_data['returns'] + obs, action, returns = [self.rollout_data[key] for key in ['observations', 'actions', 'returns']] # TODO: avoid second computation of everything because of the gradient - log_prob, entropy = self.actor.get_distribution_stats(obs, action) + log_prob, entropy = self.actor.evaluate_actions(obs, action) # Normalize returns # returns = (returns - returns.mean()) / (returns.std() + 1e-8) @@ -281,7 +282,6 @@ class TD3(BaseRLModel): gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay) - # Evaluate episode if 0 < eval_freq <= timesteps_since_eval and eval_env is not None: timesteps_since_eval %= eval_freq From f0f2f10d1e37fd18d55c767451a6660db2c61eac Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 22 Nov 2019 17:24:14 +0100 Subject: [PATCH 20/24] Change default architecture --- torchy_baselines/ppo/policies.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index e5ac7af..7eff3d5 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -16,8 +16,12 @@ class PPOPolicy(BasePolicy): ortho_init=True, use_sde=False, log_std_init=0.0): super(PPOPolicy, self).__init__(observation_space, action_space, device) self.obs_dim = self.observation_space.shape[0] + + + # Default network architecture, from stable-baselines if net_arch is None: - net_arch = [dict(pi=[64], vf=[64])] + net_arch = [dict(pi=[64, 64], vf=[64, 64])] + self.net_arch = net_arch self.activation_fn = activation_fn self.adam_epsilon = adam_epsilon From c47be0086e12bc0967aac98c8682549dec091d22 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 22 Nov 2019 17:24:47 +0100 Subject: [PATCH 21/24] Add docstrings --- torchy_baselines/common/base_class.py | 2 +- torchy_baselines/ppo/policies.py | 29 +++++++++++++++-- torchy_baselines/td3/policies.py | 45 +++++++++++++++++++++++++++ 3 files changed, 73 insertions(+), 3 deletions(-) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 9841172..bf3468c 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -23,7 +23,7 @@ class BaseRLModel(object): :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. + :param device: (str or th.device) 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: (bool) Whether the algorithm supports training diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 7eff3d5..bd85797 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -10,6 +10,20 @@ from torchy_baselines.common.distributions import make_proba_distribution,\ class PPOPolicy(BasePolicy): + """ + Policy class (with both actor and critic) for A2C and derivates (PPO). + + :param observation_space: (gym.spaces.Space) Observation space + :param action_dim: (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 device: (str or th.device) Device on which the code should run. + :param activation_fn: (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 + :param log_std_init: (float) Initial value for the log standard deviation + """ def __init__(self, observation_space, action_space, learning_rate, net_arch=None, device='cpu', activation_fn=nn.Tanh, adam_epsilon=1e-5, @@ -100,14 +114,25 @@ class PPOPolicy(BasePolicy): return action.detach().cpu().numpy() def evaluate_actions(self, obs, action, deterministic=False): + """ + Evaluate actions according to the current policy, + given the observations. + + :param obs: (th.Tensor) + :param action: (th.Tensor) + :param deterministic: (bool) + :return: (th.Tensor, th.Tensor, th.Tensor) estimated value, log likelihood of taking those actions + and entropy of the action distribution. + """ latent_pi, latent_vf = self._get_latent(obs) _, action_distribution = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic) log_prob = action_distribution.log_prob(action) value = self.value_net(latent_vf) return value, log_prob, action_distribution.entropy() - def value_forward(self): - pass + def value_forward(self, obs): + _, latent_vf = self._get_latent(obs) + return self.value_net(latent_vf) MlpPolicy = PPOPolicy diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 7ef7fd7..b883db3 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -6,6 +6,20 @@ from torchy_baselines.common.policies import BasePolicy, register_policy, create class Actor(BaseNetwork): + """ + Actor network (policy) for TD3. + + :param obs_dim: (int) Dimension of the observation + :param action_dim: (int) Dimension of the action space + :param net_arch: ([int]) Network architecture + :param activation_fn: (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 + :param lr_sde: (float) Learning rate for the standard deviation of the noise + :param full_std: (bool) Whether to use (n_features x n_actions) parameters + for the std instead of only (n_features,) + """ def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU, use_sde=False, log_std_init=-2, clip_noise=None, lr_sde=3e-4, full_std=False): @@ -43,6 +57,15 @@ class Actor(BaseNetwork): return th.ones((self.latent_dim, self.action_dim)).to(self.log_std.device) * self.log_std def evaluate_actions(self, obs, action): + """ + Evaluate actions according to the current policy, + given the observations. Only useful when using SDE. + + :param obs: (th.Tensor) + :param action: (th.Tensor) + :return: (th.Tensor, th.Tensor) log likelihood of taking those actions + and entropy of the action distribution. + """ with th.no_grad(): latent_pi = self.latent_pi(obs) mean_actions = self.actor_net(latent_pi) @@ -77,6 +100,15 @@ class Actor(BaseNetwork): class Critic(BaseNetwork): + """ + Critic network for TD3, + in fact it represents the action-state value function (Q-value function) + + :param obs_dim: (int) Dimension of the observation + :param action_dim: (int) Dimension of the action space + :param net_arch: ([int]) Network architecture + :param activation_fn: (nn.Module) Activation function + """ def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU): super(Critic, self).__init__() @@ -98,12 +130,25 @@ class Critic(BaseNetwork): class TD3Policy(BasePolicy): + """ + Policy class (with both actor and critic) for TD3. + + :param observation_space: (gym.spaces.Space) Observation space + :param action_dim: (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 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 + :param log_std_init: (float) Initial value for the log standard deviation + """ def __init__(self, observation_space, action_space, learning_rate, net_arch=None, device='cpu', activation_fn=nn.ReLU, use_sde=False, log_std_init=-2, clip_noise=None, lr_sde=3e-4): super(TD3Policy, self).__init__(observation_space, action_space, device) + # Default network architecture, from the original paper if net_arch is None: net_arch = [400, 300] From c56865e10d2b232d523e7a06a2bfcd8e4ab870f5 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 22 Nov 2019 19:02:00 +0100 Subject: [PATCH 22/24] Cleanup CEM, rename variables + add comments --- torchy_baselines/cem_rl/cem.py | 43 +++++++++++++++++++------------ torchy_baselines/cem_rl/cem_rl.py | 14 +++++----- 2 files changed, 34 insertions(+), 23 deletions(-) diff --git a/torchy_baselines/cem_rl/cem.py b/torchy_baselines/cem_rl/cem.py index bf6b6fa..ee4f484 100644 --- a/torchy_baselines/cem_rl/cem.py +++ b/torchy_baselines/cem_rl/cem.py @@ -8,36 +8,39 @@ class CEM(object): """ Cross-entropy method with diagonal covariance (separable CEM). - :param num_params: (int) + :param num_params: (int) Number of parameters per individual (dimension of the problem) :param mu_init: (np.ndarray) Initial mean of the population distribution Taken to be zero if None is passed. :param sigma_init: (float) Initial standard deviation of the population distribution :param pop_size: (int) Number of individuals in the population - :param damp: (float) Damping for preventing from early convergence. - :param damp_limit: (float) Final value of damping - :param parents: (int) - :param elitism: (bool) + :param damping_init: (float) Initial value of damping for preventing from early convergence. + :param damping_final: (float) Final value of damping + :param parents: (int) Number of parents used to compute the new distribution + of individuals. + :param elitism: (bool) Keep the best known individual in the population :param antithetic: (bool) Use a finite difference like method for sampling (mu + epsilon, mu - epsilon) """ - def __init__(self, num_params, mu_init=None, sigma_init=1e-3, - pop_size=256, damp=1e-3, damp_limit=1e-5, + pop_size=256, damping_init=1e-3, damping_final=1e-5, parents=None, elitism=False, antithetic=False): super(CEM, self).__init__() - # misc + self.num_params = num_params - # distribution parameters + # Distribution parameters if mu_init is None: self.mu = np.zeros(self.num_params) else: self.mu = np.array(mu_init) + self.sigma = sigma_init - self.damp = damp - self.damp_limit = damp_limit + # Damping parameters + self.damping = damping_init + self.damping_final = damping_final # Exponential moving average decay for damping self.tau = 0.95 + # Covariance matrix, here only the diagonal self.cov = self.sigma * np.ones(self.num_params) # elite stuff @@ -45,16 +48,20 @@ class CEM(object): self.elite = np.sqrt(self.sigma) * np.random.rand(self.num_params) self.elite_score = None - # sampling stuff + # sampling parameters self.pop_size = pop_size self.antithetic = antithetic if self.antithetic: assert (self.pop_size % 2 == 0), "Population size must be even" + if parents is None or parents <= 0: self.parents = pop_size // 2 else: self.parents = parents + + # Weighting for computing the new mean of the distributions + # from the parents. The better the individual, the higher the weight self.weights = np.array([np.log((self.parents + 1) / i) for i in range(1, self.parents + 1)]) self.weights /= self.weights.sum() @@ -69,11 +76,12 @@ class CEM(object): if self.antithetic and not pop_size % 2: epsilon_half = np.random.randn(pop_size // 2, self.num_params) epsilon = np.concatenate([epsilon_half, - epsilon_half]) - else: epsilon = np.random.randn(pop_size, self.num_params) individuals = self.mu + epsilon * np.sqrt(self.cov) + + # Keep the best known individual in the population if self.elitism: individuals[-1] = self.elite @@ -89,19 +97,22 @@ class CEM(object): # Convert rewards (we want to maximize) to cost (we want to minimize) scores = np.array(scores) scores *= -1 + # Sort the individuals by fitness idx_sorted = np.argsort(scores) old_mu = self.mu - self.damp = self.damp * self.tau + (1 - self.tau) * self.damp_limit + # Update damping using a moving average + self.damping = self.damping * self.tau + (1 - self.tau) * self.damping_final # self.mu = self.weights @ solutions[idx_sorted[:self.parents]] self.mu = self.weights.dot(solutions[idx_sorted[:self.parents]]) + # CMA-ES style would be to use the new mean here z = (solutions[idx_sorted[:self.parents]] - old_mu) - self.cov = 1 / self.parents * self.weights.dot(z * z) + self.damp * np.ones(self.num_params) + self.cov = 1 / self.parents * self.weights.dot(z * z) + self.damping * np.ones(self.num_params) + # Retrieve the best individual self.elite = solutions[idx_sorted[0]] self.elite_score = scores[idx_sorted[0]] - # print(self.cov) def get_distrib_params(self): """ diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index 0ada89c..9a432fa 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -19,9 +19,9 @@ class CEMRL(TD3): :param env: (Gym environment or str) The environment to learn from (if registered in Gym, can be str) :param sigma_init: (float) Initial standard deviation of the population distribution :param pop_size: (int) Number of individuals in the population - :param damp: (float) Damping for preventing from early convergence. - :param damp_limit: (float) Final value of damping - :param elitism: (bool) + :param damping_init: (float) Initial value of damping for preventing from early convergence. + :param damping_final: (float) Final value of damping + :param elitism: (bool) Keep the best known individual in the population :param n_grad: (int) Number of individuals that will receive a gradient update. Half of the population size in the paper. :param buffer_size: (int) size of the replay buffer @@ -48,7 +48,7 @@ class CEMRL(TD3): :param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance """ def __init__(self, policy, env, sigma_init=1e-3, pop_size=10, - damp=1e-3, damp_limit=1e-5, elitism=False, n_grad=5, + damping_init=1e-3, damping_final=1e-5, elitism=False, n_grad=5, buffer_size=int(1e6), learning_rate=1e-3, policy_delay=2, learning_starts=100, gamma=0.99, batch_size=100, tau=0.005, action_noise=None, target_policy_noise=0.2, target_noise_clip=0.5, @@ -72,8 +72,8 @@ class CEMRL(TD3): self.es = None self.sigma_init = sigma_init self.pop_size = pop_size - self.damp = damp - self.damp_limit = damp_limit + self.damping_init = damping_init + self.damping_final = damping_final self.elitism = elitism self.n_grad = n_grad self.es_params = None @@ -87,7 +87,7 @@ class CEMRL(TD3): super(CEMRL, self)._setup_model() params_vector = self.actor.parameters_to_vector() self.es = CEM(len(params_vector), mu_init=params_vector, - sigma_init=self.sigma_init, damp=self.damp, damp_limit=self.damp_limit, + sigma_init=self.sigma_init, damping_init=self.damping_init, damping_final=self.damping_final, pop_size=self.pop_size, antithetic=not self.pop_size % 2, parents=self.pop_size // 2, elitism=self.elitism) From 5bbb14188d8030fc69bca1875b9d5e919ab31599 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Mon, 25 Nov 2019 13:19:33 +0100 Subject: [PATCH 23/24] Unify A2C and TD3 SDE implementation --- torchy_baselines/common/base_class.py | 2 +- torchy_baselines/common/distributions.py | 42 +++++++++++----- torchy_baselines/ppo/policies.py | 5 +- torchy_baselines/td3/policies.py | 61 +++++++++++++----------- 4 files changed, 70 insertions(+), 40 deletions(-) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index bf3468c..40dcedb 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -488,7 +488,7 @@ class BaseRLModel(object): logger.logkv('time_elapsed', int(time.time() - self.start_time)) logger.logkv("total timesteps", num_timesteps) if self.use_sde: - logger.logkv("std", th.exp(self.actor.log_std).mean().item()) + logger.logkv("std", (self.actor.get_std()).mean().item()) logger.dumpkvs() mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0 diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py index 9d67b39..ad02734 100644 --- a/torchy_baselines/common/distributions.py +++ b/torchy_baselines/common/distributions.py @@ -234,6 +234,8 @@ class StateDependentNoiseDistribution(Distribution): compute the log probabilty of an action with that noise. :param action_dim: (int) Number of continuous actions + :param full_std: (bool) Whether to use (n_features x n_actions) parameters + for the std instead of only (n_features,) :param use_expln: (bool) Use `expln()` function instead of `exp()` to ensure a positive standard deviation (cf paper). It allows to keep variance above zero and prevent it from growing too fast. In practice, `exp()` is usually enough. @@ -241,16 +243,19 @@ class StateDependentNoiseDistribution(Distribution): this allows to ensure boundaries. :param epsilon: (float) small value to avoid NaN due to numerical imprecision. """ - def __init__(self, action_dim, use_expln=False, + def __init__(self, action_dim, full_std=True, use_expln=False, squash_output=False, epsilon=1e-6): super(StateDependentNoiseDistribution, self).__init__() self.distribution = None self.action_dim = action_dim + self.latent_dim = None self.mean_actions = None self.log_std = None self.weights_dist = None self.exploration_mat = None self.use_expln = use_expln + self.full_std = full_std + self.epsilon = epsilon if squash_output: print("== Using TanhBijector ===") self.bijector = TanhBijector(epsilon) @@ -269,12 +274,17 @@ class StateDependentNoiseDistribution(Distribution): # From SDE paper, it allows to keep variance # above zero and prevent it from growing too fast if log_std <= 0: - return th.exp(log_std) + std = th.exp(log_std) else: - return th.log(log_std + 1.0) + 1.0 + std = th.log(log_std + 1.0) + 1.0 else: # Use normal exponential - return th.exp(log_std) + std = th.exp(log_std) + + if self.full_std: + return std + # Reduce the number of parameters: + return th.ones((self.latent_dim, self.action_dim)).to(log_std.device) * std def sample_weights(self, log_std): """ @@ -283,8 +293,8 @@ class StateDependentNoiseDistribution(Distribution): :param log_std: (th.Tensor) """ - # TODO: reduce the number of learned dimensions (cf TD3) - self.weights_dist = Normal(th.zeros_like(log_std), self.get_std(log_std)) + std = self.get_std(log_std) + self.weights_dist = Normal(th.zeros_like(std), std) self.exploration_mat = self.weights_dist.rsample() def proba_distribution_net(self, latent_dim, log_std_init=0.0): @@ -297,10 +307,17 @@ class StateDependentNoiseDistribution(Distribution): :param log_std_init: (float) Initial value for the log standard deviation :return: (nn.Linear, nn.Parameter) """ - mean_actions = nn.Linear(latent_dim, self.action_dim) - log_std = nn.Parameter(th.ones(latent_dim, self.action_dim) * log_std_init) + # Network for the deterministic action, it represents the mean of the distribution + mean_actions_net = nn.Linear(latent_dim, self.action_dim) + + self.latent_dim = latent_dim + # Reduce the number of parameters if needed + log_std = th.ones(latent_dim, self.action_dim) if self.full_std else th.ones(latent_dim, 1) + # Transform it to a parameter so it can be optimized + log_std = nn.Parameter(log_std * log_std_init) + # Sample an exploration matrix self.sample_weights(log_std) - return mean_actions, log_std + return mean_actions_net, log_std def proba_distribution(self, mean_actions, log_std, latent_pi, deterministic=False): """ @@ -312,7 +329,7 @@ class StateDependentNoiseDistribution(Distribution): :return: (th.Tensor) """ variance = th.mm(latent_pi.detach() ** 2, self.get_std(log_std) ** 2) - self.distribution = Normal(mean_actions, th.sqrt(variance)) + self.distribution = Normal(mean_actions, th.sqrt(variance + self.epsilon)) if deterministic: action = self.mode() @@ -326,8 +343,11 @@ class StateDependentNoiseDistribution(Distribution): return self.bijector.forward(action) return action + def get_noise(self, latent_pi): + return th.mm(latent_pi.detach(), self.exploration_mat) + def sample(self, latent_pi): - noise = th.mm(latent_pi.detach(), self.exploration_mat) + noise = self.get_noise(latent_pi) action = self.distribution.mean + noise if self.bijector is not None: return self.bijector.forward(action) diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index bd85797..adcf239 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -58,6 +58,9 @@ class PPOPolicy(BasePolicy): self._build(learning_rate) def reset_noise_net(self): + """ + Sample new weights for the exploration matrix. + """ self.action_dist.sample_weights(self.log_std) def _build(self, learning_rate): @@ -75,7 +78,7 @@ class PPOPolicy(BasePolicy): # with small initial weight for the output if self.ortho_init: for module in [self.mlp_extractor, self.action_net, self.value_net]: - # Values from stable-baselines check why + # Values from stable-baselines, TODO: check why gain = { self.mlp_extractor: np.sqrt(2), self.action_net: 0.01, diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index b883db3..97d5862 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -1,8 +1,8 @@ import torch as th import torch.nn as nn -from torch.distributions import Normal from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork +from torchy_baselines.common.distributions import StateDependentNoiseDistribution class Actor(BaseNetwork): @@ -32,17 +32,15 @@ class Actor(BaseNetwork): self.full_std = full_std if use_sde: - latent_dim = net_arch[-1] latent_pi = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_out=False) self.latent_pi = nn.Sequential(*latent_pi) - if full_std: - self.log_std = nn.Parameter(th.ones(latent_dim, action_dim) * log_std_init, requires_grad=True) - else: - # Reduce the number of parameters: - self.log_std = nn.Parameter(th.ones(latent_dim, 1) * log_std_init, requires_grad=True) - - self.latent_dim = latent_dim - self.actor_net = nn.Sequential(nn.Linear(net_arch[-1], action_dim), nn.Tanh()) + # Create state dependent noise matrix (SDE) + self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False, + squash_output=False) + action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1], + log_std_init=log_std_init) + # Squash output + self.actor_net = nn.Sequential(action_net, nn.Tanh()) self.clip_noise = clip_noise self.sde_optimizer = th.optim.Adam([self.log_std], lr=lr_sde) self.reset_noise() @@ -50,11 +48,21 @@ class Actor(BaseNetwork): actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_out=True) self.actor_net = nn.Sequential(*actor_net) - def get_log_std(self): - if self.full_std: - return self.log_std - # Reduce the number of parameters: - return th.ones((self.latent_dim, self.action_dim)).to(self.log_std.device) * self.log_std + def get_std(self): + """ + Retrieve the standard deviation of the action distribution. + Only useful when using SDE. + It corresponds to `th.exp(log_std)` in the normal case, + but is slightly different when using `expln` function + (cf StateDependentNoiseDistribution doc). + + :return: (th.Tensor) + """ + return self.action_dist.get_std(self.log_std) + + def _get_action_dist_from_latent(self, latent_pi): + mean_actions = self.actor_net(latent_pi) + return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_pi) def evaluate_actions(self, obs, action): """ @@ -63,38 +71,37 @@ class Actor(BaseNetwork): :param obs: (th.Tensor) :param action: (th.Tensor) + :param deterministic: (bool) :return: (th.Tensor, th.Tensor) log likelihood of taking those actions and entropy of the action distribution. """ with th.no_grad(): latent_pi = self.latent_pi(obs) - mean_actions = self.actor_net(latent_pi) - - variance = th.mm(latent_pi ** 2, th.exp(self.get_log_std()) ** 2) - distribution = Normal(mean_actions, th.sqrt(variance + 1e-5)) + _, distribution = self._get_action_dist_from_latent(latent_pi) log_prob = distribution.log_prob(action) - if len(log_prob.shape) > 1: - log_prob = log_prob.sum(axis=1) - else: - log_prob = log_prob.sum() + # value = self.value_net(latent_vf) return log_prob, distribution.entropy() def reset_noise(self): - self.weights_dist = Normal(th.zeros_like(self.get_log_std()), th.exp(self.get_log_std())) - self.exploration_mat = self.weights_dist.rsample() + """ + Sample new weights for the exploration matrix. + """ + self.action_dist.sample_weights(self.log_std) def forward(self, obs, deterministic=True): if self.use_sde: latent_pi = self.latent_pi(obs) if deterministic: return self.actor_net(latent_pi) - noise = th.mm(latent_pi.detach(), self.exploration_mat) + noise = self.action_dist.get_noise(latent_pi) if self.clip_noise is not None: noise = th.clamp(noise, -self.clip_noise, self.clip_noise) # TODO: Replace with squashing -> need to account for that in the sde update - # return th.clamp(self.actor_net(latent_pi) + noise, -1, 1) + # -> set squash_out=True in the action_dist? # NOTE: the clipping is done in the rollout for now return self.actor_net(latent_pi) + noise + # action, _ = self._get_action_dist_from_latent(latent_pi) + # return action else: return self.actor_net(obs) From d0003ee4ecd7bb447ec945517d1c35e8b51abc28 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Mon, 25 Nov 2019 14:00:21 +0100 Subject: [PATCH 24/24] Enable kwargs for proba dist --- README.md | 2 -- torchy_baselines/common/distributions.py | 18 +++++++++++------- torchy_baselines/ppo/policies.py | 16 ++++++++++++++-- torchy_baselines/td3/policies.py | 2 +- 4 files changed, 26 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index cc9f015..0b19832 100644 --- a/README.md +++ b/README.md @@ -21,8 +21,6 @@ TODO: - save/load - better predict - complete logger -- 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 diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py index ad02734..f90883a 100644 --- a/torchy_baselines/common/distributions.py +++ b/torchy_baselines/common/distributions.py @@ -297,7 +297,7 @@ class StateDependentNoiseDistribution(Distribution): self.weights_dist = Normal(th.zeros_like(std), std) self.exploration_mat = self.weights_dist.rsample() - def proba_distribution_net(self, latent_dim, log_std_init=0.0): + def proba_distribution_net(self, latent_dim, log_std_init=-2.0): """ Create the layers and parameter that represent the distribution: one output will be the deterministic action, the other parameter will be the @@ -423,26 +423,30 @@ class TanhBijector(object): return th.log(1 - th.tanh(x) ** 2 + self.epsilon) -def make_proba_distribution(action_space, use_sde=False): +def make_proba_distribution(action_space, use_sde=False, dist_kwargs=None): """ Return an instance of Distribution for the correct type of action space :param action_space: (Gym Space) the input action space :param use_sde: (bool) Force the use of StateDependentNoiseDistribution instead of DiagGaussianDistribution + :param dist_kwargs: (dict) Keyword arguments to pass to the probabilty distribution :return: (Distribution) the approriate Distribution object """ + if dist_kwargs is None: + dist_kwargs = {} + if isinstance(action_space, spaces.Box): assert len(action_space.shape) == 1, "Error: the action space must be a vector" if use_sde: - return StateDependentNoiseDistribution(action_space.shape[0]) - return DiagGaussianDistribution(action_space.shape[0]) + return StateDependentNoiseDistribution(action_space.shape[0], **dist_kwargs) + return DiagGaussianDistribution(action_space.shape[0], **dist_kwargs) elif isinstance(action_space, spaces.Discrete): - return CategoricalDistribution(action_space.n) + return CategoricalDistribution(action_space.n, **dist_kwargs) # elif isinstance(action_space, spaces.MultiDiscrete): - # return MultiCategoricalDistribution(action_space.nvec) + # return MultiCategoricalDistribution(action_space.nvec, **dist_kwargs) # elif isinstance(action_space, spaces.MultiBinary): - # return BernoulliDistribution(action_space.n) + # return BernoulliDistribution(action_space.n, **dist_kwargs) else: raise NotImplementedError("Error: probability distribution, not implemented for action space of type {}." .format(type(action_space)) + diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index adcf239..46fe75c 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -23,11 +23,14 @@ class PPOPolicy(BasePolicy): :param ortho_init: (bool) Whether to use or not orthogonal initialization :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 + for the std instead of only (n_features,) when using SDE """ def __init__(self, observation_space, action_space, learning_rate, net_arch=None, device='cpu', activation_fn=nn.Tanh, adam_epsilon=1e-5, - ortho_init=True, use_sde=False, log_std_init=0.0): + ortho_init=True, use_sde=False, + log_std_init=0.0, full_std=True): super(PPOPolicy, self).__init__(observation_space, action_space, device) self.obs_dim = self.observation_space.shape[0] @@ -52,8 +55,17 @@ class PPOPolicy(BasePolicy): self.features_extractor = nn.Flatten() self.features_dim = self.obs_dim self.log_std_init = log_std_init + dist_kwargs = None + # Keyword arguments for SDE distribution + if use_sde: + dist_kwargs = { + 'full_std': full_std, + 'squash_output': False, + 'use_expln': False + } + # Action distribution - self.action_dist = make_proba_distribution(action_space, use_sde=use_sde) + self.action_dist = make_proba_distribution(action_space, use_sde=use_sde, dist_kwargs=dist_kwargs) self._build(learning_rate) diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 97d5862..cd72950 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -18,7 +18,7 @@ class Actor(BaseNetwork): :param clip_noise: (float) Clip the magnitude of the noise :param lr_sde: (float) Learning rate for the standard deviation of the noise :param full_std: (bool) Whether to use (n_features x n_actions) parameters - for the std instead of only (n_features,) + for the std instead of only (n_features,) when using SDE. """ def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU, use_sde=False, log_std_init=-2, clip_noise=None,