Flow matching for stochastic linear control systems
Yuhang Mei, Mohammad Al-Jarrah, Amirhossein Taghvaei, Yongxin Chen
TL;DR
The paper addresses steering an initial distribution $P_0$ to a target distribution $P_1$ under deterministic or stochastic linear dynamics using a flow-matching framework that operates through control channels. It constructs an interpolating probability flow $P_t$ and defines a feedback law $k(t,\xi)=\mathbb{E}\left[\frac{d X^z_t}{dt}\middle| X^z_t=\xi\right]$, enabling distributional interpolation; analytical solutions are provided for Gaussian and Gaussian-mixture targets, and a least-squares neural-network procedure approximates the law in general. The method extends flow matching to linear control systems, yielding the same control law for both deterministic and stochastic cases, with numerical experiments in 2D and higher dimensions demonstrating capability to steer mixtures and scale with dimension. This work has practical implications for robotic swarms and stochastic thermodynamics, where control inputs constrain the evolution of probability densities and efficient distribution steering is valuable.
Abstract
This paper addresses the problem of steering an initial probability distribution to a target probability distribution through a deterministic or stochastic linear control system. Our proposed approach is inspired by the flow matching methodology, with the difference that we can only affect the flow through the given control channels. The motivation comes from applications such as robotic swarms and stochastic thermodynamics, where agents or particles can only be manipulated through control actions. The feedback control law that achieves the task is characterized as the conditional expectation of the control inputs for the stochastic bridges that respect the given control system dynamics. Explicit forms are derived for special cases, and a numerical procedure is presented to approximate the control law, illustrated with examples.