HATTV20Network.py
实体编码器、CandidateEdgeDecoder、基础 actor 与 schema 投影
原文件 · 860 行 · 38376 字节
algorithms/algorithm/HATTV20Network.py
SHA-256: 7a3a5ce95102adb886b9d66ac2afd019dd5a78f999f51a605b91846aa37d1597
1"""v2.0 是后续 v2/v3 系列的主要旧基线。它的核心思想是：21. 分别编码 focal agent、其他 agent 和任务；32. 用 A-A attention 更新 focal agent 表示；43. 按当前联盟编号把自己和其他成员聚合到每个任务槽；54. 将任务编码与该任务当前联盟表示融合，再对每条 A-T 候选边打分。67注意：本版本编码的是“当前联盟状态”，没有做 v3.0 之后的候选加入前后反事实推理。8"""910import math1112import torch13import torch.nn as nn1415from algorithms.utils.distributions import FixedCategorical16from algorithms.utils.rnn import RNNLayer17from algorithms.utils.util import check181920class FeatureEncoder(nn.Module):21 """两层 MLP + LayerNorm，用于把 ego/other/task 的原始特征映射到同一隐空间。"""2223 def __init__(self, input_dim, hidden_dim):24 super().__init__()25 self.net = nn.Sequential(26 nn.Linear(input_dim, hidden_dim),27 nn.ReLU(),28 nn.Linear(hidden_dim, hidden_dim),29 nn.LayerNorm(hidden_dim),30 )3132 def forward(self, x):33 return self.net(x)343536def initialize_full_actor_from_compact_reference(37 full_actor, compact_actor, cap_dim38):39 """Make a full-schema actor an exact zero-padded extension of compact6.4041 The compact task order is ``x, y, demand, deadline, req0, type``. The42 corresponding full-schema columns are copied into the 12-D task encoder;43 all additional input columns start at zero. Every other actor parameter44 is copied verbatim. With zero-valued extra fields and restored req0 scale,45 the two actors therefore implement the same function at initialization.46 """4748 compact_state = compact_actor.state_dict()49 full_state = full_actor.state_dict()50 first_weight = "backbone.task_encoder.net.0.weight"51 compact_columns = [0, 1, 2, 3, 4, 5]52 full_columns = [0, 1, 2 + cap_dim, 3 + cap_dim, 2, 7 + cap_dim]5354 if compact_state[first_weight].shape[1] != 6:55 raise ValueError("compact reference task encoder must have six inputs")56 expected_full_dim = 2 + cap_dim + 1 + 1 + 3 + 157 if full_state[first_weight].shape[1] != expected_full_dim:58 raise ValueError(59 "full task encoder has unexpected input width: "60 f"{full_state[first_weight].shape[1]} != {expected_full_dim}"61 )6263 for name, target in full_state.items():64 source = compact_state[name]65 if name == first_weight:66 target.zero_()67 for compact_idx, full_idx in zip(compact_columns, full_columns):68 target[:, full_idx].copy_(source[:, compact_idx])69 continue70 if target.shape != source.shape:71 raise ValueError(72 f"cannot preserve actor function for {name}: "73 f"full={tuple(target.shape)}, compact={tuple(source.shape)}"74 )75 target.copy_(source)76 full_actor.load_state_dict(full_state)777879def encode_repeated_tasks(task_encoder, tasks, shared_task_obs):80 """任务观测在所有 agent 间相同时，只编码一次并广播，减少重复计算。"""8182 if shared_task_obs and tasks.ndim == 4 and tasks.shape[1] > 1:83 encoded = task_encoder(tasks[:, 0])84 return encoded.unsqueeze(1).expand(-1, tasks.shape[1], -1, -1)85 return task_encoder(tasks)868788class RelationAttention(nn.Module):89 """带边特征 bias 的稠密多头注意力，用于 A-A 关系聚合。"""9091 def __init__(self, hidden_dim, edge_dim, num_heads, dropout=0.0):92 super().__init__()93 if hidden_dim % num_heads != 0:94 raise ValueError(95 f"hidden_dim={hidden_dim} must be divisible by num_heads={num_heads}."96 )97 self.num_heads = num_heads98 self.head_dim = hidden_dim // num_heads99 self.scale = self.head_dim ** -0.5100 self.query = nn.Linear(hidden_dim, hidden_dim)101 self.key = nn.Linear(hidden_dim, hidden_dim)102 self.value = nn.Linear(hidden_dim, hidden_dim)103 self.edge_bias = nn.Sequential(104 nn.Linear(edge_dim, hidden_dim),105 nn.ReLU(),106 nn.Linear(hidden_dim, num_heads),107 )108 self.output = nn.Linear(hidden_dim, hidden_dim)109 self.dropout = nn.Dropout(dropout)110111 def forward(self, query, neighbors, edge_features, mask):112 # query: [B, A, H]；neighbors: [B, A, N, H]。113 # mask 屏蔽不存在/死亡/不可用的邻居，避免 softmax 分配注意力。114 batch_size, num_agents, num_neighbors, hidden_dim = neighbors.shape115 # 每个 observer agent 作为 query；每个邻居智能体作为 key/value。116 q = self.query(query).view(117 batch_size, num_agents, self.num_heads, self.head_dim118 )119 k = self.key(neighbors).view(120 batch_size, num_agents, num_neighbors, self.num_heads, self.head_dim121 )122 v = self.value(neighbors).view(123 batch_size, num_agents, num_neighbors, self.num_heads, self.head_dim124 )125 scores = torch.einsum("bahd,banhd->banh", q, k) * self.scale126 # 边特征直接加到每个 head 的 logit 上，用于注入相对位置和同联盟关系。127 scores = scores + self.edge_bias(edge_features)128 valid = mask.bool().unsqueeze(-1)129 scores = scores.masked_fill(~valid, torch.finfo(scores.dtype).min)130 weights = torch.softmax(scores, dim=2)131 # 全无效邻居时避免 NaN；再归一化保证有效邻居权重和为 1。132 weights = torch.where(valid, weights, torch.zeros_like(weights))133 weights = weights / weights.sum(dim=2, keepdim=True).clamp_min(1e-8)134 weights = self.dropout(weights)135 context = torch.einsum("banh,banhd->bahd", weights, v)136 context = context.reshape(batch_size, num_agents, hidden_dim)137 return self.output(context), weights138139140class HATTV20Backbone(nn.Module):141 """ 对每个观察者 agent，构建 num_tasks 个任务槽的联盟表示。142 这里的联盟表示包含“当前已经分配到该任务的成员”，并包含 focal agent 自己当前所在联盟。143 """144145 def __init__(146 self,147 ego_dim,148 other_dim,149 task_dim,150 hidden_dim,151 num_heads,152 cap_dim,153 num_agents,154 num_tasks,155 dropout=0.0,156 shared_task_obs=True,157 use_aa_attention=True,158 use_coalition_context=True,159 task_input_mode="native",160 task_requirement_scale=5.0,161 ):162 super().__init__()163 self.cap_dim = cap_dim164 self.num_agents = num_agents165 self.num_tasks = num_tasks166 self.shared_task_obs = shared_task_obs167 # Ablations preserve all parameter tensors and shapes. A disabled168 # branch contributes zeros so parameter count cannot explain a result.169 self.use_aa_attention = bool(use_aa_attention)170 self.use_coalition_context = bool(use_coalition_context)171 self.task_input_mode = str(task_input_mode)172 self.task_requirement_scale = float(task_requirement_scale)173 if self.task_input_mode not in (174 "native",175 "restore_req_scale",176 "compact_projection",177 ):178 raise ValueError(179 "task_input_mode must be native, restore_req_scale, or "180 f"compact_projection; got {self.task_input_mode!r}"181 )182 if self.task_input_mode != "native" and task_dim < 2 + cap_dim + 6:183 raise ValueError(184 f"task_input_mode={self.task_input_mode!r} requires the full "185 f"task schema, got task_dim={task_dim}"186 )187 encoded_task_dim = 6 if self.task_input_mode == "compact_projection" else task_dim188 self.self_encoder = FeatureEncoder(ego_dim, hidden_dim)189 self.other_encoder = FeatureEncoder(other_dim, hidden_dim)190 self.task_encoder = FeatureEncoder(encoded_task_dim, hidden_dim)191 self.aa_attention = RelationAttention(192 hidden_dim, edge_dim=3, num_heads=num_heads, dropout=dropout193 )194 self.agent_fusion = nn.Sequential(195 nn.Linear(hidden_dim * 2, hidden_dim),196 nn.ReLU(),197 nn.Linear(hidden_dim, hidden_dim),198 )199 self.agent_norm = nn.LayerNorm(hidden_dim)200 self.task_fusion = nn.Sequential(201 nn.Linear(hidden_dim * 3 + 1, hidden_dim),202 nn.ReLU(),203 nn.Linear(hidden_dim, hidden_dim),204 )205 self.task_norm = nn.LayerNorm(hidden_dim)206207 def _prepare_task_encoder_input(self, tasks):208 """Apply controlled schema interventions before the v2.0 task MLP.209210 The full DelivCoal schema stores req0 divided by capability_scale. The211 restore control changes only that scale. The compact projection maps212 the full schema back to the historical six semantic fields, allowing a213 width-matched actor initialization against the compact baseline.214 """215216 if self.task_input_mode == "native":217 return tasks218219 requirement = tasks[..., 2:3] * self.task_requirement_scale220 if self.task_input_mode == "restore_req_scale":221 prepared = tasks.clone()222 prepared[..., 2:3] = requirement223 return prepared224225 demand_idx = 2 + self.cap_dim226 deadline_idx = 3 + self.cap_dim227 type_idx = 7 + self.cap_dim228 return torch.cat(229 [230 tasks[..., 0:2],231 tasks[..., demand_idx:demand_idx + 1],232 tasks[..., deadline_idx:deadline_idx + 1],233 requirement,234 tasks[..., type_idx:type_idx + 1],235 ],236 dim=-1,237 )238239 def _coalition_indices(self, normalized_ids):240 # 环境中联盟编号以 [0, 1] 归一化存储，这里恢复为离散 task slot。241 return (normalized_ids * max(1, self.num_tasks - 1)).round().long().clamp(242 0, self.num_tasks - 1243 )244245 def _relation_features(self, ego, others, tasks):246 # A-A 边：相对位置 + 是否同联盟。247 # A-T 边：相对任务位置 + focal agent 当前是否分配给该任务。248 ego_position = ego[..., 0:2]249 ego_coalition = ego[..., 2 + self.cap_dim]250 other_position = others[..., 0:2]251 other_coalition = others[..., self.cap_dim + 4]252 aa_relative = other_position - ego_position.unsqueeze(2)253 same_coalition = (254 self._coalition_indices(other_coalition)255 == self._coalition_indices(ego_coalition).unsqueeze(2)256 ).float()257 aa_features = torch.cat(258 [aa_relative, same_coalition.unsqueeze(-1)], dim=-1259 )260 task_position = tasks[..., 0:2]261 at_relative = task_position - ego_position.unsqueeze(2)262 task_ids = torch.arange(self.num_tasks, device=ego.device)263 assigned = (264 self._coalition_indices(ego_coalition).unsqueeze(-1)265 == task_ids.view(1, 1, -1)266 ).float()267 at_features = torch.cat(268 [at_relative, assigned.unsqueeze(-1)], dim=-1269 )270 return aa_features, at_features271272 def _aggregate_coalitions(273 self, self_h, other_h, ego_coalition, other_coalition, other_mask274 ):275 """把 ego/other 成员按当前任务编号 scatter 到任务槽中形成联盟 sum/mean/count。"""276277 batch_size, observers, _, hidden_dim = other_h.shape278 coalition_sum = other_h.new_zeros(279 batch_size, observers, self.num_tasks, hidden_dim280 )281 coalition_count = other_h.new_zeros(282 batch_size, observers, self.num_tasks, 1283 )284 other_indices = self._coalition_indices(other_coalition)285 # scatter_add_ 需要把每个 other 的任务槽编号扩展到 hidden_dim 维。286 message_indices = other_indices.unsqueeze(-1).expand(287 -1, -1, -1, hidden_dim288 )289 # other_mask 会过滤 padding/self/dead agent，避免进入联盟聚合。290 valid_messages = other_h * other_mask.unsqueeze(-1)291 coalition_sum.scatter_add_(2, message_indices, valid_messages)292 coalition_count.scatter_add_(293 2, other_indices.unsqueeze(-1), other_mask.unsqueeze(-1)294 )295 self_indices = self._coalition_indices(ego_coalition).unsqueeze(-1)296 # v2.0 当前联盟状态包含 focal agent 自己，因此把 self_h 加回当前任务槽。297 coalition_sum.scatter_add_(298 2,299 self_indices.unsqueeze(-1).expand(-1, -1, -1, hidden_dim),300 self_h.unsqueeze(2),301 )302 coalition_count.scatter_add_(303 2,304 self_indices.unsqueeze(-1),305 torch.ones_like(self_indices).unsqueeze(-1).float(),306 )307 # mean 给任务融合提供规模无关的联盟表示；count 保留联盟人数信息。308 coalition_mean = coalition_sum / coalition_count.clamp_min(1.0)309 return coalition_sum, coalition_mean, coalition_count310311 def forward(self, obs, available_actions, return_diagnostics=True):312 del available_actions313 ego = obs["ego"]314 others = obs["others"]315 tasks = obs["tasks"]316 ego_coalition = ego[..., 2 + self.cap_dim]317 other_coalition = others[..., self.cap_dim + 4]318 # 联盟编号只用于构造关系和聚合槽，不作为连续语义特征直接送入 encoder。319 ego_input = ego.clone()320 others_input = others.clone()321 ego_input[..., 2 + self.cap_dim] = 0.0322 others_input[..., self.cap_dim + 4] = 0.0323 self_h = self.self_encoder(ego_input)324 other_h = self.other_encoder(others_input)325 task_encoder_input = self._prepare_task_encoder_input(tasks)326 task_h = encode_repeated_tasks(327 self.task_encoder, task_encoder_input, self.shared_task_obs328 )329 # relation_features 不进 encoder，而是作为 attention/decoder 的边条件。330 aa_features, at_features = self._relation_features(ego, others, tasks)331 living_index = self.cap_dim + 6332 # others 已经把对角线自身置零；这里再过滤空行和死亡 agent。333 other_mask = (334 (others.abs().sum(dim=-1) > 0)335 & (others[..., living_index] > 0.5)336 ).float()337 # A-A attention 提供“其他智能体对当前 agent 决策的上下文影响”。338 if self.use_aa_attention:339 aa_context, aa_weights = self.aa_attention(340 self_h, other_h, aa_features, other_mask.bool()341 )342 else:343 aa_context = torch.zeros_like(self_h)344 aa_weights = other_h.new_zeros(345 *other_h.shape[:-1], self.aa_attention.num_heads346 )347 agent_update = self.agent_fusion(348 torch.cat([self_h, aa_context], dim=-1)349 )350 agent_context = self.agent_norm(self_h + agent_update)351 # v2.0 的任务上下文来自当前联盟状态，不考虑“如果我加入某任务会怎样”。352 coalition_sum, coalition_mean, coalition_count = self._aggregate_coalitions(353 self_h,354 other_h,355 ego_coalition,356 other_coalition,357 other_mask,358 )359 if not self.use_coalition_context:360 coalition_sum = torch.zeros_like(coalition_sum)361 coalition_mean = torch.zeros_like(coalition_mean)362 coalition_count = torch.zeros_like(coalition_count)363 normalized_count = coalition_count / float(max(1, self.num_agents))364 task_update = self.task_fusion(365 torch.cat(366 [task_h, coalition_sum, coalition_mean, normalized_count], dim=-1367 )368 )369 task_context = self.task_norm(task_h + task_update)370 diagnostics = {}371 if return_diagnostics:372 diagnostics = {373 # diagnostics 不参与反传，只给 runner/logging 观察注意力和联盟规模。374 "aa_attention": aa_weights,375 "occupied_task_ratio": (coalition_count > 0).float(),376 "max_coalition_size": coalition_count.max(dim=2).values,377 "coalition_context_norm": coalition_sum.norm(dim=-1),378 }379 return agent_context, task_context, at_features, diagnostics380381382class CandidateEdgeDecoder(nn.Module):383 """共享 A-T 边打分器：每个候选任务动作复用同一 MLP 输出一个 logit。"""384385 def __init__(self, agent_dim, task_dim, hidden_dim, edge_dim=3):386 super().__init__()387 self.network = nn.Sequential(388 nn.Linear(agent_dim + task_dim + edge_dim, hidden_dim),389 nn.ReLU(),390 nn.Linear(hidden_dim, hidden_dim),391 nn.ReLU(),392 nn.Linear(hidden_dim, 1),393 )394 nn.init.orthogonal_(self.network[-1].weight, gain=0.01)395 nn.init.constant_(self.network[-1].bias, 0.0)396397 def forward(self, agent_context, task_context, edge_features):398 # agent_context 扩展到所有任务槽，与 task_context/edge_features 拼接后逐任务评分。399 expanded_agent = agent_context.unsqueeze(2).expand(400 -1, -1, task_context.shape[2], -1401 )402 features = torch.cat(403 [expanded_agent, task_context, edge_features], dim=-1404 )405 return self.network(features).squeeze(-1)406407408class HATTV20Actor(nn.Module):409 """MAPPO/RMAPPO actor 外壳：负责 obs 整形、RNN、动作 mask 和 Categorical 分布。"""410411 def __init__(412 self,413 args,414 ego_dim,415 other_dim,416 task_dim,417 num_agents,418 num_tasks,419 device=torch.device("cpu"),420 ):421 super().__init__()422 self.num_agents = num_agents423 self.num_tasks = num_tasks424 self.hidden_size = args.hidden_size425 self.use_policy_active_masks = args.use_policy_active_masks426 self.use_naive_recurrent_policy = args.use_naive_recurrent_policy427 self.use_recurrent_policy = args.use_recurrent_policy428 self.tpdv = dict(dtype=torch.float32, device=device)429 hidden_dim = getattr(args, "hatt_hidden_dim", args.gat_out_dim)430 num_heads = getattr(args, "hatt_num_heads", 4)431 dropout = getattr(args, "hatt_dropout", 0.0)432 self.backbone = HATTV20Backbone(433 ego_dim=ego_dim,434 other_dim=other_dim,435 task_dim=task_dim,436 hidden_dim=hidden_dim,437 num_heads=num_heads,438 cap_dim=args.cap_dim,439 num_agents=num_agents,440 num_tasks=num_tasks,441 dropout=dropout,442 shared_task_obs=getattr(args, "hatt_shared_task_obs", True),443 use_aa_attention=getattr(args, "hatt_use_aa_attention", True),444 use_coalition_context=getattr(445 args, "hatt_use_coalition_context", True446 ),447 task_input_mode=getattr(args, "hatt_v20_task_input_mode", "native"),448 task_requirement_scale=getattr(args, "hatt_capability_scale", 5.0),449 )450 self.agent_projection = nn.Sequential(451 nn.Linear(hidden_dim, self.hidden_size),452 nn.ReLU(),453 nn.LayerNorm(self.hidden_size),454 )455 if self.use_naive_recurrent_policy or self.use_recurrent_policy:456 self.rnn = RNNLayer(457 self.hidden_size,458 self.hidden_size,459 args.recurrent_N,460 args.use_orthogonal,461 )462 self.decoder = CandidateEdgeDecoder(463 agent_dim=self.hidden_size,464 task_dim=hidden_dim,465 hidden_dim=self.hidden_size,466 )467 # learned reconfiguration gate 只在环境把 agent 激活给 actor 时生效。468 # mixture 是最早的可微混合：469 # pi(a|s) = g(s) * pi_hatt(a|s) + (1-g(s)) * 1[a=current]。470 # switch/hierarchical 是后续定位实验：前者让“离开当前任务”更可辨识，471 # 后者把“是否重新决策”和“重新决策后选哪个任务”拆成显式二阶段动作。472 self.use_reconfiguration_gate = (473 getattr(args, "reconfiguration_mode", "legacy") == "learned"474 )475 self.reconfiguration_gate_floor = float(476 getattr(args, "reconfiguration_gate_floor", 0.02)477 )478 self.reconfiguration_context_features = bool(479 getattr(args, "reconfiguration_context_features", False)480 )481 self.reconfiguration_context_dim = int(482 getattr(args, "reconfiguration_context_dim", 9)483 )484 self.reconfiguration_gate_parameterization = str(485 getattr(args, "reconfiguration_gate_parameterization", "mixture")486 )487 if self.reconfiguration_gate_parameterization not in (488 "mixture", "switch", "hierarchical"489 ):490 raise ValueError(491 "reconfiguration_gate_parameterization must be "492 "'mixture', 'switch', or 'hierarchical'."493 )494 self.reconfiguration_gate_override = str(495 getattr(args, "reconfiguration_gate_override", "none")496 )497 if self.reconfiguration_gate_override not in (498 "none", "open", "closed", "event", "event_or_learned"499 ):500 raise ValueError(501 "reconfiguration_gate_override must be "502 "none/open/closed/event/event_or_learned."503 )504 if (505 self.reconfiguration_gate_override in ("event", "event_or_learned")506 and not self.reconfiguration_context_features507 ):508 raise ValueError(509 "The event gate override requires reconfiguration_context_features."510 )511 if not 0.0 <= self.reconfiguration_gate_floor < 1.0:512 raise ValueError("reconfiguration_gate_floor must be in [0, 1).")513 if self.use_reconfiguration_gate:514 self.reconfiguration_gate = nn.Sequential(515 nn.Linear(516 self.hidden_size517 + hidden_dim518 + 2519 + (520 self.reconfiguration_context_dim521 if self.reconfiguration_context_features522 else 0523 ),524 self.hidden_size,525 ),526 nn.ReLU(),527 nn.Linear(self.hidden_size, 1),528 )529 nn.init.orthogonal_(self.reconfiguration_gate[-1].weight, gain=0.01)530 initial_probability = float(531 getattr(args, "reconfiguration_gate_init", 0.8)532 )533 if not 0.0 < initial_probability < 1.0:534 raise ValueError("reconfiguration_gate_init must be in (0, 1).")535 nn.init.constant_(536 self.reconfiguration_gate[-1].bias,537 math.log(initial_probability / (1.0 - initial_probability)),538 )539 self.last_diagnostics = {}540 self.last_gate_actions = None541 self.last_gate_regularization = None542 self.to(device)543544 def _prepare_obs(self, obs):545 # HATT 路径要求 Dict 观测；flatten 只在非 HATT actor 内部使用。546 if not isinstance(obs, dict):547 raise TypeError("HATTV20Actor requires structured Dict observations.")548 result = {k: check(v).to(**self.tpdv) for k, v in obs.items()}549 if result["ego"].ndim == 2:550 # buffer 采样时常把 [B, A, ...] 压成 [B*A, ...]，这里恢复 observer 维。551 result = {552 k: v.reshape(553 v.shape[0] // self.num_agents,554 self.num_agents,555 *v.shape[1:],556 )557 for k, v in result.items()558 }559 return result560561 def _distribution(self, obs, rnn_states, masks, available_actions):562 obs = self._prepare_obs(obs)563 batch_size, num_agents, _ = obs["ego"].shape564 if num_agents != self.num_agents:565 raise ValueError(f"Expected {self.num_agents} agents, got {num_agents}.")566 # available_actions 的最后一维和任务槽一一对应，非法动作会在 logits 上置为极小值。567 available_actions = check(available_actions).to(**self.tpdv)568 available_actions = available_actions.reshape(569 batch_size, num_agents, self.num_tasks570 )571 if not torch.all(available_actions.sum(dim=-1) > 0):572 raise ValueError("Every active HATT policy row must have an available action.")573 agent_context, task_h, edge_features, diagnostics = self.backbone(574 obs, available_actions575 )576 agent_features = self.agent_projection(agent_context).reshape(577 batch_size * num_agents, self.hidden_size578 )579 # RNNLayer 沿 batch*agent 维工作，因此进入 RNN 前先摊平 agent 维。580 rnn_states = check(rnn_states).to(**self.tpdv)581 masks = check(masks).to(**self.tpdv)582 if self.use_naive_recurrent_policy or self.use_recurrent_policy:583 agent_features, rnn_states = self.rnn(584 agent_features, rnn_states, masks585 )586 agent_features = agent_features.reshape(587 batch_size, num_agents, self.hidden_size588 )589 logits = self.decoder(agent_features, task_h, edge_features)590 logits = logits.reshape(batch_size * num_agents, self.num_tasks)591 available_flat = available_actions.reshape(592 batch_size * num_agents, self.num_tasks593 )594 logits = logits.masked_fill(available_flat <= 0, -1e10)595 dist = FixedCategorical(logits=logits)596 if self.use_reconfiguration_gate:597 current_indices = (598 obs["ego"][..., 2 + self.backbone.cap_dim]599 * float(max(1, self.num_tasks - 1))600 ).round().long().clamp(0, self.num_tasks - 1)601 current_task = task_h.gather(602 2,603 current_indices[..., None, None].expand(604 -1, -1, 1, task_h.shape[-1]605 ),606 ).squeeze(2)607608 base_probs = dist.probs.reshape(609 batch_size, num_agents, self.num_tasks610 )611 current_probability = base_probs.gather(612 2, current_indices.unsqueeze(-1)613 )614 alternative_probs = base_probs.clone()615 alternative_probs.scatter_(616 2, current_indices.unsqueeze(-1), 0.0617 )618 best_alternative_probability = alternative_probs.max(619 dim=2, keepdim=True620 ).values621 # gate 输入包含当前任务上下文、当前动作概率和最优替代动作优势。622 # 若开启 reconfiguration_context_features，再拼接环境显式事件/时间特征；623 # 这样 gate 不需要从稀疏局部观测中反推出“现在是否必须决策”。624 gate_features = torch.cat(625 [626 agent_features,627 current_task,628 current_probability,629 best_alternative_probability - current_probability,630 *(631 [obs["reconfiguration"]]632 if self.reconfiguration_context_features633 else []634 ),635 ],636 dim=-1,637 )638 raw_gate_probability = torch.sigmoid(639 self.reconfiguration_gate(gate_features)640 )641 # floor 防止 gate 在早期训练中完全关死，导致 decoder 没有梯度样本。642 gate_probability = self.reconfiguration_gate_floor + (643 1.0 - self.reconfiguration_gate_floor644 ) * raw_gate_probability645 if self.reconfiguration_gate_override == "open":646 gate_probability = torch.ones_like(gate_probability)647 elif self.reconfiguration_gate_override == "closed":648 gate_probability = torch.zeros_like(gate_probability)649 elif self.reconfiguration_gate_override in ("event", "event_or_learned"):650 context = obs["reconfiguration"]651 # 事件 override 用于干预实验：新任务、任务完成、当前任务完成，652 # 以及 11 维 context 中的 mandatory 决策都会强制开门。653 event_trigger = (654 (context[..., 0:1] > 0.5)655 | (context[..., 1:2] > 0.5)656 | (context[..., 3:4] > 0.5)657 )658 if self.reconfiguration_context_dim >= 11:659 event_trigger = event_trigger | (context[..., 10:11] > 0.5)660 if self.reconfiguration_gate_override == "event":661 gate_probability = event_trigger.to(gate_probability.dtype)662 else:663 gate_probability = torch.where(664 event_trigger,665 torch.ones_like(gate_probability),666 gate_probability,667 )668 current_available = available_actions.gather(669 2, current_indices.unsqueeze(-1)670 )671 # A completed or newly unavailable current task cannot be retained.672 gate_probability = torch.where(673 current_available > 0.0,674 gate_probability,675 torch.ones_like(gate_probability),676 )677 if self.reconfiguration_gate_parameterization == "hierarchical":678 # hierarchical 不把 gate 混入任务分布，而是暂存 Bernoulli 参数。679 # forward/evaluate_actions 会单独采样/重放 gate_action，并把 joint680 # logprob 写成 log p(z|s) + z * log p(a|s,z=1)。681 self._hierarchical_gate_probability = gate_probability682 self._hierarchical_current_indices = current_indices683 if self.reconfiguration_context_features and self.reconfiguration_context_dim >= 11:684 # mandatory 位置只表示“必须让 decoder 接管”的状态；685 # 正则项不会惩罚这些由环境语义强制的重新决策。686 self._hierarchical_mandatory_mask = (687 obs["reconfiguration"][..., 10:11] > 0.5688 ).to(gate_probability.dtype)689 else:690 self._hierarchical_mandatory_mask = torch.zeros_like(691 gate_probability692 )693 mixed_probs = base_probs694 elif self.reconfiguration_gate_parameterization == "switch":695 # Make g identifiable from the executed action:696 # P(current)=1-g and P(alternative)=g*q(alternative), where q697 # has no mass on current. The historical mixture instead has698 # P(switch)=g*(1-p_current), allowing g and p_current to cancel.699 alternative_probs = base_probs.clone()700 alternative_probs.scatter_(701 2, current_indices.unsqueeze(-1), 0.0702 )703 alternative_mass = alternative_probs.sum(dim=-1, keepdim=True)704 has_alternative = alternative_mass > 1e-8705 conditional_alternatives = alternative_probs / alternative_mass.clamp_min(1e-8)706 gate_probability = torch.where(707 has_alternative,708 gate_probability,709 torch.zeros_like(gate_probability),710 )711 mixed_probs = gate_probability * conditional_alternatives712 mixed_probs.scatter_add_(713 2,714 current_indices.unsqueeze(-1),715 1.0 - gate_probability,716 )717 else:718 mixed_probs = gate_probability * base_probs719 mixed_probs.scatter_add_(720 2,721 current_indices.unsqueeze(-1),722 1.0 - gate_probability,723 )724 mixed_probs = mixed_probs / mixed_probs.sum(725 dim=-1, keepdim=True726 ).clamp_min(1e-8)727 dist = FixedCategorical(728 probs=mixed_probs.reshape(729 batch_size * num_agents, self.num_tasks730 )731 )732 diagnostics.update(733 {734 "reconfiguration_gate_probability": gate_probability,735 "reconfiguration_switch_probability": (736 gate_probability737 if self.reconfiguration_gate_parameterization == "switch"738 else gate_probability * (1.0 - current_probability)739 ),740 "reconfiguration_current_task_probability": current_probability,741 }742 )743 if self.reconfiguration_context_features:744 event_mask = (745 (obs["reconfiguration"][..., 0:1] > 0.5)746 | (obs["reconfiguration"][..., 1:2] > 0.5)747 | (obs["reconfiguration"][..., 3:4] > 0.5)748 )749 if self.reconfiguration_context_dim >= 11:750 event_mask = event_mask | (751 obs["reconfiguration"][..., 10:11] > 0.5752 )753 for suffix, selection in (754 ("on_event", event_mask),755 ("off_event", ~event_mask),756 ):757 selected = gate_probability.masked_select(selection)758 diagnostics[759 f"reconfiguration_gate_probability_{suffix}"760 ] = (761 selected.mean()762 if selected.numel()763 else torch.zeros((), **self.tpdv)764 )765 # 诊断张量只用于日志，detach 避免 runner 持有计算图。766 self.last_diagnostics = {k: v.detach() for k, v in diagnostics.items()}767 return dist, rnn_states768769 def _hierarchical_action(self, dist, deterministic):770 """采样显式二阶段动作：先决定是否重选，再在开门时采样任务动作。"""771772 gate_probability = self._hierarchical_gate_probability.reshape(-1, 1)773 gate_dist = torch.distributions.Bernoulli(probs=gate_probability)774 gate_action = (775 (gate_probability >= 0.5).to(gate_probability.dtype)776 if deterministic777 else gate_dist.sample()778 )779 task_action = dist.mode() if deterministic else dist.sample()780 current = self._hierarchical_current_indices.reshape(-1, 1)781 action = torch.where(gate_action > 0.5, task_action, current)782 # gate_action=0 时最终动作固定为 current，此时 task_action 只是未执行的783 # 内部样本，不能把它的 logprob 计入 PPO 的重要性权重。784 joint_log_prob = gate_dist.log_prob(gate_action)785 joint_log_prob = joint_log_prob + gate_action * dist.log_probs(task_action)786 self.last_gate_actions = gate_action.detach()787 return action, joint_log_prob788789 def forward(790 self,791 obs,792 rnn_states,793 masks,794 available_actions=None,795 deterministic=False,796 ):797 dist, rnn_states = self._distribution(798 obs, rnn_states, masks, available_actions799 )800 if (801 self.use_reconfiguration_gate802 and self.reconfiguration_gate_parameterization == "hierarchical"803 ):804 actions, action_log_probs = self._hierarchical_action(805 dist, deterministic806 )807 return actions, action_log_probs, rnn_states808 actions = dist.mode() if deterministic else dist.sample()809 self.last_gate_actions = None810 return actions, dist.log_probs(actions), rnn_states811812 def evaluate_actions(813 self,814 obs,815 rnn_states,816 action,817 masks,818 available_actions=None,819 active_masks=None,820 gate_actions=None,821 ):822 dist, _ = self._distribution(obs, rnn_states, masks, available_actions)823 action = check(action).to(device=self.tpdv["device"])824 if (825 self.use_reconfiguration_gate826 and self.reconfiguration_gate_parameterization == "hierarchical"827 ):828 if gate_actions is None:829 raise ValueError("hierarchical gate evaluation requires gate_actions")830 gate_actions = check(gate_actions).to(**self.tpdv).reshape(-1, 1)831 gate_probability = self._hierarchical_gate_probability.reshape(-1, 1)832 gate_dist = torch.distributions.Bernoulli(probs=gate_probability)833 # PPO update 必须重放 rollout 时实际采到的 gate_action；否则834 # old/new logprob 对不上，hierarchical gate 的梯度会变成错配样本。835 action_log_probs = gate_dist.log_prob(gate_actions)836 action_log_probs = action_log_probs + gate_actions * dist.log_probs(action)837 entropy = gate_dist.entropy().squeeze(-1) + (838 gate_probability.squeeze(-1) * dist.entropy()839 )840 learned_mask = 1.0 - self._hierarchical_mandatory_mask.reshape(-1, 1)841 if active_masks is not None:842 learned_mask = learned_mask * check(active_masks).to(**self.tpdv).reshape(-1, 1)843 # decision_coef 惩罚的是“非强制状态下主动开门”的平均概率，844 # 用来寻找更稀疏的 learned reconfiguration，而不是环境奖励项。845 self.last_gate_regularization = (846 gate_probability * learned_mask847 ).sum() / learned_mask.sum().clamp_min(1.0)848 else:849 action_log_probs = dist.log_probs(action)850 entropy = dist.entropy()851 self.last_gate_regularization = None852 if active_masks is not None and self.use_policy_active_masks:853 active_masks = check(active_masks).to(**self.tpdv).squeeze(-1)854 dist_entropy = (855 (entropy * active_masks).sum()856 / active_masks.sum().clamp_min(1.0)857 )858 else:859 dist_entropy = entropy.mean()860 return action_log_probs, dist_entropy