gaRMAPPOPolicy.py
图模型路由、actor/critic 构造与动作接口
原文件 · 870 行 · 38302 字节
algorithms/algorithm/gaRMAPPOPolicy.py
SHA-256: a7e11d5a43393b1f28a8c270b61178342e2c0a9e07a25954fa8da19bb89b6e14
1"""2# @Time : 2025/12/07 10:303# @Author : wingong4# @Email : wy1997331@126.com5# @File : gaRMAPPOPolicy.py6#7# graph_model="legacy_gat" 保留三种旧 GAT Actor 融合模式：8# "gat" — Actor 输入 = 仅 GAT 嵌入（信息瓶颈）9# "cat" — Actor 输入 = cat(obs, GAT 嵌入)10# "gated" — Actor 输入 = cat(obs, tanh(α)*GAT 嵌入)，α 可学习11# graph_model="hatt" 使用结构化 Dict 观测、关系专属注意力和任务动作解码器。12# Critic 始终使用与 NoGAT 相同的 centralized state。13"""1415import copy1617import torch18from gym import spaces19from algorithms.algorithm.r_actor_critic import R_Actor, R_Critic20from utils.util import update_linear_schedule2122from algorithms.algorithm.GATNetwork import GATNetwork23from algorithms.algorithm.HATTV10Network import HATTActor24from algorithms.algorithm.HATTV20Network import (25 HATTV20Actor,26 initialize_full_actor_from_compact_reference,27)28from algorithms.algorithm.HATTV20InfoNetwork import (29 HATTV20NoIdleActor,30 HATTV20PhysicsActor,31 HATTV20RawActor,32 HATTV20ResourceActor,33 HATTV20SetContextActor,34 HATTV20TaskAttentionActor,35)36from algorithms.algorithm.HATTV20SpecialNetwork import (37 HATTV20PlainNoIdleActor,38 HATTV20SpecialDetachedActor,39 HATTV20SpecialSplitActor,40)41from algorithms.algorithm.HATTV20TypeConditionedNetwork import (42 HATTV20TypeEmbeddingActor,43 HATTV20TypeGateActor,44)45from algorithms.algorithm.HATTV21Network import HATTV2Actor as HATTV21Actor46from algorithms.algorithm.HATTV22Network import HATTV2Actor as HATTV22Actor47from algorithms.algorithm.HATTV23Network import (48 HATTV23AActor,49 HATTV23BActor,50 HATTV23CActor,51)52from algorithms.algorithm.HATTV24Network import HATTV24Actor53from algorithms.algorithm.HATTV25Network import (54 HATTV25AActor,55 HATTV25BActor,56 HATTV25CActor,57 HATTV25FullActor,58)59from algorithms.algorithm.HATTV26Network import (60 HATTV26AActor,61 HATTV26BActor,62 HATTV26CActor,63 HATTV26FullActor,64)65from algorithms.algorithm.HATTV27Network import (66 HATTV27AActor,67 HATTV27BActor,68 HATTV27Actor,69)70from algorithms.algorithm.HATTV30Network import (71 HATTV30AActor,72 HATTV30BActor,73 HATTV30CActor,74)75from algorithms.algorithm.HATTV31Network import HATTV31AActor, HATTV31BActor76from algorithms.algorithm.HATTV32Network import (77 HATTV32AActor,78 HATTV32BActor,79 HATTV32CActor,80 HATTV32DActor,81)82from algorithms.algorithm.HATTV32RNetwork import (83 HATTV32RFullActor,84 HATTV32RNoMemberActor,85 HATTV32RNoTransitionActor,86)87from algorithms.algorithm.HATTV33AANetwork import HATTV33AAActor88from algorithms.algorithm.HATTV34Network import HATTV34Actor89from algorithms.algorithm.HATTGyMRT2ANetwork import HATTGyMRT2AActor90from algorithms.utils.graph_builder import build_graph91from utils.obs_utils import POLICY_OBS_FEATURE_RANKS, flatten_state_dict9293import numpy as np949596def _initialize_with_seed(factory, seed):97 """Build a module from a local RNG stream without shifting rollout RNG.9899 This is opt-in and used by schema attribution experiments. Separating the100 actor seed from the critic construction makes equal-shaped actors exactly101 paired even when their centralized-state dimensions differ.102 """103104 if seed is None or int(seed) < 0:105 return factory()106 with torch.random.fork_rng(devices=[]):107 torch.manual_seed(int(seed))108 return factory()109import torch.nn as nn110111from torch_geometric.data import Batch112from typing import Tuple113114115class GAMAPPOPolicy:116 """117 MAPPO Policy class. Wraps actor and critic networks to compute actions and value function predictions.118119 :param args: (argparse.Namespace) arguments containing relevant model and policy information.120 :param obs_space: (gym.Space) observation space.121 :param cent_obs_space: (gym.Space) value function input space (centralized input for MAPPO, decentralized for IPPO).122 :param action_space: (gym.Space) action space.123 :param device: (torch.device) specifies the device to run on (cpu/gpu).124 """125126 MODE_GAT = "gat"127 MODE_CAT = "cat"128 MODE_GATED = "gated"129 GRAPH_LEGACY = "legacy_gat"130 GRAPH_HATT = "hatt"131 GRAPH_HATT_V2_0 = "hatt_v2_0"132 GRAPH_HATT_V2_0_RAW = "hatt_v2_0_raw"133 GRAPH_HATT_V2_0_PHYSICS = "hatt_v2_0_physics"134 GRAPH_HATT_V2_0_TASKATTN = "hatt_v2_0_taskattn"135 GRAPH_HATT_V2_0_NOIDLE = "hatt_v2_0_noidle"136 GRAPH_HATT_V2_0_RESOURCE = "hatt_v2_0_resource"137 GRAPH_HATT_V2_0_SETCTX = "hatt_v2_0_setctx"138 GRAPH_HATT_V2_0_PLAIN_NOIDLE = "hatt_v2_0_plain_noidle"139 GRAPH_HATT_V2_0_DETACHED = "hatt_v2_0_detached"140 GRAPH_HATT_V2_0_SPLIT = "hatt_v2_0_split"141 GRAPH_HATT_V2_0_TYPE_EMBED = "hatt_v2_0_type_embed"142 GRAPH_HATT_V2_0_TYPE_GATE = "hatt_v2_0_type_gate"143 GRAPH_HATT_V2_1 = "hatt_v2_1"144 GRAPH_HATT_V2 = "hatt_v2"145 GRAPH_HATT_V2_3A = "hatt_v2_3a"146 GRAPH_HATT_V2_3B = "hatt_v2_3b"147 GRAPH_HATT_V2_3C = "hatt_v2_3c"148 GRAPH_HATT_V2_4 = "hatt_v2_4"149 GRAPH_HATT_V2_5A = "hatt_v2_5a"150 GRAPH_HATT_V2_5B = "hatt_v2_5b"151 GRAPH_HATT_V2_5C = "hatt_v2_5c"152 GRAPH_HATT_V2_5 = "hatt_v2_5"153 GRAPH_HATT_V2_6A = "hatt_v2_6a"154 GRAPH_HATT_V2_6B = "hatt_v2_6b"155 GRAPH_HATT_V2_6C = "hatt_v2_6c"156 GRAPH_HATT_V2_6 = "hatt_v2_6"157 GRAPH_HATT_V2_7A = "hatt_v2_7a"158 GRAPH_HATT_V2_7B = "hatt_v2_7b"159 GRAPH_HATT_V2_7 = "hatt_v2_7"160 GRAPH_HATT_V3_0A = "hatt_v3_0a"161 GRAPH_HATT_V3_0B = "hatt_v3_0b"162 GRAPH_HATT_V3_0C = "hatt_v3_0c"163 GRAPH_HATT_V3_1A = "hatt_v3_1a"164 GRAPH_HATT_V3_1B = "hatt_v3_1b"165 GRAPH_HATT_V3_2A = "hatt_v3_2a"166 GRAPH_HATT_V3_2B = "hatt_v3_2b"167 GRAPH_HATT_V3_2C = "hatt_v3_2c"168 GRAPH_HATT_V3_2D = "hatt_v3_2d"169 GRAPH_HATT_V3_2R = "hatt_v3_2r"170 GRAPH_HATT_V3_2R_NO_TRANSITION = "hatt_v3_2r_no_transition"171 GRAPH_HATT_V3_2R_NO_MEMBER = "hatt_v3_2r_no_member"172 GRAPH_HATT_V3_3AA = "hatt_v3_3aa"173 GRAPH_HATT_V3_4 = "hatt_v3_4"174 GRAPH_HATT_GYMRT2A = "hatt_gymrt2a"175176 @staticmethod177 def _is_hatt_graph_model(graph_model):178 """Return True for all structured Dict-observation HATT variants."""179 return graph_model.startswith("hatt")180181 def __init__(self, args, obs_space, cent_obs_space, act_space, device=torch.device("cpu")):182 self.device = device183 self.lr = args.lr184 self.critic_lr = args.critic_lr185 self.opti_eps = args.opti_eps186 self.weight_decay = args.weight_decay187188 # GAT 模块189 self.obs_space = obs_space190 self.share_obs_space = cent_obs_space191 self.act_space = act_space192193 self.graph_model = getattr(args, "graph_model", self.GRAPH_LEGACY)194 self.reconfiguration_mode = getattr(195 args, "reconfiguration_mode", "legacy"196 )197 # Checkpoint metadata normally prevents silently evaluating a policy198 # under a different coalition decision schedule. Causal evaluations199 # may opt in to crossing non-learned schedules (for example legacy vs200 # event), whose actor architectures are identical. Learned mode is201 # never covered because it adds trainable gate parameters.202 self.allow_reconfiguration_mode_mismatch = bool(203 getattr(args, "allow_reconfiguration_mode_mismatch", False)204 )205 self.reconfiguration_context_features = bool(206 getattr(args, "reconfiguration_context_features", False)207 )208 self.reconfiguration_gate_parameterization = getattr(209 args, "reconfiguration_gate_parameterization", "mixture"210 )211 graph_models = (212 self.GRAPH_LEGACY,213 self.GRAPH_HATT,214 self.GRAPH_HATT_V2_0,215 self.GRAPH_HATT_V2_0_RAW,216 self.GRAPH_HATT_V2_0_PHYSICS,217 self.GRAPH_HATT_V2_0_TASKATTN,218 self.GRAPH_HATT_V2_0_NOIDLE,219 self.GRAPH_HATT_V2_0_RESOURCE,220 self.GRAPH_HATT_V2_0_SETCTX,221 self.GRAPH_HATT_V2_0_PLAIN_NOIDLE,222 self.GRAPH_HATT_V2_0_DETACHED,223 self.GRAPH_HATT_V2_0_SPLIT,224 self.GRAPH_HATT_V2_0_TYPE_EMBED,225 self.GRAPH_HATT_V2_0_TYPE_GATE,226 self.GRAPH_HATT_V2_1,227 self.GRAPH_HATT_V2,228 self.GRAPH_HATT_V2_3A,229 self.GRAPH_HATT_V2_3B,230 self.GRAPH_HATT_V2_3C,231 self.GRAPH_HATT_V2_4,232 self.GRAPH_HATT_V2_5A,233 self.GRAPH_HATT_V2_5B,234 self.GRAPH_HATT_V2_5C,235 self.GRAPH_HATT_V2_5,236 self.GRAPH_HATT_V2_6A,237 self.GRAPH_HATT_V2_6B,238 self.GRAPH_HATT_V2_6C,239 self.GRAPH_HATT_V2_6,240 self.GRAPH_HATT_V2_7A,241 self.GRAPH_HATT_V2_7B,242 self.GRAPH_HATT_V2_7,243 self.GRAPH_HATT_V3_0A,244 self.GRAPH_HATT_V3_0B,245 self.GRAPH_HATT_V3_0C,246 self.GRAPH_HATT_V3_1A,247 self.GRAPH_HATT_V3_1B,248 self.GRAPH_HATT_V3_2A,249 self.GRAPH_HATT_V3_2B,250 self.GRAPH_HATT_V3_2C,251 self.GRAPH_HATT_V3_2D,252 self.GRAPH_HATT_V3_2R,253 self.GRAPH_HATT_V3_2R_NO_TRANSITION,254 self.GRAPH_HATT_V3_2R_NO_MEMBER,255 self.GRAPH_HATT_V3_3AA,256 self.GRAPH_HATT_V3_4,257 self.GRAPH_HATT_GYMRT2A,258 )259 if self.graph_model not in graph_models:260 raise ValueError(261 f"Unknown graph_model={self.graph_model!r}; expected one of {graph_models}."262 )263264 # Legacy GAT fusion mode. HATT has its own task-specific action decoder.265 self.mode = getattr(args, "gat_mode", self.MODE_GATED)266 if self.graph_model == self.GRAPH_LEGACY:267 assert self.mode in (self.MODE_GAT, self.MODE_CAT, self.MODE_GATED), \268 f"Unknown gat_mode: {self.mode!r}, expected 'gat', 'cat', or 'gated'"269270 obs_dim = obs_space.shape[0]271 share_obs_dim = cent_obs_space.shape[0]272 self.num_agents = args.num_agents273 self.num_tasks = act_space.n274 self.cap_dim = args.cap_dim275 spatial_hatt = self.graph_model == self.GRAPH_HATT_GYMRT2A276 self.ego_dim = int(args.gymrt2a_ego_dim) if spatial_hatt else self.cap_dim + 7277 self.other_agent_feat_dim = (278 int(args.gymrt2a_other_dim)279 if spatial_hatt280 else 2 + self.cap_dim + 1 + 1 + 1 + 1 + 1281 )282 others_dim = self.num_agents * self.other_agent_feat_dim283 self.per_task_feat = (obs_dim - self.ego_dim - others_dim) // self.num_tasks284 assert self.ego_dim + self.num_tasks * self.per_task_feat + others_dim == obs_dim, \285 f"Cannot infer per_task_feat from obs_dim={obs_dim}, ego_dim={self.ego_dim}, others_dim={others_dim}, num_tasks={self.num_tasks}"286 self.agent_node_dim = obs_dim287 self.graph_node_dim = 2 + self.cap_dim + 1 + 1 + 1288 self.gat_out = args.gat_out_dim # GAT 输出维度，例如 32289 self.hatt_v20_task_input_mode = str(290 getattr(args, "hatt_v20_task_input_mode", "native")291 )292 self.hatt_v20_init_scheme = str(293 getattr(args, "hatt_v20_init_scheme", "default")294 )295296 critic_input_dim = share_obs_dim297 fake_share_obs_space = spaces.Box(low=-np.inf, high=np.inf, shape=(critic_input_dim,), dtype=np.float32)298 paired_init_seed = int(getattr(args, "hatt_paired_init_seed", -1))299 self.critic = _initialize_with_seed(300 lambda: R_Critic(args, fake_share_obs_space, self.device),301 paired_init_seed,302 )303304 if self._is_hatt_graph_model(self.graph_model):305 actor_class = {306 self.GRAPH_HATT: HATTActor,307 self.GRAPH_HATT_V2_0: HATTV20Actor,308 self.GRAPH_HATT_V2_0_RAW: HATTV20RawActor,309 self.GRAPH_HATT_V2_0_PHYSICS: HATTV20PhysicsActor,310 self.GRAPH_HATT_V2_0_TASKATTN: HATTV20TaskAttentionActor,311 self.GRAPH_HATT_V2_0_NOIDLE: HATTV20NoIdleActor,312 self.GRAPH_HATT_V2_0_RESOURCE: HATTV20ResourceActor,313 self.GRAPH_HATT_V2_0_SETCTX: HATTV20SetContextActor,314 self.GRAPH_HATT_V2_0_PLAIN_NOIDLE: HATTV20PlainNoIdleActor,315 self.GRAPH_HATT_V2_0_DETACHED: HATTV20SpecialDetachedActor,316 self.GRAPH_HATT_V2_0_SPLIT: HATTV20SpecialSplitActor,317 self.GRAPH_HATT_V2_0_TYPE_EMBED: HATTV20TypeEmbeddingActor,318 self.GRAPH_HATT_V2_0_TYPE_GATE: HATTV20TypeGateActor,319 self.GRAPH_HATT_V2_1: HATTV21Actor,320 self.GRAPH_HATT_V2: HATTV22Actor,321 self.GRAPH_HATT_V2_3A: HATTV23AActor,322 self.GRAPH_HATT_V2_3B: HATTV23BActor,323 self.GRAPH_HATT_V2_3C: HATTV23CActor,324 self.GRAPH_HATT_V2_4: HATTV24Actor,325 self.GRAPH_HATT_V2_5A: HATTV25AActor,326 self.GRAPH_HATT_V2_5B: HATTV25BActor,327 self.GRAPH_HATT_V2_5C: HATTV25CActor,328 self.GRAPH_HATT_V2_5: HATTV25FullActor,329 self.GRAPH_HATT_V2_6A: HATTV26AActor,330 self.GRAPH_HATT_V2_6B: HATTV26BActor,331 self.GRAPH_HATT_V2_6C: HATTV26CActor,332 self.GRAPH_HATT_V2_6: HATTV26FullActor,333 self.GRAPH_HATT_V2_7A: HATTV27AActor,334 self.GRAPH_HATT_V2_7B: HATTV27BActor,335 self.GRAPH_HATT_V2_7: HATTV27Actor,336 self.GRAPH_HATT_V3_0A: HATTV30AActor,337 self.GRAPH_HATT_V3_0B: HATTV30BActor,338 self.GRAPH_HATT_V3_0C: HATTV30CActor,339 self.GRAPH_HATT_V3_1A: HATTV31AActor,340 self.GRAPH_HATT_V3_1B: HATTV31BActor,341 self.GRAPH_HATT_V3_2A: HATTV32AActor,342 self.GRAPH_HATT_V3_2B: HATTV32BActor,343 self.GRAPH_HATT_V3_2C: HATTV32CActor,344 self.GRAPH_HATT_V3_2D: HATTV32DActor,345 self.GRAPH_HATT_V3_2R: HATTV32RFullActor,346 self.GRAPH_HATT_V3_2R_NO_TRANSITION: HATTV32RNoTransitionActor,347 self.GRAPH_HATT_V3_2R_NO_MEMBER: HATTV32RNoMemberActor,348 self.GRAPH_HATT_V3_3AA: HATTV33AAActor,349 self.GRAPH_HATT_V3_4: HATTV34Actor,350 self.GRAPH_HATT_GYMRT2A: HATTGyMRT2AActor,351 }[self.graph_model]352 self.actor = _initialize_with_seed(353 lambda: actor_class(354 args=args,355 ego_dim=self.ego_dim,356 other_dim=self.other_agent_feat_dim,357 task_dim=self.per_task_feat,358 num_agents=self.num_agents,359 num_tasks=self.num_tasks,360 device=self.device,361 ),362 paired_init_seed + 1 if paired_init_seed >= 0 else -1,363 )364 if self.hatt_v20_init_scheme == "compact_function_preserving":365 if self.graph_model != self.GRAPH_HATT_V2_0:366 raise ValueError(367 "compact_function_preserving initialization is limited "368 "to graph_model=hatt_v2_0"369 )370 if self.hatt_v20_task_input_mode != "restore_req_scale":371 raise ValueError(372 "compact_function_preserving initialization requires "373 "hatt_v20_task_input_mode=restore_req_scale"374 )375 if paired_init_seed < 0:376 raise ValueError(377 "compact_function_preserving initialization requires "378 "hatt_paired_init_seed"379 )380 compact_args = copy.copy(args)381 compact_args.hatt_v20_task_input_mode = "native"382 compact_reference = _initialize_with_seed(383 lambda: HATTV20Actor(384 args=compact_args,385 ego_dim=self.ego_dim,386 other_dim=self.other_agent_feat_dim,387 task_dim=6,388 num_agents=self.num_agents,389 num_tasks=self.num_tasks,390 device=self.device,391 ),392 paired_init_seed + 1,393 )394 initialize_full_actor_from_compact_reference(395 self.actor, compact_reference, self.cap_dim396 )397 actor_params = list(self.actor.parameters())398 else:399 self.gat_encoder = GATNetwork(400 node_dim=self.graph_node_dim,401 edge_dim=self.graph_node_dim,402 hidden_dim=args.gat_hidden_dim,403 output_dim=self.gat_out404 ).to(self.device)405 if self.mode in (self.MODE_CAT, self.MODE_GATED):406 actor_input_dim = obs_dim + self.gat_out407 else:408 actor_input_dim = self.gat_out409 if self.mode == self.MODE_GATED:410 self.actor_gat_gate = nn.Parameter(torch.zeros((), device=self.device))411 fake_obs_space = spaces.Box(412 low=-np.inf, high=np.inf, shape=(actor_input_dim,), dtype=np.float32413 )414 self.actor = R_Actor(args, fake_obs_space, self.act_space, self.device)415 actor_params = list(self.actor.parameters()) + list(self.gat_encoder.parameters())416 if self.mode == self.MODE_GATED:417 actor_params += [self.actor_gat_gate]418419 self._actor_params = actor_params420 critic_params = list(self.critic.parameters())421422 self.actor_optimizer = torch.optim.Adam(actor_params,423 lr=self.lr, eps=self.opti_eps,424 weight_decay=self.weight_decay)425 self.critic_optimizer = torch.optim.Adam(critic_params,426 lr=self.critic_lr,427 eps=self.opti_eps,428 weight_decay=self.weight_decay)429430 def lr_decay(self, episode, episodes):431 """432 Decay the actor and critic learning rates.433 :param episode: (int) current training episode.434 :param episodes: (int) total number of training episodes.435 """436 update_linear_schedule(self.actor_optimizer, episode, episodes, self.lr)437 update_linear_schedule(self.critic_optimizer, episode, episodes, self.critic_lr)438439 def _reshape_batch(self, batch: np.ndarray, unit_size):440 """Reshape a batch that may be either [B, unit_size, dim] or [B*unit_size, dim]."""441 if isinstance(batch, torch.Tensor):442 batch = batch.cpu().numpy()443444 if batch.ndim == 3 and batch.shape[1] == unit_size:445 return batch446447 if batch.ndim == 2 and batch.shape[0] % unit_size == 0:448 return batch.reshape(batch.shape[0] // unit_size, unit_size, -1)449450 raise ValueError(451 f"Cannot reshape batch with shape {batch.shape} using unit_size={unit_size}."452 )453454 def _reshape_obs_dict(self, obs, unit_size):455 if not isinstance(obs, dict):456 return self._reshape_batch(obs, unit_size)457458 out = {}459 for key, value in obs.items():460 if isinstance(value, torch.Tensor):461 value = value.detach().cpu().numpy()462 feature_rank = POLICY_OBS_FEATURE_RANKS[key]463 if value.ndim == feature_rank + 2 and value.shape[1] == unit_size:464 out[key] = value465 elif value.shape[0] % unit_size == 0:466 out[key] = value.reshape(value.shape[0] // unit_size, unit_size, *value.shape[1:])467 else:468 raise ValueError(469 f"Cannot reshape obs[{key!r}] with shape {value.shape} using unit_size={unit_size}."470 )471 return out472473 def _to_tensor_obs(self, obs):474 if not isinstance(obs, dict):475 return torch.tensor(obs, dtype=torch.float32, device=self.device)476 return {477 k: torch.as_tensor(v, dtype=torch.float32, device=self.device)478 for k, v in obs.items()479 }480481 def _flatten_grouped_obs(self, obs):482 out = {}483 for k, v in obs.items():484 feature_rank = POLICY_OBS_FEATURE_RANKS[k]485 if v.ndim == feature_rank + 2 and v.shape[1] == self.num_agents:486 out[k] = v.reshape(v.shape[0] * v.shape[1], *v.shape[2:])487 else:488 out[k] = v489 return out490491 def _flat_to_grouped_obs(self, obs_flat):492 T, n_agents, _ = obs_flat.shape493 task_start = self.ego_dim494 task_end = task_start + self.num_tasks * self.per_task_feat495 return {496 "ego": obs_flat[..., :self.ego_dim],497 "tasks": obs_flat[..., task_start:task_end].view(498 T, n_agents, self.num_tasks, self.per_task_feat499 ),500 "others": obs_flat[..., task_end:].view(501 T, n_agents, n_agents, self.other_agent_feat_dim502 ),503 }504505 def _build_graph_node_features(self, obs):506 ego = obs["ego"]507 task_obs = obs["tasks"]508 others = obs["others"]509 T, n_agents, _ = ego.shape510511 # others[observer, target] stores target-agent state, with the observer's512 # own row zeroed. Aggregate over observers to recover one node per agent.513 valid_other = (others.abs().sum(dim=-1, keepdim=True) > 0).float()514 other_counts = valid_other.sum(dim=1).clamp_min(1.0)515 other_agent_state = (others * valid_other).sum(dim=1) / other_counts516517 agent_nodes = torch.zeros(518 (T, n_agents, self.graph_node_dim), dtype=torch.float32, device=self.device519 )520 resource_start = 2521 resource_end = resource_start + self.cap_dim522 amount_idx = resource_end523 type_idx = amount_idx + 1524 active_idx = type_idx + 1525526 agent_nodes[..., 0:2] = ego[..., 0:2] # position527 agent_nodes[..., resource_start:resource_end] = ego[..., 2:2 + self.cap_dim]528 agent_nodes[..., amount_idx] = ego[..., self.cap_dim + 4] # load ratio529 agent_nodes[..., type_idx] = other_agent_state[..., self.cap_dim + 5]530 agent_nodes[..., active_idx] = other_agent_state[..., self.cap_dim + 6]531532 # Task features are global for the timestep; average the duplicated533 # per-agent observations.534 task_state = task_obs.mean(dim=1)535 task_nodes = torch.zeros(536 (T, self.num_tasks, self.graph_node_dim), dtype=torch.float32, device=self.device537 )538 task_nodes[..., 0:2] = task_state[..., 0:2] # position539 task_nodes[..., resource_start:resource_end] = task_state[..., 2:2 + self.cap_dim]540 task_nodes[..., amount_idx] = task_state[..., 2 + self.cap_dim]541 task_nodes[..., type_idx] = task_state[..., 3 + self.cap_dim]542 task_nodes[..., active_idx] = (task_state.abs().sum(dim=-1) > 0).float()543544 return agent_nodes, task_nodes545546 def _gat_augment(self, obs):547 """548 使用 GAT 对每个时间步的 agent/raw task node 特征建图并编码。549550 参数551 ----552 obs : dict[str, Tensor]553 每个时间步、每个智能体的结构化观测。554 返回555 ----556 agent_emb_batch : Tensor [T, N_agents, gat_out_dim]557 每个时间步、每个智能体对应的 GAT 节点嵌入。558 global_pool_batch: Tensor [T, gat_out_dim]559 每个时间步整张图的全局池化向量（所有节点平均）。560 """561 T, n_agents, _ = obs["ego"].shape562 agent_nodes, task_nodes = self._build_graph_node_features(obs)563564 data_list = []565 for t in range(T):566 g = build_graph(agent_nodes[t], task_nodes[t])567 data_list.append(g)568569 batch = Batch.from_data_list(data_list)570571 node_emb_all = self.gat_encoder(batch.x, batch.edge_index, batch.edge_attr)572573 N = n_agents + self.num_tasks574 node_emb_split = node_emb_all.view(T, N, -1)575576 agent_emb_batch = node_emb_split[:, :n_agents, :]577 global_pool_batch = node_emb_split.mean(dim=1)578579 return agent_emb_batch, global_pool_batch580581 # ---- 输入构造辅助方法 ----582583 def _build_actor_input(self, obs, agent_emb: torch.Tensor):584 """根据模式构造 Actor 的输入张量。"""585 obs_flat_dict = self._flatten_grouped_obs(obs)586 agent_emb_flat = agent_emb.reshape(-1, self.gat_out) # [T*N, gat_out]587588 if self.mode == self.MODE_GAT:589 return {"gat": agent_emb_flat}590 elif self.mode == self.MODE_CAT:591 obs_aug = dict(obs_flat_dict)592 obs_aug["gat"] = agent_emb_flat593 return obs_aug594 else: # MODE_GATED595 alpha = torch.tanh(self.actor_gat_gate)596 obs_aug = dict(obs_flat_dict)597 obs_aug["gat"] = alpha * agent_emb_flat598 return obs_aug599600 @staticmethod601 def _build_critic_input(cent_obs):602 cent_obs = flatten_state_dict(cent_obs)603 if getattr(cent_obs, "ndim", 0) > 2:604 cent_obs = cent_obs.reshape(-1, cent_obs.shape[-1])605 return cent_obs606607 def get_gat_gate_values(self) -> Tuple[float, float]:608 """Return the actor gate and a zero compatibility value for the removed critic gate."""609 if self.graph_model != self.GRAPH_LEGACY or self.mode != self.MODE_GATED:610 return 0.0, 0.0611 actor_alpha = torch.tanh(self.actor_gat_gate).detach().cpu().item()612 return float(actor_alpha), 0.0613614 def get_actor_parameters(self):615 return self._actor_params616617 def get_actor_checkpoint(self):618 checkpoint = {619 "graph_model": self.graph_model,620 "reconfiguration_mode": self.reconfiguration_mode,621 "reconfiguration_context_features": self.reconfiguration_context_features,622 "reconfiguration_gate_parameterization": self.reconfiguration_gate_parameterization,623 "hatt_v20_task_input_mode": self.hatt_v20_task_input_mode,624 "hatt_v20_init_scheme": self.hatt_v20_init_scheme,625 "actor": self.actor.state_dict(),626 }627 if self.graph_model == self.GRAPH_LEGACY:628 checkpoint["gat_encoder"] = self.gat_encoder.state_dict()629 if self.mode == self.MODE_GATED:630 checkpoint["actor_gat_gate"] = self.actor_gat_gate.detach().cpu()631 return checkpoint632633 def load_actor_checkpoint(self, checkpoint):634 if "actor" not in checkpoint:635 # Backward compatibility with legacy actor-only checkpoints.636 self.actor.load_state_dict(checkpoint)637 return638 saved_model = checkpoint.get("graph_model", self.GRAPH_LEGACY)639 compatible_v20_checkpoint = (640 (641 self.graph_model == self.GRAPH_HATT_V2_0642 and saved_model == self.GRAPH_HATT_V2643 )644 or (645 self.graph_model == self.GRAPH_HATT_V2_0_NOIDLE646 and saved_model == self.GRAPH_HATT_V2_0_PHYSICS647 )648 )649 if saved_model != self.graph_model and not compatible_v20_checkpoint:650 raise ValueError(651 f"Checkpoint graph_model={saved_model!r} does not match current {self.graph_model!r}."652 )653 saved_reconfiguration_mode = checkpoint.get(654 "reconfiguration_mode", "legacy"655 )656 mode_mismatch = saved_reconfiguration_mode != self.reconfiguration_mode657 safe_nonlearned_intervention = (658 self.allow_reconfiguration_mode_mismatch659 and "learned" not in (660 saved_reconfiguration_mode,661 self.reconfiguration_mode,662 )663 )664 if mode_mismatch and not safe_nonlearned_intervention:665 raise ValueError(666 "Checkpoint reconfiguration_mode="667 f"{saved_reconfiguration_mode!r} does not match current "668 f"{self.reconfiguration_mode!r}."669 )670 saved_context_features = bool(671 checkpoint.get("reconfiguration_context_features", False)672 )673 if saved_context_features != self.reconfiguration_context_features:674 raise ValueError(675 "Checkpoint reconfiguration_context_features="676 f"{saved_context_features!r} does not match current "677 f"{self.reconfiguration_context_features!r}."678 )679 saved_gate_parameterization = checkpoint.get(680 "reconfiguration_gate_parameterization", "mixture"681 )682 if saved_gate_parameterization != self.reconfiguration_gate_parameterization:683 raise ValueError(684 "Checkpoint reconfiguration_gate_parameterization="685 f"{saved_gate_parameterization!r} does not match current "686 f"{self.reconfiguration_gate_parameterization!r}."687 )688 saved_task_input_mode = checkpoint.get("hatt_v20_task_input_mode", "native")689 if saved_task_input_mode != self.hatt_v20_task_input_mode:690 raise ValueError(691 "Checkpoint hatt_v20_task_input_mode="692 f"{saved_task_input_mode!r} does not match current "693 f"{self.hatt_v20_task_input_mode!r}."694 )695 self.actor.load_state_dict(checkpoint["actor"])696 if self.graph_model == self.GRAPH_LEGACY and "gat_encoder" in checkpoint:697 self.gat_encoder.load_state_dict(checkpoint["gat_encoder"])698 if (699 self.graph_model == self.GRAPH_LEGACY700 and self.mode == self.MODE_GATED701 and "actor_gat_gate" in checkpoint702 ):703 self.actor_gat_gate.data.copy_(704 checkpoint["actor_gat_gate"].to(self.device)705 )706707 def get_utility_loss(self, obs, available_actions):708 """Return auxiliary utility head loss (zero if actor has no utility head)."""709 if not hasattr(self.actor, "get_utility_loss"):710 return torch.tensor(0.0, device=self.device)711 return self.actor.get_utility_loss(obs, available_actions)712713 def get_graph_diagnostics(self):714 if not self._is_hatt_graph_model(self.graph_model):715 return {}716 result = {}717 prefix = self.graph_model718 for name, value in self.actor.last_diagnostics.items():719 if name.endswith("_attention"):720 probabilities = value.clamp_min(1e-8)721 metric = -(value * probabilities.log()).sum(dim=-2).mean()722 name = f"{name}_entropy"723 else:724 metric = value.mean()725 result[f"{prefix}/{name}"] = float(metric.cpu())726 return result727728 def get_reconfiguration_gate_regularization(self):729 """返回 actor 最近一次 evaluate_actions 产生的 gate 正则项。730731 非 hierarchical gate 或未启用 learned reconfiguration 时返回 0，732 使 trainer 可以无条件把该项接入总 loss。733 """734735 value = getattr(self.actor, "last_gate_regularization", None)736 if value is None:737 return torch.zeros((), device=self.device)738 return value739740 def get_actions(self, cent_obs, obs, rnn_states_actor, rnn_states_critic, masks, available_actions=None,741 deterministic=False):742 """743 Compute actions and value function predictions for the given inputs.744 :param cent_obs (np.ndarray): centralized input to the critic.745 :param obs (np.ndarray): local agent inputs to the actor.746 :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.747 :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.748 :param masks: (np.ndarray) denotes points at which RNN states should be reset.749 :param available_actions: (np.ndarray) denotes which actions are available to agent750 (if None, all actions available)751 :param deterministic: (bool) whether the action should be mode of distribution or should be sampled.752753 :return values: (torch.Tensor) value function predictions.754 :return actions: (torch.Tensor) actions to take.755 :return action_log_probs: (torch.Tensor) log probabilities of chosen actions.756 :return rnn_states_actor: (torch.Tensor) updated actor network RNN states.757 :return rnn_states_critic: (torch.Tensor) updated critic network RNN states.758 """759760 obs_raw = self._reshape_obs_dict(obs, self.num_agents)761 obs_t = self._to_tensor_obs(obs_raw)762763 critic_input = self._build_critic_input(cent_obs)764 if self._is_hatt_graph_model(self.graph_model):765 actor_input = obs_t766 else:767 agent_emb, _ = self._gat_augment(obs_t)768 actor_input = self._build_actor_input(obs_t, agent_emb)769770 actions, action_log_probs, rnn_states_actor = self.actor(actor_input,771 rnn_states_actor,772 masks,773 available_actions,774 deterministic)775776 values, rnn_states_critic = self.critic(critic_input, rnn_states_critic, masks)777 gate_actions = getattr(self.actor, "last_gate_actions", None)778 if gate_actions is None:779 # 普通 Categorical、mixture、switch 都没有显式 gate_action；780 # buffer 仍存一个零占位，保持所有 graph_model 的 batch 结构一致。781 gate_actions = torch.zeros_like(actions, dtype=torch.float32)782 return (783 values,784 actions,785 action_log_probs,786 gate_actions,787 rnn_states_actor,788 rnn_states_critic,789 )790791 def get_values(self, cent_obs, rnn_states_critic, masks):792 """793 Get value function predictions.794 :param cent_obs (np.ndarray): centralized input to the critic.795 :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.796 :param masks: (np.ndarray) denotes points at which RNN states should be reset.797798 :return values: (torch.Tensor) value function predictions.799 """800 values, _ = self.critic(self._build_critic_input(cent_obs), rnn_states_critic, masks)801 return values802803 def evaluate_actions(self, cent_obs, obs, rnn_states_actor, rnn_states_critic, action, masks,804 available_actions=None, active_masks=None, gate_actions=None):805 """806 Get action logprobs / entropy and value function predictions for actor update.807 :param cent_obs (np.ndarray): centralized input to the critic.808 :param obs (np.ndarray): local agent inputs to the actor.809 :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.810 :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.811 :param action: (np.ndarray) actions whose log probabilites and entropy to compute.812 :param masks: (np.ndarray) denotes points at which RNN states should be reset.813 :param available_actions: (np.ndarray) denotes which actions are available to agent814 (if None, all actions available)815 :param active_masks: (torch.Tensor) denotes whether an agent is active or dead.816 :param gate_actions: hierarchical gate rollout decisions, needed to replay817 the joint action logprob during PPO update.818819 :return values: (torch.Tensor) value function predictions.820 :return action_log_probs: (torch.Tensor) log probabilities of the input actions.821 :return dist_entropy: (torch.Tensor) action distribution entropy for the given inputs.822 """823824 obs_raw = self._reshape_obs_dict(obs, self.num_agents)825 obs_t = self._to_tensor_obs(obs_raw)826827 critic_input = self._build_critic_input(cent_obs)828 if self._is_hatt_graph_model(self.graph_model):829 actor_input = obs_t830 else:831 agent_emb, _ = self._gat_augment(obs_t)832 actor_input = self._build_actor_input(obs_t, agent_emb)833834 action_log_probs, dist_entropy = self.actor.evaluate_actions(actor_input,835 rnn_states_actor,836 action,837 masks,838 available_actions,839 active_masks,840 gate_actions)841842 values, _ = self.critic(critic_input, rnn_states_critic, masks)843 return values, action_log_probs, dist_entropy844845 def act(self, obs, rnn_states_actor, masks, available_actions=None, deterministic=False):846 """847 Compute actions using the given inputs.848 :param obs (np.ndarray): local agent inputs to the actor.849 :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.850 :param masks: (np.ndarray) denotes points at which RNN states should be reset.851 :param available_actions: (np.ndarray) denotes which actions are available to agent852 (if None, all actions available)853 :param deterministic: (bool) whether the action should be mode of distribution or should be sampled.854 """855 obs_t = self._to_tensor_obs(856 self._reshape_obs_dict(obs, self.num_agents)857 )858 if self._is_hatt_graph_model(self.graph_model):859 actor_input = obs_t860 else:861 agent_emb, _ = self._gat_augment(obs_t)862 actor_input = self._build_actor_input(obs_t, agent_emb)863 actions, _, rnn_states_actor = self.actor(864 actor_input,865 rnn_states_actor,866 masks,867 available_actions,868 deterministic,869 )870 return actions, rnn_states_actor