Trajectory-Optimized Density Control with Flow Matching

TL;DR

This work addresses trajectory-aware density control for multi-agent systems by uniting flow matching with mean-field control and Schrödinger bridge concepts, enabling explicit optimization over transport paths rather than solely endpoints. It develops a generalized SB-FBSDE framework with MF interactions, accompanied by Temporal Difference and Flow Matching objectives to train the model. An alternating optimization algorithm is proposed to jointly learn forward and backward dynamics, yielding collision-free transport plans. Empirical results on a generalized GMM with obstacles and a V-neck corridor demonstrate effective collision avoidance and robust path coordination with computational efficiency on par with standard flow-matching methods.

Abstract

Optimal transport (OT) and Schr{ö}dinger bridge (SB) problems have emerged as powerful frameworks for transferring probability distributions with minimal cost. However, existing approaches typically focus on endpoint matching while neglecting critical path-dependent properties -- particularly collision avoidance in multiagent systems -- which limits their practical applicability in robotics, economics, and other domains where inter-agent interactions are essential. Moreover, traditional density control methods often rely on independence assumptions that fail to capture swarm dynamics. We propose a novel framework that addresses these limitations by employing flow matching as the core modeling tool, where the flow model co-evolves with the control policy. Unlike prior methods that treat transport trajectories as mere interpolations between source and target distributions, our approach explicitly optimizes over the entire transport path, enabling the incorporation of trajectory-dependent costs and collision avoidance constraints. Our framework bridges optimal transport theory with mean field control, providing a principled approach to multiagent coordination problems where both endpoint alignment and path properties are critical. Experimental results demonstrate that our method successfully generates collision-free transport plans while maintaining computational efficiency comparable to standard flow matching approaches.

Figures (5)

• Figure 1: Problem setup for the Gaussian mixture model transport task. The source distribution (center) is a standard Gaussian centered at the origin, while the target distribution consists of eight equally-weighted Gaussian components arranged uniformly on a circle of radius $16$. Three circular obstacles of radius $1.5$ are strategically positioned to obstruct direct transport paths.
• Figure 2: Forward transport for the GMM problem at stage 20. Top row: Results from the method in Liu2022. Bottom row: Results from the proposed method. Left column: Agent trajectories navigating from the source to target while avoiding obstacles. Right column: Final distribution of agents, demonstrating successful coverage of all eight target components.
• Figure 3: Backward transport for the GMM problem at stage 20. Top row: Results from the method in Liu2022. Bottom row: Results from the proposed method. Left column: Agent trajectories from the distributed target back to the source. Right column: Final distribution showing successful reconcentration at the origin.
• Figure 4: Problem setup for the V-neck corridor navigation task. The source distribution (left) and target distribution (right) are connected by a V-shaped constrained passage, creating a bottleneck through which all agents must coordinate their movement.
• Figure 5: Transport results for the V-neck corridor problem at stage 20. Left: Forward transport showing agents successfully navigating through the bottleneck from left to right. Right: Backward transport demonstrating successful passage in the reverse direction.

Theorems & Definitions (1)

• Theorem 3.1: Generalized SB-FBSDEs