Nonlinear Stochastic Density Steering via Gaussian Mixture Schrodinger Bridges and Multiple Linearizations

Abstract: The paper studies the optimal density steering problem for nonlinear continuous-time stochastic systems. To accurately capture nonlinear dynamics in high-uncertainty regions that deviate significantly from a nominal linearization point, we introduce the concept of Multiple Distribution-to-Distribution Linearization. The proposed approach first approximates the boundary distributions using Gaussian Mixture Models (GMMs), and decomposes the original nonlinear problem into a collection of Gaussian-to-Gaussian Optimal Covariance Steering (OCS) subproblems between pairs of mixture components. Each elementary OCS problem is solved via local linearization around the mean trajectory connecting the corresponding initial and terminal Gaussian components. The resulting elementary policies are then combined according to their associated conditional densities. We prove that the proposed multi-linearization approach yields tighter approximation error bounds than single-linearization for a broad class of problems. The effectiveness of the approach is demonstrated through numerical experiments on an Earth-to-Mars orbit transfer scenario.

• The paper introduces a novel method that decomposes the nonlinear density steering task into tractable Gaussian-to-Gaussian optimal covariance steering subproblems.
• It leverages Gaussian Mixture Schrödinger Bridges and multiple local linearizations to accurately preserve multi-modal state distributions.
• Numerical results demonstrate significant reductions in control cost and terminal error compared to single linearization approaches.

Nonlinear Stochastic Density Steering via Gaussian Mixture Schrödinger Bridges and Multiple Linearizations

Problem Formulation and Motivation

The paper addresses nonlinear stochastic density steering, where the state of a continuous-time nonlinear SDE is transitioned from an arbitrary initial probability distribution to a target distribution under control constraints, with both distributions possibly multimodal and represented as Gaussian Mixture Models (GMMs). This problem is central in stochastic optimal control, mean-field control, and Schrödinger Bridge theory. The authors critique the limitations of existing approaches—specifically, linearization-based methods—which degrade in performance when the state distributions are distant from unimodality or global mean statistics. Such regimes arise, for example, in trajectory planning for autonomous vehicles and aerospace systems, where multi-modal uncertainty must be propagated and controlled explicitly.

The method uses a GMM-based Schrödinger Bridge (SB) construction and introduces a multiple linearization (ML) framework, wherein the nonlinear density steering problem is decomposed into tractable Gaussian-to-Gaussian Optimal Covariance Steering (OCS) subproblems. Each component-to-component subproblem is solved via local linearization of the nonlinear drift, and global control is synthesized by blending these affine subpolicies according to a transport plan.

Figure 1: The problem setting—control steers the stochastic system from an initial GMM distribution to a target GMM, with ellipses denoting high-probability regions of Gaussian mixture components.

Background: Covariance Steering, Schrödinger Bridges, and Linearization

The OCS framework allows exact steering of a linear stochastic system between Gaussian marginals using affine feedback and semidefinite programming (SDP). Nonlinear systems, however, entangle mean and covariance evolution, eliminating closed-form solutions. Iterative Covariance Steering (iCS) and single linearization (SL) methods locally approximate the system but break down under pronounced nonlinearity or multi-modality, as much of the probability mass then lies beyond the region of valid linearization.

The Schrödinger Bridge framework generalizes OCS to arbitrary marginals (with entropic regularization), naturally capturing multi-modal boundary conditions. Recent works have applied GMMs within SBs to yield efficient transport couplings, but no previous method reliably integrates nonlinear drift, GMM marginals, explicit feedback synthesis, and rigorous enforcement of terminal density constraints.

Multiple Linearizations for Nonlinear GMM Boundary Problems

The authors propose a multi-indexed collection of local OCS subproblems, each connecting an initial GMM component to a final component. For each $(i,j)$ pair, the system is linearized around a reference trajectory connecting the means, and the corresponding OCS solution is computed by SDP.

The total control at state $x$ and time $t$ is then a mixture of these affine controls, weighted by the posterior that $x$ was sampled under the $(i,j)$ bridge, with weights determined by solving an optimal transport linear program on the mixture masses. This ensures the overall control explicitly respects both the nonlinear system dynamics and the multi-modal structure of the marginals.

The paper develops formal links between this construction, Fokker-Planck evolution, and the SB dynamical system. The total mixture density evolves under an FPK equation whose drift is shown to approximate the true nonlinear drift more accurately (in expectation and under certain convexity conditions strictly) than any single global linearization.

Theoretical Error Analysis

The authors prove that, under uniform quadratic local approximation error for the drift, the expected error of function approximation using ML bounds that of SL from above. For vector fields with componentwise convexity or concavity, ML yields strictly lower expected error unless all GMM means coincide.

These results formally support the intuition that, for multi-modal distributions, local linearizations produce more accurate aggregate models of nonlinear dynamics than linearizing only at the global mean. In the density steering setting, this translates to superior terminal distribution matching and control efficiency.

Numerical Results: Stochastic Spacecraft Orbit Transfer

Comprehensive numerical results are provided using an Earth-to-Mars stochastic transfer in Sun-centered inertial coordinates, with both position and velocity subject to process noise. The system is governed by nonlinear gravity and thruster-driven dynamics, presenting highly non-Gaussian state distributions as the transfer evolves.

Case I: Multi-Modal Initial and Target Distributions

The first scenario investigates multimodal boundary distributions: three well-separated initial and two terminal Gaussian components.

Figure 2: Visualization of the multimodal initial distribution $\rho_0$ with covariance ellipses.

The SL controller, which linearizes about the aggregate mean, fails to steer the distribution mass efficiently along all high-probability trajectories, resulting in bulk misplacement and larger dispersion at the target. In contrast, the ML controller, via superposition of bridge-specific policies, preserves modal separation and tightly matches the terminal GMM marginals.

Figure 3: Phase-space evolution: SL (left) cannot match the modal structure, while ML (right) realizes bridges connecting all GMM pairs; ML achieves notably tighter sample covariance throughout transfer (zoom panels).

Figure 4: Empirical terminal distribution: ML closely matches the desired multi-modal target, while SL exhibits substantial mismatch.

Numerically, ML yields a 10.6% reduction in average control cost and a Sliced Wasserstein-2 terminal error less than one-third that of SL, confirming improved steering efficacy for multi-modal distributions.

Case II: High-Uncertainty, Broad Covariance Marginals

For broader, nearly-Gaussian marginals, artificially fitted GMMs are employed to activate ML superiority. Both controllers converge in mean, but ML realizes the distribution transfer at an 18.9% lower control cost and with improved Sliced Wasserstein-2 final matching.

Figure 5: Initial high-variance (but close to Gaussian) distribution with component ellipses.

Figure 6: Phase-space evolution in Case II, visual layout as above; ML again achieves tighter distribution control.

Figure 7: ML meets the terminal distribution constraints more faithfully than SL.

Empirical results consistently show ML dominates SL in cost and terminal matching, in some cases by a statistically significant margin.

Implications and Future Prospects

The combination of GMM Schrödinger Bridges, multiple linearization, and explicit optimal transport coupling provides a scalable method for nonlinear, distributionally robust control under complex boundary conditions. The resulting feedback law is a convex mixture of affine policies, automatically adapting to trajectory-dependent uncertainties.

Practically, this methodology has direct application in stochastic optimal control of aerospace vehicles, robotics, and any system where uncertainty propagation and nonlinear dynamics are non-negligible and must be robustly regulated to target densities.

Theoretically, this framework opens routes to hybrid SB/OCS methods for high-dimensional systems, connections to diffusion-based generative models, and structured optimal transport formulations. Future work includes extension to discrete-time systems, integration of nonparametric boundary densities, and joint learning of transport couplings in more general policy classes—directions likely to stimulate further advances in AI and control for complex stochastic dynamical systems.

Conclusion

The paper establishes a rigorous and efficient approach for nonlinear stochastic density steering with GMM marginals, demonstrating that multiple local linearizations, optimally coupled, fundamentally outperform single linearization methods on both theoretical and empirical grounds. The approach delivers both lower control effort and superior fidelity in terminal distribution matching for challenging stochastic systems, highlighting its utility for high-impact control problems with multi-modal uncertainties (2604.15576).