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=1.5, k_ref=6.0），保证不同域尺寸下每单位面积单元数一致
• 可配 exponent：^2 = P1 Helmholtz 理论最优 (污染误差 ∝ k²)，^1.5 = 工程折中。建议 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 残差 + 复数场分解 + 物理参数，节点 12 维 + 边 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 层消息传递，inner 残差 + LayerNorm）
          │    └─ MessagePassingStep × N
          │         ├─ EdgeModule:  MLP([src | dst | edge_attr])
          │         └─ NodeModule:  MLP([node | scatter(入边)])
          └─ 输出: 节点隐向量

超参数	值
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 过早收敛到下限

---

输入特征

节点特征（12 维）

维度	来源	名称	说明
1	cfg	volume	无量纲单元面积：volume / λ²
3	cfg	internal_residual / gradient_jump / sbc_residual	PDE 残差三分量（无量纲化，经 log₁₀ 压缩）： `(h_K/k_local)·√V·
1	cfg	element_penalty	单元惩罚系数 λ
1	cfg	timestep	当前 rollout 步数
1	cfg	wave_number	Helmholtz 波数 k
1	cfg	k_local_sqrt_vol	k × √体积（局域波数 × 特征长度）
1	cfg	is_sbc_boundary	是否与 SBC 吸收边界相邻 (0/1)
1	cfg	dist_to_interface	到介质圆柱边界的带符号距离，无量纲化后经 sign·ln(1+
1	fix	epsilon_r	单元中点相对介电常数（圆柱内 = εᵣ，外 = 1.0）
1	fix	total_solution_magnitude	散射场复数解的振幅

• cfg: 由 element_features 配置控制
• fix: 始终启用（Helmholtz 复数场分解，硬编码）

GNN 输入用 _compute_residual_components（k_local 无量纲化，log₁₀ 压缩）。Reward 用逐单元 η_K（_eta_indicator），与 GNN 特征公式一致但不经 log 压缩。

边特征（1 维）

维度	名称	说明
1	euclidean_distance	相邻单元中点欧几里得距离 / λ（无量纲边特征）

---

调用逻辑

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)

训练

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	被 Dörfler-P95 掩码 (η<0.05·η_P95) 滤掉的占比	因场景而异，非固定比例
sel	实际选中的细化单元数	每步最多 N_current // 4
n_budget	全局物理预算（每 episode 固定）	k=30 → ~1800

测试

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。

可视化

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_{local}} \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_{local}} \cdot \left| [[\nabla u \cdot n]] \right|^2_e} \tag{2}

r_{\text{sbc}} = \frac{h_{bnd}}{k_{local}} \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_{local} \sim [L]^{-1}，h_e \sim [L]，|\text{jump}|^2 \sim [L]^{-2}）： 三项均严格无量纲：h_e/k_{local} \cdot |\text{jump}|^2 \sim [L]^2 \cdot [L]^{-2} = 1。 细化后 h_e 缩小直接降低跳变项，为 RL agent 提供可感知的正向 reward 信号。

η_K 的计算（_compute_residual_indicator）与 GNN 输入特征（_compute_residual_components）公式完全一致，特征仅多一层 log₁₀ 压缩。关键验证点：

• 内部残差：P1 元 ∇²u_h ≡ 0，仅含反应项 k²ε_r·u + k²(ε_r-1)·u_inc，除以 k_local 后跨介质公平可比
• 梯度跳变：(h_e/k_local)·|jump|²，½ 分配给相邻左右单元；h_e 保留边积分路径，细化后自然衰减
• SBC 项在 η_K² 中为 (h_bnd²/k_local²)·|B|²，分量 r_sbc = (h_bnd/k_local)·|B|

连续尺寸场策略（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 ≥ 0.05·η_P95   // 数值底线 + Dörfler-P95
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 个单元细化（预算耗尽或 Dörfler 屏蔽所有候选）时 episode 自动结束，不再浪费 FEM 求解
• k_exponent 可配：初始网格缩放指数可通过 helmholtz.k_exponent 配置（默认 1.5），² 为 P1 Helmholtz 理论最优
• 动作掩码 (Dörfler-P95)：η_K < 0.05·η_P95 的单元移出候选池。P95 锚定物理误差尺度，免疫远场噪声稀释（与 median/mean 不同），确保只有误差达标的区域消耗细化预算

奖励计算

---

变量

符号	含义
η_K = √(r_int² + r_jump² + r_sbc²)	逐单元误差指示子，r_* 定义见式 (1)–(3)
C(i)	父单元 i 经 1-to-4 切分产生的子单元集合
M_new[j]	子单元 j 对应的父单元索引
`n_i =	C(i)
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)
// 数值底线 + Dörfler-P95 掩码
V_min_safeguard = 1e-10 × domain_area    // 纯数值安全底线，防止 FEM 退化
η_p95 = percentile(η_old, 95)
eligible = {i | V_old[i] > V_min_safeguard  AND  η_old_i ≥ 0.05·η_p95}
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）

ε_dynamic = max(0.01 × η_P95, 1e-6) // P95 锚定，免疫远场噪声稀释 ε_dynamic = max(0.05 × mean(η_new), 1e-6) // 自适应钳制，切断远场低 η 区 reward hacking spatial: r_local_i = log(η_old_i + ε_dynamic) − log( √(Σ_{j: M_new[j]=i} η_new_j²) + ε_dynamic ) spatial_max: r_local_i = log(η_old_i + ε_dynamic) − log( max_{j: M_new[j]=i} η_new_j + ε_dynamic )

> **L₂ 聚合保证 r_local ≥ 0**: 对 1-to-4 切分：
> ```
> Σ η_child² = int²/4 + jump² + sbc²   ≤   η_parent² = int² + jump² + sbc²
> → r_local = ½[log(η_parent²) − log(Σ η_child²)] ≥ 0
> ```
> - 纯 int 主导: r_local = log(2) ≈ 0.69（强正奖励）
> - 纯 jump/sbc 主导: r_local = 0（中性，不惩罚不奖励）
> - **永远不会惩罚细化**——与 L₁ sum 不同，L₂ 天然避免了对 jump/sbc 主导区的结构性负偏置。

**Step 4 — 动作惩罚**

penalty_i = λ · (n_i − 1) // λ = 0.06 + (λ_limit / N_old) · 𝟙[达到最大单元数上限] // λ_limit = 10000

r_local_i ← r_local_i − penalty_i

**Step 5 — 全局势函数塑形**（仅发给被细化的父单元）

E_global = √(Σ_K η_K²) / ||u_h||_{L₂(Ω)} global_bonus = α · [ log(E_global_old) − log(E_global_new) ] // α = 0.2

r_i = r_local_i − penalty_i + global_bonus · 𝟙[i 被细化] // 未细化的单元 reward ≈ 0

> 全局改进信号只分配给实际参与细化的单元，避免被未细化单元稀释。

---

#### 奖励标度校准（旧尺寸场下测量，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（子→父求平方和再开方） | L₂ 聚合保证 r_local ≥ 0：不惩罚任何细化，int 主导区获强正奖励 (≈+0.69)，纯 jump/sbc 区中性 |
| 动作惩罚 `λ(n_i−1)` λ=0.02 | per-parent | 轻微抑制网格膨胀（1-to-4 切分扣 0.06，仅占 r_local 的 ~16%） |
| 元素上限惩罚 | 达到 20000 上限时触发 | 极端情况兜底，λ_limit / N_old ≈ 0.05~0.5 per agent |
| 全局项 `α·ΔlogE` α=0.2 | 仅细化父单元 | L₂ 无量纲全局误差下降趋势，只发给实际参与细化的单元，避免被未细化单元稀释 |

---

## PPO 关键细节

- **单路 GAE**: 势函数塑形后的奖励已包含全局改进信号，用 `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 到 `[-4.0, -1.0]`，std ∈ [0.018, 0.368]<br>
  初始化 `-2.0` (std≈0.135)，避免 `continuous_sizing_field` 有效范围 [-3, 3] 内噪声过大
- **熵正则**: `entropy_coefficient=0.001`，防止 log_std 过早收敛