"""HATT v3.3AA 用同一套计算步骤处理智能体边和任务边。 智能体边和任务边各自保存参数。现有成员的智能体边输出会按联盟编号放入对应的 候选任务，再和当前智能体到该任务的边输出一起生成动作表示。 ``hatt_candidate_aa_mode`` 只控制智能体边的 gate：``learned`` 使用网络输出， ``uniform`` 固定为 0.5，``zero`` 将这条分支清零。 """ import torch import torch.nn as nn from algorithms.algorithm.HATTV20Network import ( HATTV20Actor, FeatureEncoder, encode_repeated_tasks, ) class TypedRelationEdge(nn.Module): """输入一组有方向的边，为每条有效边生成 gate 和消息。""" def __init__(self, hidden_dim, edge_dim, num_heads, dropout=0.0): super().__init__() # 每个 head 使用相同宽度，因此 hidden_dim 必须能被 head 数整除。 if hidden_dim % num_heads != 0: raise ValueError( f"hidden_dim={hidden_dim} must be divisible by num_heads={num_heads}." ) self.num_heads = num_heads self.head_dim = hidden_dim // num_heads # 点积除以维度平方根，避免 hidden_dim 增大后数值过大。 self.scale = self.head_dim ** -0.5 # source 生成 query，target 分别生成 key 和 value。 self.query = nn.Linear(hidden_dim, hidden_dim) self.key = nn.Linear(hidden_dim, hidden_dim) self.value = nn.Linear(hidden_dim, hidden_dim) # 相对位置和联盟关系会转换成每个 head 的额外分数。 self.edge_bias = nn.Sequential( nn.Linear(edge_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, num_heads), ) self.output = nn.Linear(hidden_dim, hidden_dim) nn.init.constant_(self.output.bias, 0.0) self.dropout = nn.Dropout(dropout) def forward(self, source, target, edge_features, valid_mask, gate_mode="learned"): # source 是边的起点，target 可以包含多个终点。每个 head 分别计算 gate， # 随后用 gate 缩放终点的 value，并合并成一条消息。 if gate_mode not in ("learned", "uniform", "zero"): raise ValueError( "gate_mode must be learned, uniform, or zero; " f"got {gate_mode!r}." ) prefix = target.shape[:-2] num_targets = target.shape[-2] hidden_dim = target.shape[-1] # 把最后一维拆成多个 head。prefix 中保留 batch、观察者等前置维度。 query = self.query(source).view(*prefix, self.num_heads, self.head_dim) key = self.key(target).view( *prefix, num_targets, self.num_heads, self.head_dim ) value = self.value(target).view( *prefix, num_targets, self.num_heads, self.head_dim ) # 每条边的分数由 query-key 点积和边特征共同决定。 logits = ( torch.einsum("...he,...nhe->...nh", query, key) * self.scale + self.edge_bias(edge_features) ) learned_gates = torch.sigmoid(logits) valid = valid_mask.bool().unsqueeze(-1) # padding、死亡智能体和不可选任务对应的 gate 固定为零。 learned_gates = torch.where( valid, learned_gates, torch.zeros_like(learned_gates) ) if gate_mode == "learned": gates = learned_gates elif gate_mode == "uniform": # uniform 用固定权重替代网络输出，供对照实验使用。 gates = valid.expand_as(learned_gates).to(learned_gates.dtype) * 0.5 else: gates = torch.zeros_like(learned_gates) # 每个 head 用自己的 gate 缩放 value，再把所有 head 拼回 hidden_dim。 gates = self.dropout(gates) message = (gates.unsqueeze(-1) * value).reshape( *prefix, num_targets, hidden_dim ) message = self.output(message) if gate_mode == "zero": # output 层可能含有 bias，这里再次清零，保证整条分支没有输出。 message = torch.zeros_like(message) # learned_gates 始终保留网络原始结果，uniform/zero 模式下也可用于诊断。 return message, gates, learned_gates class HATTV33AABackbone(nn.Module): """汇总智能体边和任务边的消息，为每个候选任务生成表示。""" 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", ): 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 # 三类实体先各自编码到相同的 hidden_dim，后续关系层才能直接组合。 self.self_encoder = FeatureEncoder(ego_dim, hidden_dim) self.other_encoder = FeatureEncoder(other_dim, hidden_dim) self.task_encoder = FeatureEncoder(task_dim, hidden_dim) # 两类边使用相同的计算步骤，参数分别保存在 aa_relation 和 at_relation 中。 self.aa_relation = TypedRelationEdge( hidden_dim, edge_dim=3, num_heads=num_heads, dropout=dropout ) self.at_relation = TypedRelationEdge( hidden_dim, edge_dim=3, num_heads=num_heads, dropout=dropout ) # 当前智能体表示会和全局 A-A、A-T 消息拼接，生成 agent_context。 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 槽。""" # 任务类型除以 3 存入最后一项，因此 Deliver 对应 1/3。 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): """生成相对位置、同配送联盟和当前任务归属特征。""" 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) # A-A 边包含对方相对位置，以及双方是否属于同一个配送联盟。 aa_features = torch.cat( [ other_position - ego_position.unsqueeze(2), same_delivery_coalition.unsqueeze(-1), ], dim=-1, ) task_position = tasks[..., 0:2] task_ids = torch.arange(self.num_tasks, device=tasks.device) # assigned 表示候选任务是否就是当前智能体正在执行的任务。 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"] # 联盟编号用于后面的任务槽查找，不作为普通连续值送入实体 encoder。 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、other_h、task_h 的最后一维都等于 hidden_dim。 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 # 全零行是 padding；living 为 0 的智能体也不能发送消息。 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", ) # 为每个其他智能体取出它当前加入的任务，只计算这一条智能体到任务的边。 # 这里复用 at_relation，无需展开“所有成员 × 所有任务”的完整组合。 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) # 最后一项固定为 1，表示这条边连接的是成员当前已经加入的任务。 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 ) # 只让存活成员向自己当前加入的 Deliver 任务发送消息。 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) # 残差保留 self_h，再用 LayerNorm 稳定合并后的数值范围。 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 的维度是“观察者 × 候选任务 × 其他智能体”。 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, ) ) # 保留原始 task_h，让新增关系消息以残差形式修改候选任务表示。 candidate_context = self.candidate_norm(task_h + candidate_update) diagnostics = {} if return_diagnostics: # eligible、nonempty、multi_member 分别记录可选配送任务、已有成员的 # 配送任务，以及至少有两名成员的配送任务。 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 只选出真正被放入某个候选联盟的 A-A gate。 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 HATTV33AAActor(HATTV20Actor): """使用 HATTV33AABackbone 的 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) # HATTV20Actor 提供动作分布和 RNN 外壳，这里只替换关系计算 backbone。 self.backbone = HATTV33AABackbone( 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"), ) self.to(device)