ppo_learner.py
MRT2A 上游 PPO learner
原文件 · 266 行 · 12641 字节
third_party/marl_mrt2a/marl_mrt2a/marl/src/learners/ppo_learner.py
SHA-256: afe1562bba86741f97e4145748fd29f6a1b10e347a448eec342c7c1c68b9fc12
1import copy2from components.episode_buffer import EpisodeBatch3from modules.critics.coma import COMACritic4from modules.critics.centralV import CentralVCritic5from utils.rl_utils import build_td_lambda_targets6import torch as th7from torch.optim import Adam8from modules.critics import REGISTRY as critic_registry9from components.standarize_stream import RunningMeanStd101112# 一个完整 episode 的栅格观测无法一次放入显存，因此按时间段前向计算并累积梯度。13# 当前策略没有使用 RNN，各时间步的损失可以分别计算后求和。14_TIME_CHUNK = 32151617class PPOLearner:18 """使用集中式 critic 更新共享的 MARL-MRT2A 策略。"""19 def __init__(self, mac, scheme, logger, args):20 self.args = args21 self.n_agents = args.n_agents22 self.n_actions = args.n_actions23 self.logger = logger2425 self.mac = mac26 # old_mac 保存本轮 PPO 更新开始前的策略，用于计算概率比值。27 self.old_mac = copy.deepcopy(mac)28 self.agent_params = list(mac.parameters())29 self.agent_optimiser = Adam(params=self.agent_params, lr=args.lr)3031 self.critic = critic_registry[args.critic_type](scheme, args)32 # target_critic 提供计算回报目标时使用的延迟参数。33 self.target_critic = copy.deepcopy(self.critic)3435 self.critic_params = list(self.critic.parameters())36 self.critic_optimiser = Adam(params=self.critic_params, lr=args.lr)3738 self.last_target_update_step = 039 self.critic_training_steps = 040 self.log_stats_t = -self.args.learner_log_interval - 14142 device = "cuda" if args.use_cuda else "cpu"43 if self.args.standardise_returns:44 self.ret_ms = RunningMeanStd(shape=(self.n_agents, ), device=device)45 if self.args.standardise_rewards:46 self.rew_ms = RunningMeanStd(shape=(1,), device=device)4748 def train(self, batch: EpisodeBatch, t_env: int, episode_num: int):49 # filled 去掉补齐的时间步，terminated 去掉 episode 结束后的数据。5051 rewards = batch["reward"][:, :-1]52 actions = batch["actions"][:, :]53 terminated = batch["terminated"][:, :-1].float()54 mask = batch["filled"][:, :-1].float()55 mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1])56 target_updates = batch["target_update"][:, :-1][...,1]57 action_exec = batch["action_exec"][:, :-1][...,0]58 # target_updates 表示本步重新选过目标，action_exec 表示环境实际执行了动作。59 actions = actions[:, :-1]60 if self.args.standardise_rewards:61 self.rew_ms.update(rewards)62 rewards = (rewards - self.rew_ms.mean) / th.sqrt(self.rew_ms.var)6364 mask = mask.repeat(1, 1, self.n_agents)65 mask = th.logical_and( mask, action_exec ).float()66 # critic 使用实际执行过动作的时间步；actor 还要求本步重新选择过目标。67 critic_mask = mask.clone()68 mask = th.logical_and( mask, target_updates ).float()69 # mask = th.logical_and( mask, action_exec ).float()7071 # 当前 batch 没有发生目标选择时，actor 没有可用样本。72 if mask.sum() == 0:73 self.logger.log_stat("Mask_Sum_Zero", 1, t_env)74 self.logger.console_logger.error("Actor Critic Learner: mask.sum() == 0 at t_env {}".format(t_env))75 return7677 with th.no_grad():78 # 整条序列先用 old_mac 计算一次，PPO 多轮更新期间保持不变。79 old_mac_out = []80 self.old_mac.init_hidden(batch.batch_size)81 for t in range(batch.max_seq_length - 1):82 agent_outs, _, _ = self.old_mac.forward(batch, t=t)83 old_mac_out.append(agent_outs)84 try:85 old_mac_out = th.stack(old_mac_out, dim=1) # Concat over time86 except:87 pass88 old_pi = old_mac_out89 old_pi[mask == 0] = 1.09091 old_pi_taken = th.gather(old_pi, dim=3, index=actions).squeeze(3)92 # gather 只取轨迹中实际选择动作的概率。93 old_log_pi_taken = th.log(old_pi_taken + 1e-10)9495 mask_sum = mask.sum()96 for k in range(self.args.epochs):97 # 每个 epoch 先更新 critic，并用最新 critic 重新计算 advantage。98 advantages, critic_train_stats = self.train_critic_sequential(self.critic, self.target_critic, batch, rewards,99 critic_mask)100 advantages = advantages.detach()101102 # 分段计算 PPO 损失，每段 backward 后保留梯度，最后统一更新参数。103 self.mac.init_hidden(batch.batch_size)104 self.agent_optimiser.zero_grad()105 pg_loss = 0.0106 pi_max_num = 0.0107 for t0 in range(0, batch.max_seq_length - 1, _TIME_CHUNK):108 t1 = min(t0 + _TIME_CHUNK, batch.max_seq_length - 1)109 pi_parts = []110 for t in range(t0, t1):111 agent_outs, _, _ = self.mac.forward(batch, t=t)112 pi_parts.append(agent_outs)113 pi = th.stack(pi_parts, dim=1) # Concat over time (chunk)114115 pi[mask[:, t0:t1] == 0] = 1.0116117 pi_taken = th.gather(pi, dim=3, index=actions[:, t0:t1]).squeeze(3)118 log_pi_taken = th.log(pi_taken + 1e-10)119120 ratios = th.exp(log_pi_taken - old_log_pi_taken[:, t0:t1].detach())121 # ratios 衡量当前策略相对更新前策略改变了多少。122 surr1 = ratios * advantages[:, t0:t1]123 surr2 = th.clamp(ratios, 1 - self.args.eps_clip, 1 + self.args.eps_clip) * advantages[:, t0:t1]124125 entropy = -th.sum(pi * th.log(pi + 1e-10), dim=-1)126 chunk_loss = -((th.min(surr1, surr2) + self.args.entropy_coef * entropy) * mask[:, t0:t1]).sum() / mask_sum127 if th.isnan(chunk_loss).any():128 raise ValueError("NaN in loss")129 chunk_loss.backward()130 pg_loss += chunk_loss.detach().item()131 pi_max_num += (pi.max(dim=-1)[0] * mask[:, t0:t1]).sum().item()132133 grad_norm = th.nn.utils.clip_grad_norm_(self.agent_params, self.args.grad_norm_clip)134 self.agent_optimiser.step()135136 self.old_mac.load_state(self.mac)137138 self.critic_training_steps += 1139 if self.args.target_update_interval_or_tau > 1 and (140 self.critic_training_steps - self.last_target_update_step) / self.args.target_update_interval_or_tau >= 1.0:141 self._update_targets_hard()142 self.last_target_update_step = self.critic_training_steps143 elif self.args.target_update_interval_or_tau <= 1.0:144 self._update_targets_soft(self.args.target_update_interval_or_tau)145146 if t_env - self.log_stats_t >= self.args.learner_log_interval:147 ts_logged = len(critic_train_stats["critic_loss"])148 for key in ["critic_loss", "critic_grad_norm", "td_error_abs", "q_taken_mean", "target_mean"]:149 self.logger.log_stat(key, sum(critic_train_stats[key]) / ts_logged, t_env)150151 self.logger.log_stat("advantage_mean", (advantages * mask).sum().item() / mask_sum.item(), t_env)152 self.logger.log_stat("pg_loss", pg_loss, t_env)153 self.logger.log_stat("agent_grad_norm", grad_norm.item(), t_env)154 self.logger.log_stat("pi_max", pi_max_num / mask_sum.item(), t_env)155 self.log_stats_t = t_env156157 def train_critic_sequential(self, critic, target_critic, batch, rewards, mask):158 # Optimise critic159 target_vals = []160 with th.no_grad():161 for t in range(batch.max_seq_length):162 target_vals.append(target_critic(batch, t=t).squeeze(3))163 target_vals = th.cat(target_vals, dim=1)164165 if self.args.standardise_returns:166 target_vals = target_vals * th.sqrt(self.ret_ms.var) + self.ret_ms.mean167168 target_returns = self.nstep_returns(rewards, mask, target_vals, self.args.q_nstep)169 if self.args.standardise_returns:170 self.ret_ms.update(target_returns)171 target_returns = (target_returns - self.ret_ms.mean) / (th.sqrt(self.ret_ms.var)+1e-5)172173 running_log = {174 "critic_loss": [],175 "critic_grad_norm": [],176 "td_error_abs": [],177 "target_mean": [],178 "q_taken_mean": [],179 }180181 # Chunked value forward with gradient accumulation so the full182 # 750-step sequence never has to be materialised on the device at once.183 mask_sum = mask.sum()184 critic_loss = 0.0185 v_parts = []186 self.critic_optimiser.zero_grad()187 for t0 in range(0, batch.max_seq_length - 1, _TIME_CHUNK):188 t1 = min(t0 + _TIME_CHUNK, batch.max_seq_length - 1)189 v = []190 for t in range(t0, t1):191 v.append(critic(batch, t=t).squeeze(3))192 v = th.cat(v, dim=1)193 masked_td_error = (target_returns[:, t0:t1].detach() - v) * mask[:, t0:t1]194 loss = (masked_td_error ** 2).sum() / mask_sum195 loss.backward()196 critic_loss += loss.detach().item()197 v_parts.append(v.detach())198199 grad_norm = th.nn.utils.clip_grad_norm_(self.critic_params, self.args.grad_norm_clip)200 self.critic_optimiser.step()201202 v = th.cat(v_parts, dim=1)203 masked_td_error = (target_returns.detach() - v) * mask204205 running_log["critic_loss"].append(critic_loss)206 running_log["critic_grad_norm"].append(grad_norm.item())207 mask_elems = mask_sum.item()208 running_log["td_error_abs"].append((masked_td_error.abs().sum().item() / mask_elems))209 running_log["q_taken_mean"].append((v * mask).sum().item() / mask_elems)210 running_log["target_mean"].append((target_returns * mask).sum().item() / mask_elems)211212 if th.isnan(masked_td_error).any():213 raise ValueError("NaN in TD error")214215 return masked_td_error, running_log216217 def nstep_returns(self, rewards, mask, values, nsteps):218 nstep_values = th.zeros_like(values[:, :-1])219 for t_start in range(rewards.size(1)):220 nstep_return_t = th.zeros_like(values[:, 0])221 for step in range(nsteps + 1):222 t = t_start + step223 if t >= rewards.size(1):224 break225 elif step == nsteps:226 nstep_return_t += self.args.gamma ** (step) * values[:, t] * mask[:, t]227 elif t == rewards.size(1) - 1 and self.args.add_value_last_step:228 nstep_return_t += self.args.gamma ** (step) * rewards[:, t] * mask[:, t]229 nstep_return_t += self.args.gamma ** (step + 1) * values[:, t + 1]230 else:231 nstep_return_t += self.args.gamma ** (step) * rewards[:, t] * mask[:, t]232 nstep_values[:, t_start, :] = nstep_return_t233 return nstep_values234235 def _update_targets(self):236 self.target_critic.load_state_dict(self.critic.state_dict())237238 def _update_targets_hard(self):239 self.target_critic.load_state_dict(self.critic.state_dict())240241 def _update_targets_soft(self, tau):242 for target_param, param in zip(self.target_critic.parameters(), self.critic.parameters()):243 target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau)244245 def cuda(self):246 self.old_mac.cuda()247 self.mac.cuda()248 self.critic.cuda()249 self.target_critic.cuda()250251 def save_models(self, path):252 self.mac.save_models(path)253 th.save(self.critic.state_dict(), "{}/critic.th".format(path))254 th.save(self.agent_optimiser.state_dict(), "{}/agent_opt.th".format(path))255 th.save(self.critic_optimiser.state_dict(), "{}/critic_opt.th".format(path))256257 def load_models(self, path, eval=False):258 self.mac.load_models(path)259 if not eval:260 self.critic.load_state_dict(th.load("{}/critic.th".format(path), map_location=lambda storage, loc: storage))261 # Not quite right but I don't want to save target networks262 self.target_critic.load_state_dict(self.critic.state_dict())263 self.critic_optimiser.load_state_dict(264 th.load("{}/critic_opt.th".format(path), map_location=lambda storage, loc: storage))265 self.agent_optimiser.load_state_dict(266 th.load("{}/agent_opt.th".format(path), map_location=lambda storage, loc: storage))