From a08382faab0d350ef1e9a627bdeb95e66164bea3 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Tue, 12 Nov 2019 18:37:13 +0100 Subject: [PATCH] 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: