您最多选择25个主题
主题必须以中文或者字母或数字开头，可以包含连字符 (-)，并且长度不得超过35个字符
326 行
13 KiB
326 行
13 KiB
# # 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.trainers.brain import BrainParameters, BrainInfo
|
|
from mlagents.trainers.tf_policy import TFPolicy
|
|
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.trainers.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_name,
|
|
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_name, 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()
|
|
self.load = load
|
|
self.multi_gpu = multi_gpu
|
|
self.seed = seed
|
|
self.policy = None
|
|
|
|
def process_experiences(
|
|
self, name_behavior_id: str, 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 name_behavior_id: string policy identifier.
|
|
: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.processing_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.processing_buffer[
|
|
agent_id
|
|
].last_brain_info
|
|
idx = bootstrapping_info.agents.index(agent_id)
|
|
else:
|
|
bootstrapping_info = next_info
|
|
idx = l
|
|
value_next = self.ppo_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.processing_buffer[agent_id][
|
|
"{}_rewards".format(name)
|
|
].get_batch()
|
|
local_value_estimates = self.processing_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.processing_buffer[agent_id]["{}_returns".format(name)].set(
|
|
local_return
|
|
)
|
|
self.processing_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, dtype=np.float32), axis=0)
|
|
)
|
|
global_returns = list(
|
|
np.mean(np.array(tmp_returns, dtype=np.float32), axis=0)
|
|
)
|
|
self.processing_buffer[agent_id]["advantages"].set(global_advantages)
|
|
self.processing_buffer[agent_id]["discounted_returns"].set(
|
|
global_returns
|
|
)
|
|
|
|
self.processing_buffer.append_to_update_buffer(
|
|
self.update_buffer,
|
|
agent_id,
|
|
batch_size=None,
|
|
training_length=self.policy.sequence_length,
|
|
)
|
|
|
|
self.processing_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.processing_buffer[agent_id]["actions_pre"].append(
|
|
actions_pre[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.processing_buffer[agent_id]["actions"].append(actions[agent_idx])
|
|
self.processing_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.processing_buffer[agent_id]["{}_rewards".format(name)].append(
|
|
reward_result.scaled_reward[agent_next_idx]
|
|
)
|
|
self.processing_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 = self.update_buffer.num_experiences
|
|
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 = self.update_buffer.num_experiences
|
|
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.update_buffer["advantages"].get_batch()
|
|
self.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.update_buffer.shuffle(sequence_length=self.policy.sequence_length)
|
|
buffer = self.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 create_policy(self, brain_parameters: BrainParameters) -> TFPolicy:
|
|
if self.multi_gpu and len(get_devices()) > 1:
|
|
policy = MultiGpuPPOPolicy(
|
|
self.seed,
|
|
brain_parameters,
|
|
self.trainer_parameters,
|
|
self.is_training,
|
|
self.load,
|
|
)
|
|
else:
|
|
policy = PPOPolicy(
|
|
self.seed,
|
|
brain_parameters,
|
|
self.trainer_parameters,
|
|
self.is_training,
|
|
self.load,
|
|
)
|
|
|
|
for _reward_signal in policy.reward_signals.keys():
|
|
self.collected_rewards[_reward_signal] = {}
|
|
|
|
return policy
|
|
|
|
|
|
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
|