256 lines
12 KiB
Markdown
256 lines
12 KiB
Markdown
# ASMR++ 网络架构与数据流 (默认配置)
|
||
|
||
> 基于 `configs/asmr_pp/asmr_default.yaml` — `value_function_aggr: spatial`, `projection_type: sum`
|
||
|
||
## 架构总览
|
||
|
||
```mermaid
|
||
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→子元素价值求和回父元素 |
|
||
|
||
## 详细数据流 (序列图)
|
||
|
||
```mermaid
|
||
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`):
|
||
|
||
```python
|
||
# _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()` 产生:
|
||
|
||
```python
|
||
# 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) 的 Σ 项
|
||
```python
|
||
# _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 时间差分反投影 — 动态规划递推
|
||
```python
|
||
# _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. **自适应细化**: 连续动作 → 概率性元素选择 → 非均匀网格，资源集中在误差大的区域
|