"""v2.0 是后续 v2/v3 系列的主要旧基线。它的核心思想是： 1. 分别编码 focal agent、其他 agent 和任务； 2. 用 A-A attention 更新 focal agent 表示； 3. 按当前联盟编号把自己和其他成员聚合到每个任务槽； 4. 将任务编码与该任务当前联盟表示融合，再对每条 A-T 候选边打分。 注意：本版本编码的是“当前联盟状态”，没有做 v3.0 之后的候选加入前后反事实推理。 """ import math import torch import torch.nn as nn from algorithms.utils.distributions import FixedCategorical from algorithms.utils.rnn import RNNLayer from algorithms.utils.util import check class FeatureEncoder(nn.Module): """两层 MLP + LayerNorm，用于把 ego/other/task 的原始特征映射到同一隐空间。""" def __init__(self, input_dim, hidden_dim): super().__init__() self.net = nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.LayerNorm(hidden_dim), ) def forward(self, x): return self.net(x) def initialize_full_actor_from_compact_reference( full_actor, compact_actor, cap_dim ): """Make a full-schema actor an exact zero-padded extension of compact6. The compact task order is ``x, y, demand, deadline, req0, type``. The corresponding full-schema columns are copied into the 12-D task encoder; all additional input columns start at zero. Every other actor parameter is copied verbatim. With zero-valued extra fields and restored req0 scale, the two actors therefore implement the same function at initialization. """ compact_state = compact_actor.state_dict() full_state = full_actor.state_dict() first_weight = "backbone.task_encoder.net.0.weight" compact_columns = [0, 1, 2, 3, 4, 5] full_columns = [0, 1, 2 + cap_dim, 3 + cap_dim, 2, 7 + cap_dim] if compact_state[first_weight].shape[1] != 6: raise ValueError("compact reference task encoder must have six inputs") expected_full_dim = 2 + cap_dim + 1 + 1 + 3 + 1 if full_state[first_weight].shape[1] != expected_full_dim: raise ValueError( "full task encoder has unexpected input width: " f"{full_state[first_weight].shape[1]} != {expected_full_dim}" ) for name, target in full_state.items(): source = compact_state[name] if name == first_weight: target.zero_() for compact_idx, full_idx in zip(compact_columns, full_columns): target[:, full_idx].copy_(source[:, compact_idx]) continue if target.shape != source.shape: raise ValueError( f"cannot preserve actor function for {name}: " f"full={tuple(target.shape)}, compact={tuple(source.shape)}" ) target.copy_(source) full_actor.load_state_dict(full_state) def encode_repeated_tasks(task_encoder, tasks, shared_task_obs): """任务观测在所有 agent 间相同时，只编码一次并广播，减少重复计算。""" if shared_task_obs and tasks.ndim == 4 and tasks.shape[1] > 1: encoded = task_encoder(tasks[:, 0]) return encoded.unsqueeze(1).expand(-1, tasks.shape[1], -1, -1) return task_encoder(tasks) class RelationAttention(nn.Module): """带边特征 bias 的稠密多头注意力，用于 A-A 关系聚合。""" def __init__(self, hidden_dim, edge_dim, num_heads, dropout=0.0): super().__init__() 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 self.scale = self.head_dim ** -0.5 self.query = nn.Linear(hidden_dim, hidden_dim) self.key = nn.Linear(hidden_dim, hidden_dim) self.value = nn.Linear(hidden_dim, hidden_dim) 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) self.dropout = nn.Dropout(dropout) def forward(self, query, neighbors, edge_features, mask): # query: [B, A, H]；neighbors: [B, A, N, H]。 # mask 屏蔽不存在/死亡/不可用的邻居，避免 softmax 分配注意力。 batch_size, num_agents, num_neighbors, hidden_dim = neighbors.shape # 每个 observer agent 作为 query；每个邻居智能体作为 key/value。 q = self.query(query).view( batch_size, num_agents, self.num_heads, self.head_dim ) k = self.key(neighbors).view( batch_size, num_agents, num_neighbors, self.num_heads, self.head_dim ) v = self.value(neighbors).view( batch_size, num_agents, num_neighbors, self.num_heads, self.head_dim ) scores = torch.einsum("bahd,banhd->banh", q, k) * self.scale # 边特征直接加到每个 head 的 logit 上，用于注入相对位置和同联盟关系。 scores = scores + self.edge_bias(edge_features) valid = mask.bool().unsqueeze(-1) scores = scores.masked_fill(~valid, torch.finfo(scores.dtype).min) weights = torch.softmax(scores, dim=2) # 全无效邻居时避免 NaN；再归一化保证有效邻居权重和为 1。 weights = torch.where(valid, weights, torch.zeros_like(weights)) weights = weights / weights.sum(dim=2, keepdim=True).clamp_min(1e-8) weights = self.dropout(weights) context = torch.einsum("banh,banhd->bahd", weights, v) context = context.reshape(batch_size, num_agents, hidden_dim) return self.output(context), weights class HATTV20Backbone(nn.Module): """ 对每个观察者 agent，构建 num_tasks 个任务槽的联盟表示。 这里的联盟表示包含“当前已经分配到该任务的成员”，并包含 focal agent 自己当前所在联盟。 """ 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, use_aa_attention=True, use_coalition_context=True, task_input_mode="native", task_requirement_scale=5.0, ): super().__init__() self.cap_dim = cap_dim self.num_agents = num_agents self.num_tasks = num_tasks self.shared_task_obs = shared_task_obs # Ablations preserve all parameter tensors and shapes. A disabled # branch contributes zeros so parameter count cannot explain a result. self.use_aa_attention = bool(use_aa_attention) self.use_coalition_context = bool(use_coalition_context) self.task_input_mode = str(task_input_mode) self.task_requirement_scale = float(task_requirement_scale) if self.task_input_mode not in ( "native", "restore_req_scale", "compact_projection", ): raise ValueError( "task_input_mode must be native, restore_req_scale, or " f"compact_projection; got {self.task_input_mode!r}" ) if self.task_input_mode != "native" and task_dim < 2 + cap_dim + 6: raise ValueError( f"task_input_mode={self.task_input_mode!r} requires the full " f"task schema, got task_dim={task_dim}" ) encoded_task_dim = 6 if self.task_input_mode == "compact_projection" else task_dim self.self_encoder = FeatureEncoder(ego_dim, hidden_dim) self.other_encoder = FeatureEncoder(other_dim, hidden_dim) self.task_encoder = FeatureEncoder(encoded_task_dim, hidden_dim) self.aa_attention = RelationAttention( hidden_dim, edge_dim=3, num_heads=num_heads, dropout=dropout ) self.agent_fusion = nn.Sequential( nn.Linear(hidden_dim * 2, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), ) self.agent_norm = nn.LayerNorm(hidden_dim) self.task_fusion = nn.Sequential( nn.Linear(hidden_dim * 3 + 1, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), ) self.task_norm = nn.LayerNorm(hidden_dim) def _prepare_task_encoder_input(self, tasks): """Apply controlled schema interventions before the v2.0 task MLP. The full DelivCoal schema stores req0 divided by capability_scale. The restore control changes only that scale. The compact projection maps the full schema back to the historical six semantic fields, allowing a width-matched actor initialization against the compact baseline. """ if self.task_input_mode == "native": return tasks requirement = tasks[..., 2:3] * self.task_requirement_scale if self.task_input_mode == "restore_req_scale": prepared = tasks.clone() prepared[..., 2:3] = requirement return prepared demand_idx = 2 + self.cap_dim deadline_idx = 3 + self.cap_dim type_idx = 7 + self.cap_dim return torch.cat( [ tasks[..., 0:2], tasks[..., demand_idx:demand_idx + 1], tasks[..., deadline_idx:deadline_idx + 1], requirement, tasks[..., type_idx:type_idx + 1], ], dim=-1, ) def _coalition_indices(self, normalized_ids): # 环境中联盟编号以 [0, 1] 归一化存储，这里恢复为离散 task slot。 return (normalized_ids * max(1, self.num_tasks - 1)).round().long().clamp( 0, self.num_tasks - 1 ) def _relation_features(self, ego, others, tasks): # A-A 边：相对位置 + 是否同联盟。 # A-T 边：相对任务位置 + focal agent 当前是否分配给该任务。 ego_position = ego[..., 0:2] ego_coalition = ego[..., 2 + self.cap_dim] other_position = others[..., 0:2] other_coalition = others[..., self.cap_dim + 4] aa_relative = other_position - ego_position.unsqueeze(2) same_coalition = ( self._coalition_indices(other_coalition) == self._coalition_indices(ego_coalition).unsqueeze(2) ).float() aa_features = torch.cat( [aa_relative, same_coalition.unsqueeze(-1)], dim=-1 ) task_position = tasks[..., 0:2] at_relative = task_position - ego_position.unsqueeze(2) task_ids = torch.arange(self.num_tasks, device=ego.device) assigned = ( self._coalition_indices(ego_coalition).unsqueeze(-1) == task_ids.view(1, 1, -1) ).float() at_features = torch.cat( [at_relative, assigned.unsqueeze(-1)], dim=-1 ) return aa_features, at_features def _aggregate_coalitions( self, self_h, other_h, ego_coalition, other_coalition, other_mask ): """把 ego/other 成员按当前任务编号 scatter 到任务槽中形成联盟 sum/mean/count。""" batch_size, observers, _, hidden_dim = other_h.shape coalition_sum = other_h.new_zeros( batch_size, observers, self.num_tasks, hidden_dim ) coalition_count = other_h.new_zeros( batch_size, observers, self.num_tasks, 1 ) other_indices = self._coalition_indices(other_coalition) # scatter_add_ 需要把每个 other 的任务槽编号扩展到 hidden_dim 维。 message_indices = other_indices.unsqueeze(-1).expand( -1, -1, -1, hidden_dim ) # other_mask 会过滤 padding/self/dead agent，避免进入联盟聚合。 valid_messages = other_h * other_mask.unsqueeze(-1) coalition_sum.scatter_add_(2, message_indices, valid_messages) coalition_count.scatter_add_( 2, other_indices.unsqueeze(-1), other_mask.unsqueeze(-1) ) self_indices = self._coalition_indices(ego_coalition).unsqueeze(-1) # v2.0 当前联盟状态包含 focal agent 自己，因此把 self_h 加回当前任务槽。 coalition_sum.scatter_add_( 2, self_indices.unsqueeze(-1).expand(-1, -1, -1, hidden_dim), self_h.unsqueeze(2), ) coalition_count.scatter_add_( 2, self_indices.unsqueeze(-1), torch.ones_like(self_indices).unsqueeze(-1).float(), ) # mean 给任务融合提供规模无关的联盟表示；count 保留联盟人数信息。 coalition_mean = coalition_sum / coalition_count.clamp_min(1.0) return coalition_sum, coalition_mean, coalition_count def forward(self, obs, available_actions, return_diagnostics=True): del available_actions ego = obs["ego"] others = obs["others"] tasks = obs["tasks"] ego_coalition = ego[..., 2 + self.cap_dim] other_coalition = others[..., self.cap_dim + 4] # 联盟编号只用于构造关系和聚合槽，不作为连续语义特征直接送入 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 = self.self_encoder(ego_input) other_h = self.other_encoder(others_input) task_encoder_input = self._prepare_task_encoder_input(tasks) task_h = encode_repeated_tasks( self.task_encoder, task_encoder_input, self.shared_task_obs ) # relation_features 不进 encoder，而是作为 attention/decoder 的边条件。 aa_features, at_features = self._relation_features(ego, others, tasks) living_index = self.cap_dim + 6 # others 已经把对角线自身置零；这里再过滤空行和死亡 agent。 other_mask = ( (others.abs().sum(dim=-1) > 0) & (others[..., living_index] > 0.5) ).float() # A-A attention 提供“其他智能体对当前 agent 决策的上下文影响”。 if self.use_aa_attention: aa_context, aa_weights = self.aa_attention( self_h, other_h, aa_features, other_mask.bool() ) else: aa_context = torch.zeros_like(self_h) aa_weights = other_h.new_zeros( *other_h.shape[:-1], self.aa_attention.num_heads ) agent_update = self.agent_fusion( torch.cat([self_h, aa_context], dim=-1) ) agent_context = self.agent_norm(self_h + agent_update) # v2.0 的任务上下文来自当前联盟状态，不考虑“如果我加入某任务会怎样”。 coalition_sum, coalition_mean, coalition_count = self._aggregate_coalitions( self_h, other_h, ego_coalition, other_coalition, other_mask, ) if not self.use_coalition_context: coalition_sum = torch.zeros_like(coalition_sum) coalition_mean = torch.zeros_like(coalition_mean) coalition_count = torch.zeros_like(coalition_count) normalized_count = coalition_count / float(max(1, self.num_agents)) task_update = self.task_fusion( torch.cat( [task_h, coalition_sum, coalition_mean, normalized_count], dim=-1 ) ) task_context = self.task_norm(task_h + task_update) diagnostics = {} if return_diagnostics: diagnostics = { # diagnostics 不参与反传，只给 runner/logging 观察注意力和联盟规模。 "aa_attention": aa_weights, "occupied_task_ratio": (coalition_count > 0).float(), "max_coalition_size": coalition_count.max(dim=2).values, "coalition_context_norm": coalition_sum.norm(dim=-1), } return agent_context, task_context, at_features, diagnostics class CandidateEdgeDecoder(nn.Module): """共享 A-T 边打分器：每个候选任务动作复用同一 MLP 输出一个 logit。""" def __init__(self, agent_dim, task_dim, hidden_dim, edge_dim=3): super().__init__() self.network = nn.Sequential( nn.Linear(agent_dim + task_dim + edge_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 1), ) nn.init.orthogonal_(self.network[-1].weight, gain=0.01) nn.init.constant_(self.network[-1].bias, 0.0) def forward(self, agent_context, task_context, edge_features): # agent_context 扩展到所有任务槽，与 task_context/edge_features 拼接后逐任务评分。 expanded_agent = agent_context.unsqueeze(2).expand( -1, -1, task_context.shape[2], -1 ) features = torch.cat( [expanded_agent, task_context, edge_features], dim=-1 ) return self.network(features).squeeze(-1) class HATTV20Actor(nn.Module): """MAPPO/RMAPPO actor 外壳：负责 obs 整形、RNN、动作 mask 和 Categorical 分布。""" def __init__( self, args, ego_dim, other_dim, task_dim, num_agents, num_tasks, device=torch.device("cpu"), ): super().__init__() self.num_agents = num_agents self.num_tasks = num_tasks self.hidden_size = args.hidden_size self.use_policy_active_masks = args.use_policy_active_masks self.use_naive_recurrent_policy = args.use_naive_recurrent_policy self.use_recurrent_policy = args.use_recurrent_policy self.tpdv = dict(dtype=torch.float32, device=device) hidden_dim = getattr(args, "hatt_hidden_dim", args.gat_out_dim) num_heads = getattr(args, "hatt_num_heads", 4) dropout = getattr(args, "hatt_dropout", 0.0) self.backbone = HATTV20Backbone( ego_dim=ego_dim, other_dim=other_dim, task_dim=task_dim, hidden_dim=hidden_dim, num_heads=num_heads, cap_dim=args.cap_dim, num_agents=num_agents, num_tasks=num_tasks, dropout=dropout, shared_task_obs=getattr(args, "hatt_shared_task_obs", True), use_aa_attention=getattr(args, "hatt_use_aa_attention", True), use_coalition_context=getattr( args, "hatt_use_coalition_context", True ), task_input_mode=getattr(args, "hatt_v20_task_input_mode", "native"), task_requirement_scale=getattr(args, "hatt_capability_scale", 5.0), ) self.agent_projection = nn.Sequential( nn.Linear(hidden_dim, self.hidden_size), nn.ReLU(), nn.LayerNorm(self.hidden_size), ) if self.use_naive_recurrent_policy or self.use_recurrent_policy: self.rnn = RNNLayer( self.hidden_size, self.hidden_size, args.recurrent_N, args.use_orthogonal, ) self.decoder = CandidateEdgeDecoder( agent_dim=self.hidden_size, task_dim=hidden_dim, hidden_dim=self.hidden_size, ) # learned reconfiguration gate 只在环境把 agent 激活给 actor 时生效。 # mixture 是最早的可微混合： # pi(a|s) = g(s) * pi_hatt(a|s) + (1-g(s)) * 1[a=current]。 # switch/hierarchical 是后续定位实验：前者让“离开当前任务”更可辨识， # 后者把“是否重新决策”和“重新决策后选哪个任务”拆成显式二阶段动作。 self.use_reconfiguration_gate = ( getattr(args, "reconfiguration_mode", "legacy") == "learned" ) self.reconfiguration_gate_floor = float( getattr(args, "reconfiguration_gate_floor", 0.02) ) self.reconfiguration_context_features = bool( getattr(args, "reconfiguration_context_features", False) ) self.reconfiguration_context_dim = int( getattr(args, "reconfiguration_context_dim", 9) ) self.reconfiguration_gate_parameterization = str( getattr(args, "reconfiguration_gate_parameterization", "mixture") ) if self.reconfiguration_gate_parameterization not in ( "mixture", "switch", "hierarchical" ): raise ValueError( "reconfiguration_gate_parameterization must be " "'mixture', 'switch', or 'hierarchical'." ) self.reconfiguration_gate_override = str( getattr(args, "reconfiguration_gate_override", "none") ) if self.reconfiguration_gate_override not in ( "none", "open", "closed", "event", "event_or_learned" ): raise ValueError( "reconfiguration_gate_override must be " "none/open/closed/event/event_or_learned." ) if ( self.reconfiguration_gate_override in ("event", "event_or_learned") and not self.reconfiguration_context_features ): raise ValueError( "The event gate override requires reconfiguration_context_features." ) if not 0.0 <= self.reconfiguration_gate_floor < 1.0: raise ValueError("reconfiguration_gate_floor must be in [0, 1).") if self.use_reconfiguration_gate: self.reconfiguration_gate = nn.Sequential( nn.Linear( self.hidden_size + hidden_dim + 2 + ( self.reconfiguration_context_dim if self.reconfiguration_context_features else 0 ), self.hidden_size, ), nn.ReLU(), nn.Linear(self.hidden_size, 1), ) nn.init.orthogonal_(self.reconfiguration_gate[-1].weight, gain=0.01) initial_probability = float( getattr(args, "reconfiguration_gate_init", 0.8) ) if not 0.0 < initial_probability < 1.0: raise ValueError("reconfiguration_gate_init must be in (0, 1).") nn.init.constant_( self.reconfiguration_gate[-1].bias, math.log(initial_probability / (1.0 - initial_probability)), ) self.last_diagnostics = {} self.last_gate_actions = None self.last_gate_regularization = None self.to(device) def _prepare_obs(self, obs): # HATT 路径要求 Dict 观测；flatten 只在非 HATT actor 内部使用。 if not isinstance(obs, dict): raise TypeError("HATTV20Actor requires structured Dict observations.") result = {k: check(v).to(**self.tpdv) for k, v in obs.items()} if result["ego"].ndim == 2: # buffer 采样时常把 [B, A, ...] 压成 [B*A, ...]，这里恢复 observer 维。 result = { k: v.reshape( v.shape[0] // self.num_agents, self.num_agents, *v.shape[1:], ) for k, v in result.items() } return result def _distribution(self, obs, rnn_states, masks, available_actions): obs = self._prepare_obs(obs) batch_size, num_agents, _ = obs["ego"].shape if num_agents != self.num_agents: raise ValueError(f"Expected {self.num_agents} agents, got {num_agents}.") # available_actions 的最后一维和任务槽一一对应，非法动作会在 logits 上置为极小值。 available_actions = check(available_actions).to(**self.tpdv) available_actions = available_actions.reshape( batch_size, num_agents, self.num_tasks ) if not torch.all(available_actions.sum(dim=-1) > 0): raise ValueError("Every active HATT policy row must have an available action.") agent_context, task_h, edge_features, diagnostics = self.backbone( obs, available_actions ) agent_features = self.agent_projection(agent_context).reshape( batch_size * num_agents, self.hidden_size ) # RNNLayer 沿 batch*agent 维工作，因此进入 RNN 前先摊平 agent 维。 rnn_states = check(rnn_states).to(**self.tpdv) masks = check(masks).to(**self.tpdv) if self.use_naive_recurrent_policy or self.use_recurrent_policy: agent_features, rnn_states = self.rnn( agent_features, rnn_states, masks ) agent_features = agent_features.reshape( batch_size, num_agents, self.hidden_size ) logits = self.decoder(agent_features, task_h, edge_features) logits = logits.reshape(batch_size * num_agents, self.num_tasks) available_flat = available_actions.reshape( batch_size * num_agents, self.num_tasks ) logits = logits.masked_fill(available_flat <= 0, -1e10) dist = FixedCategorical(logits=logits) if self.use_reconfiguration_gate: current_indices = ( obs["ego"][..., 2 + self.backbone.cap_dim] * float(max(1, self.num_tasks - 1)) ).round().long().clamp(0, self.num_tasks - 1) current_task = task_h.gather( 2, current_indices[..., None, None].expand( -1, -1, 1, task_h.shape[-1] ), ).squeeze(2) base_probs = dist.probs.reshape( batch_size, num_agents, self.num_tasks ) current_probability = base_probs.gather( 2, current_indices.unsqueeze(-1) ) alternative_probs = base_probs.clone() alternative_probs.scatter_( 2, current_indices.unsqueeze(-1), 0.0 ) best_alternative_probability = alternative_probs.max( dim=2, keepdim=True ).values # gate 输入包含当前任务上下文、当前动作概率和最优替代动作优势。 # 若开启 reconfiguration_context_features，再拼接环境显式事件/时间特征； # 这样 gate 不需要从稀疏局部观测中反推出“现在是否必须决策”。 gate_features = torch.cat( [ agent_features, current_task, current_probability, best_alternative_probability - current_probability, *( [obs["reconfiguration"]] if self.reconfiguration_context_features else [] ), ], dim=-1, ) raw_gate_probability = torch.sigmoid( self.reconfiguration_gate(gate_features) ) # floor 防止 gate 在早期训练中完全关死，导致 decoder 没有梯度样本。 gate_probability = self.reconfiguration_gate_floor + ( 1.0 - self.reconfiguration_gate_floor ) * raw_gate_probability if self.reconfiguration_gate_override == "open": gate_probability = torch.ones_like(gate_probability) elif self.reconfiguration_gate_override == "closed": gate_probability = torch.zeros_like(gate_probability) elif self.reconfiguration_gate_override in ("event", "event_or_learned"): context = obs["reconfiguration"] # 事件 override 用于干预实验：新任务、任务完成、当前任务完成， # 以及 11 维 context 中的 mandatory 决策都会强制开门。 event_trigger = ( (context[..., 0:1] > 0.5) | (context[..., 1:2] > 0.5) | (context[..., 3:4] > 0.5) ) if self.reconfiguration_context_dim >= 11: event_trigger = event_trigger | (context[..., 10:11] > 0.5) if self.reconfiguration_gate_override == "event": gate_probability = event_trigger.to(gate_probability.dtype) else: gate_probability = torch.where( event_trigger, torch.ones_like(gate_probability), gate_probability, ) current_available = available_actions.gather( 2, current_indices.unsqueeze(-1) ) # A completed or newly unavailable current task cannot be retained. gate_probability = torch.where( current_available > 0.0, gate_probability, torch.ones_like(gate_probability), ) if self.reconfiguration_gate_parameterization == "hierarchical": # hierarchical 不把 gate 混入任务分布，而是暂存 Bernoulli 参数。 # forward/evaluate_actions 会单独采样/重放 gate_action，并把 joint # logprob 写成 log p(z|s) + z * log p(a|s,z=1)。 self._hierarchical_gate_probability = gate_probability self._hierarchical_current_indices = current_indices if self.reconfiguration_context_features and self.reconfiguration_context_dim >= 11: # mandatory 位置只表示“必须让 decoder 接管”的状态； # 正则项不会惩罚这些由环境语义强制的重新决策。 self._hierarchical_mandatory_mask = ( obs["reconfiguration"][..., 10:11] > 0.5 ).to(gate_probability.dtype) else: self._hierarchical_mandatory_mask = torch.zeros_like( gate_probability ) mixed_probs = base_probs elif self.reconfiguration_gate_parameterization == "switch": # Make g identifiable from the executed action: # P(current)=1-g and P(alternative)=g*q(alternative), where q # has no mass on current. The historical mixture instead has # P(switch)=g*(1-p_current), allowing g and p_current to cancel. alternative_probs = base_probs.clone() alternative_probs.scatter_( 2, current_indices.unsqueeze(-1), 0.0 ) alternative_mass = alternative_probs.sum(dim=-1, keepdim=True) has_alternative = alternative_mass > 1e-8 conditional_alternatives = alternative_probs / alternative_mass.clamp_min(1e-8) gate_probability = torch.where( has_alternative, gate_probability, torch.zeros_like(gate_probability), ) mixed_probs = gate_probability * conditional_alternatives mixed_probs.scatter_add_( 2, current_indices.unsqueeze(-1), 1.0 - gate_probability, ) else: mixed_probs = gate_probability * base_probs mixed_probs.scatter_add_( 2, current_indices.unsqueeze(-1), 1.0 - gate_probability, ) mixed_probs = mixed_probs / mixed_probs.sum( dim=-1, keepdim=True ).clamp_min(1e-8) dist = FixedCategorical( probs=mixed_probs.reshape( batch_size * num_agents, self.num_tasks ) ) diagnostics.update( { "reconfiguration_gate_probability": gate_probability, "reconfiguration_switch_probability": ( gate_probability if self.reconfiguration_gate_parameterization == "switch" else gate_probability * (1.0 - current_probability) ), "reconfiguration_current_task_probability": current_probability, } ) if self.reconfiguration_context_features: event_mask = ( (obs["reconfiguration"][..., 0:1] > 0.5) | (obs["reconfiguration"][..., 1:2] > 0.5) | (obs["reconfiguration"][..., 3:4] > 0.5) ) if self.reconfiguration_context_dim >= 11: event_mask = event_mask | ( obs["reconfiguration"][..., 10:11] > 0.5 ) for suffix, selection in ( ("on_event", event_mask), ("off_event", ~event_mask), ): selected = gate_probability.masked_select(selection) diagnostics[ f"reconfiguration_gate_probability_{suffix}" ] = ( selected.mean() if selected.numel() else torch.zeros((), **self.tpdv) ) # 诊断张量只用于日志，detach 避免 runner 持有计算图。 self.last_diagnostics = {k: v.detach() for k, v in diagnostics.items()} return dist, rnn_states def _hierarchical_action(self, dist, deterministic): """采样显式二阶段动作：先决定是否重选，再在开门时采样任务动作。""" gate_probability = self._hierarchical_gate_probability.reshape(-1, 1) gate_dist = torch.distributions.Bernoulli(probs=gate_probability) gate_action = ( (gate_probability >= 0.5).to(gate_probability.dtype) if deterministic else gate_dist.sample() ) task_action = dist.mode() if deterministic else dist.sample() current = self._hierarchical_current_indices.reshape(-1, 1) action = torch.where(gate_action > 0.5, task_action, current) # gate_action=0 时最终动作固定为 current，此时 task_action 只是未执行的 # 内部样本，不能把它的 logprob 计入 PPO 的重要性权重。 joint_log_prob = gate_dist.log_prob(gate_action) joint_log_prob = joint_log_prob + gate_action * dist.log_probs(task_action) self.last_gate_actions = gate_action.detach() return action, joint_log_prob def forward( self, obs, rnn_states, masks, available_actions=None, deterministic=False, ): dist, rnn_states = self._distribution( obs, rnn_states, masks, available_actions ) if ( self.use_reconfiguration_gate and self.reconfiguration_gate_parameterization == "hierarchical" ): actions, action_log_probs = self._hierarchical_action( dist, deterministic ) return actions, action_log_probs, rnn_states actions = dist.mode() if deterministic else dist.sample() self.last_gate_actions = None return actions, dist.log_probs(actions), rnn_states def evaluate_actions( self, obs, rnn_states, action, masks, available_actions=None, active_masks=None, gate_actions=None, ): dist, _ = self._distribution(obs, rnn_states, masks, available_actions) action = check(action).to(device=self.tpdv["device"]) if ( self.use_reconfiguration_gate and self.reconfiguration_gate_parameterization == "hierarchical" ): if gate_actions is None: raise ValueError("hierarchical gate evaluation requires gate_actions") gate_actions = check(gate_actions).to(**self.tpdv).reshape(-1, 1) gate_probability = self._hierarchical_gate_probability.reshape(-1, 1) gate_dist = torch.distributions.Bernoulli(probs=gate_probability) # PPO update 必须重放 rollout 时实际采到的 gate_action；否则 # old/new logprob 对不上，hierarchical gate 的梯度会变成错配样本。 action_log_probs = gate_dist.log_prob(gate_actions) action_log_probs = action_log_probs + gate_actions * dist.log_probs(action) entropy = gate_dist.entropy().squeeze(-1) + ( gate_probability.squeeze(-1) * dist.entropy() ) learned_mask = 1.0 - self._hierarchical_mandatory_mask.reshape(-1, 1) if active_masks is not None: learned_mask = learned_mask * check(active_masks).to(**self.tpdv).reshape(-1, 1) # decision_coef 惩罚的是“非强制状态下主动开门”的平均概率， # 用来寻找更稀疏的 learned reconfiguration，而不是环境奖励项。 self.last_gate_regularization = ( gate_probability * learned_mask ).sum() / learned_mask.sum().clamp_min(1.0) else: action_log_probs = dist.log_probs(action) entropy = dist.entropy() self.last_gate_regularization = None if active_masks is not None and self.use_policy_active_masks: active_masks = check(active_masks).to(**self.tpdv).squeeze(-1) dist_entropy = ( (entropy * active_masks).sum() / active_masks.sum().clamp_min(1.0) ) else: dist_entropy = entropy.mean() return action_log_probs, dist_entropy