""" # @Time : 2025/12/07 10:30 # @Author : wingong # @Email : wy1997331@126.com # @File : gaRMAPPOPolicy.py # # graph_model="legacy_gat" 保留三种旧 GAT Actor 融合模式： # "gat" — Actor 输入 = 仅 GAT 嵌入（信息瓶颈） # "cat" — Actor 输入 = cat(obs, GAT 嵌入) # "gated" — Actor 输入 = cat(obs, tanh(α)*GAT 嵌入)，α 可学习 # graph_model="hatt" 使用结构化 Dict 观测、关系专属注意力和任务动作解码器。 # Critic 始终使用与 NoGAT 相同的 centralized state。 """ import copy import torch from gym import spaces from algorithms.algorithm.r_actor_critic import R_Actor, R_Critic from utils.util import update_linear_schedule from algorithms.algorithm.GATNetwork import GATNetwork from algorithms.algorithm.HATTV10Network import HATTActor from algorithms.algorithm.HATTV20Network import ( HATTV20Actor, initialize_full_actor_from_compact_reference, ) from algorithms.algorithm.HATTV20InfoNetwork import ( HATTV20NoIdleActor, HATTV20PhysicsActor, HATTV20RawActor, HATTV20ResourceActor, HATTV20SetContextActor, HATTV20TaskAttentionActor, ) from algorithms.algorithm.HATTV20SpecialNetwork import ( HATTV20PlainNoIdleActor, HATTV20SpecialDetachedActor, HATTV20SpecialSplitActor, ) from algorithms.algorithm.HATTV20TypeConditionedNetwork import ( HATTV20TypeEmbeddingActor, HATTV20TypeGateActor, ) from algorithms.algorithm.HATTV21Network import HATTV2Actor as HATTV21Actor from algorithms.algorithm.HATTV22Network import HATTV2Actor as HATTV22Actor from algorithms.algorithm.HATTV23Network import ( HATTV23AActor, HATTV23BActor, HATTV23CActor, ) from algorithms.algorithm.HATTV24Network import HATTV24Actor from algorithms.algorithm.HATTV25Network import ( HATTV25AActor, HATTV25BActor, HATTV25CActor, HATTV25FullActor, ) from algorithms.algorithm.HATTV26Network import ( HATTV26AActor, HATTV26BActor, HATTV26CActor, HATTV26FullActor, ) from algorithms.algorithm.HATTV27Network import ( HATTV27AActor, HATTV27BActor, HATTV27Actor, ) from algorithms.algorithm.HATTV30Network import ( HATTV30AActor, HATTV30BActor, HATTV30CActor, ) from algorithms.algorithm.HATTV31Network import HATTV31AActor, HATTV31BActor from algorithms.algorithm.HATTV32Network import ( HATTV32AActor, HATTV32BActor, HATTV32CActor, HATTV32DActor, ) from algorithms.algorithm.HATTV32RNetwork import ( HATTV32RFullActor, HATTV32RNoMemberActor, HATTV32RNoTransitionActor, ) from algorithms.algorithm.HATTV33AANetwork import HATTV33AAActor from algorithms.algorithm.HATTV34Network import HATTV34Actor from algorithms.algorithm.HATTGyMRT2ANetwork import HATTGyMRT2AActor from algorithms.utils.graph_builder import build_graph from utils.obs_utils import POLICY_OBS_FEATURE_RANKS, flatten_state_dict import numpy as np def _initialize_with_seed(factory, seed): """Build a module from a local RNG stream without shifting rollout RNG. This is opt-in and used by schema attribution experiments. Separating the actor seed from the critic construction makes equal-shaped actors exactly paired even when their centralized-state dimensions differ. """ if seed is None or int(seed) < 0: return factory() with torch.random.fork_rng(devices=[]): torch.manual_seed(int(seed)) return factory() import torch.nn as nn from torch_geometric.data import Batch from typing import Tuple class GAMAPPOPolicy: """ MAPPO Policy class. Wraps actor and critic networks to compute actions and value function predictions. :param args: (argparse.Namespace) arguments containing relevant model and policy information. :param obs_space: (gym.Space) observation space. :param cent_obs_space: (gym.Space) value function input space (centralized input for MAPPO, decentralized for IPPO). :param action_space: (gym.Space) action space. :param device: (torch.device) specifies the device to run on (cpu/gpu). """ MODE_GAT = "gat" MODE_CAT = "cat" MODE_GATED = "gated" GRAPH_LEGACY = "legacy_gat" GRAPH_HATT = "hatt" GRAPH_HATT_V2_0 = "hatt_v2_0" GRAPH_HATT_V2_0_RAW = "hatt_v2_0_raw" GRAPH_HATT_V2_0_PHYSICS = "hatt_v2_0_physics" GRAPH_HATT_V2_0_TASKATTN = "hatt_v2_0_taskattn" GRAPH_HATT_V2_0_NOIDLE = "hatt_v2_0_noidle" GRAPH_HATT_V2_0_RESOURCE = "hatt_v2_0_resource" GRAPH_HATT_V2_0_SETCTX = "hatt_v2_0_setctx" GRAPH_HATT_V2_0_PLAIN_NOIDLE = "hatt_v2_0_plain_noidle" GRAPH_HATT_V2_0_DETACHED = "hatt_v2_0_detached" GRAPH_HATT_V2_0_SPLIT = "hatt_v2_0_split" GRAPH_HATT_V2_0_TYPE_EMBED = "hatt_v2_0_type_embed" GRAPH_HATT_V2_0_TYPE_GATE = "hatt_v2_0_type_gate" GRAPH_HATT_V2_1 = "hatt_v2_1" GRAPH_HATT_V2 = "hatt_v2" GRAPH_HATT_V2_3A = "hatt_v2_3a" GRAPH_HATT_V2_3B = "hatt_v2_3b" GRAPH_HATT_V2_3C = "hatt_v2_3c" GRAPH_HATT_V2_4 = "hatt_v2_4" GRAPH_HATT_V2_5A = "hatt_v2_5a" GRAPH_HATT_V2_5B = "hatt_v2_5b" GRAPH_HATT_V2_5C = "hatt_v2_5c" GRAPH_HATT_V2_5 = "hatt_v2_5" GRAPH_HATT_V2_6A = "hatt_v2_6a" GRAPH_HATT_V2_6B = "hatt_v2_6b" GRAPH_HATT_V2_6C = "hatt_v2_6c" GRAPH_HATT_V2_6 = "hatt_v2_6" GRAPH_HATT_V2_7A = "hatt_v2_7a" GRAPH_HATT_V2_7B = "hatt_v2_7b" GRAPH_HATT_V2_7 = "hatt_v2_7" GRAPH_HATT_V3_0A = "hatt_v3_0a" GRAPH_HATT_V3_0B = "hatt_v3_0b" GRAPH_HATT_V3_0C = "hatt_v3_0c" GRAPH_HATT_V3_1A = "hatt_v3_1a" GRAPH_HATT_V3_1B = "hatt_v3_1b" GRAPH_HATT_V3_2A = "hatt_v3_2a" GRAPH_HATT_V3_2B = "hatt_v3_2b" GRAPH_HATT_V3_2C = "hatt_v3_2c" GRAPH_HATT_V3_2D = "hatt_v3_2d" GRAPH_HATT_V3_2R = "hatt_v3_2r" GRAPH_HATT_V3_2R_NO_TRANSITION = "hatt_v3_2r_no_transition" GRAPH_HATT_V3_2R_NO_MEMBER = "hatt_v3_2r_no_member" GRAPH_HATT_V3_3AA = "hatt_v3_3aa" GRAPH_HATT_V3_4 = "hatt_v3_4" GRAPH_HATT_GYMRT2A = "hatt_gymrt2a" @staticmethod def _is_hatt_graph_model(graph_model): """Return True for all structured Dict-observation HATT variants.""" return graph_model.startswith("hatt") def __init__(self, args, obs_space, cent_obs_space, act_space, device=torch.device("cpu")): self.device = device self.lr = args.lr self.critic_lr = args.critic_lr self.opti_eps = args.opti_eps self.weight_decay = args.weight_decay # GAT 模块 self.obs_space = obs_space self.share_obs_space = cent_obs_space self.act_space = act_space self.graph_model = getattr(args, "graph_model", self.GRAPH_LEGACY) self.reconfiguration_mode = getattr( args, "reconfiguration_mode", "legacy" ) # Checkpoint metadata normally prevents silently evaluating a policy # under a different coalition decision schedule. Causal evaluations # may opt in to crossing non-learned schedules (for example legacy vs # event), whose actor architectures are identical. Learned mode is # never covered because it adds trainable gate parameters. self.allow_reconfiguration_mode_mismatch = bool( getattr(args, "allow_reconfiguration_mode_mismatch", False) ) self.reconfiguration_context_features = bool( getattr(args, "reconfiguration_context_features", False) ) self.reconfiguration_gate_parameterization = getattr( args, "reconfiguration_gate_parameterization", "mixture" ) graph_models = ( self.GRAPH_LEGACY, self.GRAPH_HATT, self.GRAPH_HATT_V2_0, self.GRAPH_HATT_V2_0_RAW, self.GRAPH_HATT_V2_0_PHYSICS, self.GRAPH_HATT_V2_0_TASKATTN, self.GRAPH_HATT_V2_0_NOIDLE, self.GRAPH_HATT_V2_0_RESOURCE, self.GRAPH_HATT_V2_0_SETCTX, self.GRAPH_HATT_V2_0_PLAIN_NOIDLE, self.GRAPH_HATT_V2_0_DETACHED, self.GRAPH_HATT_V2_0_SPLIT, self.GRAPH_HATT_V2_0_TYPE_EMBED, self.GRAPH_HATT_V2_0_TYPE_GATE, self.GRAPH_HATT_V2_1, self.GRAPH_HATT_V2, self.GRAPH_HATT_V2_3A, self.GRAPH_HATT_V2_3B, self.GRAPH_HATT_V2_3C, self.GRAPH_HATT_V2_4, self.GRAPH_HATT_V2_5A, self.GRAPH_HATT_V2_5B, self.GRAPH_HATT_V2_5C, self.GRAPH_HATT_V2_5, self.GRAPH_HATT_V2_6A, self.GRAPH_HATT_V2_6B, self.GRAPH_HATT_V2_6C, self.GRAPH_HATT_V2_6, self.GRAPH_HATT_V2_7A, self.GRAPH_HATT_V2_7B, self.GRAPH_HATT_V2_7, self.GRAPH_HATT_V3_0A, self.GRAPH_HATT_V3_0B, self.GRAPH_HATT_V3_0C, self.GRAPH_HATT_V3_1A, self.GRAPH_HATT_V3_1B, self.GRAPH_HATT_V3_2A, self.GRAPH_HATT_V3_2B, self.GRAPH_HATT_V3_2C, self.GRAPH_HATT_V3_2D, self.GRAPH_HATT_V3_2R, self.GRAPH_HATT_V3_2R_NO_TRANSITION, self.GRAPH_HATT_V3_2R_NO_MEMBER, self.GRAPH_HATT_V3_3AA, self.GRAPH_HATT_V3_4, self.GRAPH_HATT_GYMRT2A, ) if self.graph_model not in graph_models: raise ValueError( f"Unknown graph_model={self.graph_model!r}; expected one of {graph_models}." ) # Legacy GAT fusion mode. HATT has its own task-specific action decoder. self.mode = getattr(args, "gat_mode", self.MODE_GATED) if self.graph_model == self.GRAPH_LEGACY: assert self.mode in (self.MODE_GAT, self.MODE_CAT, self.MODE_GATED), \ f"Unknown gat_mode: {self.mode!r}, expected 'gat', 'cat', or 'gated'" obs_dim = obs_space.shape[0] share_obs_dim = cent_obs_space.shape[0] self.num_agents = args.num_agents self.num_tasks = act_space.n self.cap_dim = args.cap_dim spatial_hatt = self.graph_model == self.GRAPH_HATT_GYMRT2A self.ego_dim = int(args.gymrt2a_ego_dim) if spatial_hatt else self.cap_dim + 7 self.other_agent_feat_dim = ( int(args.gymrt2a_other_dim) if spatial_hatt else 2 + self.cap_dim + 1 + 1 + 1 + 1 + 1 ) others_dim = self.num_agents * self.other_agent_feat_dim self.per_task_feat = (obs_dim - self.ego_dim - others_dim) // self.num_tasks assert self.ego_dim + self.num_tasks * self.per_task_feat + others_dim == obs_dim, \ 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}" self.agent_node_dim = obs_dim self.graph_node_dim = 2 + self.cap_dim + 1 + 1 + 1 self.gat_out = args.gat_out_dim # GAT 输出维度，例如 32 self.hatt_v20_task_input_mode = str( getattr(args, "hatt_v20_task_input_mode", "native") ) self.hatt_v20_init_scheme = str( getattr(args, "hatt_v20_init_scheme", "default") ) critic_input_dim = share_obs_dim fake_share_obs_space = spaces.Box(low=-np.inf, high=np.inf, shape=(critic_input_dim,), dtype=np.float32) paired_init_seed = int(getattr(args, "hatt_paired_init_seed", -1)) self.critic = _initialize_with_seed( lambda: R_Critic(args, fake_share_obs_space, self.device), paired_init_seed, ) if self._is_hatt_graph_model(self.graph_model): actor_class = { self.GRAPH_HATT: HATTActor, self.GRAPH_HATT_V2_0: HATTV20Actor, self.GRAPH_HATT_V2_0_RAW: HATTV20RawActor, self.GRAPH_HATT_V2_0_PHYSICS: HATTV20PhysicsActor, self.GRAPH_HATT_V2_0_TASKATTN: HATTV20TaskAttentionActor, self.GRAPH_HATT_V2_0_NOIDLE: HATTV20NoIdleActor, self.GRAPH_HATT_V2_0_RESOURCE: HATTV20ResourceActor, self.GRAPH_HATT_V2_0_SETCTX: HATTV20SetContextActor, self.GRAPH_HATT_V2_0_PLAIN_NOIDLE: HATTV20PlainNoIdleActor, self.GRAPH_HATT_V2_0_DETACHED: HATTV20SpecialDetachedActor, self.GRAPH_HATT_V2_0_SPLIT: HATTV20SpecialSplitActor, self.GRAPH_HATT_V2_0_TYPE_EMBED: HATTV20TypeEmbeddingActor, self.GRAPH_HATT_V2_0_TYPE_GATE: HATTV20TypeGateActor, self.GRAPH_HATT_V2_1: HATTV21Actor, self.GRAPH_HATT_V2: HATTV22Actor, self.GRAPH_HATT_V2_3A: HATTV23AActor, self.GRAPH_HATT_V2_3B: HATTV23BActor, self.GRAPH_HATT_V2_3C: HATTV23CActor, self.GRAPH_HATT_V2_4: HATTV24Actor, self.GRAPH_HATT_V2_5A: HATTV25AActor, self.GRAPH_HATT_V2_5B: HATTV25BActor, self.GRAPH_HATT_V2_5C: HATTV25CActor, self.GRAPH_HATT_V2_5: HATTV25FullActor, self.GRAPH_HATT_V2_6A: HATTV26AActor, self.GRAPH_HATT_V2_6B: HATTV26BActor, self.GRAPH_HATT_V2_6C: HATTV26CActor, self.GRAPH_HATT_V2_6: HATTV26FullActor, self.GRAPH_HATT_V2_7A: HATTV27AActor, self.GRAPH_HATT_V2_7B: HATTV27BActor, self.GRAPH_HATT_V2_7: HATTV27Actor, self.GRAPH_HATT_V3_0A: HATTV30AActor, self.GRAPH_HATT_V3_0B: HATTV30BActor, self.GRAPH_HATT_V3_0C: HATTV30CActor, self.GRAPH_HATT_V3_1A: HATTV31AActor, self.GRAPH_HATT_V3_1B: HATTV31BActor, self.GRAPH_HATT_V3_2A: HATTV32AActor, self.GRAPH_HATT_V3_2B: HATTV32BActor, self.GRAPH_HATT_V3_2C: HATTV32CActor, self.GRAPH_HATT_V3_2D: HATTV32DActor, self.GRAPH_HATT_V3_2R: HATTV32RFullActor, self.GRAPH_HATT_V3_2R_NO_TRANSITION: HATTV32RNoTransitionActor, self.GRAPH_HATT_V3_2R_NO_MEMBER: HATTV32RNoMemberActor, self.GRAPH_HATT_V3_3AA: HATTV33AAActor, self.GRAPH_HATT_V3_4: HATTV34Actor, self.GRAPH_HATT_GYMRT2A: HATTGyMRT2AActor, }[self.graph_model] self.actor = _initialize_with_seed( lambda: actor_class( args=args, ego_dim=self.ego_dim, other_dim=self.other_agent_feat_dim, task_dim=self.per_task_feat, num_agents=self.num_agents, num_tasks=self.num_tasks, device=self.device, ), paired_init_seed + 1 if paired_init_seed >= 0 else -1, ) if self.hatt_v20_init_scheme == "compact_function_preserving": if self.graph_model != self.GRAPH_HATT_V2_0: raise ValueError( "compact_function_preserving initialization is limited " "to graph_model=hatt_v2_0" ) if self.hatt_v20_task_input_mode != "restore_req_scale": raise ValueError( "compact_function_preserving initialization requires " "hatt_v20_task_input_mode=restore_req_scale" ) if paired_init_seed < 0: raise ValueError( "compact_function_preserving initialization requires " "hatt_paired_init_seed" ) compact_args = copy.copy(args) compact_args.hatt_v20_task_input_mode = "native" compact_reference = _initialize_with_seed( lambda: HATTV20Actor( args=compact_args, ego_dim=self.ego_dim, other_dim=self.other_agent_feat_dim, task_dim=6, num_agents=self.num_agents, num_tasks=self.num_tasks, device=self.device, ), paired_init_seed + 1, ) initialize_full_actor_from_compact_reference( self.actor, compact_reference, self.cap_dim ) actor_params = list(self.actor.parameters()) else: self.gat_encoder = GATNetwork( node_dim=self.graph_node_dim, edge_dim=self.graph_node_dim, hidden_dim=args.gat_hidden_dim, output_dim=self.gat_out ).to(self.device) if self.mode in (self.MODE_CAT, self.MODE_GATED): actor_input_dim = obs_dim + self.gat_out else: actor_input_dim = self.gat_out if self.mode == self.MODE_GATED: self.actor_gat_gate = nn.Parameter(torch.zeros((), device=self.device)) fake_obs_space = spaces.Box( low=-np.inf, high=np.inf, shape=(actor_input_dim,), dtype=np.float32 ) self.actor = R_Actor(args, fake_obs_space, self.act_space, self.device) actor_params = list(self.actor.parameters()) + list(self.gat_encoder.parameters()) if self.mode == self.MODE_GATED: actor_params += [self.actor_gat_gate] self._actor_params = actor_params critic_params = list(self.critic.parameters()) self.actor_optimizer = torch.optim.Adam(actor_params, lr=self.lr, eps=self.opti_eps, weight_decay=self.weight_decay) self.critic_optimizer = torch.optim.Adam(critic_params, lr=self.critic_lr, eps=self.opti_eps, weight_decay=self.weight_decay) def lr_decay(self, episode, episodes): """ Decay the actor and critic learning rates. :param episode: (int) current training episode. :param episodes: (int) total number of training episodes. """ update_linear_schedule(self.actor_optimizer, episode, episodes, self.lr) update_linear_schedule(self.critic_optimizer, episode, episodes, self.critic_lr) def _reshape_batch(self, batch: np.ndarray, unit_size): """Reshape a batch that may be either [B, unit_size, dim] or [B*unit_size, dim].""" if isinstance(batch, torch.Tensor): batch = batch.cpu().numpy() if batch.ndim == 3 and batch.shape[1] == unit_size: return batch if batch.ndim == 2 and batch.shape[0] % unit_size == 0: return batch.reshape(batch.shape[0] // unit_size, unit_size, -1) raise ValueError( f"Cannot reshape batch with shape {batch.shape} using unit_size={unit_size}." ) def _reshape_obs_dict(self, obs, unit_size): if not isinstance(obs, dict): return self._reshape_batch(obs, unit_size) out = {} for key, value in obs.items(): if isinstance(value, torch.Tensor): value = value.detach().cpu().numpy() feature_rank = POLICY_OBS_FEATURE_RANKS[key] if value.ndim == feature_rank + 2 and value.shape[1] == unit_size: out[key] = value elif value.shape[0] % unit_size == 0: out[key] = value.reshape(value.shape[0] // unit_size, unit_size, *value.shape[1:]) else: raise ValueError( f"Cannot reshape obs[{key!r}] with shape {value.shape} using unit_size={unit_size}." ) return out def _to_tensor_obs(self, obs): if not isinstance(obs, dict): return torch.tensor(obs, dtype=torch.float32, device=self.device) return { k: torch.as_tensor(v, dtype=torch.float32, device=self.device) for k, v in obs.items() } def _flatten_grouped_obs(self, obs): out = {} for k, v in obs.items(): feature_rank = POLICY_OBS_FEATURE_RANKS[k] if v.ndim == feature_rank + 2 and v.shape[1] == self.num_agents: out[k] = v.reshape(v.shape[0] * v.shape[1], *v.shape[2:]) else: out[k] = v return out def _flat_to_grouped_obs(self, obs_flat): T, n_agents, _ = obs_flat.shape task_start = self.ego_dim task_end = task_start + self.num_tasks * self.per_task_feat return { "ego": obs_flat[..., :self.ego_dim], "tasks": obs_flat[..., task_start:task_end].view( T, n_agents, self.num_tasks, self.per_task_feat ), "others": obs_flat[..., task_end:].view( T, n_agents, n_agents, self.other_agent_feat_dim ), } def _build_graph_node_features(self, obs): ego = obs["ego"] task_obs = obs["tasks"] others = obs["others"] T, n_agents, _ = ego.shape # others[observer, target] stores target-agent state, with the observer's # own row zeroed. Aggregate over observers to recover one node per agent. valid_other = (others.abs().sum(dim=-1, keepdim=True) > 0).float() other_counts = valid_other.sum(dim=1).clamp_min(1.0) other_agent_state = (others * valid_other).sum(dim=1) / other_counts agent_nodes = torch.zeros( (T, n_agents, self.graph_node_dim), dtype=torch.float32, device=self.device ) resource_start = 2 resource_end = resource_start + self.cap_dim amount_idx = resource_end type_idx = amount_idx + 1 active_idx = type_idx + 1 agent_nodes[..., 0:2] = ego[..., 0:2] # position agent_nodes[..., resource_start:resource_end] = ego[..., 2:2 + self.cap_dim] agent_nodes[..., amount_idx] = ego[..., self.cap_dim + 4] # load ratio agent_nodes[..., type_idx] = other_agent_state[..., self.cap_dim + 5] agent_nodes[..., active_idx] = other_agent_state[..., self.cap_dim + 6] # Task features are global for the timestep; average the duplicated # per-agent observations. task_state = task_obs.mean(dim=1) task_nodes = torch.zeros( (T, self.num_tasks, self.graph_node_dim), dtype=torch.float32, device=self.device ) task_nodes[..., 0:2] = task_state[..., 0:2] # position task_nodes[..., resource_start:resource_end] = task_state[..., 2:2 + self.cap_dim] task_nodes[..., amount_idx] = task_state[..., 2 + self.cap_dim] task_nodes[..., type_idx] = task_state[..., 3 + self.cap_dim] task_nodes[..., active_idx] = (task_state.abs().sum(dim=-1) > 0).float() return agent_nodes, task_nodes def _gat_augment(self, obs): """ 使用 GAT 对每个时间步的 agent/raw task node 特征建图并编码。 参数 ---- obs : dict[str, Tensor] 每个时间步、每个智能体的结构化观测。 返回 ---- agent_emb_batch : Tensor [T, N_agents, gat_out_dim] 每个时间步、每个智能体对应的 GAT 节点嵌入。 global_pool_batch: Tensor [T, gat_out_dim] 每个时间步整张图的全局池化向量（所有节点平均）。 """ T, n_agents, _ = obs["ego"].shape agent_nodes, task_nodes = self._build_graph_node_features(obs) data_list = [] for t in range(T): g = build_graph(agent_nodes[t], task_nodes[t]) data_list.append(g) batch = Batch.from_data_list(data_list) node_emb_all = self.gat_encoder(batch.x, batch.edge_index, batch.edge_attr) N = n_agents + self.num_tasks node_emb_split = node_emb_all.view(T, N, -1) agent_emb_batch = node_emb_split[:, :n_agents, :] global_pool_batch = node_emb_split.mean(dim=1) return agent_emb_batch, global_pool_batch # ---- 输入构造辅助方法 ---- def _build_actor_input(self, obs, agent_emb: torch.Tensor): """根据模式构造 Actor 的输入张量。""" obs_flat_dict = self._flatten_grouped_obs(obs) agent_emb_flat = agent_emb.reshape(-1, self.gat_out) # [T*N, gat_out] if self.mode == self.MODE_GAT: return {"gat": agent_emb_flat} elif self.mode == self.MODE_CAT: obs_aug = dict(obs_flat_dict) obs_aug["gat"] = agent_emb_flat return obs_aug else: # MODE_GATED alpha = torch.tanh(self.actor_gat_gate) obs_aug = dict(obs_flat_dict) obs_aug["gat"] = alpha * agent_emb_flat return obs_aug @staticmethod def _build_critic_input(cent_obs): cent_obs = flatten_state_dict(cent_obs) if getattr(cent_obs, "ndim", 0) > 2: cent_obs = cent_obs.reshape(-1, cent_obs.shape[-1]) return cent_obs def get_gat_gate_values(self) -> Tuple[float, float]: """Return the actor gate and a zero compatibility value for the removed critic gate.""" if self.graph_model != self.GRAPH_LEGACY or self.mode != self.MODE_GATED: return 0.0, 0.0 actor_alpha = torch.tanh(self.actor_gat_gate).detach().cpu().item() return float(actor_alpha), 0.0 def get_actor_parameters(self): return self._actor_params def get_actor_checkpoint(self): checkpoint = { "graph_model": self.graph_model, "reconfiguration_mode": self.reconfiguration_mode, "reconfiguration_context_features": self.reconfiguration_context_features, "reconfiguration_gate_parameterization": self.reconfiguration_gate_parameterization, "hatt_v20_task_input_mode": self.hatt_v20_task_input_mode, "hatt_v20_init_scheme": self.hatt_v20_init_scheme, "actor": self.actor.state_dict(), } if self.graph_model == self.GRAPH_LEGACY: checkpoint["gat_encoder"] = self.gat_encoder.state_dict() if self.mode == self.MODE_GATED: checkpoint["actor_gat_gate"] = self.actor_gat_gate.detach().cpu() return checkpoint def load_actor_checkpoint(self, checkpoint): if "actor" not in checkpoint: # Backward compatibility with legacy actor-only checkpoints. self.actor.load_state_dict(checkpoint) return saved_model = checkpoint.get("graph_model", self.GRAPH_LEGACY) compatible_v20_checkpoint = ( ( self.graph_model == self.GRAPH_HATT_V2_0 and saved_model == self.GRAPH_HATT_V2 ) or ( self.graph_model == self.GRAPH_HATT_V2_0_NOIDLE and saved_model == self.GRAPH_HATT_V2_0_PHYSICS ) ) if saved_model != self.graph_model and not compatible_v20_checkpoint: raise ValueError( f"Checkpoint graph_model={saved_model!r} does not match current {self.graph_model!r}." ) saved_reconfiguration_mode = checkpoint.get( "reconfiguration_mode", "legacy" ) mode_mismatch = saved_reconfiguration_mode != self.reconfiguration_mode safe_nonlearned_intervention = ( self.allow_reconfiguration_mode_mismatch and "learned" not in ( saved_reconfiguration_mode, self.reconfiguration_mode, ) ) if mode_mismatch and not safe_nonlearned_intervention: raise ValueError( "Checkpoint reconfiguration_mode=" f"{saved_reconfiguration_mode!r} does not match current " f"{self.reconfiguration_mode!r}." ) saved_context_features = bool( checkpoint.get("reconfiguration_context_features", False) ) if saved_context_features != self.reconfiguration_context_features: raise ValueError( "Checkpoint reconfiguration_context_features=" f"{saved_context_features!r} does not match current " f"{self.reconfiguration_context_features!r}." ) saved_gate_parameterization = checkpoint.get( "reconfiguration_gate_parameterization", "mixture" ) if saved_gate_parameterization != self.reconfiguration_gate_parameterization: raise ValueError( "Checkpoint reconfiguration_gate_parameterization=" f"{saved_gate_parameterization!r} does not match current " f"{self.reconfiguration_gate_parameterization!r}." ) saved_task_input_mode = checkpoint.get("hatt_v20_task_input_mode", "native") if saved_task_input_mode != self.hatt_v20_task_input_mode: raise ValueError( "Checkpoint hatt_v20_task_input_mode=" f"{saved_task_input_mode!r} does not match current " f"{self.hatt_v20_task_input_mode!r}." ) self.actor.load_state_dict(checkpoint["actor"]) if self.graph_model == self.GRAPH_LEGACY and "gat_encoder" in checkpoint: self.gat_encoder.load_state_dict(checkpoint["gat_encoder"]) if ( self.graph_model == self.GRAPH_LEGACY and self.mode == self.MODE_GATED and "actor_gat_gate" in checkpoint ): self.actor_gat_gate.data.copy_( checkpoint["actor_gat_gate"].to(self.device) ) def get_utility_loss(self, obs, available_actions): """Return auxiliary utility head loss (zero if actor has no utility head).""" if not hasattr(self.actor, "get_utility_loss"): return torch.tensor(0.0, device=self.device) return self.actor.get_utility_loss(obs, available_actions) def get_graph_diagnostics(self): if not self._is_hatt_graph_model(self.graph_model): return {} result = {} prefix = self.graph_model for name, value in self.actor.last_diagnostics.items(): if name.endswith("_attention"): probabilities = value.clamp_min(1e-8) metric = -(value * probabilities.log()).sum(dim=-2).mean() name = f"{name}_entropy" else: metric = value.mean() result[f"{prefix}/{name}"] = float(metric.cpu()) return result def get_reconfiguration_gate_regularization(self): """返回 actor 最近一次 evaluate_actions 产生的 gate 正则项。 非 hierarchical gate 或未启用 learned reconfiguration 时返回 0， 使 trainer 可以无条件把该项接入总 loss。 """ value = getattr(self.actor, "last_gate_regularization", None) if value is None: return torch.zeros((), device=self.device) return value def get_actions(self, cent_obs, obs, rnn_states_actor, rnn_states_critic, masks, available_actions=None, deterministic=False): """ Compute actions and value function predictions for the given inputs. :param cent_obs (np.ndarray): centralized input to the critic. :param obs (np.ndarray): local agent inputs to the actor. :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor. :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic. :param masks: (np.ndarray) denotes points at which RNN states should be reset. :param available_actions: (np.ndarray) denotes which actions are available to agent (if None, all actions available) :param deterministic: (bool) whether the action should be mode of distribution or should be sampled. :return values: (torch.Tensor) value function predictions. :return actions: (torch.Tensor) actions to take. :return action_log_probs: (torch.Tensor) log probabilities of chosen actions. :return rnn_states_actor: (torch.Tensor) updated actor network RNN states. :return rnn_states_critic: (torch.Tensor) updated critic network RNN states. """ obs_raw = self._reshape_obs_dict(obs, self.num_agents) obs_t = self._to_tensor_obs(obs_raw) critic_input = self._build_critic_input(cent_obs) if self._is_hatt_graph_model(self.graph_model): actor_input = obs_t else: agent_emb, _ = self._gat_augment(obs_t) actor_input = self._build_actor_input(obs_t, agent_emb) actions, action_log_probs, rnn_states_actor = self.actor(actor_input, rnn_states_actor, masks, available_actions, deterministic) values, rnn_states_critic = self.critic(critic_input, rnn_states_critic, masks) gate_actions = getattr(self.actor, "last_gate_actions", None) if gate_actions is None: # 普通 Categorical、mixture、switch 都没有显式 gate_action； # buffer 仍存一个零占位，保持所有 graph_model 的 batch 结构一致。 gate_actions = torch.zeros_like(actions, dtype=torch.float32) return ( values, actions, action_log_probs, gate_actions, rnn_states_actor, rnn_states_critic, ) def get_values(self, cent_obs, rnn_states_critic, masks): """ Get value function predictions. :param cent_obs (np.ndarray): centralized input to the critic. :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic. :param masks: (np.ndarray) denotes points at which RNN states should be reset. :return values: (torch.Tensor) value function predictions. """ values, _ = self.critic(self._build_critic_input(cent_obs), rnn_states_critic, masks) return values def evaluate_actions(self, cent_obs, obs, rnn_states_actor, rnn_states_critic, action, masks, available_actions=None, active_masks=None, gate_actions=None): """ Get action logprobs / entropy and value function predictions for actor update. :param cent_obs (np.ndarray): centralized input to the critic. :param obs (np.ndarray): local agent inputs to the actor. :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor. :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic. :param action: (np.ndarray) actions whose log probabilites and entropy to compute. :param masks: (np.ndarray) denotes points at which RNN states should be reset. :param available_actions: (np.ndarray) denotes which actions are available to agent (if None, all actions available) :param active_masks: (torch.Tensor) denotes whether an agent is active or dead. :param gate_actions: hierarchical gate rollout decisions, needed to replay the joint action logprob during PPO update. :return values: (torch.Tensor) value function predictions. :return action_log_probs: (torch.Tensor) log probabilities of the input actions. :return dist_entropy: (torch.Tensor) action distribution entropy for the given inputs. """ obs_raw = self._reshape_obs_dict(obs, self.num_agents) obs_t = self._to_tensor_obs(obs_raw) critic_input = self._build_critic_input(cent_obs) if self._is_hatt_graph_model(self.graph_model): actor_input = obs_t else: agent_emb, _ = self._gat_augment(obs_t) actor_input = self._build_actor_input(obs_t, agent_emb) action_log_probs, dist_entropy = self.actor.evaluate_actions(actor_input, rnn_states_actor, action, masks, available_actions, active_masks, gate_actions) values, _ = self.critic(critic_input, rnn_states_critic, masks) return values, action_log_probs, dist_entropy def act(self, obs, rnn_states_actor, masks, available_actions=None, deterministic=False): """ Compute actions using the given inputs. :param obs (np.ndarray): local agent inputs to the actor. :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor. :param masks: (np.ndarray) denotes points at which RNN states should be reset. :param available_actions: (np.ndarray) denotes which actions are available to agent (if None, all actions available) :param deterministic: (bool) whether the action should be mode of distribution or should be sampled. """ obs_t = self._to_tensor_obs( self._reshape_obs_dict(obs, self.num_agents) ) if self._is_hatt_graph_model(self.graph_model): actor_input = obs_t else: agent_emb, _ = self._gat_augment(obs_t) actor_input = self._build_actor_input(obs_t, agent_emb) actions, _, rnn_states_actor = self.actor( actor_input, rnn_states_actor, masks, available_actions, deterministic, ) return actions, rnn_states_actor