Distribution Steering via Sliced Optimal Transport Control
Kaito Ito, Anqi Dong
The University of Tokyo · KTH Royal Institute of Technology
math.OC, cs.LG, cs.SY, eess.SY, stat.ML
Submitted: 2026-08-13
Updated: 2026-08-14
Comments: 39 pages
License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
Importance score: 75/100
The gist: Distribution steering seeks feedback laws that drive the state law of a dynamical system between prescribed initial and terminal distributions.
Terminology
Summary
Distribution steering seeks feedback laws that drive the state law of a dynamical system between prescribed initial and terminal distributions. Optimal transport provides a natural geometric approach, but its implementation generally requires a transport map or coupling in the full state space. Sliced optimal transport avoids this full-dimensional construction through one-dimensional projections. Yet, the resulting projected maps specify only directional displacements and do not by themselves prescribe a realizable feedback law. To this end, we develop a finite-horizon control framework based on sliced optimal transport. At each sampling instant, a projected optimal transport map defines a directional terminal condition, whose minimum-energy realization yields a randomized single-direction controller. Averaging over projection directions gives a deterministic sliced feedback. For the single-integrator dynamics, the averaged feedback makes the sliced Wasserstein distance to the target non-increasing. For Gaussian endpoint laws, it is affine, preserves Gaussianity, and steers the mean and covariance to their prescribed terminal values. We further identify a law-dependent gain that yields linear decay of the sliced Wasserstein distance together with an explicit characterization of the control energy. We also prove that the randomized controller converges to the averaged sliced flow as the sampling period vanishes. Finally, we extend the construction to linear dynamical systems. Reachability-normalized coordinates allow instantaneous realization of the sliced velocity for uniformly fully actuated systems, while local controllability Gramians provide exact finite-step realization for general controllable systems. Numerical examples illustrate the resulting distributional flows.
Improvements for AI systems
Improvements to AI Systems:
-
Sample-Efficient Distributional Control for Robotics: Use the sliced optimal transport (SOT) feedback law to design controllers that steer a robot swarm’s spatial distribution (e.g., from a Gaussian start to a target formation) without needing full-state coupling maps. The AI system can compute real-time, deterministic feedback from 1D projections, reducing computational cost by orders of magnitude compared to full optimal transport solvers.
-
Stochastic Model Predictive Control (MPC) with Guaranteed Convergence: Integrate the averaged sliced feedback into MPC frameworks for linear systems. The AI system can now guarantee non-increasing sliced Wasserstein distance to a target distribution at each step, enabling safe, provably convergent planning under stochastic disturbances—useful for autonomous vehicles or drones in uncertain environments.
-
Gaussian Mean-Covariance Steering for Kalman Filters: The affine, Gaussianity-preserving property allows AI systems to directly steer the mean and covariance of a belief state (e.g., in a Kalman filter) to prescribed values. This improves active perception systems—e.g., a robot can move to minimize uncertainty (covariance) or match a desired belief distribution for target tracking.
-
Energy-Aware Path Planning: The explicit characterization of control energy and linear decay rate lets AI systems trade off convergence speed vs. energy consumption. An improved planner can select a law-dependent gain to achieve a desired decay rate while minimizing fuel or battery use—critical for long-duration autonomous missions.
-
Real-Time Randomized-to-Deterministic Conversion: The proof that the randomized controller converges to the averaged flow as sampling period vanishes enables AI systems to use cheap, randomized single-direction controllers in high-frequency loops, then switch to the deterministic averaged law for stability—ideal for embedded systems with limited compute.
-
General Linear System Reachability Normalization: For fully actuated or controllable linear systems, the AI system can precompute reachability-normalized coordinates or local controllability Gramians. This allows exact finite-step realization of the sliced velocity, enabling distributional control for non-trivial dynamics (e.g., quadrotors, manipulators) without iterative optimization.
-
Distributional Reinforcement Learning (RL) Prior: The SOT-based feedback can serve as a model-based prior or shaping reward in RL. The AI system can initialize policies to match target state distributions, accelerating learning and providing a theoretical convergence guarantee during early exploration—reducing sample complexity in tasks like robotic assembly or navigation.
What the Improved AI System Can Do:
-
Steer a swarm or a single stochastic system to a target distribution in real time with minimal computation.
-
Guarantee monotonic convergence in distributional distance, with explicit energy bounds.
-
Handle Gaussian and non-Gaussian endpoints without full-dimensional transport maps.
-
Operate on linear dynamics with exact finite-step reachability, not just integrators.
-
Provide both randomized (cheap) and deterministic (smooth) control modes with proven equivalence in the limit.
-
Enable energy-aware, distributionally safe planning for autonomous systems under uncertainty.
Sources
Related papers
- Lions and Muons: Optimization via Stochastic Frank-Wolfe under Heavy-Tailed Noise
- Adam-HNAG: A Convergent Reformulation of Adam with Accelerated Rate
- Incremental Learning in Mirror Flows
- Online Control via Counterfactual Tracking
- Asynchronous Replanning in Two Population Linear Quadratic Mean Field Games: Information Requirements and Stability
- Petrov-Galerkin operator inference with application to stability-encouraging identification