ml-agents/ml-agents/mlagents/trainers/ppo/trainer.py

# # 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 AllBrainInfo
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: AllBrainInfo, next_info: AllBrainInfo
    ) -> 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: Dictionary of all current brains and corresponding BrainInfo.
        :param next_info: Dictionary of all next brains and corresponding BrainInfo.
        """
        info = next_info[self.brain_name]
        # if self.is_training:
        #     self.policy.update_normalization(info.vector_observations)
        for l in range(len(info.agents)):
            agent_actions = self.training_buffer[info.agents[l]]["actions"]
            if (
                info.local_done[l]
                or len(agent_actions) > self.trainer_parameters["time_horizon"]
            ) and len(agent_actions) > 0:
                agent_id = info.agents[l]
                if info.max_reached[l]:
                    bootstrapping_info = self.training_buffer[agent_id].last_brain_info
                    idx = bootstrapping_info.agents.index(agent_id)
                else:
                    bootstrapping_info = info
                    idx = l
                value_next = self.policy.get_value_estimates(
                    bootstrapping_info,
                    idx,
                    info.local_done[l] and not 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 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