HATTV34Network.py
v3.4 服务任务上下文、AA/AT 与候选路由；方法章主要实现
原文件 · 386 行 · 15273 字节
algorithms/algorithm/HATTV34Network.py
SHA-256: 4f6e181d1cccb1044c6a16787d567abfc60caa8a47f945fde3062e87a97f4f88
1"""HATT v3.4 在 v3.3AA 基础上，把成员所服务任务的条件上下文送入 A-A 边 gate。23v3.3AA 的 A-A 边只看到 ``[相对位置, 是否同联盟]``，因此 gate 是"任务盲"的：它无法4判断 k 所服务的联盟还剩多少需求、已有多少承诺运力。v3.4 在每条 (i, k) 边追加 k 所5服务任务的 4 维上下文：剩余需求、优先级、已承诺载荷比、成员数。边数量不变，只有6边特征从 3 维扩到 7 维。78``hatt_aa_task_context`` 控制是否启用这 4 维：9- ``true`` → 边特征 7 维（本版本主方法）；10- ``false`` → 边特征 3 维，计算图与 v3.3AA 完全一致（作为消融对照）。11"""1213import torch14import torch.nn as nn1516from algorithms.algorithm.HATTV20Network import (17 HATTV20Actor,18 FeatureEncoder,19 encode_repeated_tasks,20)21from algorithms.algorithm.HATTV33AANetwork import TypedRelationEdge222324class HATTV34Backbone(nn.Module):25 """在 v3.3AA 的 A-A 边上叠加任务条件上下文。"""2627 AA_MODES = ("learned", "uniform", "zero")2829 def __init__(30 self,31 ego_dim,32 other_dim,33 task_dim,34 hidden_dim,35 num_heads,36 cap_dim,37 num_agents,38 num_tasks,39 dropout=0.0,40 shared_task_obs=True,41 candidate_aa_mode="learned",42 task_context=True,43 coalition_task_feature_mode="compact",44 ):45 super().__init__()46 if candidate_aa_mode not in self.AA_MODES:47 raise ValueError(48 f"hatt_candidate_aa_mode must be one of {self.AA_MODES}, "49 f"got {candidate_aa_mode!r}."50 )51 self.cap_dim = cap_dim52 self.num_agents = num_agents53 self.num_tasks = num_tasks54 self.shared_task_obs = shared_task_obs55 self.candidate_aa_mode = candidate_aa_mode56 self.task_context = bool(task_context)5758 # 任务特征里 demand / priority 的列下标随 schema 变化。59 # compact: [pos(2), demand, deadline, priority, type]60 # full: [pos(2), req_vec(cap_dim), demand, deadline, ..., type]61 if coalition_task_feature_mode == "full":62 self._demand_idx = 2 + cap_dim63 self._priority_idx = 2 # full 模式下 priority 即 req_vec[0] / scale64 else:65 self._demand_idx = 266 self._priority_idx = 46768 self.self_encoder = FeatureEncoder(ego_dim, hidden_dim)69 self.other_encoder = FeatureEncoder(other_dim, hidden_dim)70 self.task_encoder = FeatureEncoder(task_dim, hidden_dim)71 # A-A 边特征维度：基础 3 维，task_context 开启时再加 4 维任务上下文。72 aa_edge_dim = 3 + (4 if self.task_context else 0)73 self.aa_relation = TypedRelationEdge(74 hidden_dim, edge_dim=aa_edge_dim, num_heads=num_heads, dropout=dropout75 )76 self.at_relation = TypedRelationEdge(77 hidden_dim, edge_dim=3, num_heads=num_heads, dropout=dropout78 )79 self.agent_fusion = nn.Sequential(80 nn.Linear(hidden_dim * 3, hidden_dim),81 nn.ReLU(),82 nn.Linear(hidden_dim, hidden_dim),83 )84 self.agent_norm = nn.LayerNorm(hidden_dim)85 self.candidate_fusion = nn.Sequential(86 nn.Linear(hidden_dim * 4 + 1, hidden_dim),87 nn.ReLU(),88 nn.Linear(hidden_dim, hidden_dim),89 )90 self.candidate_norm = nn.LayerNorm(hidden_dim)9192 def _coalition_indices(self, normalized_ids):93 """把观测中的归一化联盟编号恢复成任务槽下标。"""94 return (normalized_ids * max(1, self.num_tasks - 1)).round().long().clamp(95 0, self.num_tasks - 196 )9798 @staticmethod99 def _delivery_mask(tasks):100 """标出任务列表中的 Deliver 槽。"""101 return torch.isclose(102 tasks[..., -1],103 tasks.new_tensor(1.0 / 3.0),104 atol=1e-5,105 rtol=0.0,106 )107108 @staticmethod109 def _masked_mean(value, mask):110 """只计算 mask 选中元素的均值；没有元素时返回零。"""111 selected = value.masked_select(mask)112 return selected.mean() if selected.numel() else value.new_zeros(())113114 def _relation_features(self, ego, others, tasks):115 """生成相对位置、同配送联盟，以及（可选）k 所服务任务的条件上下文。"""116 ego_position = ego[..., 0:2]117 ego_coalition = ego[..., 2 + self.cap_dim]118 other_position = others[..., 0:2]119 other_coalition = others[..., self.cap_dim + 4]120 ego_indices = self._coalition_indices(ego_coalition)121 other_indices = self._coalition_indices(other_coalition)122123 # 先查出当前智能体所处任务的类型，控制动作不产生"同配送联盟"关系。124 current_task_type = tasks[..., -1].gather(2, ego_indices.unsqueeze(-1))125 current_is_delivery = torch.isclose(126 current_task_type,127 tasks.new_tensor(1.0 / 3.0),128 atol=1e-5,129 rtol=0.0,130 )131 same_delivery_coalition = (132 (other_indices == ego_indices.unsqueeze(2)) & current_is_delivery133 ).to(ego.dtype)134 base_aa_features = torch.cat(135 [136 other_position - ego_position.unsqueeze(2),137 same_delivery_coalition.unsqueeze(-1),138 ],139 dim=-1,140 )141142 if not self.task_context:143 aa_features = base_aa_features144 else:145 # k 所服务任务的条件上下文：剩余需求、优先级、已承诺载荷比、成员数。146 serving_is_delivery = self._delivery_mask(tasks).gather(2, other_indices)147 serving_is_delivery = serving_is_delivery.to(ego.dtype)148 serving_demand = tasks[..., self._demand_idx].gather(2, other_indices)149 serving_priority = tasks[..., self._priority_idx].gather(2, other_indices)150 # 同联盟聚合：k 与 m 处于同一任务槽，且 m 存活。151 living = (others[..., self.cap_dim + 6] > 0.5).to(ego.dtype)152 same_coalition_mask = (153 other_indices.unsqueeze(-1) == other_indices.unsqueeze(-2)154 ).to(ego.dtype)155 member_mask = same_coalition_mask * living.unsqueeze(-2)156 other_load = others[..., self.cap_dim + 2]157 serving_committed_load = (158 (member_mask * other_load.unsqueeze(-2)).sum(dim=-1)159 / max(1, self.num_agents)160 )161 serving_member_count = member_mask.sum(dim=-1) / max(1, self.num_agents)162 task_context = torch.cat(163 [164 serving_demand.unsqueeze(-1),165 serving_priority.unsqueeze(-1),166 serving_committed_load.unsqueeze(-1),167 serving_member_count.unsqueeze(-1),168 ],169 dim=-1,170 )171 # 非配送任务上下文置零，与 member_task_is_delivery 掩码保持一致。172 task_context = task_context * serving_is_delivery.unsqueeze(-1)173 aa_features = torch.cat([base_aa_features, task_context], dim=-1)174175 task_position = tasks[..., 0:2]176 task_ids = torch.arange(self.num_tasks, device=tasks.device)177 assigned = (ego_indices.unsqueeze(-1) == task_ids.view(1, 1, -1)).to(178 ego.dtype179 )180 at_features = torch.cat(181 [182 task_position - ego_position.unsqueeze(2),183 assigned.unsqueeze(-1),184 ],185 dim=-1,186 )187 return aa_features, at_features, other_indices188189 @staticmethod190 def _pool_messages(messages, valid_mask):191 """平均有效边的消息，用于生成当前智能体的整体关系信息。"""192 valid = valid_mask.to(messages.dtype).unsqueeze(-1)193 count = valid.sum(dim=-2).clamp_min(1.0)194 return (messages * valid).sum(dim=-2) / count195196 def forward(self, obs, available_actions, return_diagnostics=True):197 """生成当前智能体表示、各候选任务表示、A-T 边特征和诊断数据。"""198 ego = obs["ego"]199 others = obs["others"]200 tasks = obs["tasks"]201 ego_input = ego.clone()202 others_input = others.clone()203 ego_input[..., 2 + self.cap_dim] = 0.0204 others_input[..., self.cap_dim + 4] = 0.0205 self_h = self.self_encoder(ego_input)206 other_h = self.other_encoder(others_input)207 task_h = encode_repeated_tasks(208 self.task_encoder, tasks, self.shared_task_obs209 )210 aa_features, at_features, other_indices = self._relation_features(211 ego, others, tasks212 )213214 living_index = self.cap_dim + 6215 other_mask = (216 (others.abs().sum(dim=-1) > 0)217 & (others[..., living_index] > 0.5)218 )219 available_mask = (220 available_actions > 0221 if available_actions is not None222 else torch.ones(223 *tasks.shape[:-1], dtype=torch.bool, device=tasks.device224 )225 )226 aa_messages, aa_gates, learned_aa_gates = self.aa_relation(227 self_h,228 other_h,229 aa_features,230 other_mask,231 gate_mode=self.candidate_aa_mode,232 )233 at_messages, at_gates, _ = self.at_relation(234 self_h,235 task_h,236 at_features,237 available_mask,238 gate_mode="learned",239 )240241 member_task_h = task_h.unsqueeze(2).expand(242 -1, -1, self.num_agents, -1, -1243 ).gather(244 3,245 other_indices[..., None, None].expand(246 -1, -1, -1, 1, task_h.shape[-1]247 ),248 )249 member_task_position = tasks[..., 0:2].unsqueeze(2).expand(250 -1, -1, self.num_agents, -1, -1251 ).gather(252 3,253 other_indices[..., None, None].expand(-1, -1, -1, 1, 2),254 ).squeeze(3)255 member_at_features = torch.cat(256 [257 member_task_position - others[..., 0:2],258 torch.ones_like(other_indices, dtype=tasks.dtype).unsqueeze(-1),259 ],260 dim=-1,261 ).unsqueeze(3)262 member_task_is_delivery = self._delivery_mask(tasks).gather(263 2, other_indices264 )265 member_at_valid = (other_mask & member_task_is_delivery).unsqueeze(-1)266 member_at_messages, member_at_gates, _ = self.at_relation(267 other_h,268 member_task_h,269 member_at_features,270 member_at_valid,271 gate_mode="learned",272 )273 member_at_messages = member_at_messages.squeeze(3)274275 global_aa = self._pool_messages(aa_messages, other_mask)276 global_at = self._pool_messages(at_messages, available_mask)277 agent_update = self.agent_fusion(278 torch.cat([self_h, global_aa, global_at], dim=-1)279 )280 agent_context = self.agent_norm(self_h + agent_update)281282 task_ids = torch.arange(self.num_tasks, device=tasks.device).view(1, 1, -1)283 delivery_mask = self._delivery_mask(tasks)284 member_mask = (285 (other_indices.unsqueeze(2) == task_ids.unsqueeze(-1))286 & other_mask.unsqueeze(2)287 & delivery_mask.unsqueeze(-1)288 )289 routed_aa = aa_messages.unsqueeze(2) * member_mask.unsqueeze(-1).to(290 aa_messages.dtype291 )292 routed_member_at = member_at_messages.unsqueeze(2) * member_mask.unsqueeze(293 -1294 ).to(member_at_messages.dtype)295 member_count = member_mask.sum(dim=3, keepdim=True).to(task_h.dtype)296 coalition_aa = routed_aa.sum(dim=3) / member_count.clamp_min(1.0).sqrt()297 coalition_member_at = (298 routed_member_at.sum(dim=3) / member_count.clamp_min(1.0).sqrt()299 )300 candidate_update = self.candidate_fusion(301 torch.cat(302 [303 task_h,304 at_messages,305 coalition_member_at,306 coalition_aa,307 member_count / float(max(1, self.num_agents)),308 ],309 dim=-1,310 )311 )312 candidate_context = self.candidate_norm(task_h + candidate_update)313314 diagnostics = {}315 if return_diagnostics:316 eligible = delivery_mask & available_mask317 nonempty = eligible & (member_count[..., 0] > 0)318 multi_member = eligible & (member_count[..., 0] > 1)319 routed_gates = aa_gates.unsqueeze(2).expand(320 -1, -1, self.num_tasks, -1, -1321 )322 routed_learned = learned_aa_gates.unsqueeze(2).expand_as(routed_gates)323 route_mask = member_mask.unsqueeze(-1).expand_as(routed_gates)324 diagnostics = {325 "aa_relation_gate": aa_gates,326 "at_relation_gate": at_gates,327 "candidate_aa_context_norm": coalition_aa.norm(dim=-1),328 "candidate_member_at_context_norm": coalition_member_at.norm(dim=-1),329 "candidate_aa_member_count": member_count,330 "candidate_aa_eligible_ratio": eligible.to(task_h.dtype),331 "candidate_aa_nonempty_ratio": nonempty.to(task_h.dtype),332 "candidate_aa_multi_member_ratio": multi_member.to(task_h.dtype),333 "candidate_aa_routed_gate": self._masked_mean(334 routed_gates, route_mask335 ),336 "candidate_aa_learned_uniform_gate_delta": self._masked_mean(337 (routed_learned - 0.5).abs(), route_mask338 ),339 "global_aa_context_norm": global_aa.norm(dim=-1),340 "global_at_context_norm": global_at.norm(dim=-1),341 "member_at_relation_gate": member_at_gates,342 }343 return agent_context, candidate_context, at_features, diagnostics344345346class HATTV34Actor(HATTV20Actor):347 """使用 HATTV34Backbone 的 MAPPO actor。"""348349 def __init__(350 self,351 args,352 ego_dim,353 other_dim,354 task_dim,355 num_agents,356 num_tasks,357 device=torch.device("cpu"),358 ):359 super().__init__(360 args,361 ego_dim,362 other_dim,363 task_dim,364 num_agents,365 num_tasks,366 device,367 )368 hidden_dim = getattr(args, "hatt_hidden_dim", args.gat_out_dim)369 self.backbone = HATTV34Backbone(370 ego_dim=ego_dim,371 other_dim=other_dim,372 task_dim=task_dim,373 hidden_dim=hidden_dim,374 num_heads=getattr(args, "hatt_num_heads", 4),375 cap_dim=args.cap_dim,376 num_agents=num_agents,377 num_tasks=num_tasks,378 dropout=getattr(args, "hatt_dropout", 0.0),379 shared_task_obs=getattr(args, "hatt_shared_task_obs", True),380 candidate_aa_mode=getattr(args, "hatt_candidate_aa_mode", "learned"),381 task_context=getattr(args, "hatt_aa_task_context", True),382 coalition_task_feature_mode=getattr(383 args, "coalition_task_feature_mode", "compact"384 ),385 )386 self.to(device)