HATTV33AANetwork.py
TypedRelationEdge 与 v3.3AA 基础；learned/uniform/zero 语义
原文件 · 436 行 · 18592 字节
algorithms/algorithm/HATTV33AANetwork.py
SHA-256: 6b0ee905eff4b1d5fcd203023c5aa0d00a256b39cd0aef791fa0fc3b6f13dbdb
1"""HATT v3.3AA 用同一套计算步骤处理智能体边和任务边。23智能体边和任务边各自保存参数。现有成员的智能体边输出会按联盟编号放入对应的4候选任务，再和当前智能体到该任务的边输出一起生成动作表示。56``hatt_candidate_aa_mode`` 只控制智能体边的 gate：``learned`` 使用网络输出，7``uniform`` 固定为 0.5，``zero`` 将这条分支清零。8"""910import torch11import torch.nn as nn1213from algorithms.algorithm.HATTV20Network import (14 HATTV20Actor,15 FeatureEncoder,16 encode_repeated_tasks,17)181920class TypedRelationEdge(nn.Module):21 """输入一组有方向的边，为每条有效边生成 gate 和消息。"""2223 def __init__(self, hidden_dim, edge_dim, num_heads, dropout=0.0):24 super().__init__()25 # 每个 head 使用相同宽度，因此 hidden_dim 必须能被 head 数整除。26 if hidden_dim % num_heads != 0:27 raise ValueError(28 f"hidden_dim={hidden_dim} must be divisible by num_heads={num_heads}."29 )30 self.num_heads = num_heads31 self.head_dim = hidden_dim // num_heads32 # 点积除以维度平方根，避免 hidden_dim 增大后数值过大。33 self.scale = self.head_dim ** -0.534 # source 生成 query，target 分别生成 key 和 value。35 self.query = nn.Linear(hidden_dim, hidden_dim)36 self.key = nn.Linear(hidden_dim, hidden_dim)37 self.value = nn.Linear(hidden_dim, hidden_dim)38 # 相对位置和联盟关系会转换成每个 head 的额外分数。39 self.edge_bias = nn.Sequential(40 nn.Linear(edge_dim, hidden_dim),41 nn.ReLU(),42 nn.Linear(hidden_dim, num_heads),43 )44 self.output = nn.Linear(hidden_dim, hidden_dim)45 nn.init.constant_(self.output.bias, 0.0)46 self.dropout = nn.Dropout(dropout)4748 def forward(self, source, target, edge_features, valid_mask, gate_mode="learned"):49 # source 是边的起点，target 可以包含多个终点。每个 head 分别计算 gate，50 # 随后用 gate 缩放终点的 value，并合并成一条消息。51 if gate_mode not in ("learned", "uniform", "zero"):52 raise ValueError(53 "gate_mode must be learned, uniform, or zero; "54 f"got {gate_mode!r}."55 )56 prefix = target.shape[:-2]57 num_targets = target.shape[-2]58 hidden_dim = target.shape[-1]59 # 把最后一维拆成多个 head。prefix 中保留 batch、观察者等前置维度。60 query = self.query(source).view(*prefix, self.num_heads, self.head_dim)61 key = self.key(target).view(62 *prefix, num_targets, self.num_heads, self.head_dim63 )64 value = self.value(target).view(65 *prefix, num_targets, self.num_heads, self.head_dim66 )67 # 每条边的分数由 query-key 点积和边特征共同决定。68 logits = (69 torch.einsum("...he,...nhe->...nh", query, key) * self.scale70 + self.edge_bias(edge_features)71 )72 learned_gates = torch.sigmoid(logits)73 valid = valid_mask.bool().unsqueeze(-1)74 # padding、死亡智能体和不可选任务对应的 gate 固定为零。75 learned_gates = torch.where(76 valid, learned_gates, torch.zeros_like(learned_gates)77 )78 if gate_mode == "learned":79 gates = learned_gates80 elif gate_mode == "uniform":81 # uniform 用固定权重替代网络输出，供对照实验使用。82 gates = valid.expand_as(learned_gates).to(learned_gates.dtype) * 0.583 else:84 gates = torch.zeros_like(learned_gates)85 # 每个 head 用自己的 gate 缩放 value，再把所有 head 拼回 hidden_dim。86 gates = self.dropout(gates)87 message = (gates.unsqueeze(-1) * value).reshape(88 *prefix, num_targets, hidden_dim89 )90 message = self.output(message)91 if gate_mode == "zero":92 # output 层可能含有 bias，这里再次清零，保证整条分支没有输出。93 message = torch.zeros_like(message)94 # learned_gates 始终保留网络原始结果，uniform/zero 模式下也可用于诊断。95 return message, gates, learned_gates969798class HATTV33AABackbone(nn.Module):99 """汇总智能体边和任务边的消息，为每个候选任务生成表示。"""100101 AA_MODES = ("learned", "uniform", "zero")102103 def __init__(104 self,105 ego_dim,106 other_dim,107 task_dim,108 hidden_dim,109 num_heads,110 cap_dim,111 num_agents,112 num_tasks,113 dropout=0.0,114 shared_task_obs=True,115 candidate_aa_mode="learned",116 ):117 super().__init__()118 # 配置错误在初始化时直接报出，避免训练过程中静默切换计算方式。119 if candidate_aa_mode not in self.AA_MODES:120 raise ValueError(121 f"hatt_candidate_aa_mode must be one of {self.AA_MODES}, "122 f"got {candidate_aa_mode!r}."123 )124 self.cap_dim = cap_dim125 self.num_agents = num_agents126 self.num_tasks = num_tasks127 self.shared_task_obs = shared_task_obs128 self.candidate_aa_mode = candidate_aa_mode129130 # 三类实体先各自编码到相同的 hidden_dim，后续关系层才能直接组合。131 self.self_encoder = FeatureEncoder(ego_dim, hidden_dim)132 self.other_encoder = FeatureEncoder(other_dim, hidden_dim)133 self.task_encoder = FeatureEncoder(task_dim, hidden_dim)134 # 两类边使用相同的计算步骤，参数分别保存在 aa_relation 和 at_relation 中。135 self.aa_relation = TypedRelationEdge(136 hidden_dim, edge_dim=3, num_heads=num_heads, dropout=dropout137 )138 self.at_relation = TypedRelationEdge(139 hidden_dim, edge_dim=3, num_heads=num_heads, dropout=dropout140 )141 # 当前智能体表示会和全局 A-A、A-T 消息拼接，生成 agent_context。142 self.agent_fusion = nn.Sequential(143 nn.Linear(hidden_dim * 3, hidden_dim),144 nn.ReLU(),145 nn.Linear(hidden_dim, hidden_dim),146 )147 self.agent_norm = nn.LayerNorm(hidden_dim)148 # 每个候选任务合并任务本身、当前智能体到任务的消息、联盟成员消息和人数。149 self.candidate_fusion = nn.Sequential(150 nn.Linear(hidden_dim * 4 + 1, hidden_dim),151 nn.ReLU(),152 nn.Linear(hidden_dim, hidden_dim),153 )154 self.candidate_norm = nn.LayerNorm(hidden_dim)155156 def _coalition_indices(self, normalized_ids):157 """把观测中的归一化联盟编号恢复成任务槽下标。"""158 return (normalized_ids * max(1, self.num_tasks - 1)).round().long().clamp(159 0, self.num_tasks - 1160 )161162 @staticmethod163 def _delivery_mask(tasks):164 """标出任务列表中的 Deliver 槽。"""165 # 任务类型除以 3 存入最后一项，因此 Deliver 对应 1/3。166 return torch.isclose(167 tasks[..., -1],168 tasks.new_tensor(1.0 / 3.0),169 atol=1e-5,170 rtol=0.0,171 )172173 @staticmethod174 def _masked_mean(value, mask):175 """只计算 mask 选中元素的均值；没有元素时返回零。"""176 selected = value.masked_select(mask)177 return selected.mean() if selected.numel() else value.new_zeros(())178179 def _relation_features(self, ego, others, tasks):180 """生成相对位置、同配送联盟和当前任务归属特征。"""181 ego_position = ego[..., 0:2]182 ego_coalition = ego[..., 2 + self.cap_dim]183 other_position = others[..., 0:2]184 other_coalition = others[..., self.cap_dim + 4]185 ego_indices = self._coalition_indices(ego_coalition)186 other_indices = self._coalition_indices(other_coalition)187188 # 先查出当前智能体所处任务的类型，控制动作不产生“同配送联盟”关系。189 current_task_type = tasks[..., -1].gather(2, ego_indices.unsqueeze(-1))190 current_is_delivery = torch.isclose(191 current_task_type,192 tasks.new_tensor(1.0 / 3.0),193 atol=1e-5,194 rtol=0.0,195 )196 same_delivery_coalition = (197 (other_indices == ego_indices.unsqueeze(2)) & current_is_delivery198 ).to(ego.dtype)199 # A-A 边包含对方相对位置，以及双方是否属于同一个配送联盟。200 aa_features = torch.cat(201 [202 other_position - ego_position.unsqueeze(2),203 same_delivery_coalition.unsqueeze(-1),204 ],205 dim=-1,206 )207208 task_position = tasks[..., 0:2]209 task_ids = torch.arange(self.num_tasks, device=tasks.device)210 # assigned 表示候选任务是否就是当前智能体正在执行的任务。211 assigned = (ego_indices.unsqueeze(-1) == task_ids.view(1, 1, -1)).to(212 ego.dtype213 )214 at_features = torch.cat(215 [216 task_position - ego_position.unsqueeze(2),217 assigned.unsqueeze(-1),218 ],219 dim=-1,220 )221 return aa_features, at_features, other_indices222223 @staticmethod224 def _pool_messages(messages, valid_mask):225 """平均有效边的消息，用于生成当前智能体的整体关系信息。"""226 valid = valid_mask.to(messages.dtype).unsqueeze(-1)227 count = valid.sum(dim=-2).clamp_min(1.0)228 return (messages * valid).sum(dim=-2) / count229230 def forward(self, obs, available_actions, return_diagnostics=True):231 """生成当前智能体表示、各候选任务表示、A-T 边特征和诊断数据。"""232 ego = obs["ego"]233 others = obs["others"]234 tasks = obs["tasks"]235 # 联盟编号用于后面的任务槽查找，不作为普通连续值送入实体 encoder。236 ego_input = ego.clone()237 others_input = others.clone()238 ego_input[..., 2 + self.cap_dim] = 0.0239 others_input[..., self.cap_dim + 4] = 0.0240 # self_h、other_h、task_h 的最后一维都等于 hidden_dim。241 self_h = self.self_encoder(ego_input)242 other_h = self.other_encoder(others_input)243 task_h = encode_repeated_tasks(244 self.task_encoder, tasks, self.shared_task_obs245 )246 aa_features, at_features, other_indices = self._relation_features(247 ego, others, tasks248 )249250 living_index = self.cap_dim + 6251 # 全零行是 padding；living 为 0 的智能体也不能发送消息。252 other_mask = (253 (others.abs().sum(dim=-1) > 0)254 & (others[..., living_index] > 0.5)255 )256 available_mask = (257 available_actions > 0258 if available_actions is not None259 else torch.ones(260 *tasks.shape[:-1], dtype=torch.bool, device=tasks.device261 )262 )263 # 当前智能体分别向所有其他智能体和所有可选任务建立有方向的边。264 aa_messages, aa_gates, learned_aa_gates = self.aa_relation(265 self_h,266 other_h,267 aa_features,268 other_mask,269 gate_mode=self.candidate_aa_mode,270 )271 at_messages, at_gates, _ = self.at_relation(272 self_h,273 task_h,274 at_features,275 available_mask,276 gate_mode="learned",277 )278279 # 为每个其他智能体取出它当前加入的任务，只计算这一条智能体到任务的边。280 # 这里复用 at_relation，无需展开“所有成员 × 所有任务”的完整组合。281 member_task_h = task_h.unsqueeze(2).expand(282 -1, -1, self.num_agents, -1, -1283 ).gather(284 3,285 other_indices[..., None, None].expand(286 -1, -1, -1, 1, task_h.shape[-1]287 ),288 )289 # 同样取出该任务的位置，用于计算成员到其当前任务的相对位置。290 member_task_position = tasks[..., 0:2].unsqueeze(2).expand(291 -1, -1, self.num_agents, -1, -1292 ).gather(293 3,294 other_indices[..., None, None].expand(-1, -1, -1, 1, 2),295 ).squeeze(3)296 # 最后一项固定为 1，表示这条边连接的是成员当前已经加入的任务。297 member_at_features = torch.cat(298 [299 member_task_position - others[..., 0:2],300 torch.ones_like(other_indices, dtype=tasks.dtype).unsqueeze(-1),301 ],302 dim=-1,303 ).unsqueeze(3)304 member_task_is_delivery = self._delivery_mask(tasks).gather(305 2, other_indices306 )307 # 只让存活成员向自己当前加入的 Deliver 任务发送消息。308 member_at_valid = (other_mask & member_task_is_delivery).unsqueeze(-1)309 member_at_messages, member_at_gates, _ = self.at_relation(310 other_h,311 member_task_h,312 member_at_features,313 member_at_valid,314 gate_mode="learned",315 )316 member_at_messages = member_at_messages.squeeze(3)317318 # 分别平均所有有效的智能体边消息和任务边消息，生成当前智能体的整体表示。319 global_aa = self._pool_messages(aa_messages, other_mask)320 global_at = self._pool_messages(at_messages, available_mask)321 # 残差保留 self_h，再用 LayerNorm 稳定合并后的数值范围。322 agent_update = self.agent_fusion(323 torch.cat([self_h, global_aa, global_at], dim=-1)324 )325 agent_context = self.agent_norm(self_h + agent_update)326327 task_ids = torch.arange(self.num_tasks, device=tasks.device).view(1, 1, -1)328 delivery_mask = self._delivery_mask(tasks)329 # member_mask 的维度是“观察者 × 候选任务 × 其他智能体”。330 member_mask = (331 (other_indices.unsqueeze(2) == task_ids.unsqueeze(-1))332 & other_mask.unsqueeze(2)333 & delivery_mask.unsqueeze(-1)334 )335 # 根据联盟编号，把已有的智能体边消息放到对应的候选任务中。336 routed_aa = aa_messages.unsqueeze(2) * member_mask.unsqueeze(-1).to(337 aa_messages.dtype338 )339 routed_member_at = member_at_messages.unsqueeze(2) * member_mask.unsqueeze(340 -1341 ).to(member_at_messages.dtype)342 # 每个候选任务分别统计当前联盟人数。343 member_count = member_mask.sum(dim=3, keepdim=True).to(task_h.dtype)344 # 除以人数平方根，减小大联盟简单累加带来的数值增长。345 coalition_aa = routed_aa.sum(dim=3) / member_count.clamp_min(1.0).sqrt()346 coalition_member_at = (347 routed_member_at.sum(dim=3) / member_count.clamp_min(1.0).sqrt()348 )349 candidate_update = self.candidate_fusion(350 torch.cat(351 [352 task_h,353 at_messages,354 coalition_member_at,355 coalition_aa,356 member_count / float(max(1, self.num_agents)),357 ],358 dim=-1,359 )360 )361 # 保留原始 task_h，让新增关系消息以残差形式修改候选任务表示。362 candidate_context = self.candidate_norm(task_h + candidate_update)363364 diagnostics = {}365 if return_diagnostics:366 # eligible、nonempty、multi_member 分别记录可选配送任务、已有成员的367 # 配送任务，以及至少有两名成员的配送任务。368 eligible = delivery_mask & available_mask369 nonempty = eligible & (member_count[..., 0] > 0)370 multi_member = eligible & (member_count[..., 0] > 1)371 routed_gates = aa_gates.unsqueeze(2).expand(372 -1, -1, self.num_tasks, -1, -1373 )374 routed_learned = learned_aa_gates.unsqueeze(2).expand_as(routed_gates)375 # route_mask 只选出真正被放入某个候选联盟的 A-A gate。376 route_mask = member_mask.unsqueeze(-1).expand_as(routed_gates)377 diagnostics = {378 "aa_relation_gate": aa_gates,379 "at_relation_gate": at_gates,380 "candidate_aa_context_norm": coalition_aa.norm(dim=-1),381 "candidate_member_at_context_norm": coalition_member_at.norm(dim=-1),382 "candidate_aa_member_count": member_count,383 "candidate_aa_eligible_ratio": eligible.to(task_h.dtype),384 "candidate_aa_nonempty_ratio": nonempty.to(task_h.dtype),385 "candidate_aa_multi_member_ratio": multi_member.to(task_h.dtype),386 "candidate_aa_routed_gate": self._masked_mean(387 routed_gates, route_mask388 ),389 "candidate_aa_learned_uniform_gate_delta": self._masked_mean(390 (routed_learned - 0.5).abs(), route_mask391 ),392 "global_aa_context_norm": global_aa.norm(dim=-1),393 "global_at_context_norm": global_at.norm(dim=-1),394 "member_at_relation_gate": member_at_gates,395 }396 return agent_context, candidate_context, at_features, diagnostics397398399class HATTV33AAActor(HATTV20Actor):400 """使用 HATTV33AABackbone 的 MAPPO actor。"""401402 def __init__(403 self,404 args,405 ego_dim,406 other_dim,407 task_dim,408 num_agents,409 num_tasks,410 device=torch.device("cpu"),411 ):412 super().__init__(413 args,414 ego_dim,415 other_dim,416 task_dim,417 num_agents,418 num_tasks,419 device,420 )421 hidden_dim = getattr(args, "hatt_hidden_dim", args.gat_out_dim)422 # HATTV20Actor 提供动作分布和 RNN 外壳，这里只替换关系计算 backbone。423 self.backbone = HATTV33AABackbone(424 ego_dim=ego_dim,425 other_dim=other_dim,426 task_dim=task_dim,427 hidden_dim=hidden_dim,428 num_heads=getattr(args, "hatt_num_heads", 4),429 cap_dim=args.cap_dim,430 num_agents=num_agents,431 num_tasks=num_tasks,432 dropout=getattr(args, "hatt_dropout", 0.0),433 shared_task_obs=getattr(args, "hatt_shared_task_obs", True),434 candidate_aa_mode=getattr(args, "hatt_candidate_aa_mode", "learned"),435 )436 self.to(device)