def _update_learning_rate(self, optimizers: Union[List[th.optim.Optimizer], th.optim.Optimizer]) -> None: """ Update the optimizers learning rate using the current learning rate schedule and the current progress (from 1 to 0). :param optimizers: (Union[List[th.optim.Optimizer], th.optim.Optimizer]) An optimizer or a list of optimizers. """ # Log the current learning rate logger.logkv("learning_rate", self.lr_schedule(self._current_progress)) if not isinstance(optimizers, list): optimizers = [optimizers] for optimizer in optimizers: update_learning_rate(optimizer, self.lr_schedule(self._current_progress))
def test_main(): """ tests for the logger module """ info("hi") debug("shouldn't appear") set_level(DEBUG) debug("should appear") folder = "/tmp/testlogging" if os.path.exists(folder): shutil.rmtree(folder) configure(folder=folder) logkv("a", 3) logkv("b", 2.5) dumpkvs() logkv("b", -2.5) logkv("a", 5.5) dumpkvs() info("^^^ should see a = 5.5") logkv_mean("b", -22.5) logkv_mean("b", -44.4) logkv("a", 5.5) dumpkvs() with ScopedConfigure(None, None): info("^^^ should see b = 33.3") with ScopedConfigure("/tmp/test-logger/", ["json"]): logkv("b", -2.5) dumpkvs() reset() logkv("a", "longasslongasslongasslongasslongasslongassvalue") dumpkvs() warn("hey") error("oh") logkvs({"test": 1})
def train(self, gradient_steps: int, batch_size: Optional[int] = None) -> None: # Update optimizer learning rate self._update_learning_rate(self.policy.optimizer) # A2C with gradient_steps > 1 does not make sense assert gradient_steps == 1, "A2C does not support multiple gradient steps" # We do not use minibatches for A2C assert batch_size is None, "A2C does not support minibatch" for rollout_data in self.rollout_buffer.get(batch_size=None): actions = rollout_data.actions if isinstance(self.action_space, spaces.Discrete): # Convert discrete action from float to long actions = actions.long().flatten() # TODO: avoid second computation of everything because of the gradient values, log_prob, entropy = self.policy.evaluate_actions( rollout_data.observations, actions) values = values.flatten() # Normalize advantage (not present in the original implementation) advantages = rollout_data.advantages if self.normalize_advantage: advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) # Policy gradient loss policy_loss = -(advantages * log_prob).mean() # Value loss using the TD(gae_lambda) target value_loss = F.mse_loss(rollout_data.returns, values) # Entropy loss favor exploration if entropy is None: # Approximate entropy when no analytical form entropy_loss = -log_prob.mean() else: entropy_loss = -th.mean(entropy) loss = policy_loss + self.ent_coef * entropy_loss + self.vf_coef * value_loss # # MSA debugging learning # self.loss_hist.append([loss.item(), policy_loss.item(), value_loss.item(), entropy_loss.item() ]) # if len (self.loss_hist) == 25: # import matplotlib.pyplot as plt # l = [] # pl = [] # vl = [] # el = [] # for losses in self.loss_hist: # l.append (losses[0]) # pl.append (losses[1]) # vl.append (losses[2]) # el.append (losses[3]) # plt.plot (l, marker="o") # plt.plot (pl, marker="o") # plt.plot (vl, marker="o") # plt.plot (el, marker="o") # plt.title ('Losses') # plt.legend (['loss', 'policy loss', 'value loss', 'ent loss']) # filename = "RL_detailed_plots/2/losses.png" # plt.savefig (filename) # plt.close() # Optimization step self.policy.optimizer.zero_grad() loss.backward() # Clip grad norm th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) self.policy.optimizer.step() explained_var = explained_variance( self.rollout_buffer.returns.flatten(), self.rollout_buffer.values.flatten()) self._n_updates += 1 logger.logkv("n_updates", self._n_updates) logger.logkv("explained_variance", explained_var) logger.logkv("entropy_loss", entropy_loss.item()) logger.logkv("policy_loss", policy_loss.item()) logger.logkv("value_loss", value_loss.item()) if hasattr(self.policy, 'log_std'): logger.logkv("std", th.exp(self.policy.log_std).mean().item())
def learn(self, total_timesteps: int, callback: MaybeCallback = None, log_interval: int = 1, eval_env: Optional[GymEnv] = None, eval_freq: int = -1, n_eval_episodes: int = 5, tb_log_name: str = "PPO", eval_log_path: Optional[str] = None, reset_num_timesteps: bool = True) -> 'PPO': iteration = 0 callback = self._setup_learn(eval_env, callback, eval_freq, n_eval_episodes, eval_log_path, reset_num_timesteps) # if self.tensorboard_log is not None and SummaryWriter is not None: # self.tb_writer = SummaryWriter(log_dir=os.path.join(self.tensorboard_log, tb_log_name)) callback.on_training_start(locals(), globals()) while self.num_timesteps < total_timesteps: continue_training = self.collect_rollouts(self.env, callback, self.rollout_buffer, n_rollout_steps=self.n_steps) if continue_training is False: break iteration += 1 self._update_current_progress(self.num_timesteps, total_timesteps) # Display training infos if self.verbose >= 1 and log_interval is not None and iteration % log_interval == 0: fps = int(self.num_timesteps / (time.time() - self.start_time)) logger.logkv("iterations", iteration) 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("fps", fps) logger.logkv('time_elapsed', int(time.time() - self.start_time)) logger.logkv("total timesteps", self.num_timesteps) logger.dumpkvs() self.train(self.n_epochs, batch_size=self.batch_size) # For tensorboard integration # if self.tb_writer is not None: # self.tb_writer.add_scalar('Eval/reward', mean_reward, self.num_timesteps) callback.on_training_end() return self
def train(self, n_epochs: int, batch_size: int = 64) -> None: # Update optimizer learning rate self._update_learning_rate(self.policy.optimizer) # Compute current clip range clip_range = self.clip_range(self._current_progress) # Optional: clip range for the value function if self.clip_range_vf is not None: clip_range_vf = self.clip_range_vf(self._current_progress) entropy_losses, all_kl_divs = [], [] pg_losses, value_losses = [], [] clip_fractions = [] # train for gradient_steps epochs for epoch in range(n_epochs): approx_kl_divs = [] # Do a complete pass on the rollout buffer for rollout_data in self.rollout_buffer.get(batch_size): actions = rollout_data.actions if isinstance(self.action_space, spaces.Discrete): # Convert discrete action from float to long actions = rollout_data.actions.long().flatten() # Re-sample the noise matrix because the log_std has changed # TODO: investigate why there is no issue with the gradient # if that line is commented (as in SAC) if self.use_sde: self.policy.reset_noise(batch_size) values, log_prob, entropy = self.policy.evaluate_actions(rollout_data.observations, actions) values = values.flatten() # Normalize advantage advantages = rollout_data.advantages advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) # ratio between old and new policy, should be one at the first iteration ratio = th.exp(log_prob - rollout_data.old_log_prob) # clipped surrogate loss policy_loss_1 = advantages * ratio policy_loss_2 = advantages * th.clamp(ratio, 1 - clip_range, 1 + clip_range) policy_loss = -th.min(policy_loss_1, policy_loss_2).mean() # Logging pg_losses.append(policy_loss.item()) clip_fraction = th.mean((th.abs(ratio - 1) > clip_range).float()).item() clip_fractions.append(clip_fraction) if self.clip_range_vf is None: # No clipping values_pred = values else: # Clip the different between old and new value # NOTE: this depends on the reward scaling values_pred = rollout_data.old_values + th.clamp(values - rollout_data.old_values, -clip_range_vf, clip_range_vf) # Value loss using the TD(gae_lambda) target value_loss = F.mse_loss(rollout_data.returns, values_pred) value_losses.append(value_loss.item()) # Entropy loss favor exploration if entropy is None: # Approximate entropy when no analytical form entropy_loss = -log_prob.mean() else: entropy_loss = -th.mean(entropy) entropy_losses.append(entropy_loss.item()) loss = policy_loss + self.ent_coef * entropy_loss + self.vf_coef * value_loss # Optimization step self.policy.optimizer.zero_grad() loss.backward() # Clip grad norm th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) self.policy.optimizer.step() approx_kl_divs.append(th.mean(rollout_data.old_log_prob - log_prob).detach().cpu().numpy()) all_kl_divs.append(np.mean(approx_kl_divs)) if self.target_kl is not None and np.mean(approx_kl_divs) > 1.5 * self.target_kl: print(f"Early stopping at step {epoch} due to reaching max kl: {np.mean(approx_kl_divs):.2f}") break self._n_updates += n_epochs explained_var = explained_variance(self.rollout_buffer.returns.flatten(), self.rollout_buffer.values.flatten()) logger.logkv("n_updates", self._n_updates) logger.logkv("clip_fraction", np.mean(clip_fraction)) logger.logkv("clip_range", clip_range) if self.clip_range_vf is not None: logger.logkv("clip_range_vf", clip_range_vf) logger.logkv("approx_kl", np.mean(approx_kl_divs)) logger.logkv("explained_variance", explained_var) logger.logkv("entropy_loss", np.mean(entropy_losses)) logger.logkv("policy_gradient_loss", np.mean(pg_losses)) logger.logkv("value_loss", np.mean(value_losses)) if hasattr(self.policy, 'log_std'): logger.logkv("std", th.exp(self.policy.log_std).mean().item())
def collect_rollouts( self, env: VecEnv, # Type hint as string to avoid circular import callback: 'BaseCallback', n_episodes: int = 1, n_steps: int = -1, action_noise: Optional[ActionNoise] = None, learning_starts: int = 0, replay_buffer: Optional[ReplayBuffer] = None, log_interval: Optional[int] = None) -> RolloutReturn: """ Collect rollout using the current policy (and possibly fill the replay buffer) :param env: (VecEnv) The training environment :param n_episodes: (int) Number of episodes to use to collect rollout data You can also specify a ``n_steps`` instead :param n_steps: (int) Number of steps to use to collect rollout data You can also specify a ``n_episodes`` instead. :param action_noise: (Optional[ActionNoise]) Action noise that will be used for exploration Required for deterministic policy (e.g. TD3). This can also be used in addition to the stochastic policy for SAC. :param callback: (BaseCallback) Callback that will be called at each step (and at the beginning and end of the rollout) :param learning_starts: (int) Number of steps before learning for the warm-up phase. :param replay_buffer: (ReplayBuffer) :param log_interval: (int) Log data every ``log_interval`` episodes :return: (RolloutReturn) """ episode_rewards, total_timesteps = [], [] total_steps, total_episodes = 0, 0 assert isinstance(env, VecEnv), "You must pass a VecEnv" assert env.num_envs == 1, "OffPolicyRLModel only support single environment" if self.use_sde: self.actor.reset_noise() callback.on_rollout_start() continue_training = True while total_steps < n_steps or total_episodes < n_episodes: done = False episode_reward, episode_timesteps = 0.0, 0 while not done: if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0: # Sample a new noise matrix self.actor.reset_noise() # Select action randomly or according to policy if self.num_timesteps < learning_starts and not ( self.use_sde and self.use_sde_at_warmup): # Warmup phase unscaled_action = np.array([self.action_space.sample()]) else: # Note: we assume that the policy uses tanh to scale the action # We use non-deterministic action in the case of SAC, for TD3, it does not matter unscaled_action, _ = self.predict(self._last_obs, deterministic=False) # Rescale the action from [low, high] to [-1, 1] if isinstance(self.action_space, gym.spaces.Box): scaled_action = self.policy.scale_action(unscaled_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 scaled_action = np.clip(scaled_action + action_noise(), -1, 1) # We store the scaled action in the buffer buffer_action = scaled_action action = self.policy.unscale_action(scaled_action) else: # Discrete case, no need to normalize or clip buffer_action = unscaled_action action = buffer_action # Rescale and perform action new_obs, reward, done, infos = env.step(action) #env.render() # global reward_history_disp #reward_history_disp += [reward] # Only stop training if return value is False, not when it is None. if callback.on_step() is False: return RolloutReturn(0.0, total_steps, total_episodes, continue_training=False) episode_reward += reward #print("Last reward: ", reward, " Episode reward", reward) # Retrieve reward and episode length if using Monitor wrapper self._update_info_buffer(infos, done) # Store data in replay buffer if replay_buffer is not None: # Store only the unnormalized version if self._vec_normalize_env is not None: new_obs_ = self._vec_normalize_env.get_original_obs() reward_ = self._vec_normalize_env.get_original_reward() else: # Avoid changing the original ones self._last_original_obs, new_obs_, reward_ = self._last_obs, new_obs, reward replay_buffer.add(self._last_original_obs, new_obs_, buffer_action, reward_, done) self._last_obs = new_obs # Save the unnormalized observation if self._vec_normalize_env is not None: self._last_original_obs = new_obs_ self.num_timesteps += 1 episode_timesteps += 1 total_steps += 1 if 0 < n_steps <= total_steps: break if done: total_episodes += 1 self._episode_num += 1 episode_rewards.append(episode_reward) total_timesteps.append(episode_timesteps) #print("DONE: Last reward: ", episode_rewards[-1]) if action_noise is not None: action_noise.reset() # Display training infos if self.verbose >= 1 and log_interval is not None and self._episode_num % log_interval == 0: fps = int(self.num_timesteps / (time.time() - self.start_time)) logger.logkv("episodes", self._episode_num) 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("fps", fps) logger.logkv('time_elapsed', int(time.time() - self.start_time)) logger.logkv("total timesteps", self.num_timesteps) if self.use_sde: logger.logkv("std", (self.actor.get_std()).mean().item()) if len(self.ep_success_buffer) > 0: logger.logkv('success rate', self.safe_mean(self.ep_success_buffer)) logger.dumpkvs() mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0 #print("Mean reward: ", mean_reward) callback.on_rollout_end() return RolloutReturn(mean_reward, total_steps, total_episodes, continue_training)
def train(self, gradient_steps: int, batch_size: int = 100, policy_delay: int = 2) -> None: # Update learning rate according to lr schedule self._update_learning_rate( [self.actor.optimizer, self.critic.optimizer]) for gradient_step in range(gradient_steps): # Sample replay buffer replay_data = self.replay_buffer.sample( batch_size, env=self._vec_normalize_env) with th.no_grad(): # Select action according to policy and add clipped noise noise = replay_data.actions.clone().data.normal_( 0, self.target_policy_noise) noise = noise.clamp(-self.target_noise_clip, self.target_noise_clip) next_actions = ( self.actor_target(replay_data.next_observations) + noise).clamp(-1, 1) # Compute the target Q value target_q1, target_q2 = self.critic_target( replay_data.next_observations, next_actions) target_q = th.min(target_q1, target_q2) target_q = replay_data.rewards + ( 1 - replay_data.dones) * self.gamma * target_q # Get current Q estimates current_q1, current_q2 = self.critic(replay_data.observations, replay_data.actions) # Compute critic loss critic_loss = F.mse_loss(current_q1, target_q) + F.mse_loss( current_q2, target_q) # Optimize the critic self.critic.optimizer.zero_grad() critic_loss.backward() self.critic.optimizer.step() # Delayed policy updates if gradient_step % policy_delay == 0: # Compute actor loss actor_loss = -self.critic.q1_forward( replay_data.observations, self.actor(replay_data.observations)).mean() # Optimize the actor self.actor.optimizer.zero_grad() actor_loss.backward() self.actor.optimizer.step() # Update the frozen target networks for param, target_param in zip( self.critic.parameters(), self.critic_target.parameters()): target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data) for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()): target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data) self._n_updates += gradient_steps logger.logkv("n_updates", self._n_updates)
def train(self, gradient_steps: int, batch_size: Optional[int] = None) -> None: # Update optimizer learning rate self._update_learning_rate(self.policy.optimizer) # A2C with gradient_steps > 1 does not make sense assert gradient_steps == 1, "A2C does not support multiple gradient steps" # We do not use minibatches for A2C assert batch_size is None, "A2C does not support minibatch" for rollout_data in self.rollout_buffer.get(batch_size=None): actions = rollout_data.actions if isinstance(self.action_space, spaces.Discrete): # Convert discrete action from float to long actions = actions.long().flatten() # TODO: avoid second computation of everything because of the gradient values, log_prob, entropy = self.policy.evaluate_actions( rollout_data.observations, actions) values = values.flatten() # Normalize advantage (not present in the original implementation) advantages = rollout_data.advantages if self.normalize_advantage: advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) # Policy gradient loss policy_loss = -(advantages * log_prob).mean() # Value loss using the TD(gae_lambda) target value_loss = F.mse_loss(rollout_data.returns, values) # Entropy loss favor exploration if entropy is None: # Approximate entropy when no analytical form entropy_loss = -log_prob.mean() else: entropy_loss = -th.mean(entropy) loss = policy_loss + self.ent_coef * entropy_loss + self.vf_coef * value_loss # Optimization step self.policy.optimizer.zero_grad() loss.backward() # Clip grad norm th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) self.policy.optimizer.step() explained_var = explained_variance( self.rollout_buffer.returns.flatten(), self.rollout_buffer.values.flatten()) self._n_updates += 1 logger.logkv("n_updates", self._n_updates) logger.logkv("explained_variance", explained_var) logger.logkv("entropy_loss", entropy_loss.item()) logger.logkv("policy_loss", policy_loss.item()) logger.logkv("value_loss", value_loss.item()) if hasattr(self.policy, 'log_std'): logger.logkv("std", th.exp(self.policy.log_std).mean().item())
def train(self, gradient_steps: int, batch_size: int = 64) -> None: # Update optimizers learning rate optimizers = [self.actor.optimizer, self.critic.optimizer] if self.ent_coef_optimizer is not None: optimizers += [self.ent_coef_optimizer] # Update learning rate according to lr schedule self._update_learning_rate(optimizers) ent_coef_losses, ent_coefs = [], [] actor_losses, critic_losses = [], [] for gradient_step in range(gradient_steps): # Sample replay buffer replay_data = self.replay_buffer.sample( batch_size, env=self._vec_normalize_env) # We need to sample because `log_std` may have changed between two gradient steps if self.use_sde: self.actor.reset_noise() # Action by the current actor for the sampled state actions_pi, log_prob = self.actor.action_log_prob( replay_data.observations) log_prob = log_prob.reshape(-1, 1) ent_coef_loss = None if self.ent_coef_optimizer is not None: # Important: detach the variable from the graph # so we don't change it with other losses # see https://github.com/rail-berkeley/softlearning/issues/60 ent_coef = th.exp(self.log_ent_coef.detach()) ent_coef_loss = -( self.log_ent_coef * (log_prob + self.target_entropy).detach()).mean() ent_coef_losses.append(ent_coef_loss.item()) else: ent_coef = self.ent_coef_tensor ent_coefs.append(ent_coef.item()) # Optimize entropy coefficient, also called # entropy temperature or alpha in the paper if ent_coef_loss is not None: self.ent_coef_optimizer.zero_grad() ent_coef_loss.backward() self.ent_coef_optimizer.step() with th.no_grad(): # Select action according to policy next_actions, next_log_prob = self.actor.action_log_prob( replay_data.next_observations) # Compute the target Q value target_q1, target_q2 = self.critic_target( replay_data.next_observations, next_actions) target_q = th.min(target_q1, target_q2) target_q = replay_data.rewards + ( 1 - replay_data.dones) * self.gamma * target_q # td error + entropy term q_backup = target_q - ent_coef * next_log_prob.reshape(-1, 1) # Get current Q estimates # using action from the replay buffer current_q1, current_q2 = self.critic(replay_data.observations, replay_data.actions) # Compute critic loss critic_loss = 0.5 * (F.mse_loss(current_q1, q_backup) + F.mse_loss(current_q2, q_backup)) critic_losses.append(critic_loss.item()) # Optimize the critic self.critic.optimizer.zero_grad() critic_loss.backward() self.critic.optimizer.step() # Compute actor loss # Alternative: actor_loss = th.mean(log_prob - qf1_pi) qf1_pi, qf2_pi = self.critic.forward(replay_data.observations, actions_pi) min_qf_pi = th.min(qf1_pi, qf2_pi) actor_loss = (ent_coef * log_prob - min_qf_pi).mean() actor_losses.append(actor_loss.item()) # Optimize the actor self.actor.optimizer.zero_grad() actor_loss.backward() self.actor.optimizer.step() # Update target networks if gradient_step % self.target_update_interval == 0: for param, target_param in zip( self.critic.parameters(), self.critic_target.parameters()): target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data) self._n_updates += gradient_steps logger.logkv("n_updates", self._n_updates) logger.logkv("ent_coef", np.mean(ent_coefs)) logger.logkv("actor_loss", np.mean(actor_losses)) logger.logkv("critic_loss", np.mean(critic_losses)) if len(ent_coef_losses) > 0: logger.logkv("ent_coef_loss", np.mean(ent_coef_losses))