ga_rmappo.py
PPO 更新、active masks、损失与优化停止条件
原文件 · 343 行 · 16370 字节
algorithms/algorithm/ga_rmappo.py
SHA-256: c68bf923dfbca391d4ec7138c2e5e9064cfa0061d3aae231390c84c34d2797e7
1"""2# @Time : 2021/7/1 6:52 下午3# @Author : hezhiqiang014# @Email : hezhiqiang01@baidu.com5# @File : r_mappo.py6"""78import numpy as np9import torch10import torch.nn as nn11from utils.util import get_gard_norm, huber_loss, mse_loss12from utils.valuenorm import ValueNorm13from algorithms.utils.util import check141516class GARMAPPO():17 """18 Trainer class for MAPPO to update policies.19 :param args: (argparse.Namespace) arguments containing relevant model, policy, and env information.20 :param policy: (R_MAPPO_Policy) policy to update.21 :param device: (torch.device) specifies the device to run on (cpu/gpu).22 """2324 def __init__(self,25 args,26 policy,27 device=torch.device("cpu")):2829 self.device = device30 self.tpdv = dict(dtype=torch.float32, device=device)31 self.policy = policy3233 self.clip_param = args.clip_param34 self.target_kl = float(getattr(args, "ppo_target_kl", 0.0) or 0.0)35 self.max_clip_fraction = float(getattr(args, "ppo_max_clip_fraction", 0.0) or 0.0)36 self.ppo_epoch = args.ppo_epoch37 self.num_mini_batch = args.num_mini_batch38 self.data_chunk_length = args.data_chunk_length39 self.value_loss_coef = args.value_loss_coef40 self.utility_loss_coef = getattr(args, "utility_loss_coef", 0.0)41 self.entropy_coef = args.entropy_coef42 self.max_grad_norm = args.max_grad_norm43 self.huber_delta = args.huber_delta4445 self._use_recurrent_policy = args.use_recurrent_policy46 self._use_naive_recurrent = args.use_naive_recurrent_policy47 self._use_max_grad_norm = args.use_max_grad_norm48 self._use_clipped_value_loss = args.use_clipped_value_loss49 self._use_huber_loss = args.use_huber_loss50 self._use_popart = args.use_popart51 self._use_valuenorm = args.use_valuenorm52 self._use_value_active_masks = args.use_value_active_masks53 self._use_policy_active_masks = args.use_policy_active_masks54 self.reconfiguration_decision_coef = float(55 getattr(args, "reconfiguration_decision_coef", 0.0)56 )5758 assert (self._use_popart and self._use_valuenorm) == False, (59 "self._use_popart and self._use_valuenorm can not be set True simultaneously")6061 if self._use_popart:62 self.value_normalizer = self.policy.critic.v_out63 elif self._use_valuenorm:64 self.value_normalizer = ValueNorm(1, device=self.device)65 else:66 self.value_normalizer = None6768 @staticmethod69 def _scalar(value):70 if isinstance(value, torch.Tensor):71 return float(value.detach().cpu().item())72 return float(value)7374 def cal_value_loss(self, values, value_preds_batch, return_batch, active_masks_batch):75 """76 Calculate value function loss.77 :param values: (torch.Tensor) value function predictions.78 :param value_preds_batch: (torch.Tensor) "old" value predictions from data batch (used for value clip loss)79 :param return_batch: (torch.Tensor) reward to go returns.80 :param active_masks_batch: (torch.Tensor) denotes if agent is active or dead at a given timesep.8182 :return value_loss: (torch.Tensor) value function loss.83 """84 value_pred_clipped = value_preds_batch + (values - value_preds_batch).clamp(-self.clip_param,85 self.clip_param)86 if self._use_popart or self._use_valuenorm:87 self.value_normalizer.update(return_batch)88 error_clipped = self.value_normalizer.normalize(return_batch) - value_pred_clipped89 error_original = self.value_normalizer.normalize(return_batch) - values90 else:91 error_clipped = return_batch - value_pred_clipped92 error_original = return_batch - values9394 if self._use_huber_loss:95 value_loss_clipped = huber_loss(error_clipped, self.huber_delta)96 value_loss_original = huber_loss(error_original, self.huber_delta)97 else:98 value_loss_clipped = mse_loss(error_clipped)99 value_loss_original = mse_loss(error_original)100101 if self._use_clipped_value_loss:102 value_loss = torch.max(value_loss_original, value_loss_clipped)103 else:104 value_loss = value_loss_original105106 if self._use_value_active_masks:107 value_loss = (value_loss * active_masks_batch).sum() / active_masks_batch.sum()108 else:109 value_loss = value_loss.mean()110111 return value_loss112113 def ppo_update(self, sample, update_actor=True):114 """115 Update actor and critic networks.116 :param sample: (Tuple) contains data batch with which to update networks.117 :update_actor: (bool) whether to update actor network.118119 :return value_loss: (torch.Tensor) value function loss.120 :return critic_grad_norm: (torch.Tensor) gradient norm from critic up9date.121 ;return policy_loss: (torch.Tensor) actor(policy) loss value.122 :return dist_entropy: (torch.Tensor) action entropies.123 :return actor_grad_norm: (torch.Tensor) gradient norm from actor update.124 :return imp_weights: (torch.Tensor) importance sampling weights.125 """126 share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch, \127 value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch, \128 adv_targ, available_actions_batch, gate_actions_batch = sample129130 old_action_log_probs_batch = check(old_action_log_probs_batch).to(**self.tpdv)131 adv_targ = check(adv_targ).to(**self.tpdv)132 value_preds_batch = check(value_preds_batch).to(**self.tpdv)133 return_batch = check(return_batch).to(**self.tpdv)134 active_masks_batch = check(active_masks_batch).to(**self.tpdv)135 # Reshape to do in a single forward pass for all steps136 values, action_log_probs, dist_entropy = self.policy.evaluate_actions(share_obs_batch,137 obs_batch,138 rnn_states_batch,139 rnn_states_critic_batch,140 actions_batch,141 masks_batch,142 available_actions_batch,143 active_masks_batch,144 gate_actions_batch)145 imp_weights = torch.exp(action_log_probs - old_action_log_probs_batch)146 clipped = (imp_weights - 1.0).abs() > self.clip_param147 with torch.no_grad():148 log_ratio = action_log_probs - old_action_log_probs_batch149 approx_kl_per_sample = (imp_weights - 1.0) - log_ratio150 if self._use_policy_active_masks:151 denom = active_masks_batch.sum().clamp_min(1.0)152 approx_kl = (approx_kl_per_sample * active_masks_batch).sum() / denom153 clip_fraction = (clipped.float() * active_masks_batch).sum() / denom154 ratio_mean = (imp_weights * active_masks_batch).sum() / denom155 else:156 approx_kl = approx_kl_per_sample.mean()157 clip_fraction = clipped.float().mean()158 ratio_mean = imp_weights.mean()159 if self._use_policy_active_masks:160 active_ratios = imp_weights[active_masks_batch > 0.5]161 else:162 active_ratios = imp_weights.reshape(-1)163 if active_ratios.numel() == 0:164 active_ratios = imp_weights.new_ones(1)165 ratio_min = active_ratios.min()166 ratio_max = active_ratios.max()167 ratio_p01 = torch.quantile(active_ratios, 0.01)168 ratio_p99 = torch.quantile(active_ratios, 0.99)169170 surr1 = imp_weights * adv_targ171 surr2 = torch.clamp(imp_weights, 1.0 - self.clip_param, 1.0 + self.clip_param) * adv_targ172173 if self._use_policy_active_masks:174 policy_action_loss = (-torch.sum(torch.min(surr1, surr2),175 dim=-1,176 keepdim=True) * active_masks_batch).sum() / active_masks_batch.sum()177 else:178 policy_action_loss = -torch.sum(torch.min(surr1, surr2), dim=-1, keepdim=True).mean()179180 policy_loss = policy_action_loss181 gate_regularization = self.policy.get_reconfiguration_gate_regularization()182183 # critic update184 # 2. 再算 critic 的 value_loss（保持原来的写法）185 value_loss = self.cal_value_loss(values, value_preds_batch, return_batch, active_masks_batch)186187 # 3. auxiliary utility-head loss (non-zero only for HATT v3.2d+)188 utility_loss = self.policy.get_utility_loss(obs_batch, available_actions_batch)189 if isinstance(utility_loss, torch.Tensor):190 utility_loss = utility_loss.to(**self.tpdv)191192 # 4. 一次性清两个优化器的梯度193 self.policy.actor_optimizer.zero_grad()194 self.policy.critic_optimizer.zero_grad()195196 # 5. 组合一个总 loss197 if update_actor:198 # reconfiguration_decision_coef 只作用于 hierarchical learned gate。199 # 它压低非强制状态下的开门概率，避免模型退化成 always-redecide。200 total_loss = policy_loss \201 - dist_entropy * self.entropy_coef \202 + self.reconfiguration_decision_coef * gate_regularization \203 + self.value_loss_coef * value_loss \204 + self.utility_loss_coef * utility_loss205 else:206 # 只更新 critic（以及和 critic 共享的 GAT），不更新 actor207 total_loss = self.value_loss_coef * value_loss \208 + self.utility_loss_coef * utility_loss209210 # 5. 一次 backward，图只走一遍211 total_loss.backward()212213 actor_parameters = self.policy.get_actor_parameters()214 if self._use_max_grad_norm:215 actor_grad_norm = nn.utils.clip_grad_norm_(actor_parameters,216 self.max_grad_norm)217 critic_grad_norm = nn.utils.clip_grad_norm_(self.policy.critic.parameters(),218 self.max_grad_norm)219 else:220 actor_grad_norm = get_gard_norm(actor_parameters)221 critic_grad_norm = get_gard_norm(self.policy.critic.parameters())222223 if update_actor:224 self.policy.actor_optimizer.step()225 self.policy.critic_optimizer.step()226227 stop_update = False228 if update_actor:229 if self.target_kl > 0.0 and float(approx_kl.item()) > self.target_kl:230 stop_update = True231 if self.max_clip_fraction > 0.0 and float(clip_fraction.item()) > self.max_clip_fraction:232 stop_update = True233234 return {235 "value_loss": value_loss,236 "critic_grad_norm": critic_grad_norm,237 "policy_loss": policy_loss,238 "dist_entropy": dist_entropy,239 "actor_grad_norm": actor_grad_norm,240 "ratio": ratio_mean,241 "ratio_min": ratio_min,242 "ratio_max": ratio_max,243 "ratio_p01": ratio_p01,244 "ratio_p99": ratio_p99,245 "approx_kl": approx_kl,246 "clip_fraction": clip_fraction,247 "utility_loss": utility_loss,248 "gate_regularization": gate_regularization,249 "stop_update": stop_update,250 }251252 def train(self, buffer, update_actor=True):253 """254 Perform a training update using minibatch GD.255 :param buffer: (SharedReplayBuffer) buffer containing training data.256 :param update_actor: (bool) whether to update actor network.257258 :return train_info: (dict) contains information regarding training update (e.g. loss, grad norms, etc).259 """260 if self._use_popart or self._use_valuenorm:261 advantages = buffer.returns[:-1] - self.value_normalizer.denormalize(buffer.value_preds[:-1])262 else:263 advantages = buffer.returns[:-1] - buffer.value_preds[:-1]264 advantages_copy = advantages.copy()265 advantages_copy[buffer.active_masks[:-1] == 0.0] = np.nan266 mean_advantages = np.nanmean(advantages_copy)267 std_advantages = np.nanstd(advantages_copy)268 advantages = (advantages - mean_advantages) / (std_advantages + 1e-5)269270 train_info = {}271272 train_info['value_loss'] = 0273 train_info['policy_loss'] = 0274 train_info['dist_entropy'] = 0275 train_info['actor_grad_norm'] = 0276 train_info['critic_grad_norm'] = 0277 train_info['ratio'] = 0278 train_info['ratio_min'] = 0279 train_info['ratio_max'] = 0280 train_info['ratio_p01'] = 0281 train_info['ratio_p99'] = 0282 train_info['approx_kl'] = 0283 train_info['clip_fraction'] = 0284 train_info['ppo_updates'] = 0285 train_info['ppo_early_stop'] = 0286 train_info['utility_loss'] = 0287 train_info['gate_regularization'] = 0288289 num_updates = 0290 early_stop = False291 for _ in range(self.ppo_epoch):292 if self._use_recurrent_policy:293 data_generator = buffer.recurrent_generator(advantages, self.num_mini_batch, self.data_chunk_length)294 elif self._use_naive_recurrent:295 data_generator = buffer.naive_recurrent_generator(advantages, self.num_mini_batch)296 else:297 data_generator = buffer.feed_forward_generator(advantages, self.num_mini_batch)298299 for sample in data_generator:300 update_info = self.ppo_update(sample, update_actor)301302 train_info['value_loss'] += update_info["value_loss"].item()303 train_info['policy_loss'] += update_info["policy_loss"].item()304 train_info['dist_entropy'] += update_info["dist_entropy"].item()305 train_info['actor_grad_norm'] += self._scalar(update_info["actor_grad_norm"])306 train_info['critic_grad_norm'] += self._scalar(update_info["critic_grad_norm"])307 train_info['ratio'] += self._scalar(update_info["ratio"])308 train_info['ratio_min'] += self._scalar(update_info["ratio_min"])309 train_info['ratio_max'] += self._scalar(update_info["ratio_max"])310 train_info['ratio_p01'] += self._scalar(update_info["ratio_p01"])311 train_info['ratio_p99'] += self._scalar(update_info["ratio_p99"])312 train_info['approx_kl'] += self._scalar(update_info["approx_kl"])313 train_info['clip_fraction'] += self._scalar(update_info["clip_fraction"])314 utility_loss = update_info["utility_loss"]315 train_info['utility_loss'] += utility_loss.item() if isinstance(utility_loss, torch.Tensor) else 0.0316 train_info['gate_regularization'] += self._scalar(317 update_info["gate_regularization"]318 )319 num_updates += 1320321 if update_info["stop_update"]:322 early_stop = True323 break324 if early_stop:325 break326327 for k in train_info.keys():328 if k in ("ppo_updates", "ppo_early_stop"):329 continue330 train_info[k] /= max(1, num_updates)331 train_info["ppo_updates"] = num_updates332 train_info["ppo_early_stop"] = float(early_stop)333334 return train_info335336 def prep_training(self):337 self.policy.actor.train()338 self.policy.critic.train()339340 def prep_rollout(self):341 self.policy.actor.eval()342 self.policy.critic.eval()343