面向目标无关的偏微分方程联合嵌入预测控制

# 迈向偏微分方程的目标无关联合嵌入预测控制
来源：https://arxiv.org/html/2607.21644
Roberto Guglielmi 滑铁卢大学 roberto\.guglielmi@uwaterloo\.ca †通讯作者

###### 摘要

我们提出了一种针对偏微分方程\(PDE\)的目标无关控制框架，该框架围绕端到端的联合嵌入预测架构\(JEPA\)构建。一个轻量级的2D ViT编码器和动作条件潜变量动力学模型在离线状态下训练，无需奖励或下游目标，训练后冻结，并由模型预测路径积分\(MPPI\)控制器重复使用。我们发现，在可用的情况下，控制目标最好应用于显式的物理可观测量\(若满足单射性\)，而不是最小化学习到的潜空间中的原始欧几里得距离\(L^2\)。对于冻结潜变量轨迹上学习到的线性动能\(KE\)探测器，我们可以再现对保留轨迹的控制，\(R^2=0.989\)，同时不需要改变底层世界模型。在PDE控制Gym 2D Navier–Stokes基准测试中，使用KE探测规划将匹配的50集原生奖励从潜\(L^2\)规划的\(-12.08\pm0.86\)提升到\(-10.90\pm0.91\)（95%置信区间），同时将最后一季度速度场RMSE从0.0765降低到0.0692。在三个有意保留的、不同、非周期目标上，KE规划将后期场RMSE相对于潜\(L^2\)规划降低了53%（0.0220对比0.0469），在所有30个配对情节中获胜。相同的冻结模型还支持通过直接调控KE实现对稳态配置稳定的控制目标，达到2.7%的平均相对误差。虽然潜变量探测器对测量噪声和缺失像素较为脆弱，但我们相信这些结果支持以下主张：潜变量动力学可以保持动态性和目标无关性，同时校准后的可观测变量（只要保证唯一延拓）可能是状态控制的更好目标。

## 1 引言

偏微分方程\(PDE\)出现在众多应用领域，从气候建模、生物过程到化学过程监控。为偏微分方程系统制定控制策略可能具有挑战性，因为状态是高维的，且量之间的关系通常是非线性的。因此，候选控制序列的评估通常需要计算代价高昂的数值模拟，这通常阻碍了在线规划和控制。现有的偏微分方程控制方法，如基于模型的控制、学习型代理控制或强化学习方法，已取得不同程度的成功。传统的基于模型控制器需要针对特定问题的分析，或以其他方式巧妙构建降阶模型\(ROMs\)[6 (https://arxiv.org/html/2607.21644#bib.bib43)]。学习型代理方法引入神经网络作为系统动力学演化算子的近似[18 (https://arxiv.org/html/2607.21644#bib.bib25),19 (https://arxiv.org/html/2607.21644#bib.bib26),26 (https://arxiv.org/html/2607.21644#bib.bib27)]，但这些代理通常是在状态空间前向模拟任务上训练，并针对特定目标进行优化，继承了全状态重构的挑战，并且在目标改变时需要重新训练[35 (https://arxiv.org/html/2607.21644#bib.bib28),37 (https://arxiv.org/html/2607.21644#bib.bib6)]。强化学习在定义明确的任务上表现出色[7 (https://arxiv.org/html/2607.21644#bib.bib40),10 (https://arxiv.org/html/2607.21644#bib.bib41),28 (https://arxiv.org/html/2607.21644#bib.bib38),32 (https://arxiv.org/html/2607.21644#bib.bib39)]，但需要仔细选择奖励函数，这限制了泛化能力。

联合嵌入预测架构\(JEPAs\)则专注于学习如何将应用的动作映射到学习到的潜空间或“世界模型”的元素之间。这是在不进行任何类型的状态重构的情况下完成的[2 (https://arxiv.org/html/2607.21644#bib.bib2)]。最近在潜世界模型方面的工作已使用这种免重构特征在主要视觉领域进行零样本规划[39 (https://arxiv.org/html/2607.21644#bib.bib5)]，而现有的潜变量PDE工作表明，压缩后的动力学可以加速模拟和逆优化[37 (https://arxiv.org/html/2607.21644#bib.bib6)]。我们相信JEPA风格的潜世界模型为PDE问题提供了独特的优势。使用随机掩码进行训练可能使编码器能够同时处理完整和不完整的场，同时映射到相同的潜表示。同时，动作条件潜变量预测器能够实现目标无关的模型预测控制。学习到的动力学模型对运行时的目标保持无关性，因为在部署时更改此目标不需要重新训练编码器或预测器。

潜规划因多种原因而有吸引力：其计算效率高和目标无关性是最突出的两点。不幸的是，为可预测性和动作敏感性训练的潜表示并非自动是一个校准的度量空间。众所周知，嵌入之间的\(L^2\)距离会扭曲物理场误差的尺度和排序。因此，我们使用时间拉直来改善潜变量轨迹的几何结构[33 (https://arxiv.org/html/2607.21644#bib.bib9)]，但我们的实验表明，仅靠正则化不足以完全解决此问题。相反，我们为物理状态可观测变量拟合一个轻量级的读出器，并在该读出器之后以物理单位形成跟踪误差。我们仔细选择可观测变量以保证唯一延拓，并观察到该控制器在不重新训练或解码世界模型的情况下显著改善了干净的闭环跟踪。

最近使用联合嵌入世界模型的工作已探索了视觉操作、导航、运动控制和空中领域的规划[1 (https://arxiv.org/html/2607.21644#bib.bib21),39 (https://arxiv.org/html/2607.21644#bib.bib5),24 (https://arxiv.org/html/2607.21644#bib.bib8),38 (https://arxiv.org/html/2607.21644#bib.bib7),29 (https://arxiv.org/html/2607.21644#bib.bib11)]。据我们所知，类似JEPA的联合嵌入世界模型尚未应用于PDE系统的闭环控制。本工作旨在展示联合嵌入架构在PDE控制中的潜力，表明一个目标无关的联合嵌入与下游预测器可以成为一个独特的动态控制器，能够原生支持多个PDE控制目标。

我们在PDE控制Gym二维Navier–Stokes环境中验证了这个框架[4 (https://arxiv.org/html/2607.21644#bib.bib1)]。我们训练掩码联合嵌入动力学，使用VICReg防止坍塌[3 (https://arxiv.org/html/2607.21644#bib.bib3)]，添加时间拉直和类似Delta-JEPA的潜动作解码[38 (https://arxiv.org/html/2607.21644#bib.bib7)]，并使用带有冻结骨干网络的MPPI[36 (https://arxiv.org/html/2607.21644#bib.bib4)]。

我们的贡献：

1. 据我们所知，我们首次实现了使用联合嵌入潜世界模型的闭环PDE控制，该模型能够独立于下游控制任务学习动作条件动力学。
2. 我们引入了一个可观测变量对齐控制器：一个训练成本低的线性探测器从预测的潜变量中估计动能，并且MPPI在该物理标量上评估跟踪误差，而不是使用原始潜距离。
3. 在PDE控制Gym基准测试中，与使用潜\(L^2\)代价的控制器相比，这同时提高了原生奖励和速度场RMSE。在三个不同的非周期信号上，它将后期场RMSE降低了53%，在所有30个配对情节中都有改进。
4. 我们重复使用一个冻结的检查点和固定的五步潜变量规划视界，用于时间依赖跟踪、固定场稳定和动能控制，始终在状态空间中工作或进行任务特定的微调。

## 2 目标无关控制公式

### 2.1 无需目标学习的动力学

让一个有控制的场按照 \(U_{t+1}=F(U_t,a_t)\) 演化。我们学习一个编码器 \(E_\theta\) 和一个动作条件潜变量预测器 \(g_\phi\)，完全从离线状态转换中学习，不提供奖励、参考轨迹、平衡态或物理可观测量。训练后，编码器和动作条件预测器保持冻结，而MPPI规划器提供候选动作序列集，其效果可通过 \(g_\phi\) 探查到潜状态。这允许在部署时在MPPI规划器内指定一个目标，而不触及 \(E_\theta\) 或 \(g_\phi\) 的权重。

### 2.2 部署时目标接口

规划机制的灵活性支持任意的控制目标。我们研究了三种：

1. **控制Gym动作惩罚跟踪**。作为编码状态 \(z^\star_{t+1:t+H}\) 或物理观测 \(q^\star_{t+1:t+H}\) 提供的目标序列可以被跟踪，同时进行动作惩罚权衡，在未以代价感知方式训练的情况下达到与RL基线可比的性能。
2. **时间依赖跟踪**。目标序列可以作为编码状态 \(z^\star_{t+1:t+H}\) 或作为时间依赖的物理可观测量 \(q^\star_{t+1:t+H}\) 提供。
3. **固定目标稳定**。指定一个状态配置并仅编码一次。编码后的潜变量 \(z^\star\) 被视为一个常数目标，据此评估控制作用，从而得到一个平衡目标。

只有用户指定的代价在这些任务之间变化，这意味着动作条件潜变量动力学模型和MPPI优化器在其他方面保持不变。

### 2.3 PDE控制Gym基准测试

我们使用PDE控制Gym二维Navier–Stokes环境于其发布示例所用的配置[4 (https://arxiv.org/html/2607.21644#bib.bib1)]来评估我们的模型。域为一个单位正方形，离散化为 \(\Delta x=\Delta y=0.05\) 的 \(21 \times 21\) 网格。模拟器使用 \(\Delta t=10^{-3}\)，规划视界 \(T=0.2\)，粘度 \(\nu=0.1\)，密度 \(1\)，以及2000次压力迭代。观测为速度场 \(U_t=(u_t,v_t) \in \mathbb{R}^{21\times21\times2}\)。一个标量动作 \(a_t\in[0,4]\) 沿上边界设置一个可控的Dirichlet条件（在正x方向）。

参考基准奖励为

\[
r_t = -\frac{1}{2HW}\lVert U_t - U_t^\star \rVert_2^2 - \frac{\gamma}{2}(a_t - a_t^{\mathrm{ref}})^2,
\] (1)

其中 \(H=W=21\) 为分辨率，\(\gamma=0.1\) 为动作惩罚，\(a_t^{\mathrm{ref}}=2\) 为测量动作的参考速度。我们报告累积原生奖励、状态空间RMSE和动作总变差\(TV\)。我们的控制器每两个原生求解器步骤动作一次，因此99个宏动作覆盖基准测试199个得分原生步骤中的198个。

### 2.4 掩码观测

训练我们的编码器（轻量级ViT）时，状态观测被随机掩码。任何状态都有0.15的概率被完全观测；否则，掩码比例在范围 \([0.3,0.9]\) 内均匀采样。掩码分别以0.5、0.3和0.2的概率从连续矩形、随机像素或行条中选择。评估时，当前测量要么完全观测，被零均值高斯噪声（标准差等于每个速度通道近期RMS的5–15%）破坏，要么随机缺失5–25%的像素。目标场总是完全指定的，其中任何目标都被解释为指定的目标，而不是传感器测量。

掩码协议的一个示例显示在图1 (https://arxiv.org/html/2607.21644#S2.F1) 中。

参见图注图1：掩码协议：连续矩形、随机像素和行条，每张图都大约隐藏了60%的场。灰色表示掩码（零填充）像素。

## 3 目标无关JEPA控制框架

### 3.1 掩码潜变量动力学

设 \(E_\theta\) 为一个掩码场编码器，并设 \(z_t = E_\theta(U_t, m_t)\)。编码器是一个紧凑的二维视觉变换器，输出256维，\(3\times3\) 图像块，四个变换器块和四个注意力头。预测器 \(g_\phi\) 消耗四个潜状态及其动作的历史，然后以步长2自回归地滚动出十二个潜步骤。预测器由一个具有每层512个隐藏单元的两层MLP组成，训练用于预测潜残差而非完整状态。

对于给定的训练样本，编码器有两个任务。首先，它必须对掩码历史进行编码，然后必须对包含未来场的第二个未掩码序列进行编码。

\[
z^\mathrm{ctx}_{t-K+1:t} = E_\theta(U_{t-K+1:t}, m),\qquad z^\mathrm{real}_{t:t+12} = E_\theta(U_{t:t+12}, \mathbf{1}).
\]

然后预测器的任务是从 \(z^\mathrm{ctx}\) 向前滚动，其中未掩码的 \(z^\mathrm{real}\) 值是潜目标。在这种设置中，滚动梯度可以流过预测器、掩码历史编码器分支和未掩码未来编码器分支。因此损失为

\[
\mathcal{L}_{\mathrm{roll}} = \frac{1}{12}\sum_{h=1}^{12}\lVert\hat{z}_{t+h} - z^\mathrm{real}_{t+h}\rVert_2^2.
\] (2)

我们同时使用VICReg方差和协方差惩罚来防止坍塌[3 (https://arxiv.org/html/2607.21644#bib.bib3)]，以及一个时间拉直损失，该损失阻止潜轨迹不必要地绕弯。一个辅助解码器 \(d_\psi\) 的任务是从未掩码、编码的真实序列的相邻位移（而不是从预测的滚动位移）重建应用的动作：

\[
\mathcal{L}_{\Delta a} = \frac{1}{12}\sum_{h=1}^{12}\lVert d_\psi(z^\mathrm{real}_{t+h} - z^\mathrm{real}_{t+h-1}) - a_{t+h-1}\rVert_2^2,
\] (3)

遵循Delta-JEPA的动作敏感性动机[38 (https://arxiv.org/html/2607.21644#bib.bib7)]。时间拉直惩罚同样在未掩码、编码的真实序列上进行评估，而不是在预测的滚动路径上。完整的复合损失为

\[
\mathcal{L}_{\mathrm{total}} = \mathcal{L}_{\mathrm{roll}} + \lambda_{\mathrm{var}}\,\mathcal{L}_{\mathrm{var}} + \lambda_{\mathrm{cov}}\,\mathcal{L}_{\mathrm{cov}} + \lambda_{\mathrm{str}}\,\mathcal{L}_{\mathrm{str}} + \lambda_{\Delta a}\,\mathcal{L}_{\Delta a}
\] (4)

权重分别为 \(\lambda_{\mathrm{var}}=1.0\)，\(\lambda_{\mathrm{cov}}=1.0\)，\(\lambda_{\mathrm{str}}=0.1\)，\(\lambda_{\Delta a}=10\)。没有解码器将潜状态 \(z\) 映射回场 \(U\)。完整的训练过程总结在图2 (https://arxiv.org/html/2607.21644#S3.F2) 中。