大气传输模型机器学习代理的长程轨迹优化

传输计算是大气传输过程的重要组成部分，总体上可以分为垂直扩散、垂直对流、水平扩散、水平对流四个部分，而四个部分计算范式截然不同，在 CMAQ 当中是污染物传输和扩散的重要求解过程，其求解精度影响多个下游模块。优化计算速度的同时，保证传输部分高精度以不影响下游模块计算是大气模拟的重要条件。

而水平传输是传输模块中的计算密集型部分，其使用双向交替的二遍算子分裂，每一个子时间步都运行 X 和 Y 方向的分段抛物线计算。同时其嵌套使用同步步长和子步步长，受 Courant 数约束，大部分子步被切得很小，在一个同步步内要循环很多次，每一步都要进行 X 和 Y 方向的分段抛物线计算。所以单位时间（每小时）的计算量不是恒定的，而是随天气强度（风速）动态增长的——风越大，算得越慢、越费时。

同时计算过程会产生大量 I/O 操作，读取逐步长、逐层的输入数据，缺乏必要的算子融合。PPM 抛物线算法也缺乏向量化（SIMD 范式计算），编译器基本上只能生成标量代码，而且访问模式是那种很吃内存带宽的 gather/scatter（聚集/散射）模式。整个传输计算过程高度碎片化，存在大量条件判断、子步迭代、非向量化操作，以此用统一的机器学习模型对传输模型进行替代，将迭代求解转换成单次前向传播，配合 GPU 的高速张量计算，我们可以把传输部分的同步时间步的单次计算耗时降低至原来的 1/10。

模型与数据构造

我们用一个简单的 3D CNN 网络学习 CMAQ 输运算子（transport-only，不含化学/排放），通过 rollout training 解决一个关键问题：单步精度好但自回归链式推理时误差累积爆炸。替代目标是传输模块当中的 HADV、ZADV、HDIFF 三个模块，将其融合为一个统一的机器学习模型。通过选取真实 CMAQ 运行案例，在模块开始与结束位置进行插桩，转储输入输出数据，构建真实运行场景数据集。输入变量包含 U、V 方向风场（来自 MCIP 数据）、浓度张量、层数、Sigma Layer Thickness 等变量，输出 Arakawa-C 网格边界面物质传输通量，通过将物种作为 Batch 维度，让所有物种共享同一套权重，从而实现物种并行（Species Parallel）。

对于回归链式推理时误差累积爆炸，我们给出的答案是大规模长程轨迹合成数据。通过将目标模块单独隔离出来，用 Python 重写成一个独立运行的模块，采样转储的输入数据，测试重写的部分能与原代码保持 bit 级一致，通过真实运行案例的传输数据作为输入，定义为 $\mathrm{input}_{t0}$，由此得到 $\mathrm{output}_{t0}$，将 $\mathrm{output}_{t0}$ 作为下一步迭代的输入 $\mathrm{input}_{t1}$，由此不断循环，最终我们能得到一段 $K=N$ 的高保真自迭代轨迹数据，起始点由真实环境采样得到，同时通过拉丁超立方采样，我们可以探索输入向量空间的更多可能性，从而形成拓展合成轨迹数据（Extended synthetic trajectory data）。

损失函数与 rollout training

对于 loss 函数来说：

$$ L_{\mathrm{ss}} = \frac{1}{N}\sum_{k=0}^{N-1}\left\Vert NN\!\left(C_b^{(k)}\right) - C_a^{(k)} \right\Vert_2^2 \tag{1} $$

加入 rollout 项之后：

$$ L = L_{\mathrm{ss}} + \lambda_{\mathrm{roll}} \cdot L_{\mathrm{roll}}, \qquad \lambda_{\mathrm{roll}} = 0.5 \tag{2} $$

同时在 rollout training 当中，关键的训练配方是同步计算多步自迭代预测的损失，而只针对第一步的梯度来更新权重，这样模型既能学习到第一步的预测精度，同时又对多步预测做好准备，综合损失函数在多步损失当中会表现平滑。需要注意的是：单步的高精度未必是多步预测稳定的必要条件，机器学习模型对输入存在不同程度的敏感性，在自迭代链式预测过程中，微小误差会随着多步累积而放大，这在控制论当中被称之为自激震荡。犹如一个高保真的麦克风，无论多还原音色，靠近音响的时候也难免产生刺耳的啸叫声。

所以在长程合成轨迹模型训练当中，一个在第一步 loss 很高的模型未必会在多步预测的最后一步 loss 很高，模型会逐渐学会多步轨迹预测损失的平滑值，从而避免因为自激震荡导致数值爆炸的问题。我们采用混合步长作为合成数据配方，从 $K=3$ 到 $K=20$ 不等，在训练过程中混合调配样本比例，避免在长程监督学习中产生"灾难性遗忘"。下图很好地解释了不同训练步长导致的 loss 变化：

如果我们不选择使用 rollout training，对于以 NO2 为例的物种，其相对误差在自迭代 4 个时间步后便急剧放大，以模型 Sync Step = 600 seconds 的条件下，最多只能撑过 40 分钟便会产生灾难性的数值爆炸现象。这里为方便说明起见，我们选择一个不采用混合步长训练的 $K=8$ 的小参数模型作为训练目标，可以看到 $K=8$ 的训练效果能让多步误差累积更加平滑，差不多能抑制约 5 倍的相对误差。而 $K=8$ 的 v6 版本在第一个时间步的相对误差甚至更高！关于 $K=N$ 的步长数量和样本比例的训练配方，并没有一个比较科学的依据，受训练时间限制，最终产出模型采用了从 $K=5$ 到 $K=20$ 的混合比例，通过 Top_K 随机采样在训练集当中进行训练。

我们以一块 NVIDIA A800 GPU + 4 个英特尔 Intel® Xeon® Gold 6326 CPU 核心作为测试硬件基准，于在线测试中测试模型在传输模块的墙钟模拟时间，可以发现在单模块实现了约 8.9 倍的加速。我们的测试案例采用 $292 \times 273 \times 14$ 约 116 万个网格的模拟配置，总有效物种数 867 个，模拟网格分辨率 27 公里。通过有效的机器学习做算子融合和 GPU 并行推理加速，将物种放入 Batch 维度，实现了高效的物种并行，单次前向传播计算可同时预测 867 个物种的界面通量。配合尾端的物种质量守恒性校验，可以更加高效地验证物理约束是否吻合。

启发式强化学习：如何改进大气机器学习代理的自迭代链式推理？

我们的目标是构建这样一个链条：NN surrogate → physics diagnosis → correct/retry/fallback → continue rollout。

短期预测看似准确，但在 7 天至 1 个月的运行过程中，误差会逐渐累积，尤其对于 NOx、铵、硝酸盐、PM 等化学类别的预测更为明显。以传输模块为例，顺序是 HADV → ZADV → HDIFF，state（状态）→ action（动作）→ criticize（评价，奖励或惩罚，这里用物理信息引导，可以参考的标准有原子守恒、前体物依赖顺序、散度/梯度等）→ strategy update（策略更新，我们的案例中用 PPO）→ next state（下一步的状态）。

对于一个长程轨迹优化问题来说，传统的监督学习（supervised learning）能给模型一个比较好的监督信号，这种信号很密集、通过优化器我们能知道梯度变化方向，从而优化模型权重来使得 loss 函数更小，这是一个机器学习的入门知识。但是监督学习不能解决的问题是：优化器结合当前梯度优化的方向，是否真的符合真实大气环境的物理规律？ 权重迭代的方向会使得模型预测趋于平滑，这样会让全局损失函数更小，但这犹如农场里的火鸡，能观察到"农场主打开门 → 会有食物"这样一个前后关系，但是却没有真的学会物理规律——有食物吃的原因是每天农场主都要养着它们，而总有一天打开门的不是食物，而是屠刀。机器学习从海量的数据当中学会了统计规律，却不能真的理解背后的逻辑和因果关系。

由此：机器学习模型能够在我们的大气传输模型中进行长程训练，但是仅仅有监督信号是不够的，还要有物理信息引导的强化学习，按照物理规律设置奖励/惩罚函数。

一个具体的问题：物理引导的 PPO（近端策略优化）如何处理

也就是说，模型从当前状态 $s_t$ 出发，根据策略 $\pi_\theta$ 产生连续动作 $a_t$（连续动作空间，因为大气内部的变量都是连续的，我们这里的动作空间输出的是高斯分布），经过数值演化得到下一状态 $s_{t+1}$，并将预测误差与质量守恒、原子守恒、梯度一致性和散度等物理约束共同构造成 physics-guided reward，即：

$$ r_t = -L_{\mathrm{pred}} - \lambda L_{\mathrm{physics}} \tag{3} $$

与传统单步监督学习不同，模型会利用自身预测结果不断向前推进，形成完整的 rollout trajectory：

$$ \tau = (s_0, a_0, s_1, a_1, \ldots, s_T) \tag{4} $$

因此训练目标关注的是整个长期轨迹的累计回报，而不仅是下一时间步的预测精度。最后，PPO 根据 rollout 得到的 reward 和 advantage 判断哪些动作有利于长期预测准确性和物理一致性，提高这些动作在相似状态下出现的概率，同时通过 clipping 限制每次策略更新幅度，从而逐步训练出一个在长期积分过程中既保持预测精度，又满足关键物理约束并减少误差累积的 surrogate model。

对于 PPO 来说，核心目标函数是 $L_{\mathrm{PPO}}(\theta)=\mathbb{E}_t\left[\min\left(r_t(\theta)\hat{A}_t,\ \operatorname{clip}\!\left(r_t(\theta),1-\epsilon,1+\epsilon\right)\hat{A}_t\right)\right]$，其中最关键的是两个量。

首先是新旧策略的概率比 $r_t(\theta)=\pi_\theta(a_t\,|\,s_t)\,/\,\pi_{\theta_{\mathrm{old}}}(a_t\,|\,s_t)$，例如某个连续 correction $a_t$ 在 rollout 中表现很好，新策略就希望增加产生这个 action 的概率，因此 $r_t(\theta)>1$，反之则降低其概率。

第二个量是 advantage $\hat{A}_t$，它回答的是：在状态 $s_t$ 下采取 $a_t$，从长期 rollout 的结果来看，比平均水平好还是差？框架里，这个"好坏"最终来自前面定义的 physics-guided reward $R_t=-\mathcal{L}_{\mathrm{pred}}-\lambda\mathcal{L}_{\mathrm{physics}}$，因此可以把整个逻辑理解成 Physics Reward → Rollout Return → $\hat{A}_t$ → PPO Update → $\pi_{\theta_{\mathrm{new}}}$：例如一个 action 当前时刻预测误差很小，而且使后续几十、几百步的质量守恒和原子守恒都更好，那么它的长期 advantage 可能满足 $\hat{A}_t>0$，PPO 就会提高类似 action 的概率；如果某个 action 单步看起来不错，但最终导致 NOy 在长期 rollout 中持续漂移，那么它可能得到 $\hat{A}_t<0$，PPO 就会降低类似 action 的概率。

而公式里的 $\operatorname{clip}$ 是 PPO 很重要的一点：即使发现某个 action 很好或很差，也不允许策略一次改变得太剧烈，例如 $\epsilon=0.2$，就限制每轮优化对策略进行相对保守的调整，从而提高训练稳定性。