"""HATT v3.4 在 v3.3AA 基础上，把成员所服务任务的条件上下文送入 A-A 边 gate。 v3.3AA 的 A-A 边只看到 ``[相对位置, 是否同联盟]``，因此 gate 是"任务盲"的：它无法 判断 k 所服务的联盟还剩多少需求、已有多少承诺运力。v3.4 在每条 (i, k) 边追加 k 所 服务任务的 4 维上下文：剩余需求、优先级、已承诺载荷比、成员数。边数量不变，只有 边特征从 3 维扩到 7 维。 ``hatt_aa_task_context`` 控制是否启用这 4 维： - ``true`` → 边特征 7 维（本版本主方法）； - ``false`` → 边特征 3 维，计算图与 v3.3AA 完全一致（作为消融对照）。 """ import torch import torch.nn as nn from algorithms.algorithm.HATTV20Network import ( HATTV20Actor, FeatureEncoder, encode_repeated_tasks, ) from algorithms.algorithm.HATTV33AANetwork import TypedRelationEdge class HATTV34Backbone(nn.Module): """在 v3.3AA 的 A-A 边上叠加任务条件上下文。""" AA_MODES = ("learned", "uniform", "zero") def __init__( self, ego_dim, other_dim, task_dim, hidden_dim, num_heads, cap_dim, num_agents, num_tasks, dropout=0.0, shared_task_obs=True, candidate_aa_mode="learned", task_context=True, coalition_task_feature_mode="compact", ): super().__init__() if candidate_aa_mode not in self.AA_MODES: raise ValueError( f"hatt_candidate_aa_mode must be one of {self.AA_MODES}, " f"got {candidate_aa_mode!r}." ) self.cap_dim = cap_dim self.num_agents = num_agents self.num_tasks = num_tasks self.shared_task_obs = shared_task_obs self.candidate_aa_mode = candidate_aa_mode self.task_context = bool(task_context) # 任务特征里 demand / priority 的列下标随 schema 变化。 # compact: [pos(2), demand, deadline, priority, type] # full: [pos(2), req_vec(cap_dim), demand, deadline, ..., type] if coalition_task_feature_mode == "full": self._demand_idx = 2 + cap_dim self._priority_idx = 2 # full 模式下 priority 即 req_vec[0] / scale else: self._demand_idx = 2 self._priority_idx = 4 self.self_encoder = FeatureEncoder(ego_dim, hidden_dim) self.other_encoder = FeatureEncoder(other_dim, hidden_dim) self.task_encoder = FeatureEncoder(task_dim, hidden_dim) # A-A 边特征维度：基础 3 维，task_context 开启时再加 4 维任务上下文。 aa_edge_dim = 3 + (4 if self.task_context else 0) self.aa_relation = TypedRelationEdge( hidden_dim, edge_dim=aa_edge_dim, num_heads=num_heads, dropout=dropout ) self.at_relation = TypedRelationEdge( hidden_dim, edge_dim=3, num_heads=num_heads, dropout=dropout ) self.agent_fusion = nn.Sequential( nn.Linear(hidden_dim * 3, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), ) self.agent_norm = nn.LayerNorm(hidden_dim) self.candidate_fusion = nn.Sequential( nn.Linear(hidden_dim * 4 + 1, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), ) self.candidate_norm = nn.LayerNorm(hidden_dim) def _coalition_indices(self, normalized_ids): """把观测中的归一化联盟编号恢复成任务槽下标。""" return (normalized_ids * max(1, self.num_tasks - 1)).round().long().clamp( 0, self.num_tasks - 1 ) @staticmethod def _delivery_mask(tasks): """标出任务列表中的 Deliver 槽。""" return torch.isclose( tasks[..., -1], tasks.new_tensor(1.0 / 3.0), atol=1e-5, rtol=0.0, ) @staticmethod def _masked_mean(value, mask): """只计算 mask 选中元素的均值；没有元素时返回零。""" selected = value.masked_select(mask) return selected.mean() if selected.numel() else value.new_zeros(()) def _relation_features(self, ego, others, tasks): """生成相对位置、同配送联盟，以及（可选）k 所服务任务的条件上下文。""" ego_position = ego[..., 0:2] ego_coalition = ego[..., 2 + self.cap_dim] other_position = others[..., 0:2] other_coalition = others[..., self.cap_dim + 4] ego_indices = self._coalition_indices(ego_coalition) other_indices = self._coalition_indices(other_coalition) # 先查出当前智能体所处任务的类型，控制动作不产生"同配送联盟"关系。 current_task_type = tasks[..., -1].gather(2, ego_indices.unsqueeze(-1)) current_is_delivery = torch.isclose( current_task_type, tasks.new_tensor(1.0 / 3.0), atol=1e-5, rtol=0.0, ) same_delivery_coalition = ( (other_indices == ego_indices.unsqueeze(2)) & current_is_delivery ).to(ego.dtype) base_aa_features = torch.cat( [ other_position - ego_position.unsqueeze(2), same_delivery_coalition.unsqueeze(-1), ], dim=-1, ) if not self.task_context: aa_features = base_aa_features else: # k 所服务任务的条件上下文：剩余需求、优先级、已承诺载荷比、成员数。 serving_is_delivery = self._delivery_mask(tasks).gather(2, other_indices) serving_is_delivery = serving_is_delivery.to(ego.dtype) serving_demand = tasks[..., self._demand_idx].gather(2, other_indices) serving_priority = tasks[..., self._priority_idx].gather(2, other_indices) # 同联盟聚合：k 与 m 处于同一任务槽，且 m 存活。 living = (others[..., self.cap_dim + 6] > 0.5).to(ego.dtype) same_coalition_mask = ( other_indices.unsqueeze(-1) == other_indices.unsqueeze(-2) ).to(ego.dtype) member_mask = same_coalition_mask * living.unsqueeze(-2) other_load = others[..., self.cap_dim + 2] serving_committed_load = ( (member_mask * other_load.unsqueeze(-2)).sum(dim=-1) / max(1, self.num_agents) ) serving_member_count = member_mask.sum(dim=-1) / max(1, self.num_agents) task_context = torch.cat( [ serving_demand.unsqueeze(-1), serving_priority.unsqueeze(-1), serving_committed_load.unsqueeze(-1), serving_member_count.unsqueeze(-1), ], dim=-1, ) # 非配送任务上下文置零，与 member_task_is_delivery 掩码保持一致。 task_context = task_context * serving_is_delivery.unsqueeze(-1) aa_features = torch.cat([base_aa_features, task_context], dim=-1) task_position = tasks[..., 0:2] task_ids = torch.arange(self.num_tasks, device=tasks.device) assigned = (ego_indices.unsqueeze(-1) == task_ids.view(1, 1, -1)).to( ego.dtype ) at_features = torch.cat( [ task_position - ego_position.unsqueeze(2), assigned.unsqueeze(-1), ], dim=-1, ) return aa_features, at_features, other_indices @staticmethod def _pool_messages(messages, valid_mask): """平均有效边的消息，用于生成当前智能体的整体关系信息。""" valid = valid_mask.to(messages.dtype).unsqueeze(-1) count = valid.sum(dim=-2).clamp_min(1.0) return (messages * valid).sum(dim=-2) / count def forward(self, obs, available_actions, return_diagnostics=True): """生成当前智能体表示、各候选任务表示、A-T 边特征和诊断数据。""" ego = obs["ego"] others = obs["others"] tasks = obs["tasks"] ego_input = ego.clone() others_input = others.clone() ego_input[..., 2 + self.cap_dim] = 0.0 others_input[..., self.cap_dim + 4] = 0.0 self_h = self.self_encoder(ego_input) other_h = self.other_encoder(others_input) task_h = encode_repeated_tasks( self.task_encoder, tasks, self.shared_task_obs ) aa_features, at_features, other_indices = self._relation_features( ego, others, tasks ) living_index = self.cap_dim + 6 other_mask = ( (others.abs().sum(dim=-1) > 0) & (others[..., living_index] > 0.5) ) available_mask = ( available_actions > 0 if available_actions is not None else torch.ones( *tasks.shape[:-1], dtype=torch.bool, device=tasks.device ) ) aa_messages, aa_gates, learned_aa_gates = self.aa_relation( self_h, other_h, aa_features, other_mask, gate_mode=self.candidate_aa_mode, ) at_messages, at_gates, _ = self.at_relation( self_h, task_h, at_features, available_mask, gate_mode="learned", ) member_task_h = task_h.unsqueeze(2).expand( -1, -1, self.num_agents, -1, -1 ).gather( 3, other_indices[..., None, None].expand( -1, -1, -1, 1, task_h.shape[-1] ), ) member_task_position = tasks[..., 0:2].unsqueeze(2).expand( -1, -1, self.num_agents, -1, -1 ).gather( 3, other_indices[..., None, None].expand(-1, -1, -1, 1, 2), ).squeeze(3) member_at_features = torch.cat( [ member_task_position - others[..., 0:2], torch.ones_like(other_indices, dtype=tasks.dtype).unsqueeze(-1), ], dim=-1, ).unsqueeze(3) member_task_is_delivery = self._delivery_mask(tasks).gather( 2, other_indices ) member_at_valid = (other_mask & member_task_is_delivery).unsqueeze(-1) member_at_messages, member_at_gates, _ = self.at_relation( other_h, member_task_h, member_at_features, member_at_valid, gate_mode="learned", ) member_at_messages = member_at_messages.squeeze(3) global_aa = self._pool_messages(aa_messages, other_mask) global_at = self._pool_messages(at_messages, available_mask) agent_update = self.agent_fusion( torch.cat([self_h, global_aa, global_at], dim=-1) ) agent_context = self.agent_norm(self_h + agent_update) task_ids = torch.arange(self.num_tasks, device=tasks.device).view(1, 1, -1) delivery_mask = self._delivery_mask(tasks) member_mask = ( (other_indices.unsqueeze(2) == task_ids.unsqueeze(-1)) & other_mask.unsqueeze(2) & delivery_mask.unsqueeze(-1) ) routed_aa = aa_messages.unsqueeze(2) * member_mask.unsqueeze(-1).to( aa_messages.dtype ) routed_member_at = member_at_messages.unsqueeze(2) * member_mask.unsqueeze( -1 ).to(member_at_messages.dtype) member_count = member_mask.sum(dim=3, keepdim=True).to(task_h.dtype) coalition_aa = routed_aa.sum(dim=3) / member_count.clamp_min(1.0).sqrt() coalition_member_at = ( routed_member_at.sum(dim=3) / member_count.clamp_min(1.0).sqrt() ) candidate_update = self.candidate_fusion( torch.cat( [ task_h, at_messages, coalition_member_at, coalition_aa, member_count / float(max(1, self.num_agents)), ], dim=-1, ) ) candidate_context = self.candidate_norm(task_h + candidate_update) diagnostics = {} if return_diagnostics: eligible = delivery_mask & available_mask nonempty = eligible & (member_count[..., 0] > 0) multi_member = eligible & (member_count[..., 0] > 1) routed_gates = aa_gates.unsqueeze(2).expand( -1, -1, self.num_tasks, -1, -1 ) routed_learned = learned_aa_gates.unsqueeze(2).expand_as(routed_gates) route_mask = member_mask.unsqueeze(-1).expand_as(routed_gates) diagnostics = { "aa_relation_gate": aa_gates, "at_relation_gate": at_gates, "candidate_aa_context_norm": coalition_aa.norm(dim=-1), "candidate_member_at_context_norm": coalition_member_at.norm(dim=-1), "candidate_aa_member_count": member_count, "candidate_aa_eligible_ratio": eligible.to(task_h.dtype), "candidate_aa_nonempty_ratio": nonempty.to(task_h.dtype), "candidate_aa_multi_member_ratio": multi_member.to(task_h.dtype), "candidate_aa_routed_gate": self._masked_mean( routed_gates, route_mask ), "candidate_aa_learned_uniform_gate_delta": self._masked_mean( (routed_learned - 0.5).abs(), route_mask ), "global_aa_context_norm": global_aa.norm(dim=-1), "global_at_context_norm": global_at.norm(dim=-1), "member_at_relation_gate": member_at_gates, } return agent_context, candidate_context, at_features, diagnostics class HATTV34Actor(HATTV20Actor): """使用 HATTV34Backbone 的 MAPPO actor。""" def __init__( self, args, ego_dim, other_dim, task_dim, num_agents, num_tasks, device=torch.device("cpu"), ): super().__init__( args, ego_dim, other_dim, task_dim, num_agents, num_tasks, device, ) hidden_dim = getattr(args, "hatt_hidden_dim", args.gat_out_dim) self.backbone = HATTV34Backbone( ego_dim=ego_dim, other_dim=other_dim, task_dim=task_dim, hidden_dim=hidden_dim, num_heads=getattr(args, "hatt_num_heads", 4), cap_dim=args.cap_dim, num_agents=num_agents, num_tasks=num_tasks, dropout=getattr(args, "hatt_dropout", 0.0), shared_task_obs=getattr(args, "hatt_shared_task_obs", True), candidate_aa_mode=getattr(args, "hatt_candidate_aa_mode", "learned"), task_context=getattr(args, "hatt_aa_task_context", True), coalition_task_feature_mode=getattr( args, "coalition_task_feature_mode", "compact" ), ) self.to(device)