r_mappo.py
标准 PPO trainer
原文件 · 300 行 · 13217 字节
algorithms/algorithm/r_mappo.py
SHA-256: 0f430148b51eea83dbbc4124b54f0fcd29155d729d84ce797ddfb052195efca2
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 RMAPPO():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.ppo_epoch = args.ppo_epoch35 self.num_mini_batch = args.num_mini_batch36 self.data_chunk_length = args.data_chunk_length37 self.value_loss_coef = args.value_loss_coef38 self.entropy_coef = args.entropy_coef39 self.max_grad_norm = args.max_grad_norm40 self.huber_delta = args.huber_delta4142 self._use_recurrent_policy = args.use_recurrent_policy43 self._use_naive_recurrent = args.use_naive_recurrent_policy44 self._use_max_grad_norm = args.use_max_grad_norm45 self._use_clipped_value_loss = args.use_clipped_value_loss46 self._use_huber_loss = args.use_huber_loss47 self._use_popart = args.use_popart48 self._use_valuenorm = args.use_valuenorm49 self._use_value_active_masks = args.use_value_active_masks50 self._use_policy_active_masks = args.use_policy_active_masks5152 assert (self._use_popart and self._use_valuenorm) == False, (53 "self._use_popart and self._use_valuenorm can not be set True simultaneously")5455 if self._use_popart:56 self.value_normalizer = self.policy.critic.v_out57 elif self._use_valuenorm:58 self.value_normalizer = ValueNorm(1, device=self.device)59 else:60 self.value_normalizer = None6162 def cal_value_loss(self, values, value_preds_batch, return_batch, active_masks_batch):63 """64 Calculate value function loss.65 :param values: (torch.Tensor) value function predictions.66 :param value_preds_batch: (torch.Tensor) "old" value predictions from data batch (used for value clip loss)67 :param return_batch: (torch.Tensor) reward to go returns.68 :param active_masks_batch: (torch.Tensor) denotes if agent is active or dead at a given timesep.6970 :return value_loss: (torch.Tensor) value function loss.71 """72 value_pred_clipped = value_preds_batch + (values - value_preds_batch).clamp(-self.clip_param,73 self.clip_param)74 if self._use_popart or self._use_valuenorm:75 self.value_normalizer.update(return_batch)76 error_clipped = self.value_normalizer.normalize(return_batch) - value_pred_clipped77 error_original = self.value_normalizer.normalize(return_batch) - values78 else:79 error_clipped = return_batch - value_pred_clipped80 error_original = return_batch - values8182 if self._use_huber_loss:83 value_loss_clipped = huber_loss(error_clipped, self.huber_delta)84 value_loss_original = huber_loss(error_original, self.huber_delta)85 else:86 value_loss_clipped = mse_loss(error_clipped)87 value_loss_original = mse_loss(error_original)8889 if self._use_clipped_value_loss:90 value_loss = torch.max(value_loss_original, value_loss_clipped)91 else:92 value_loss = value_loss_original9394 if self._use_value_active_masks:95 value_loss = (value_loss * active_masks_batch).sum() / active_masks_batch.sum()96 else:97 value_loss = value_loss.mean()9899 return value_loss100101 def ppo_update(self, sample, update_actor=True):102 """103 Update actor and critic networks.104 :param sample: (Tuple) contains data batch with which to update networks.105 :update_actor: (bool) whether to update actor network.106107 :return value_loss: (torch.Tensor) value function loss.108 :return critic_grad_norm: (torch.Tensor) gradient norm from critic up9date.109 ;return policy_loss: (torch.Tensor) actor(policy) loss value.110 :return dist_entropy: (torch.Tensor) action entropies.111 :return actor_grad_norm: (torch.Tensor) gradient norm from actor update.112 :return imp_weights: (torch.Tensor) importance sampling weights.113 """114 share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch, \115 value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch, \116 adv_targ, available_actions_batch = sample117118 old_action_log_probs_batch = check(old_action_log_probs_batch).to(**self.tpdv)119 adv_targ = check(adv_targ).to(**self.tpdv)120 value_preds_batch = check(value_preds_batch).to(**self.tpdv)121 return_batch = check(return_batch).to(**self.tpdv)122 active_masks_batch = check(active_masks_batch).to(**self.tpdv)123124 # Reshape to do in a single forward pass for all steps125 values, action_log_probs, dist_entropy = self.policy.evaluate_actions(share_obs_batch,126 obs_batch,127 rnn_states_batch,128 rnn_states_critic_batch,129 actions_batch,130 masks_batch,131 available_actions_batch,132 active_masks_batch)133 # actor update134 imp_weights = torch.exp(action_log_probs - old_action_log_probs_batch)135 clipped = (imp_weights - 1.0).abs() > self.clip_param136 with torch.no_grad():137 log_ratio = action_log_probs - old_action_log_probs_batch138 approx_kl_per_sample = (imp_weights - 1.0) - log_ratio139 if self._use_policy_active_masks:140 denominator = active_masks_batch.sum().clamp_min(1.0)141 approx_kl = (142 approx_kl_per_sample * active_masks_batch143 ).sum() / denominator144 clip_fraction = (145 clipped.float() * active_masks_batch146 ).sum() / denominator147 ratio_mean = (148 imp_weights * active_masks_batch149 ).sum() / denominator150 active_ratios = imp_weights[active_masks_batch > 0.5]151 else:152 approx_kl = approx_kl_per_sample.mean()153 clip_fraction = clipped.float().mean()154 ratio_mean = imp_weights.mean()155 active_ratios = imp_weights.reshape(-1)156 if active_ratios.numel() == 0:157 active_ratios = imp_weights.new_ones(1)158 ratio_min = active_ratios.min()159 ratio_max = active_ratios.max()160 ratio_p01 = torch.quantile(active_ratios, 0.01)161 ratio_p99 = torch.quantile(active_ratios, 0.99)162163 surr1 = imp_weights * adv_targ164 surr2 = torch.clamp(imp_weights, 1.0 - self.clip_param, 1.0 + self.clip_param) * adv_targ165166 if self._use_policy_active_masks:167 policy_action_loss = (-torch.sum(torch.min(surr1, surr2),168 dim=-1,169 keepdim=True) * active_masks_batch).sum() / active_masks_batch.sum()170 else:171 policy_action_loss = -torch.sum(torch.min(surr1, surr2), dim=-1, keepdim=True).mean()172173 policy_loss = policy_action_loss174175 self.policy.actor_optimizer.zero_grad()176177 if update_actor:178 (policy_loss - dist_entropy * self.entropy_coef).backward()179180 if self._use_max_grad_norm:181 actor_grad_norm = nn.utils.clip_grad_norm_(self.policy.actor.parameters(), self.max_grad_norm)182 else:183 actor_grad_norm = get_gard_norm(self.policy.actor.parameters())184185 self.policy.actor_optimizer.step()186187 # critic update188 value_loss = self.cal_value_loss(values, value_preds_batch, return_batch, active_masks_batch)189190 self.policy.critic_optimizer.zero_grad()191192 (value_loss * self.value_loss_coef).backward()193194 if self._use_max_grad_norm:195 critic_grad_norm = nn.utils.clip_grad_norm_(self.policy.critic.parameters(), self.max_grad_norm)196 else:197 critic_grad_norm = get_gard_norm(self.policy.critic.parameters())198199 self.policy.critic_optimizer.step()200201 return (202 value_loss,203 critic_grad_norm,204 policy_loss,205 dist_entropy,206 actor_grad_norm,207 ratio_mean,208 ratio_min,209 ratio_max,210 ratio_p01,211 ratio_p99,212 approx_kl,213 clip_fraction,214 )215216 def train(self, buffer, update_actor=True):217 """218 Perform a training update using minibatch GD.219 :param buffer: (SharedReplayBuffer) buffer containing training data.220 :param update_actor: (bool) whether to update actor network.221222 :return train_info: (dict) contains information regarding training update (e.g. loss, grad norms, etc).223 """224 if self._use_popart or self._use_valuenorm:225 advantages = buffer.returns[:-1] - self.value_normalizer.denormalize(buffer.value_preds[:-1])226 else:227 advantages = buffer.returns[:-1] - buffer.value_preds[:-1]228 advantages_copy = advantages.copy()229 advantages_copy[buffer.active_masks[:-1] == 0.0] = np.nan230 mean_advantages = np.nanmean(advantages_copy)231 std_advantages = np.nanstd(advantages_copy)232 advantages = (advantages - mean_advantages) / (std_advantages + 1e-5)233234 train_info = {}235236 train_info['value_loss'] = 0237 train_info['policy_loss'] = 0238 train_info['dist_entropy'] = 0239 train_info['actor_grad_norm'] = 0240 train_info['critic_grad_norm'] = 0241 train_info['ratio'] = 0242 train_info['ratio_min'] = 0243 train_info['ratio_max'] = 0244 train_info['ratio_p01'] = 0245 train_info['ratio_p99'] = 0246 train_info['approx_kl'] = 0247 train_info['clip_fraction'] = 0248249 for _ in range(self.ppo_epoch):250 if self._use_recurrent_policy:251 data_generator = buffer.recurrent_generator(advantages, self.num_mini_batch, self.data_chunk_length)252 elif self._use_naive_recurrent:253 data_generator = buffer.naive_recurrent_generator(advantages, self.num_mini_batch)254 else:255 data_generator = buffer.feed_forward_generator(advantages, self.num_mini_batch)256257 for sample in data_generator:258 (259 value_loss,260 critic_grad_norm,261 policy_loss,262 dist_entropy,263 actor_grad_norm,264 ratio_mean,265 ratio_min,266 ratio_max,267 ratio_p01,268 ratio_p99,269 approx_kl,270 clip_fraction,271 ) = self.ppo_update(sample, update_actor)272273 train_info['value_loss'] += value_loss.item()274 train_info['policy_loss'] += policy_loss.item()275 train_info['dist_entropy'] += dist_entropy.item()276 train_info['actor_grad_norm'] += actor_grad_norm277 train_info['critic_grad_norm'] += critic_grad_norm278 train_info['ratio'] += ratio_mean.item()279 train_info['ratio_min'] += ratio_min.item()280 train_info['ratio_max'] += ratio_max.item()281 train_info['ratio_p01'] += ratio_p01.item()282 train_info['ratio_p99'] += ratio_p99.item()283 train_info['approx_kl'] += approx_kl.item()284 train_info['clip_fraction'] += clip_fraction.item()285286 num_updates = self.ppo_epoch * self.num_mini_batch287288 for k in train_info.keys():289 train_info[k] /= num_updates290291 return train_info292293 def prep_training(self):294 self.policy.actor.train()295 self.policy.critic.train()296297 def prep_rollout(self):298 self.policy.actor.eval()299 self.policy.critic.eval()300