""" # @Time : 2021/7/1 6:52 下午 # @Author : hezhiqiang01 # @Email : hezhiqiang01@baidu.com # @File : r_mappo.py """ import numpy as np import torch import torch.nn as nn from utils.util import get_gard_norm, huber_loss, mse_loss from utils.valuenorm import ValueNorm from algorithms.utils.util import check class GARMAPPO(): """ Trainer class for MAPPO to update policies. :param args: (argparse.Namespace) arguments containing relevant model, policy, and env information. :param policy: (R_MAPPO_Policy) policy to update. :param device: (torch.device) specifies the device to run on (cpu/gpu). """ def __init__(self, args, policy, device=torch.device("cpu")): self.device = device self.tpdv = dict(dtype=torch.float32, device=device) self.policy = policy self.clip_param = args.clip_param self.target_kl = float(getattr(args, "ppo_target_kl", 0.0) or 0.0) self.max_clip_fraction = float(getattr(args, "ppo_max_clip_fraction", 0.0) or 0.0) self.ppo_epoch = args.ppo_epoch self.num_mini_batch = args.num_mini_batch self.data_chunk_length = args.data_chunk_length self.value_loss_coef = args.value_loss_coef self.utility_loss_coef = getattr(args, "utility_loss_coef", 0.0) self.entropy_coef = args.entropy_coef self.max_grad_norm = args.max_grad_norm self.huber_delta = args.huber_delta self._use_recurrent_policy = args.use_recurrent_policy self._use_naive_recurrent = args.use_naive_recurrent_policy self._use_max_grad_norm = args.use_max_grad_norm self._use_clipped_value_loss = args.use_clipped_value_loss self._use_huber_loss = args.use_huber_loss self._use_popart = args.use_popart self._use_valuenorm = args.use_valuenorm self._use_value_active_masks = args.use_value_active_masks self._use_policy_active_masks = args.use_policy_active_masks self.reconfiguration_decision_coef = float( getattr(args, "reconfiguration_decision_coef", 0.0) ) assert (self._use_popart and self._use_valuenorm) == False, ( "self._use_popart and self._use_valuenorm can not be set True simultaneously") if self._use_popart: self.value_normalizer = self.policy.critic.v_out elif self._use_valuenorm: self.value_normalizer = ValueNorm(1, device=self.device) else: self.value_normalizer = None @staticmethod def _scalar(value): if isinstance(value, torch.Tensor): return float(value.detach().cpu().item()) return float(value) def cal_value_loss(self, values, value_preds_batch, return_batch, active_masks_batch): """ Calculate value function loss. :param values: (torch.Tensor) value function predictions. :param value_preds_batch: (torch.Tensor) "old" value predictions from data batch (used for value clip loss) :param return_batch: (torch.Tensor) reward to go returns. :param active_masks_batch: (torch.Tensor) denotes if agent is active or dead at a given timesep. :return value_loss: (torch.Tensor) value function loss. """ value_pred_clipped = value_preds_batch + (values - value_preds_batch).clamp(-self.clip_param, self.clip_param) if self._use_popart or self._use_valuenorm: self.value_normalizer.update(return_batch) error_clipped = self.value_normalizer.normalize(return_batch) - value_pred_clipped error_original = self.value_normalizer.normalize(return_batch) - values else: error_clipped = return_batch - value_pred_clipped error_original = return_batch - values if self._use_huber_loss: value_loss_clipped = huber_loss(error_clipped, self.huber_delta) value_loss_original = huber_loss(error_original, self.huber_delta) else: value_loss_clipped = mse_loss(error_clipped) value_loss_original = mse_loss(error_original) if self._use_clipped_value_loss: value_loss = torch.max(value_loss_original, value_loss_clipped) else: value_loss = value_loss_original if self._use_value_active_masks: value_loss = (value_loss * active_masks_batch).sum() / active_masks_batch.sum() else: value_loss = value_loss.mean() return value_loss def ppo_update(self, sample, update_actor=True): """ Update actor and critic networks. :param sample: (Tuple) contains data batch with which to update networks. :update_actor: (bool) whether to update actor network. :return value_loss: (torch.Tensor) value function loss. :return critic_grad_norm: (torch.Tensor) gradient norm from critic up9date. ;return policy_loss: (torch.Tensor) actor(policy) loss value. :return dist_entropy: (torch.Tensor) action entropies. :return actor_grad_norm: (torch.Tensor) gradient norm from actor update. :return imp_weights: (torch.Tensor) importance sampling weights. """ share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch, \ value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch, \ adv_targ, available_actions_batch, gate_actions_batch = sample old_action_log_probs_batch = check(old_action_log_probs_batch).to(**self.tpdv) adv_targ = check(adv_targ).to(**self.tpdv) value_preds_batch = check(value_preds_batch).to(**self.tpdv) return_batch = check(return_batch).to(**self.tpdv) active_masks_batch = check(active_masks_batch).to(**self.tpdv) # Reshape to do in a single forward pass for all steps values, action_log_probs, dist_entropy = self.policy.evaluate_actions(share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch, masks_batch, available_actions_batch, active_masks_batch, gate_actions_batch) imp_weights = torch.exp(action_log_probs - old_action_log_probs_batch) clipped = (imp_weights - 1.0).abs() > self.clip_param with torch.no_grad(): log_ratio = action_log_probs - old_action_log_probs_batch approx_kl_per_sample = (imp_weights - 1.0) - log_ratio if self._use_policy_active_masks: denom = active_masks_batch.sum().clamp_min(1.0) approx_kl = (approx_kl_per_sample * active_masks_batch).sum() / denom clip_fraction = (clipped.float() * active_masks_batch).sum() / denom ratio_mean = (imp_weights * active_masks_batch).sum() / denom else: approx_kl = approx_kl_per_sample.mean() clip_fraction = clipped.float().mean() ratio_mean = imp_weights.mean() if self._use_policy_active_masks: active_ratios = imp_weights[active_masks_batch > 0.5] else: active_ratios = imp_weights.reshape(-1) if active_ratios.numel() == 0: active_ratios = imp_weights.new_ones(1) ratio_min = active_ratios.min() ratio_max = active_ratios.max() ratio_p01 = torch.quantile(active_ratios, 0.01) ratio_p99 = torch.quantile(active_ratios, 0.99) surr1 = imp_weights * adv_targ surr2 = torch.clamp(imp_weights, 1.0 - self.clip_param, 1.0 + self.clip_param) * adv_targ if self._use_policy_active_masks: policy_action_loss = (-torch.sum(torch.min(surr1, surr2), dim=-1, keepdim=True) * active_masks_batch).sum() / active_masks_batch.sum() else: policy_action_loss = -torch.sum(torch.min(surr1, surr2), dim=-1, keepdim=True).mean() policy_loss = policy_action_loss gate_regularization = self.policy.get_reconfiguration_gate_regularization() # critic update # 2. 再算 critic 的 value_loss（保持原来的写法） value_loss = self.cal_value_loss(values, value_preds_batch, return_batch, active_masks_batch) # 3. auxiliary utility-head loss (non-zero only for HATT v3.2d+) utility_loss = self.policy.get_utility_loss(obs_batch, available_actions_batch) if isinstance(utility_loss, torch.Tensor): utility_loss = utility_loss.to(**self.tpdv) # 4. 一次性清两个优化器的梯度 self.policy.actor_optimizer.zero_grad() self.policy.critic_optimizer.zero_grad() # 5. 组合一个总 loss if update_actor: # reconfiguration_decision_coef 只作用于 hierarchical learned gate。 # 它压低非强制状态下的开门概率，避免模型退化成 always-redecide。 total_loss = policy_loss \ - dist_entropy * self.entropy_coef \ + self.reconfiguration_decision_coef * gate_regularization \ + self.value_loss_coef * value_loss \ + self.utility_loss_coef * utility_loss else: # 只更新 critic（以及和 critic 共享的 GAT），不更新 actor total_loss = self.value_loss_coef * value_loss \ + self.utility_loss_coef * utility_loss # 5. 一次 backward，图只走一遍 total_loss.backward() actor_parameters = self.policy.get_actor_parameters() if self._use_max_grad_norm: actor_grad_norm = nn.utils.clip_grad_norm_(actor_parameters, self.max_grad_norm) critic_grad_norm = nn.utils.clip_grad_norm_(self.policy.critic.parameters(), self.max_grad_norm) else: actor_grad_norm = get_gard_norm(actor_parameters) critic_grad_norm = get_gard_norm(self.policy.critic.parameters()) if update_actor: self.policy.actor_optimizer.step() self.policy.critic_optimizer.step() stop_update = False if update_actor: if self.target_kl > 0.0 and float(approx_kl.item()) > self.target_kl: stop_update = True if self.max_clip_fraction > 0.0 and float(clip_fraction.item()) > self.max_clip_fraction: stop_update = True return { "value_loss": value_loss, "critic_grad_norm": critic_grad_norm, "policy_loss": policy_loss, "dist_entropy": dist_entropy, "actor_grad_norm": actor_grad_norm, "ratio": ratio_mean, "ratio_min": ratio_min, "ratio_max": ratio_max, "ratio_p01": ratio_p01, "ratio_p99": ratio_p99, "approx_kl": approx_kl, "clip_fraction": clip_fraction, "utility_loss": utility_loss, "gate_regularization": gate_regularization, "stop_update": stop_update, } def train(self, buffer, update_actor=True): """ Perform a training update using minibatch GD. :param buffer: (SharedReplayBuffer) buffer containing training data. :param update_actor: (bool) whether to update actor network. :return train_info: (dict) contains information regarding training update (e.g. loss, grad norms, etc). """ if self._use_popart or self._use_valuenorm: advantages = buffer.returns[:-1] - self.value_normalizer.denormalize(buffer.value_preds[:-1]) else: advantages = buffer.returns[:-1] - buffer.value_preds[:-1] advantages_copy = advantages.copy() advantages_copy[buffer.active_masks[:-1] == 0.0] = np.nan mean_advantages = np.nanmean(advantages_copy) std_advantages = np.nanstd(advantages_copy) advantages = (advantages - mean_advantages) / (std_advantages + 1e-5) train_info = {} train_info['value_loss'] = 0 train_info['policy_loss'] = 0 train_info['dist_entropy'] = 0 train_info['actor_grad_norm'] = 0 train_info['critic_grad_norm'] = 0 train_info['ratio'] = 0 train_info['ratio_min'] = 0 train_info['ratio_max'] = 0 train_info['ratio_p01'] = 0 train_info['ratio_p99'] = 0 train_info['approx_kl'] = 0 train_info['clip_fraction'] = 0 train_info['ppo_updates'] = 0 train_info['ppo_early_stop'] = 0 train_info['utility_loss'] = 0 train_info['gate_regularization'] = 0 num_updates = 0 early_stop = False for _ in range(self.ppo_epoch): if self._use_recurrent_policy: data_generator = buffer.recurrent_generator(advantages, self.num_mini_batch, self.data_chunk_length) elif self._use_naive_recurrent: data_generator = buffer.naive_recurrent_generator(advantages, self.num_mini_batch) else: data_generator = buffer.feed_forward_generator(advantages, self.num_mini_batch) for sample in data_generator: update_info = self.ppo_update(sample, update_actor) train_info['value_loss'] += update_info["value_loss"].item() train_info['policy_loss'] += update_info["policy_loss"].item() train_info['dist_entropy'] += update_info["dist_entropy"].item() train_info['actor_grad_norm'] += self._scalar(update_info["actor_grad_norm"]) train_info['critic_grad_norm'] += self._scalar(update_info["critic_grad_norm"]) train_info['ratio'] += self._scalar(update_info["ratio"]) train_info['ratio_min'] += self._scalar(update_info["ratio_min"]) train_info['ratio_max'] += self._scalar(update_info["ratio_max"]) train_info['ratio_p01'] += self._scalar(update_info["ratio_p01"]) train_info['ratio_p99'] += self._scalar(update_info["ratio_p99"]) train_info['approx_kl'] += self._scalar(update_info["approx_kl"]) train_info['clip_fraction'] += self._scalar(update_info["clip_fraction"]) utility_loss = update_info["utility_loss"] train_info['utility_loss'] += utility_loss.item() if isinstance(utility_loss, torch.Tensor) else 0.0 train_info['gate_regularization'] += self._scalar( update_info["gate_regularization"] ) num_updates += 1 if update_info["stop_update"]: early_stop = True break if early_stop: break for k in train_info.keys(): if k in ("ppo_updates", "ppo_early_stop"): continue train_info[k] /= max(1, num_updates) train_info["ppo_updates"] = num_updates train_info["ppo_early_stop"] = float(early_stop) return train_info def prep_training(self): self.policy.actor.train() self.policy.critic.train() def prep_rollout(self): self.policy.actor.eval() self.policy.critic.eval()