MeanFlow Meets Control: Scaling Sampled-Data Control for Swarms

1. Problem Statement

The paper addresses the challenge of steering large-scale swarms (e.g., robots, autonomous vehicles) from an initial probability distribution ($\rho_0$) to a target distribution ($\rho_1$) using sampled-data control.

• The Core Conflict: Traditional generative models (like Flow Matching) learn instantaneous velocity fields ($v_t$) assuming infinitesimal time steps. However, real-world control systems operate in sampled-data regimes: control inputs are updated at discrete intervals and held constant (or follow a specific profile) over finite windows.
• The Consequence: In "few-step" control scenarios (where large time intervals exist between updates), the mismatch between learning an instantaneous velocity and applying a finite-window control leads to significant discretization errors and instability.
• The Goal: Develop a learning framework that respects the sampled-data structure, learning a quantity that is naturally compatible with finite-horizon control updates rather than infinitesimal steps.

2. Methodology

The authors propose a Control-Space MeanFlow framework adapted for Linear Time-Invariant (LTI) swarm dynamics.

A. System Dynamics

The swarm is modeled by a linear system:
$$\dot{z}(\tau) = Az(\tau) + Bu(\tau)$$
where $z$ is the state, $u$ is the control input, and $(A, B)$ are fixed matrices. The system is assumed to be controllable.

B. The Learned Object: Interval Coefficient

Instead of learning an instantaneous velocity $v(z, t)$, the model learns a finite-horizon minimum-energy steering coefficient, denoted as $c_\theta(z_t, t, r)$.

• Definition: For a time interval $[t, r]$, the optimal control that moves the system from $z_t$ to $z_r$ with minimum energy is given by:
  $$u^*(\tau) = B^\top \Phi(r, \tau)^\top c(z_t, t, r)$$
  where $\Phi$ is the state transition matrix.
• The Coefficient: The vector $c$ is derived from the Controllability Gramian $W(t, r)$:
  $$c(z_t, t, r) = W(t, r)^{-1} (z_r - \Phi(r, t)z_t)$$
  This single vector $c$ fully parameterizes the optimal control profile over the entire interval $[t, r]$.

C. Training via Bridge Supervision

The framework uses a "bridge" construction connecting samples from $\rho_0$ and $\rho_1$.

1. Bridge Trajectory: A minimum-energy path is constructed between sampled endpoints $z_0$ and $z_1$ over the horizon $[0, 1]$.
2. Local Differential Identity: The authors derive a key identity relating the interval coefficient $c$ to local quantities along the bridge. By differentiating the integral representation of $c$ with respect to the start time $t$, they obtain a residual equation:
   $$\frac{d}{dt}c(z_t, t, r) \approx \text{function of local velocity } v(z_t, t)$$
3. Loss Function: The model $c_\theta$ is trained to satisfy this differential identity in a least-squares sense. Crucially, the training target involves a stop-gradient operator on the target term, ensuring stable one-sided regression similar to MeanFlow.
   $$\mathcal{L}(\theta) = \mathbb{E} \left\| R_\theta(z_t, t, r) \right\|^2$$
   where $R_\theta$ is the residual between the predicted coefficient's evolution and the local bridge dynamics.

D. Implementation (Sampling)

During inference, the bridge is discarded. The controller operates as a sampled-data state-feedback:

1. At discrete time $t_k$, observe the current state $z_k$.
2. Query the network to get the coefficient $c_k = c_\theta(z_k, t_k, t_{k+1})$.
3. Apply the control signal $u(\tau) = B^\top \Phi(t_{k+1}, \tau)^\top c_k$ for $\tau \in [t_k, t_{k+1}]$.
4. Propagate the state analytically to $z_{k+1}$.

3. Key Contributions

1. Control-Space MeanFlow: The paper successfully adapts the MeanFlow concept (learning window-level quantities) to the specific constraints of linear control theory. It shifts the learning target from state-space velocity to the control-space coefficient.
2. Theoretical Consistency:
   • Lemma 1: Proves that when $A=0, B=I$, the formulation reduces exactly to standard MeanFlow.
   • Lemma 2 (Additivity): Demonstrates that the steering action over a long interval can be decomposed into successive shorter intervals without losing dynamic consistency, ensuring the method works for multi-step propagation.
3. Sampled-Data Native Design: The learned coefficient is not just a training artifact; it is the exact quantity used in the implementation. This guarantees that the prescribed dynamics and actuation maps are respected by construction, eliminating the "mismatch" error common in standard flow matching applied to control.
4. Robustness: Because the coefficient depends on the current state $z_k$, the controller naturally adapts to perturbations or particle loss during propagation (closed-loop behavior), even though the training uses open-loop bridges.

4. Results

The authors validate the framework through numerical experiments in PyTorch:

• 2D Planar Navigation:
  • Task: Transforming a swarm shaped like "AYKJ" to "DCJK".
  • Findings: The method successfully handles both drift-free dynamics and systems with rotational drift ($A \neq 0$). In the rotational case, the learned transport plan correctly adapts the particle correspondence to account for the 90° rotation induced by the drift, whereas a naive approach would fail.
• 3D Spatial Maneuvering:
  • Task: Transforming a pyramid distribution to a torus distribution in 3D space.
  • Findings: The model successfully navigates the complex geometric reorganization (carving out a hollow center) under various drift conditions (planar rotation vs. 3D rotation). The 3D rotational drift significantly alters the intermediate geometry, and the model correctly learns the non-trivial paths required to reach the target.
• Performance: The method achieves high-quality distribution matching using only 16 discrete control steps, demonstrating scalability for few-step control.

5. Significance

This paper bridges a critical gap between generative AI (Flow Matching) and control theory (Sampled-Data Control).

• Practical Relevance: It provides a scalable solution for real-world swarm control where high-frequency updates are impossible or inefficient.
• Theoretical Insight: It establishes that for linear systems, the "natural" object to learn is not a velocity field, but the coefficient of the minimum-energy control over a finite window.
• Future Directions: The framework opens avenues for extending MeanFlow to nonlinear dynamics, constrained control, and stability analysis under coarse-time execution, moving beyond the current LTI assumption.

In summary, MeanFlow Meets Control offers a principled, mathematically rigorous, and practically implementable framework for steering large swarms with minimal control updates, ensuring that the learned policy is inherently compatible with the discrete nature of real-world actuators.