Kolmogorov Regression for Robust Diffusion Policies

Lekan Molu

cs.LGJun 16, 2026

Kolmogorov Regression for Robust Diffusion Policies

Authors: Lekan Molu

Organizations: Bala Cynwyd, PA 19004

Abstract

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

Explore similar work

Jul 8, 2026stat.ML

Expressivity and Statistical Trade-offs in Diffusion Policy Learning

Diffusion-based policies have recently emerged as powerful policy parameterizations for reinforcement learning, representing state-conditioned action distributions as terminal laws of diffusion processes with parameterized drifts. This terminal-law representation has shown substantial expressive flexibility in practice, enabling diffusion policies to model complex, multimodal, and highly non-Gaussian action distributions; however, it remains unclear what mathematically drives this expressivity and how to fully exploit it when the policy is learned from finite data. In this paper, we identify the drift Lipschitz budget

K

as a central quantity governing the expressivity and statistical behavior of diffusion policies. We quantify expressivity through approximation: diffusion policies with

K

-Lipschitz drifts can concentrate near optimal deterministic policies and achieve value approximation error of order

1/K

; moreover, we prove a matching lower bound under nondegenerate diffusion noise. This increased expressivity comes with a statistical cost. When the drift is parameterized by neural networks, increasing

K

improves approximation but increases statistical complexity. Balancing these two terms yields a finite-sample performance gap of order

\tilde{O}(n^{-2/(m+6)})

for generic neural-network drifts, and a sharper rate

\tilde{O}(n^{-2/(m+4)})

for one-sided dissipative drift classes, where

n

is the sample size and

m

is the dimension of the state space. Numerical experiments provide empirical evidence for the sample-dependent trade-off in

K

, supporting both theoretical regimes. Our framework also suggests a practical implementation principle: choose the diffusion budget

K

according to the available sample size, and then select a neural-network architecture with the corresponding fixed Lipschitz coefficient.

Viet Vu, Renyuan Xu, Jiacheng Zhang +1

Diffusion Policies Lipschitz Constant Expressivity Trade-Off

Ask PDF

Jun 19, 2026cs.RO

BayesFP: Posterior Estimation for Flow-Based Policies via Feynman-Kac Sampling

Robots must generate trajectories that remain faithful to learned expert behavior while satisfying safety constraints and task-specific objectives specified only at inference time. We formulate constrained trajectory generation for pretrained diffusion and flow-matching policies as Bayesian posterior sampling, with the learned demonstration distribution as a prior and an inference-time, cost-derived likelihood tilting it toward feasible, optimal trajectories. To sample from this posterior without any retraining of the base policy, we leverage the Feynman--Kac corrector framework, originally formulated for diffusion models, and extend it to deterministic flow-matching policies. The result is a unified, inference-time, retraining-free sampler for diffusion and flow policies. We validate the approach on pretrained Diffusion Policy, GR00T-N1.6, and

π_{0.5}

checkpoints across simulated and real-world manipulation tasks, including planning around non-convex obstacles introduced at inference time, and show improvements over the base

π_{0.5}

on zero-shot tasks.

Sreevardhan Sirigiri, Weiming Zhi, Fabio Ramos

Flow-Based Policy Diffusion Policies Posterior Sampling Feynman-Kac Latent Trajectories Posterior+1

Ask PDF

May 21, 2026cs.LG

Factored Diffusion Policies:Compositionally Generalized Robot Control with a Single Score Network

Robotic tasks are typically specified by a tuple of factors, such as the object to be grasped, the obstacles to be avoided, the color of the target, and so on. Collecting expert demonstrations for every combination of factor values grows combinatorially. We present factored diffusion policies: a single shared diffusion network trained with per-factor null-token dropout, whose score decomposes additively across factors at inference. Under approximate conditional independence between factors given the action-observation pair, this composition approximates the true joint score with a bounded uniform error, reducing the training-task budget from a product of factor cardinalities to a sum. A trajectory-tube certificate chains this score-level bound through the reverse-time sampling ODE and a contracting tracking controller into a closed-loop state-trajectory tube whose radius factors into an ODE-sensitivity constant and a per-factor score-error budget. Unlike compositional-diffusion methods for control that combine separately trained networks, we use one shared network. Drone racing experiments confirm both the generalization bound and the certificate. On state-based multi-gate racing, the factored policy passes 90% of held-out gates -- matching an oracle -- while a K-network composition baseline collapses to 3%; on vision-based single-gate traversal, it transfers zero-shot to an unseen venue with +11.7pp success-rate gain and 2.4X crash-rate reduction.

Sayan Mitra, Ege Yuceel, Noah Giles +1

Diffusion Policies Robot Policies Autonomous Racing

Ask PDF