# AFEM — 自适应网格细化的 GNN + PPO 强化学习
## 项目架构
```
afem/
├── src/ # 应用层
│ ├── config.yaml # 配置文件
│ ├── main.py # 入口：解析命令行 → train / test / viz
│ ├── network.py # GNN + Actor-Critic 完整网络定义
│ ├── ppo.py # RolloutBuffer + PPOTrainer
│ ├── utils.py # 读配置、保存/加载 checkpoint
│ └── visualize.py # viz 模式：加载模型 → 推理 → 存 PNG
│
├── environment/ # 仿真环境层
│ ├── mesh_refinement.py # ★ 核心：网格细化 RL 环境
│ │ # - GNN 图观测构建（节点 + 边特征）
│ │ # - continuous_sizing_field (score-based + budget) 细化策略
│ │ # - spatial 奖励
│ ├── helmholtz.py # Helmholtz FEM 求解器 + 残差误差估计
│ ├── fem_problem.py # FEM 问题封装 + PDE 循环缓冲区
│ ├── fem_util.py # 三角形面积、中点、随机采样、尺寸场函数
│ ├── domain.py # 计算域：meshpy 三角剖分
│ ├── utils.py # 数组拼接、随机索引采样
│ └── visualization.py # plotly 网格渲染（RL 环境用）
│
├── checkpoints/ # 模型保存
├── result/ # 可视化输出
└── README.md
```
---
## 项目简介
### 物理场景
二维 Helmholtz 电磁散射：
```
∇²u_scat + k²·ε_r·u_scat = k²·(1-ε_r)·u_inc
```
- **入射波**: 沿 -x 方向的平面波 `u_inc = exp(i·k·x)`
- **散射体**: 圆形介质柱（ε_r 随机采样），位置和半径可配
- **边界条件**: SBC (Sommerfeld) `∂u/∂n = i·k·u`
- **域**: 可配矩形域，初始网格密度自适应 + domain area 线性缩放：`N_init = N_base × (k/k_ref)^k_exponent × domain_area`。k_ref 和 k_exponent 均可通过 helmholtz config 配置（默认 k_exponent=2.0, k_ref=6.0），保证不同域尺寸下每单位面积单元数一致
- 可配 exponent：^2 = P1 Helmholtz 理论最优 (污染误差 ∝ k²)。建议 N_base 配合 exponent 调整，使 N_init 约为 COMSOL 目标 (λ/10√ε_r) 的 30-50%，为 RL agent 留出充分细化空间
- **介质区前渐近区边缘约束**: 介质内 λ_d = 2π/(k√ε_r) 更短，强制迭代细化至 h ≤ λ_d/N（默认 N=1.5，helmholtz.pre_asymptotic_N 可配）。约 1.5 点/波长，刚好跨过渐近区门槛，赋予初始网格基本相位解析能力但不过度消耗物理预算，为 RL agent 留出充分的选择性细化空间
- **后验误差**: 残差型 indicator（Ainsworth & Oden 风格），含单元内部残差 + 梯度跳变 + SBC 边界残差
### 强化学习建模
| 概念 | 对应实体 |
|------|---------|
| **智能体** | 每个三角形网格单元 |
| **状态** | GNN 节点特征（几何 + PDE 残差 + 振幅 + 相位方向 + 物理参数，节点 13 维 + 边 1 维） |
| **动作** | 1 维连续标量 x_i → score = -x_i 排序，在物理预算内 top-k 选细化单元（x 越小优先级越高） |
| **奖励** | 局部子单元 η 的 log-ratio 改善（spatial: sum 聚合 / spatial_max: max 聚合）+ α 衰减全局 η log-ratio shaping |
| **终止** | 达到最大步数或超过最大单元数 |
---
## 网络架构
双 GNN 架构（policy / value 各自独立基座）：
```
图观测 → MessagePassingBase → MLP → 动作分布 / value 标量
├─ nn.Linear（嵌入）
├─ MessagePassingStack（2 层消息传递 + GVN 全局广播，inner 残差 + LayerNorm）
│ ├─ MessagePassingStep × N
│ │ ├─ EdgeModule: MLP([src | dst | edge_attr])
│ │ └─ NodeModule: MLP([node | scatter(入边)])
│ └─ GlobalVirtualNode (GVN): η_K 加权注意力池化 → 注意力门控广播
│ h_V = Σ(η_v/Ση)·h_v，α_v = σ(W_att[h_v || h_V])，h_v ← h_v + α_v ⊙ W_V·h_V
└─ 输出: 节点隐向量
```
| 超参数 | 值 |
|--------|-----|
| latent_dim | 64 |
| 消息传递层数 | 2 |
| 残差连接 | inner |
| 归一化 | inner LayerNorm |
| 边 dropout | 0.1 |
| Actor MLP | 2 层 tanh |
| Critic MLP | 2 层 tanh |
| Optimizer | Adam, lr=3e-4, lr_decay=0.995 |
| **动作分布** | `DiagGaussianDistribution`（连续 Box 动作空间），`log_std` 可学习，clamp 在 [-4.0, -1.0] |
| **log_std 策略** | 初始化 -2.0（std≈0.135），每步 optimizer.step() 后 clamp 到 [-4.0, -1.0]（std ∈ [0.018, 0.368]），熵系数 0.001 |
### 动作分布策略说明
环境定义的是 `_action_space`（下划线前缀），网络初始化时必须用 `environment._action_space` 而非 `environment.action_space`（后者默认为 None，会错误回退到 `CategoricalDistribution(1)`，导致 policy gradient 恒为零）。
`continuous_sizing_field`（score-based）的动作有效范围约 [-3, 3]：
- score = -x_i，x 越小 ⇒ 优先级越高（纯排序，不设正负门槛）
- `initial_log_std=-2.0`（std≈0.135），clamp 在 [-4.0, -1.0]（std ∈ [0.018, 0.368]）
- 加 `entropy_coefficient=0.001` 提供微弱探索压力，避免 log_std 过早收敛到下限
---
## 输入特征
### 节点特征（13 维）
| 维度 | 来源 | 名称 | 说明 |
|------|------|------|------|
| 1 | cfg | `volume` | 无量纲单元面积：volume / λ² |
| 3 | cfg | `internal_residual` / `gradient_jump` / `sbc_residual` | PDE 残差三分量（真空波数 k 归一化，经 log₁₀ 压缩）：
`(h_K/k)·√V·|r|` / `√(½Σ h_e·\|jump\|²/k)` / `(h_bnd/k)·\|SBC\|` |
| 1 | cfg | `element_penalty` | 单元惩罚系数 λ |
| 1 | cfg | `timestep` | 当前 rollout 步数 |
| 1 | cfg | `k_local_sqrt_vol` | k × √(ε_r) × √(V)（局域波数 × 特征长度） |
| 1 | cfg | `is_sbc_boundary` | 是否与 SBC 吸收边界相邻 (0/1) |
| 1 | cfg | `dist_to_interface` | 到介质圆柱边界的带符号距离，无量纲化后经 sign·ln(1+|d|) 压缩：`sign(d)·ln(1+|(dist-radius)/λ|)` — 近场近似线性保留分辨力，远场对数压缩避免 OOD，与残差 log₁₀ 风格一致 |
| 1 | fix | `epsilon_r` | 单元中点相对介电常数（圆柱内 = εᵣ，外 = 1.0） |
| 1 | fix | `total_solution_magnitude` | 散射场振幅 \|u_scat\|（per-element 均值） |
| 1 | fix | `cos_phase` | Re(u) / (\|u\| + 1e-8)，相位方向余弦，∈ [−1, 1]，无分支切割 |
| 1 | fix | `sin_phase` | Im(u) / (\|u\| + 1e-8)，相位方向正弦，与 cos 联合编码相位 |
> - **cfg**: 由 `element_features` 配置控制
> - **fix**: 始终启用（Helmholtz 振幅 + 相位方向，硬编码）
>
> GNN 输入用 `_compute_residual_components`（真空波数 k 归一化，log₁₀ 压缩）。Reward 用逐单元 η_K（`_eta_indicator`），与 GNN 特征公式一致但不经 log 压缩。SBC 边界条件保留 `k_local`。
### 边特征（1 维）
| 维度 | 名称 | 说明 |
|------|------|------|
| 1 | `phase_distance` | 相邻单元中点相位距离 = d × √(k_local_src·k_local_dst) / 2π — 介质内短波长自然放大，赋予 GNN k 不变性 |
---
## 调用逻辑
```
main.py --mode train/test/viz
│
├─→ utils.load_config() # 读 YAML
├─→ environment.MeshRefinement # 创建 RL 环境
│ └─→ FEMProblemCircularQueue # 管理 N 个随机 PDE 实例
│ └─→ HelmholtzProblem # FEM 求解 + 残差误差
├─→ network.create_model() # 创建 ActorCritic
│
└─ [train] → ppo.PPOTrainer.fit_iteration() 循环
├─ collect_rollouts() # 256 步 rollout
│ └─ buffer.compute_returns_and_advantage()
│ └─ 单路 GAE # 逐 agent 时序差分（scatter_add 处理网格细化）
│ └─ Return / value 归一化
└─ train_step() # 多 epoch PPO 更新
├─ policy_loss() # Clipped PPO
├─ value_loss() # Clipped value loss
└─ entropy_loss() # 熵正则
```
### 环境内部调用
```
MeshRefinement.reset()
└─→ FEMProblemWrapper.reset()
└─→ initial_mesh (meshpy → 介质内 前渐近区边缘迭代细化)
MeshRefinement.step(action)
├─→ score = -x 排序 + 物理预算约束 → top-k 细化单元
├─→ FEMProblemWrapper.refine_mesh() # scikit-fem refine
├─→ calculate_solution_and_get_error()
│ ├─→ HelmholtzProblem.calculate_solution() # FEM 求解
│ └─→ _compute_residual_indicator() # 残差误差
├─→ _get_reward_by_type() # spatial 奖励
└─→ last_observation # 构建 Data(x, edge_index, edge_attr)
```
### 训练
```bash
CUDA_VISIBLE_DEVICES=7 python src/main.py --mode train --config src/config.yaml
```
首次迭代需收集 256 步 rollout（含 FEM 求解），后续打印：
```
it | loss ev agents reward x<0 elig sel time
```
| 字段 | 含义 | 健康范围 |
|------|------|---------|
| `x<0` | `mean(x_i < 0)`，负值动作比例（纯诊断） | 越负的单元优先级越高 |
| `elig` | 通过双过滤器的候选占比 | 排除数值退化 + 低误差的单元 |
| `mask` | 被 Reverse Dörfler 剔除的噪声尾部占比（累积能量 <1% 总误差的底部单元） | 因场景而异，非固定比例 |
| `sel` | 实际选中的细化单元数 | 每步最多 N_current // 4 |
| `n_budget` | 全局物理预算（每 episode 固定） | k=30 → ~1800 |
### 测试
```bash
python src/main.py --mode test --checkpoint checkpoints/model_final.pt --k-test 6.0
python src/main.py --mode test --checkpoint checkpoints/model_final.pt \
--k-test 6.0 --center 0.3,0.6 --radius 0.15
```
输出：
```
Step 0: reward=--- error=1.0000 elements=174 budget=1885
Step 1: reward=+12.345 error=0.7160 elements=618 x<0=0.45 sel=87
...
```
每步打印 `reward error elements x<0 sel`，第 0 步额外显示 `N_budget`。
### 可视化
```bash
python src/main.py --mode viz --checkpoint checkpoints/model_final.pt --k-test 30.0
```
输出: `result/visualization.png`（总览）+ `result/visualization_steps/step*.png`（逐步对比）。
---
## 后验误差估计
### 残差 indicator 公式（无量纲化）
引入局部波数 $k_{local} = k\sqrt{\max(\varepsilon_r, 1.0)}$，消除纯几何尺度 $h$ 带来的特征偏差，
使误差指示子反映"相位分辨率残差"而非"网格粗疏程度"。
对 P1 三角单元 K，三项残差分量为：
$$r_{\text{int}} = \frac{h_K}{k} \sqrt{V_K} \cdot \left| k^2\varepsilon_r u + k^2(\varepsilon_r-1)u_{inc} \right|_K \tag{1}$$
$$r_{\text{jump}} = \sqrt{\frac{1}{2}\sum_{e\in\partial K} \frac{h_e}{k} \cdot \left| [[\nabla u \cdot n]] \right|^2_e} \tag{2}$$
$$r_{\text{sbc}} = \frac{h_{bnd}}{k} \cdot \left| \frac{\partial u}{\partial n} - ik_{local}u \right| \tag{3}$$
**逐单元误差指示子**：
$$\eta_K = \sqrt{r_{\text{int}}^2 + r_{\text{jump}}^2 + r_{\text{sbc}}^2}$$
量纲分析（$k \sim [L]^{-1}$，$h_e \sim [L]$，$|\text{jump}|^2 \sim [L]^{-2}$）：
三项均严格无量纲：$h_e/k \cdot |\text{jump}|^2 \sim [L]^2 \cdot [L]^{-2} = 1$。
SBC 边界条件仍用 $k_{local}$（物理正确），仅归一化因子改用 $k$。
介质内残差不再被 $\sqrt{\varepsilon_r}$ 压低，Agent 获得正确的介质内/外优先级信号。
`η_K` 的计算（`_compute_residual_indicator`）与 GNN 输入特征（`_compute_residual_components`）公式完全一致，特征仅多一层 log₁₀ 压缩。关键验证点：
- 内部残差：P1 元 ∇²u_h ≡ 0，仅含反应项 `k²ε_r·u + k²(ε_r-1)·u_inc`，真空波数 k 归一化
- 梯度跳变：`(h_e/k)·|jump|²`，½ 分配给相邻左右单元；$h_e$ 保留边积分路径，细化后自然衰减
- SBC 项归一化用 k，物理条件保留 k_local：`(h_bnd²/k²)·|∂u/∂n − i·k_local·u|²`
### 连续尺寸场策略（score-based + 物理预算约束 + 动作掩码）
Actor 输出标量 x_i → score = -x_i 直接排序，在预算和上限内选 top-k：
```
A_budget_i = ½(λ_local_i / 6)² // 每局部波长方向 ~6 尺度点（仅用于 N_budget 计算）
λ_local_i = 2π / (k · √ε_r_i)
N_budget = max(N_phys, ⌈5·N_init⌉) // rho_min=5.0，至少 5× 初始单元数，保证 RL 多步细化空间
N_phys = ⌈ Σ |K_i| / A_budget_i ⌉ // 全局物理预算（k=30 真空 ~1800）
remaining = N_budget − N_current
V_min_safeguard = 1e-10 × domain_area // 纯数值底线（防止 FEM 求解器退化）
eligible: area > V_min_safeguard AND η_K ∈ Reverse Dörfler 保留集 // 数值底线 + 能量尾部淘汰 (ε_noise=0.01, ≥20% floor)
num = min(|eligible|, N_current//4, remaining//3)
selected = top-k by score = -x_i → 1-to-4 切分
```
- score = -x_i：x 越小 ⇒ 优先级越高（纯排序，不设正负门槛）
- 不再使用 `0.25·A_budget` 启发式面积地板：RL 应自主学会"细化到多细"，而非被人类经验 (12 点/波长) 限制。仅保留数值底线 V_min_safeguard = 1e-10 × domain_area 防止浮点精度问题。
- per-step cap 从固定 200 改为自适应 `N_current // 4`，随网格规模缩放但增速更缓，避免大网格时单步消耗过多预算。rho_min 从 3.0 提升到 5.0，赋予更多预算余量。
- **sel=0 提前终止**：当 agent 选中 0 个单元细化（预算耗尽或 Reverse Dörfler 屏蔽所有候选）时 episode 自动结束，不再浪费 FEM 求解
- **k_exponent 可配**：初始网格缩放指数可通过 `helmholtz.k_exponent` 配置（默认 2.0），² 为 P1 Helmholtz 理论最优；对 k=30 的 $N_{init}$ 为 k=6 的 25× 倍
- **动作掩码 (Reverse Dörfler)**：按 η_K 升序排列，剔除累积平方误差贡献 < ε_noise·Ση² 的底部单元（数值噪声/已收敛区）。基于能量分布而非密度分位数，在重尾和均匀误差分布下均自适应。保留率不低于 20% 确保 Agent 始终有充分的选择空间
### 奖励计算
---
#### 变量
| 符号 | 含义 |
|------|------|
| `η_K = √(r_int² + r_jump² + r_sbc²)` | 逐单元误差指示子，`r_*` 定义见式 (1)–(3) |
| `C(i)` | 父单元 i 经 1-to-4 切分产生的子单元集合 |
| `M_new[j]` | 子单元 j 对应的父单元索引 |
| `n_i = |C(i)|` | 父单元 i 的子单元数（1 表示未切分） |
| `E_global = √(Σ η_K²) / \|\|u_h\|\|_{L₂(Ω)}` | 全局无量纲误差 |
---
#### 算法
**Step 0 — 保存旧状态** (`_set_previous_step`)
```
η_old ← 旧逐单元 η_K
||u_h_old|| ← 旧解 L₂ 范数 (≈ √(Σ |ū_K|² · area_K))
```
**Step 1 — 网格细化** (`_refine_mesh`)
```
x = action.flatten()
score = -x // x 越小 ⇒ 优先级越高
remaining = N_budget − N_old
max_by_budget = max(0, remaining // 3)
// 数值底线 + Reverse Dörfler 能量尾部淘汰
V_min_safeguard = 1e-10 × domain_area // 纯数值安全底线，防止 FEM 退化
η_sq = η_old²; total_energy = Σ η_sq
k_dorfler = searchsorted(cumsum(sort_asc(η_sq)), ε_noise·total_energy) // ε_noise=0.01
k = min(k_dorfler, N − max(1, N//5)) // ≥20% floor
eligible = {i | V_old[i] > V_min_safeguard AND i ∈ sort_asc_idx[k:] }
num = min(|eligible|, N_old//3, max_by_budget)
elements_to_refine = top-k of eligible by score
M_new[j] ∈ {0,…,N_old-1} // 子→父映射
```
**Step 2 — FEM 求解 + 误差估计**
```
η_new ← 新逐单元 η_K
||u_h_new|| ← 新解 L₂ 范数
```
**Step 3 — 因果奖励**（零和预算审查）
ε_dynamic = max(0.01 × η_P95, 1e-6)
// Refined parents: r_local + zero-sum bonus − penalty
if i ∈ refined_parents:
r_i = log(η_old + ε) − log(√(Σ η_child²) + ε) // r_local ≥ 0 (L₂ 聚合)
+ 0.3 × (η_old / μ − 1.0) // zero-sum bonus (Σ = 0)
− 0.06 // action penalty
// Unrefined parents: causal isolation
else:
r_i = 0
> **零和奖金**：α·(η/μ−1) 全场求和为零。细化高于均值的单元得正奖金，低于均值的倒扣。
> 这是 Dörfler 准则的 RL 对偶：Agent 必须选出误差超过全均水平的单元。
> **因果隔离**：未细化单元 r ≡ 0。零和奖金本身足够强（介质内 +0.51）、
> 不再需要忽视惩罚的推力，排序机制自动淘汰不划算的单元。
> **L₂ 聚合**：√(Σ η_child²) ≤ η_parent 天然成立，r_local ≥ 0 永不惩罚细化。
**Step 4 — 全局误差（仅诊断）**
global_bonus = α·[log(E_old) − log(E_new)]，α = 0.5
不注入 Actor reward。Helmholtz 污染误差可使 E_new > E_old 在正确细化后发生，
注入 global_bonus 导致因果断裂。Actor 仅优化 Step 3 的 per-element reward。
---
#### 奖励标度校准（旧尺寸场下测量，score-based 后需重新标定）
在随机策略下实测各分量量级（1321 个 refined-parent 样本）：
| 分量 | 均值 | 占 r_local 比例 |
|------|------|:---:|
| `r_local` (仅 refined parents) | +0.364 | — |
| `penalty` λ·(n−1), λ=0.02 | +0.045 | 1/8 |
| `α·ΔlogE` α=0.2 | +0.069 | 1/5 |
| **net** | **+0.387** | |
满足 `r_local ≫ penalty` 且 `α·ΔlogE ≈ r_local / 5`，局部 credit assignment 不被全局信号淹没。
---
#### 设计要点
| 组件 | 聚合 | 作用 |
|------|------|------|
| 局部项 `log(η_old / √(Σ η_child²))` | scatter_add，仅 refined parents | L₂ 保证 r_local ≥ 0；int 主导 +0.69 |
| 零和奖金 `0.3×(η/μ−1)` | 仅 refined parents | Σ=0；高于 μ 得正奖，低于 μ 倒扣 (Dörfler 准则的 RL 对偶) |
| 动作惩罚 `λ=0.06` | per-refined-parent | 轻微抑制网格膨胀（1-to-4 扣 0.06） |
| 因果隔离 `r=0` | unrefined parents | 零和奖金足够强，不需额外推力 |
| 全局项 `α·ΔlogE` α=0.5 | 仅诊断 | 不注入 Actor，避免污染误差因果断裂 |
---
## PPO 关键细节
- **单路 GAE**: r_local 自身已闭合因果（细化单元的局部误差改善），不需要势函数塑形。用 `scatter_add` 将细化后的子单元值聚合回父单元，单路 GAE 即可
- **奖励归一化**: rollout 内 reward 做 z-score 归一化（std < 1e-8 则跳过）
- **Value clipping**: 默认 clip_range=0.2
- **梯度裁剪**: max_grad_norm=0.5
- **log_std clamp**: 每步 `optimizer.step()` 后将 `log_std` clamp 到 `[-3.0, -1.0]`，σ ∈ [0.05, 0.37]
初始化 `-2.0` (σ≈0.135)，放宽下限防止策略过早确定化
- **熵正则**: `entropy_coefficient=0.005`，施加有意义的探索压力防止 x<0 崩塌
- **epochs_per_iteration**: 3，减少对同一 rollout 的过拟合