text: # # Unity ML-Agents Toolkit
  # ## ML-Agent Learning (PPO)
  # Contains an implementation of PPO as described in: https://arxiv.org/abs/1707.06347
  
  import logging
  from collections import defaultdict
  from typing import Dict
  
  import numpy as np
  
  from mlagents.envs.brain import BrainInfo
  from mlagents.trainers.ppo.policy import PPOPolicy
  from mlagents.trainers.ppo.multi_gpu_policy import MultiGpuPPOPolicy, get_devices
  from mlagents.trainers.rl_trainer import RLTrainer, AllRewardsOutput
  from mlagents.envs.action_info import ActionInfoOutputs
  
  logger = logging.getLogger("mlagents.trainers")
  
  
  class PPOTrainer(RLTrainer):
      """The PPOTrainer is an implementation of the PPO algorithm."""
  
      def __init__(
          self,
          brain,
          reward_buff_cap,
          trainer_parameters,
          training,
          load,
          seed,
          run_id,
          multi_gpu,
      ):
          """
          Responsible for collecting experiences and training PPO model.
          :param trainer_parameters: The parameters for the trainer (dictionary).
          :param reward_buff_cap: Max reward history to track in the reward buffer
          :param training: Whether the trainer is set for training.
          :param load: Whether the model should be loaded.
          :param seed: The seed the model will be initialized with
          :param run_id: The identifier of the current run
          """
          super(PPOTrainer, self).__init__(
              brain, trainer_parameters, training, run_id, reward_buff_cap
          )
          self.param_keys = [
              "batch_size",
              "beta",
              "buffer_size",
              "epsilon",
              "hidden_units",
              "lambd",
              "learning_rate",
              "max_steps",
              "normalize",
              "num_epoch",
              "num_layers",
              "time_horizon",
              "sequence_length",
              "summary_freq",
              "use_recurrent",
              "summary_path",
              "memory_size",
              "model_path",
              "reward_signals",
          ]
          self.check_param_keys()
  
          if multi_gpu and len(get_devices()) > 1:
              self.policy = MultiGpuPPOPolicy(
                  seed, brain, trainer_parameters, self.is_training, load
              )
          else:
              self.policy = PPOPolicy(
                  seed, brain, trainer_parameters, self.is_training, load
              )
  
          for _reward_signal in self.policy.reward_signals.keys():
              self.collected_rewards[_reward_signal] = {}
  
      def process_experiences(
          self, current_info: BrainInfo, next_info: BrainInfo
      ) -> None:
          """
          Checks agent histories for processing condition, and processes them as necessary.
          Processing involves calculating value and advantage targets for model updating step.
          :param current_info: current BrainInfo.
          :param next_info: next BrainInfo.
          """
          if self.is_training:
              self.policy.update_normalization(next_info.vector_observations)
          for l in range(len(next_info.agents)):
              agent_actions = self.training_buffer[next_info.agents[l]]["actions"]
              if (
                  next_info.local_done[l]
                  or len(agent_actions) > self.trainer_parameters["time_horizon"]
              ) and len(agent_actions) > 0:
                  agent_id = next_info.agents[l]
                  if next_info.max_reached[l]:
                      bootstrapping_info = self.training_buffer[agent_id].last_brain_info
                      idx = bootstrapping_info.agents.index(agent_id)
                  else:
                      bootstrapping_info = next_info
                      idx = l
                  value_next = self.policy.get_value_estimates(
                      bootstrapping_info,
                      idx,
                      next_info.local_done[l] and not next_info.max_reached[l],
                  )
  
                  tmp_advantages = []
                  tmp_returns = []
                  for name in self.policy.reward_signals:
                      bootstrap_value = value_next[name]
  
                      local_rewards = self.training_buffer[agent_id][
                          "{}_rewards".format(name)
                      ].get_batch()
                      local_value_estimates = self.training_buffer[agent_id][
                          "{}_value_estimates".format(name)
                      ].get_batch()
                      local_advantage = get_gae(
                          rewards=local_rewards,
                          value_estimates=local_value_estimates,
                          value_next=bootstrap_value,
                          gamma=self.policy.reward_signals[name].gamma,
                          lambd=self.trainer_parameters["lambd"],
                      )
                      local_return = local_advantage + local_value_estimates
                      # This is later use as target for the different value estimates
                      self.training_buffer[agent_id]["{}_returns".format(name)].set(
                          local_return
                      )
                      self.training_buffer[agent_id]["{}_advantage".format(name)].set(
                          local_advantage
                      )
                      tmp_advantages.append(local_advantage)
                      tmp_returns.append(local_return)
  
                  global_advantages = list(np.mean(np.array(tmp_advantages), axis=0))
                  global_returns = list(np.mean(np.array(tmp_returns), axis=0))
                  self.training_buffer[agent_id]["advantages"].set(global_advantages)
                  self.training_buffer[agent_id]["discounted_returns"].set(global_returns)
  
                  self.training_buffer.append_update_buffer(
                      agent_id,
                      batch_size=None,
                      training_length=self.policy.sequence_length,
                  )
  
                  self.training_buffer[agent_id].reset_agent()
                  if next_info.local_done[l]:
                      self.stats["Environment/Episode Length"].append(
                          self.episode_steps.get(agent_id, 0)
                      )
                      self.episode_steps[agent_id] = 0
                      for name, rewards in self.collected_rewards.items():
                          if name == "environment":
                              self.cumulative_returns_since_policy_update.append(
                                  rewards.get(agent_id, 0)
                              )
                              self.stats["Environment/Cumulative Reward"].append(
                                  rewards.get(agent_id, 0)
                              )
                              self.reward_buffer.appendleft(rewards.get(agent_id, 0))
                              rewards[agent_id] = 0
                          else:
                              self.stats[
                                  self.policy.reward_signals[name].stat_name
                              ].append(rewards.get(agent_id, 0))
                              rewards[agent_id] = 0
  
      def add_policy_outputs(
          self, take_action_outputs: ActionInfoOutputs, agent_id: str, agent_idx: int
      ) -> None:
          """
          Takes the output of the last action and store it into the training buffer.
          """
          actions = take_action_outputs["action"]
          if self.policy.use_continuous_act:
              actions_pre = take_action_outputs["pre_action"]
              self.training_buffer[agent_id]["actions_pre"].append(actions_pre[agent_idx])
              epsilons = take_action_outputs["random_normal_epsilon"]
              self.training_buffer[agent_id]["random_normal_epsilon"].append(
                  epsilons[agent_idx]
              )
          a_dist = take_action_outputs["log_probs"]
          # value is a dictionary from name of reward to value estimate of the value head
          self.training_buffer[agent_id]["actions"].append(actions[agent_idx])
          self.training_buffer[agent_id]["action_probs"].append(a_dist[agent_idx])
  
      def add_rewards_outputs(
          self,
          rewards_out: AllRewardsOutput,
          values: Dict[str, np.ndarray],
          agent_id: str,
          agent_idx: int,
          agent_next_idx: int,
      ) -> None:
          """
          Takes the value output of the last action and store it into the training buffer.
          """
          for name, reward_result in rewards_out.reward_signals.items():
              # 0 because we use the scaled reward to train the agent
              self.training_buffer[agent_id]["{}_rewards".format(name)].append(
                  reward_result.scaled_reward[agent_next_idx]
              )
              self.training_buffer[agent_id]["{}_value_estimates".format(name)].append(
                  values[name][agent_idx][0]
              )
  
      def is_ready_update(self):
          """
          Returns whether or not the trainer has enough elements to run update model
          :return: A boolean corresponding to whether or not update_model() can be run
          """
          size_of_buffer = len(self.training_buffer.update_buffer["actions"])
          return size_of_buffer > self.trainer_parameters["buffer_size"]
  
      def update_policy(self):
          """
          Uses demonstration_buffer to update the policy.
          The reward signal generators must be updated in this method at their own pace.
          """
          buffer_length = len(self.training_buffer.update_buffer["actions"])
          self.trainer_metrics.start_policy_update_timer(
              number_experiences=buffer_length,
              mean_return=float(np.mean(self.cumulative_returns_since_policy_update)),
          )
          self.cumulative_returns_since_policy_update.clear()
  
          # Make sure batch_size is a multiple of sequence length. During training, we
          # will need to reshape the data into a batch_size x sequence_length tensor.
          batch_size = (
              self.trainer_parameters["batch_size"]
              - self.trainer_parameters["batch_size"] % self.policy.sequence_length
          )
          # Make sure there is at least one sequence
          batch_size = max(batch_size, self.policy.sequence_length)
  
          n_sequences = max(
              int(self.trainer_parameters["batch_size"] / self.policy.sequence_length), 1
          )
  
          advantages = self.training_buffer.update_buffer["advantages"].get_batch()
          self.training_buffer.update_buffer["advantages"].set(
              (advantages - advantages.mean()) / (advantages.std() + 1e-10)
          )
          num_epoch = self.trainer_parameters["num_epoch"]
          batch_update_stats = defaultdict(list)
          for _ in range(num_epoch):
              self.training_buffer.update_buffer.shuffle(
                  sequence_length=self.policy.sequence_length
              )
              buffer = self.training_buffer.update_buffer
              max_num_batch = buffer_length // batch_size
              for l in range(0, max_num_batch * batch_size, batch_size):
                  update_stats = self.policy.update(
                      buffer.make_mini_batch(l, l + batch_size), n_sequences
                  )
                  for stat_name, value in update_stats.items():
                      batch_update_stats[stat_name].append(value)
  
          for stat, stat_list in batch_update_stats.items():
              self.stats[stat].append(np.mean(stat_list))
  
          if self.policy.bc_module:
              update_stats = self.policy.bc_module.update()
              for stat, val in update_stats.items():
                  self.stats[stat].append(val)
          self.clear_update_buffer()
          self.trainer_metrics.end_policy_update()
  
  
  def discount_rewards(r, gamma=0.99, value_next=0.0):
      """
      Computes discounted sum of future rewards for use in updating value estimate.
      :param r: List of rewards.
      :param gamma: Discount factor.
      :param value_next: T+1 value estimate for returns calculation.
      :return: discounted sum of future rewards as list.
      """
      discounted_r = np.zeros_like(r)
      running_add = value_next
      for t in reversed(range(0, r.size)):
          running_add = running_add * gamma + r[t]
          discounted_r[t] = running_add
      return discounted_r
  
  
  def get_gae(rewards, value_estimates, value_next=0.0, gamma=0.99, lambd=0.95):
      """
      Computes generalized advantage estimate for use in updating policy.
      :param rewards: list of rewards for time-steps t to T.
      :param value_next: Value estimate for time-step T+1.
      :param value_estimates: list of value estimates for time-steps t to T.
      :param gamma: Discount factor.
      :param lambd: GAE weighing factor.
      :return: list of advantage estimates for time-steps t to T.
      """
      value_estimates = np.append(value_estimates, value_next)
      delta_t = rewards + gamma * value_estimates[1:] - value_estimates[:-1]
      advantage = discount_rewards(r=delta_t, gamma=gamma * lambd)
      return advantage
  