Kolmogorov Regression for Robust Diffusion Policies

📄 arXiv: 2606.18186v1 📥 PDF

作者: Lekan Molu

分类: cs.LG, cs.AI

发布日期: 2026-06-16

💡 一句话要点

提出Kolmogorov回归以解决扩散策略的鲁棒性问题

🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱二：RL算法与架构 (RL & Architecture)

关键词: 扩散策略 Kolmogorov方程 Cameron-Martin空间 高斯测度 鲁棒性 动态系统 制造控制

📋 核心要点

  1. 现有的有限维扩散策略在物理系统中表现出时间漂移，影响长期性能。
  2. 本文提出通过反向Kolmogorov方程将扩散策略提升至Cameron-Martin空间，替代随机得分匹配为确定性PDE问题。
  3. 实验表明，Cameron-Martin损失在PushT基准上提升17%最大回合奖励，并在制造线控制中RMSE降低28.4%。

📝 摘要（中文）

有限维扩散策略因离散化伪影而导致时间漂移，降低了在物理系统上的长期性能。本文引入了一种向Cameron-Martin空间提升扩散策略的反向Kolmogorov方程，实质上将随机得分匹配替换为确定性边值PDE问题。核心创新基于高斯测度理论，通过从有色噪声分布中实现扩散噪声协方差算子，规定了推理时模型样本的规律性。我们使用推导出的精度加权Cameron-Martin损失训练扩散模型，并在推理过程中引入Kolmogorov残差作为PDE诊断。这些替代方案带来了收敛保证、改进的轨迹规律性和无奖励信号的确定性故障检测。验证结果显示，在PushT操作基准上，Cameron-Martin损失在最大回合奖励上提高了17%，在制造线控制中RMSE降低了28.4%。

🔬 方法详解

问题定义：本文旨在解决有限维扩散策略在物理系统中因离散化伪影导致的时间漂移问题，这种漂移会显著降低长期性能。

核心思路：通过引入反向Kolmogorov方程，将扩散策略提升至Cameron-Martin空间，替代传统的随机得分匹配方法，转而使用确定性边值PDE问题来增强模型的稳定性和鲁棒性。

技术框架：整体架构包括扩散模型的训练和推理两个阶段。在训练阶段，使用精度加权的Cameron-Martin损失函数；在推理阶段，引入Kolmogorov残差作为PDE诊断工具，以提高模型的推理质量。

关键创新：最重要的创新在于通过高斯测度理论实现扩散噪声协方差算子的构建，从而在推理时对模型样本的规律性进行有效控制。这一方法与现有的随机得分匹配方法有本质区别。

关键设计：在损失函数设计上，采用了精度加权的Cameron-Martin损失，并在推理过程中引入Kolmogorov残差。此外，模型的训练和推理过程中均考虑了噪声的颜色特性，以增强模型的适应性和鲁棒性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，在PushT操作基准上，Cameron-Martin损失实现了17%的最大回合奖励提升（0.95对比0.78），并在推理过程中减少了67.6%的步间漂移。在6站制造线控制中，RMSE比传统LSTM基线降低了28.4%，并实现了高召回率和有效瓶颈识别。

🎯 应用场景

该研究的潜在应用领域包括机器人控制、自动化制造和其他需要高鲁棒性的动态系统。通过提高扩散策略的稳定性和性能，该方法能够在实际应用中显著减少故障和提升效率，具有重要的实际价值和未来影响。

📄 摘要（原文）

Finite-dimensional (FD) diffusion policies exhibit temporal drift owing to discretization artifacts that degrade long-horizon performance (when deployed on physical systems). We introduce a backward Kolmogorov equation that lifts diffusion policies to a Cameron-Martin space -- a subset of the Hilbert space. Essentially, replacing stochastic score matching with a deterministic boundary-value PDE problem. Our core innovation thrives on Gaussian measure theory whereupon the diffusion noise covariance operator is realized from a colored noise distribution which prescribes a notion of regularity on samples from the model at inference time. We train the diffusion model with a derived precision-weighted Cameron- Martin loss and a Kolmogorov residual is introduced as a PDE diagnostic during inference. These substitutions yield (i) convergence guarantees where the bound's constants depend on the effective rank of the kernel rather than action dimension, (ii) improved trajectory regularity via spectral weighting, and (iii) a deterministic failure detector without reward signals. Validation across two application domains demonstrates substantial improvements: on the PushT manipulation benchmark, the Cameron-Martin loss achieves a 17% improvement in maximum episode reward (0.95 vs. 0.78 for MSE) and 67.6% reduction in inter-step drifts during inference via the introduced residual magnitude. Similarly, on a 6-station manufacturing line with constant work-in-process (CONWIP) flow control, we achieve 28.4% lower RMSE than classical LSTM baselines; a high starvation-event recall (1.0 in test cycles), and effective bottleneck identification (Precision@1 = 1.0 in test set, 13x signal-to-noise ratio). We then certify the dispatch policies with Hamilton-Jacobi reachability theory which reduces deadlock events by 96% compared to uncontrolled dispatch over 100 simulated runs (351 events prevented).