afem/asmr++_architecture.md

ASMR++ 网络架构与数据流 (默认配置)

基于 configs/asmr_pp/asmr_default.yaml — value_function_aggr: spatial, projection_type: sum

架构总览

flowchart TD
    subgraph ENV["♻️ 环境: MeshRefinement"]
        A1["FEMProblemCircularQueue<br/>随机采样 PDE 问题"]
        A2["生成初始粗网格<br/>(meshpy, 2D 三角剖分)"]
        A3["FEM 求解器<br/>计算 PDE 解和逐单元误差"]
        A4["构建观测图<br/>(节点=单元, 边=邻接关系)"]
        A1 --> A2 --> A3 --> A4
    end

    subgraph GRAPH["📊 观测图 (torch_geometric Data)"]
        B1["<b>节点特征 (x)</b><br/>━━━━━━━━━━━━━━━━<br/>solution_mean / solution_std<br/>volume / timestep<br/>element_penalty<br/>source_term (PDE 特征)<br/>共 ~10-15 维"]
        B2["<b>边特征 (edge_attr)</b><br/>━━━━━━━━━━━━━━━━<br/>euclidean_distance<br/>共 1 维"]
        B3["<b>边索引 (edge_index)</b><br/>━━━━━━━━━━━━━━━━<br/>双向邻接 + 自环"]
    end

    subgraph NORM["📏 观测归一化器"]
        C1["node.x: running mean/std"]
        C2["edge_attr: running mean/std"]
    end

    subgraph HMPN["🧠 HMPN 基础网络 (HomogeneousMessagePassingBase)"]
        subgraph EMBED["输入嵌入"]
            D1["节点嵌入: Linear(in→64)"]
            D2["边嵌入: Linear(in→64)"]
        end
        subgraph STACK["消息传递堆栈 (num_steps=2, residual=inner, layernorm=inner)"]
            subgraph STEP1["Step 1/2"]
                E1["<b>边更新</b> HomogeneousEdgeModule<br/>concat[src(64), dst(64), edge(64)]<br/>→ LatentMLP(192→64, 2层, LeakyReLU)<br/>→ LayerNorm → +inner residual"]
                E2["<b>节点更新</b> HomogeneousMessagePassingNodeModule<br/>scatter_mean(edge→dest) → concat[node(64), agg(64)]<br/>→ LatentMLP(128→64, 2层, LeakyReLU)<br/>→ LayerNorm → +inner residual"]
                E1 --> E2
            end
            subgraph STEP2["Step 2/2"]
                F1["边更新 (同上)"]
                F2["节点更新 (同上)"]
                F1 --> F2
            end
            STEP1 --> STEP2
        end
        D1 --> STEP1
        D2 --> STEP1
        STEP2 --> G["输出: 节点潜在特征 (num_nodes, 64)"]
    end

    subgraph HEADS["🎯 策略与价值头 (share_base=False, 各自独立 GNN)"]
        subgraph ACTOR["Actor 头"]
            H1["Policy MLP<br/>2层, Tanh<br/>64→64→64"]
            H2["Linear(64→action_dim)"]
            H3["log_std (可学习)"]
            H4["DiagGaussian(μ, σ)<br/>每节点输出独立动作"]
            H1 --> H2 --> H4
            H3 --> H4
        end
        subgraph CRITIC["Critic 头 — 逐节点价值，不做 scatter 聚合"]
            I1["Value MLP<br/>2层, Tanh<br/>64→64→1"]
            I2["输出形状: (num_agents, 1)<br/>每个 agent 独立 V_i(s)<br/><b>value_function_aggr=spatial<br/>→ 不聚合，保持逐节点</b>"]
            I1 --> I2
        end
        G --> H1
        G --> I1
    end

    subgraph BUFFER["🗃️ MixedOnPolicyBuffer (global_weight=0.5)"]
        J1["<b>局部 GAE (逐节点)</b><br/>δ_i = r_i + γ·Σ_j φ_ij·V_j(s') - V_i(s)<br/>projection_type='sum': Σ 通过 agent_mapping 反投影"]
        J2["<b>全局 GAE (图级别)</b><br/>δ_global = r_global + γ·V_mean(s') - V_mean(s)"]
        J3["<b>混合 Advantage</b><br/>A_i = (1-0.5)·A_i_local + 0.5·A_global"]
        J1 --> J3
        J2 --> J3
    end

    subgraph PPO["🔄 PPO 训练"]
        K1["256 步 Rollout"]
        K2["5 Epochs, batch_size=32"]
        K3["policy_loss + 0.5·value_loss<br/>clip_range=0.2"]
        K4["梯度裁剪 0.5, Adam lr=3e-4"]
        K1 --> K2 --> K3 --> K4
    end

    ENV --> GRAPH --> NORM --> HMPN
    HMPN --> HEADS
    ACTOR -->|动作| ENV
    CRITIC -->|"V_i(s) 逐节点"| BUFFER
    ENV -->|"r_i, agent_mapping φ"| BUFFER
    BUFFER --> PPO
    PPO -->|更新参数| HMPN
    PPO -->|更新参数| HEADS

核心纠正: projection_type 的真实作用

之前的错误理解:

• Critic 输出 scatter_sum → 图级别价值 ❌

正确理解:

• value_function_aggr: "spatial" → Critic 不做任何聚合，输出 (num_agents, 1) 逐节点价值 ✅
• projection_type: "sum" → 在 Buffer 中通过 agent_mapping 反投影下一步价值时使用 ✅

两个参数作用于完全不同的位置:

参数	作用位置	作用
value_function_aggr: "spatial"	SwarmPPOActorCritic._get_values_and_distribution()	控制 Critic 输出是否聚合: "spatial" → 保持逐节点
projection_type: "sum"	SpatialOnPolicyBuffer._project_to_previous_step()	控制 agent_mapping 反投影方式: sum→子元素价值求和回父元素

详细数据流 (序列图)

sequenceDiagram
    actor Trainer
    participant Env as MeshRefinement
    participant Norm as Normalizer
    participant GNN as HMPN Base
    participant Actor as Policy Head
    participant Critic as Value Head
    participant Buffer as MixedOnPolicyBuffer

    Note over Trainer,Buffer: === Rollout (256 步) ===

    Trainer->>Env: reset()
    Env->>Env: 随机 Poisson PDE + 随机域 + GMM 负载
    Env->>Env: 初始粗网格 → FEM 求解 → 构建观测图

    loop 256 步
        Env-->>Norm: 观测图 (原始 node.x, edge_attr)
        Norm-->>GNN: 归一化后图
        GNN->>GNN: Edge Dropout (0.1, 仅训练)
        GNN->>GNN: 嵌入 → MP Step1 → MP Step2
        GNN-->>Actor: node_features (num_nodes, 64)
        GNN-->>Critic: node_features (num_nodes, 64)

        Actor->>Actor: MLP → μ, σ → 采样动作
        Critic->>Critic: MLP(64→1) → <b>V_i(s): (num_agents, 1) 逐节点</b>

        Actor-->>Env: actions (num_agents, 1)
        Env->>Env: 元素选择 → 网格细分
        Env->>Env: FEM 求解 → 计算空间奖励 r_i
        Env-->>Buffer: (obs, a, r_i, V_i, log_prob, agent_mapping φ)
    end

    Note over Buffer: === GAE 计算 (逐节点 + 混合奖励) ===

    Buffer->>Buffer: <b>局部 δ_i(t) = r_i + γ·Σ_j φ_ij·V_j(t+1) - V_i(t)</b>
    Buffer->>Buffer: projection_type='sum': Σ_j 通过 agent_mapping 反投影
    Buffer->>Buffer: 局部 GAE → A_local_i (逐节点)
    Buffer->>Buffer: 全局 GAE → A_global (图级, 用 mean(V_i) 算)
    Buffer->>Buffer: <b>A_i = 0.5·A_local_i + 0.5·A_global</b>
    Buffer->>Buffer: R_i = A_i + V_i(s)

    Note over Trainer,Buffer: === 训练 (5 Epochs × batch 32) ===

    loop 5 Epochs
        Buffer-->>Trainer: (obs, a, old_log_prob, old_V_i, A_i, R_i)
        Trainer->>GNN: 重新前向传播
        GNN-->>Actor: node_features
        GNN-->>Critic: node_features
        Actor->>Actor: 新 log_prob
        Critic->>Critic: 新 V_i (逐节点)
        Trainer->>Trainer: ratio = exp(log_prob_new - log_prob_old)
        Trainer->>Trainer: policy_loss = -min(ratio·A_i, clip(ratio,0.8,1.2)·A_i)
        Trainer->>Trainer: value_loss = 0.5·clip(V_new, V_old±0.2) vs R_i
        Trainer->>Trainer: backward() + grad_clip(0.5) + Adam.step()
    end

论文公式 (3) 与代码对应

论文中的 TD 误差公式:

\delta^t_i = r(s^t, a^t)_i + \gamma \sum_j \phi_{ij}^t V_j(s^{t+1}) - V_i(s^t)

在代码中的实现路径 (spatial_on_policy_buffer.py:174-178):

# _get_agent_wise_advantages_and_returns()
for step in range(self.buffer_size):
    if self.dones[step]:
        delta = self.rewards[step] - self.values[step]  # r_i - V_i(s)
    else:
        delta = self.rewards[step] \
              + self.discount_factor * projected_next_values[step] \  # + γ·Σ_j φ_ij·V_j(s')
              - self.values[step]  # - V_i(s)

其中 projected_next_values[step] 由 _project_to_previous_step() 产生:

# projection_type='sum'
projected_value = scatter_sum(values[step], index=agent_mappings[step], dim=0)
#                                    ^^^^^^^^^^        ^^^^^^^^^^^^^^^^^^^^^
#                                    V_j(s_{t+1})      φ_ij: 新agent j → 旧agent i

关键默认参数

参数	值	代码位置
算法	PPO	config["algorithm"]["name"]
网络骨架	Homogeneous MPN	config["network"]["type_of_base"]
GNN 架构	mpn (message passing)	config["network"]["base"]["architecture"]
潜在维度	64	config["network"]["latent_dimension"]
MP 步数	2	config["network"]["base"]["stack"]["num_steps"]
残差	inner	config["network"]["base"]["stack"]["residual_connections"]
层归一化	inner	config["network"]["base"]["stack"]["layer_norm"]
边→节点聚合	mean	config["network"]["base"]["scatter_reduce"]
Base MLP	2层, LeakyReLU	config["network"]["base"]["stack"]["mlp"]
Actor MLP	2层, Tanh	config["network"]["actor"]["mlp"]
Critic MLP	2层, Tanh	config["network"]["critic"]["mlp"]
价值函数范围	spatial (逐节点, 不聚合)	config["algorithm"]["ppo"]["value_function_aggr"]
价值投影方式	sum (agent_mapping 反投影用)	config["algorithm"]["ppo"]["projection_type"]
混合奖励权重	0.5	config["algorithm"]["mixed_return"]["global_weight"]
共享 Base	False (Actor/Critic 各自独立 GNN)	config["network"]["share_base"]
动作分布	DiagGaussian (连续)	动作空间为 gym.spaces.Box
Rollout 步数	256	config["algorithm"]["ppo"]["num_rollout_steps"]
训练轮次	5	config["algorithm"]["ppo"]["epochs_per_iteration"]
Batch 大小	32	config["algorithm"]["batch_size"]
GAE λ	0.95	config["algorithm"]["ppo"]["gae_lambda"]
折现 γ	1.0	config["algorithm"]["discount_factor"]
PPO clip	0.2	config["algorithm"]["ppo"]["clip_range"]
梯度裁剪	0.5	config["algorithm"]["ppo"]["max_grad_norm"]
学习率	3e-4	config["network"]["training"]["learning_rate"]
边 Dropout	0.1	config["network"]["base"]["edge_dropout"]
Episode 步数	6	config["environment"]["mesh_refinement"]["num_timesteps"]
PDE	Poisson (GMM 负载, zero Dirichlet)	config["environment"]["mesh_refinement"]["fem"]["pde_type"]

projection_type 的两种职责

projection_type 在 Buffer 中有两处使用，都是通过 agent_mapping 做跨时间步的 agent 反投影:

1. 价值反投影 — 公式 (3) 的 Σ 项

# _project_to_previous_step() — spatial_on_policy_buffer.py:33
projected_value = scatter_sum(values[step], index=agent_mappings[step], dim=0)
# 下一步的 V_j(s_{t+1}) 按 agent_mapping φ_ij 求和回当前步的 agent i

2. GAE 时间差分反投影 — 动态规划递推

# _get_agent_wise_advantages_and_returns() — spatial_on_policy_buffer.py:169
projected_last_gae = scatter_sum(last_gae, index=self._agent_mappings[step], dim=0)
# 上一步累积的 GAE 按 agent_mapping 反投影

核心创新点

1. Swarm 视角 + 变长 Agent: 每个网格元素是一个 agent，元素分裂后 agent 数量动态增长
2. 空间奖励 + agent_mapping: 通过 agent_mapping φ_ij 追踪父→子关系，支持逐节点的 TD 误差计算（公式 3）
3. 混合奖励学习: 局部逐节点 Advantage + 全局图级 Advantage 加权混合 (0.5:0.5)
4. MPN 通信: 边更新 + 节点更新的消息传递，元素通过共享三角形边交换 PDE 解信息
5. 自适应细化: 连续动作 → 概率性元素选择 → 非均匀网格，资源集中在误差大的区域