HATTGyMRT2ANetwork.py
GymRT2A 空间候选 HATT actor
原文件 · 135 行 · 6795 字节
algorithms/algorithm/HATTGyMRT2ANetwork.py
SHA-256: dd9e27c9f677dee4686c5af0180d7a5cc227f05b516f2e8d7309c85b5736da2f
1"""用于 GyMRT2A 空间动作网格的 HATT actor。23动作网格中的每个格子都作为一个候选位置，并共用同一个 decoder。当前可见的其他4智能体先通过 A-A attention 汇总，再参与候选位置打分。5"""67import torch8import torch.nn as nn910from algorithms.algorithm.HATTV20Network import FeatureEncoder, RelationAttention11from algorithms.utils.distributions import FixedCategorical12from algorithms.utils.rnn import RNNLayer13from algorithms.utils.util import check141516class HATTGyMRT2AActor(nn.Module):17 """将空间观测编码成网格中每个位置的动作分数。"""18 def __init__(self, args, ego_dim, other_dim, task_dim, num_agents, num_tasks, device):19 super().__init__()20 self.num_agents = num_agents21 self.num_tasks = num_tasks22 self.hidden_size = args.hidden_size23 self.use_policy_active_masks = args.use_policy_active_masks24 self.use_recurrent = args.use_naive_recurrent_policy or args.use_recurrent_policy25 self.tpdv = dict(dtype=torch.float32, device=device)26 hidden = args.hatt_hidden_dim2728 self.ego_encoder = FeatureEncoder(ego_dim, hidden)29 # 当前智能体、其他智能体和候选位置的字段分别编码。30 self.other_encoder = FeatureEncoder(other_dim, hidden)31 self.candidate_encoder = FeatureEncoder(task_dim, hidden)32 self.aa_attention = RelationAttention(hidden, edge_dim=3, num_heads=args.hatt_num_heads)33 # A-A 输出与当前智能体表示合并后，再转换到 MAPPO 的 hidden_size。34 self.agent_fusion = nn.Sequential(nn.Linear(hidden * 2, hidden), nn.ReLU(), nn.Linear(hidden, hidden))35 self.agent_norm = nn.LayerNorm(hidden)36 self.agent_projection = nn.Sequential(37 nn.Linear(hidden, self.hidden_size), nn.ReLU(), nn.LayerNorm(self.hidden_size)38 )39 if self.use_recurrent:40 self.rnn = RNNLayer(41 self.hidden_size, self.hidden_size, args.recurrent_N, args.use_orthogonal42 )43 self.decoder = nn.Sequential(44 # 每个格子的输入包含智能体表示、候选位置表示和二维相对坐标。45 nn.Linear(self.hidden_size + hidden + 2, self.hidden_size),46 nn.ReLU(),47 nn.Linear(self.hidden_size, self.hidden_size),48 nn.ReLU(),49 nn.Linear(self.hidden_size, 1),50 )51 nn.init.orthogonal_(self.decoder[-1].weight, gain=0.01)52 nn.init.zeros_(self.decoder[-1].bias)53 self.last_diagnostics = {}54 self.to(device)5556 def _prepare_obs(self, obs):57 """整理观测张量，并记录调用方是否只传入了一个 batch。"""58 if not isinstance(obs, dict):59 raise TypeError("HATTGyMRT2AActor requires Dict observations.")60 result = {key: check(value).to(**self.tpdv) for key, value in obs.items()}61 if result["ego"].ndim == 2:62 # buffer 会把 batch 和智能体维压平，这里恢复成两维。63 result = {64 key: value.reshape(value.shape[0] // self.num_agents, self.num_agents, *value.shape[1:])65 for key, value in result.items()66 }67 return result6869 def _distribution(self, obs, rnn_states, masks, available_actions):70 # HATT 先输出任务槽分数，环境适配层已将任务槽排成空间动作网格。71 # RNN 保存跨时间步的信息，动作掩码负责去掉当前不可达的位置。72 obs = self._prepare_obs(obs)73 batch, agents, _ = obs["ego"].shape74 available = check(available_actions).to(**self.tpdv).reshape(batch, agents, self.num_tasks)7576 ego_h = self.ego_encoder(obs["ego"])77 other_h = self.other_encoder(obs["others"])78 candidate_h = self.candidate_encoder(obs["tasks"])79 # others 的前两项是相对位置，第 3 项表示该智能体是否可见。80 relative = obs["others"][..., :2]81 visible = obs["others"][..., 2] > 0.582 edge = torch.cat([relative, visible.float().unsqueeze(-1)], dim=-1)83 aa_context, aa_weights = self.aa_attention(ego_h, other_h, edge, visible)84 # 残差保留当前智能体自己的编码。85 agent_h = self.agent_norm(ego_h + self.agent_fusion(torch.cat([ego_h, aa_context], dim=-1)))86 agent_h = self.agent_projection(agent_h).reshape(batch * agents, self.hidden_size)8788 rnn_states = check(rnn_states).to(**self.tpdv)89 masks = check(masks).to(**self.tpdv)90 if self.use_recurrent:91 # RNN 的输入按 batch × 智能体展开。92 agent_h, rnn_states = self.rnn(agent_h, rnn_states, masks)93 agent_h = agent_h.reshape(batch, agents, self.hidden_size)9495 expanded = agent_h.unsqueeze(2).expand(-1, -1, self.num_tasks, -1)96 # 将当前智能体表示复制到每个网格位置，再逐格计算 logit。97 candidate_xy = obs["tasks"][..., :2]98 logits = self.decoder(torch.cat([expanded, candidate_h, candidate_xy], dim=-1)).squeeze(-1)99 logits = logits.reshape(batch * agents, self.num_tasks)100 logits = logits.masked_fill(available.reshape(batch * agents, self.num_tasks) <= 0, -1e10)101 # 诊断项只用于日志，全部从反向传播图中分离。102 self.last_diagnostics = {103 "aa_attention": aa_weights.detach(),104 "visible_agents": visible.float().sum(dim=-1).detach(),105 "candidate_context_norm": candidate_h.norm(dim=-1).detach(),106 }107 return FixedCategorical(logits=logits), rnn_states108109 def forward(self, obs, rnn_states, masks, available_actions=None, deterministic=False):110 dist, rnn_states = self._distribution(obs, rnn_states, masks, available_actions)111 actions = dist.mode() if deterministic else dist.sample()112 return actions, dist.log_probs(actions), rnn_states113114 def evaluate_actions(115 self,116 obs,117 rnn_states,118 action,119 masks,120 available_actions=None,121 active_masks=None,122 gate_actions=None,123 ):124 del gate_actions # GyMRT2A 只有一个离散动作输出，不使用额外 gate 动作。125 dist, _ = self._distribution(obs, rnn_states, masks, available_actions)126 action = check(action).to(device=self.tpdv["device"])127 log_probs = dist.log_probs(action)128 entropy = dist.entropy()129 # 未激活的智能体不参与平均熵计算。130 if active_masks is not None and self.use_policy_active_masks:131 active_masks = check(active_masks).to(**self.tpdv).squeeze(-1)132 entropy = (entropy * active_masks).sum() / active_masks.sum().clamp_min(1.0)133 else:134 entropy = entropy.mean()135 return log_probs, entropy