Continuity-Regularized Flow Matching for Offline Reinforcement Learning

Flow-matching policies have recently emerged as a powerful class of generative models for offline reinforcement learning (RL), capable of capturing complex, multi-modal action distributions from static datasets. However, standard training objectives are largely agnostic to the global properties of the generative path, permitting learned vector fields that are irregular and unstable, which can hinder performance. In this work, we introduce PDE-regularized Q-Learning (PQL), a novel algorithm that addresses this limitation by imposing a principled structure on the entire probability flow. PQL makes two synergistic contributions: first, a partial differential equation based regularizer derived from the continuity equation promotes global smoothness and stability on the flow. Second, to solve the complex optimization problem introduced by this regularizer, we propose a Beta-distributed timestep sampling strategy that focuses learning on the critical trajectory segments where the trade-off between imitation and smoothness is most acute. Through extensive experiments, we demonstrate that by structuring the generative journey and not just its destination, PQL achieves state-of-the-art performance on a wide range of challenging offline RL tasks.

How can AI systems learn good behaviors purely from past recordings, without risky real-world trial and error? This matters for robotics, self-driving cars, and any setting where on-the-fly experimentation is expensive or dangerous. A popular new approach teaches AI to generate skilled actions by smoothly transforming random starting points into useful answers, then steers this transformation toward whatever looks most valuable based on past experience. The catch: the AI's sense of "value" is only reliable for actions resembling the recordings. Pushing into unfamiliar territory turns these judgments into wild guesses, and the smooth transformation becomes erratic, with the AI recommending actions that look great on paper but fail in practice. We propose PQL, a method that adds a smoothness rule keeping the transformation well-behaved even under value-driven pressure. We also focus training on the part of the process where erratic behavior is most likely. Across major benchmarks for AI decision-making, PQL outperforms previous methods, with the largest gains on the hardest problems. This brings us closer to AI that can be safely deployed in physical environments without extensive trial and error.