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

1. Problem Statement

The paper addresses the Stochastic Optimal Control (SOC) problem of steering a nonlinear dynamical system from a prescribed initial probability distribution to a desired terminal probability distribution.

• System Dynamics: The system is modeled by an Itô stochastic differential equation (SDE) with nonlinear drift $f(x_t)$, linear control input, and additive Brownian motion noise.
• Objective: Minimize the expected control effort (quadratic cost) while satisfying explicit boundary constraints on the state distribution at $t=0$ and $t=T$.
• The Challenge:
  • Standard SOC problems with terminal distribution constraints are analytically intractable for nonlinear systems due to the Hamilton–Jacobi–Bellman (HJB) equation's complexity and the "curse of dimensionality."
  • Existing methods often rely on single linearization (linearizing the nonlinear dynamics around a single global mean trajectory). This approach fails when the state distribution is multi-modal or exhibits large uncertainty, as significant portions of the probability mass may lie outside the region where the single linear approximation is valid.
  • While Gaussian Mixture Models (GMMs) can represent complex, multi-modal distributions, applying them to nonlinear density steering requires a framework that handles the coupling between mean and covariance dynamics in a nonlinear setting.

2. Methodology

The authors propose a Multiple Distribution-to-Distribution Linearization framework that decomposes the global nonlinear problem into a set of local, solvable subproblems.

A. Core Concept: Decomposition into GMMs

Instead of treating the initial ($\rho_0$) and terminal ($\rho_T$) distributions as single Gaussians, the method approximates them as Gaussian Mixture Models (GMMs) with $N_0$ and $N_1$ components, respectively.

• The problem is decomposed into $N_0 \times N_1$ subproblems, where each subproblem steers a specific initial Gaussian component ($i$) to a specific terminal Gaussian component ($j$).

B. Multiple Linearization Strategy

For each pair $(i, j)$:

1. Reference Trajectory: Solve a deterministic optimal control problem to find the optimal open-loop trajectory $\mu_{t|ij}$ connecting the mean of the $i$-th initial component to the mean of the $j$-th terminal component.
2. Local Linearization: Linearize the nonlinear drift function $f(x)$ around this specific reference trajectory $\mu_{t|ij}$. This yields a local linear time-varying system with Jacobian $A_{t|ij}$.
3. Local Optimal Covariance Steering (OCS): Solve the linear OCS problem for each $(i, j)$ pair using Semi-Definite Programming (SDP). This generates an optimal affine feedback control law $u_{t|ij}(x)$ and a conditional Gaussian density $\rho_{t|ij}$.

C. Global Control Synthesis

The final nonlinear control policy $u_t(x)$ is constructed as a weighted mixture of these local affine policies:
$$u_t(x) = \sum_{i,j} u_{t|ij}(x) \cdot w_{ij}(x)$$
Where the weights $w_{ij}(x)$ are determined by the posterior probability that the current state $x$ belongs to the $(i, j)$-th Gaussian component, derived from the Schrödinger Bridge framework. The transport plan (weights $\lambda_{ij}$) between components is optimized via a Linear Program (LP) to minimize total transport cost.

D. Theoretical Validation

The authors provide theoretical error analysis comparing Single Linearization (SL) vs. Multiple Linearization (ML):

• Theorem 1: Under uniform quadratic bounds, the upper bound on the expected linearization error for ML is strictly lower than or equal to that of SL.
• Theorem 2: If the system dynamics are component-wise convex or concave, ML provides a strictly better approximation of the nonlinear function than SL, provided the GMM component means are distinct.
• FPK Consistency: The authors show that the resulting mixture density satisfies the Fokker-Planck-Kolmogorov (FPK) equation with a drift term that is a better approximation of the true nonlinear drift than a single linearization.

3. Key Contributions

1. Novel Framework: Introduction of a multiple linearization framework that bridges Gaussian Mixture Schrödinger Bridges with Optimal Covariance Steering (OCS) for nonlinear systems.
2. Handling Multi-Modality: The method explicitly handles multi-modal uncertainty and large deviations from the mean, which single-linearization methods fail to capture.
3. Computational Tractability: By decomposing the problem, the authors reduce a complex nonlinear stochastic control problem to a series of convex SDP problems and a single LP, making the solution computationally viable.
4. Theoretical Guarantees: Proof that multiple linearizations yield tighter error bounds and better approximation accuracy than single linearization for multi-modal distributions.

4. Numerical Results

The method was validated using an Earth-to-Mars orbit transfer scenario under nonlinear orbital dynamics (Sun-centered inertial frame) and non-Gaussian uncertainty. Two cases were tested:

• Case I (Multi-Modal Boundaries): The initial and terminal distributions consisted of separated Gaussian clusters (3 initial, 2 terminal).
  • Result: The Single Linearization (SL) approach failed to capture the separation of probability mass, leading to a poor terminal distribution match. The Multiple Linearization (ML) approach successfully steered each cluster to its target, maintaining tight covariance control.
• Case II (Large Covariance): Distributions with large variances were approximated by GMMs to leverage multiple linearizations.
  • Result: The ML approach achieved the target distribution with significantly lower control cost and higher accuracy compared to SL.

Quantitative Metrics (Table 2):

• Control Cost ($J_{ctrl}$): ML reduced control cost by ~10.6% (Case I) and ~18.9% (Case II) compared to SL.
• Distribution Matching (Sliced $W_2$): ML achieved a drastically lower Wasserstein distance (e.g., 1.25 vs. 4.45 in Case I), indicating a much closer match to the desired terminal distribution.

5. Significance

This work fills a critical gap in stochastic control for aerospace and autonomous systems where:

• Nonlinearity is dominant: Standard linearization techniques break down under high uncertainty or complex dynamics (e.g., orbital mechanics).
• Distribution constraints are explicit: Applications require precise control over the shape of the state distribution (e.g., avoiding collision zones, ensuring multi-target rendezvous), not just the mean.
• Scalability is required: The decomposition into convex subproblems allows for the application of rigorous optimal control to complex, high-dimensional nonlinear systems without resorting to intractable numerical PDE solvers.

The proposed framework offers a computationally efficient, theoretically grounded, and empirically superior alternative to existing linearization-based methods for nonlinear stochastic density steering.