# Kolmogorov Regression for Robust Diffusion Policies

arXiv: https://arxiv.org/abs/2606.18186 · HTML: https://mlsanity.com/paper/2606.18186/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models.
- Builds on 15 of its 45 references that are in the map.
- First posted on arXiv 16 Jun 2026.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)

## Builds on (most cited first, 15 of 15)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [OpenVLA: An Open-Source Vision-Language-Action Model](https://mlsanity.com/paper/2406.09246/) (2024)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Improved Techniques for Training Score-Based Generative Models](https://mlsanity.com/paper/2006.09011/) (2020)
- [Implicit Behavioral Cloning](https://mlsanity.com/paper/2109.00137/) (2021)
- [Convergence of denoising diffusion models under the manifold hypothesis](https://mlsanity.com/paper/2208.05314/) (2022)
- [Improved Analysis of Score-based Generative Modeling: User-Friendly Bounds under Minimal Smoothness Assumptions](https://mlsanity.com/paper/2211.01916/) (2022)
- [BeyondMimic: From Motion Tracking to Versatile Humanoid Control via Guided Diffusion](https://mlsanity.com/paper/2508.08241/) (2025)
- [A Careful Examination of Large Behavior Models for Multitask Dexterous Manipulation](https://mlsanity.com/paper/2507.05331/) (2025)
- [The Principles of Diffusion Models](https://mlsanity.com/paper/2510.21890/) (2025)
- [Infinite-Dimensional Diffusion Models](https://mlsanity.com/paper/2302.10130/) (2023)
- [The Python LevelSet Toolbox (LevelSetPy)](https://mlsanity.com/paper/2411.03501/) (2024)

## Abstract (arXiv)

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).