Schrodinger Bridges and Density Steering Problems for Gaussian Mixtures Models in Discrete-Time

1. Problem Statement

The paper addresses the challenge of controlling the probability density of discrete-time dynamical systems to steer them from a prescribed initial distribution ($\rho_0$) to a desired terminal distribution ($\rho_N$) within a finite time horizon. Specifically, the authors focus on Gaussian Mixture Models (GMMs) as the boundary distributions, which are more expressive and realistic than single Gaussian distributions.

The paper tackles two distinct but related formulations:

• Discrete-Time Schrödinger Bridge (dtSB): Finding the stochastic process closest to a reference dynamics (minimizing Kullback-Leibler divergence) that satisfies the marginal distribution constraints.
• Discrete-Time Density Steering (DS): Finding a control policy that minimizes the expected control effort (e.g., $\sum E[\|u_k\|^2]$) while satisfying the same marginal constraints.

The core difficulty lies in the fact that while analytical solutions exist for single Gaussian boundaries, the non-linearity introduced by mixing multiple Gaussian components makes the general GMM problem intractable without approximation.

2. Methodology

The authors propose a randomized Markovian feedback policy that constructs a feasible transport plan by mixing elementary "component-to-component" optimal policies.

A. Feasible Set Construction

Instead of solving the high-dimensional optimization directly, the authors construct a family of feasible processes based on the individual components of the GMMs:

1. Decomposition: The initial GMM $\rho_0$ is decomposed into $N_1$ components ($\mu_i$), and the terminal GMM $\rho_N$ into $N_2$ components ($\nu_j$).
2. Elementary Bridges: For every pair $(i, j)$, an optimal transport policy (either a dtSB or a minimum-effort steering policy) is computed to move mass from component $\mu_i$ to $\nu_j$. Let these policies be denoted as $p_{ij}$.
3. Mixing Weights: A set of scalar mixing weights $\lambda_{ij} \geq 0$ is introduced, subject to marginal constraints ($\sum_i \lambda_{ij} = \beta_j$ and $\sum_j \lambda_{ij} = \alpha_i$), representing the probability of a sample starting in component $i$ and ending in component $j$.

B. The Proposed Policy (Randomization at Every Step)

The key innovation is the construction of the global transition kernel $p_{k+1|k}$ by mixing the elementary kernels $p_{ij}$ dynamically at every time step, rather than fixing the component assignment at the start.
The transition probability is defined as:
$$p_{k+1|k}(x_{k+1}|x_k) = \sum_{i,j} p_{ij}^{k+1|k}(x_{k+1}|x_k) \frac{\lambda_{ij} p_{ij}^k(x_k)}{p_k(x_k)}$$
where $p_k(x_k) = \sum_{i,j} \lambda_{ij} p_{ij}^k(x_k)$ is the current state density.

This results in a Markovian policy where the control action at time $k$ is selected randomly based on the current state $x_k$ and the posterior probability of which component pair $(i,j)$ the state belongs to.

C. Optimization

• For Density Steering: The total control cost is shown to be a linear combination of the costs of the elementary policies: $J = \sum_{i,j} \lambda_{ij} J_{ij}^{DS}$. The optimal mixing weights $\lambda_{ij}$ are found by solving a linear programming problem.
• For Schrödinger Bridges: The KL-divergence cost is upper-bounded by the linear combination of elementary costs: $D_{KL}(p\|q) \leq \sum_{i,j} \lambda_{ij} J_{ij}^{SB}$. The authors optimize $\lambda_{ij}$ to minimize this upper bound, providing a tractable approximation.

3. Key Contributions

1. Markovian vs. Memory-Based Control:
   The paper contrasts their approach with previous work (e.g., [14]) that randomizes the component assignment once at the beginning of the horizon. The proposed method randomizes at every time step.

   • Result: The proposed method yields a Markovian policy (state-feedback only), whereas the "randomize once" approach yields a process with memory (requiring knowledge of the initial latent variable).
   • Significance: Markovian policies are easier to implement in real-time digital control systems and do not require storing latent variables.

2. Cost Superiority in SB Problems:
   The authors prove that for Schrödinger Bridge problems, the Markovian policy (randomizing at every step) yields a cost less than or equal to the cost of the memory-based policy.

   • Reasoning: Randomizing at every step maximizes the differential entropy of the process compared to randomizing once, leading to a lower KL-divergence relative to the reference process. This provides a less conservative approximation.

3. Exactness in Density Steering:
   For minimum-effort density steering, the proposed method achieves the exact same control cost as the memory-based approach, but with the added benefit of being a state-feedback law.

4. Continuous-Time Convergence:
   The paper demonstrates that as the time step $\Delta T \to 0$, the proposed discrete-time policy converges to the continuous-time Schrödinger Bridge approximation for GMMs recently proposed in the literature [4]. This validates the discrete-time formulation as a consistent discretization of the continuous problem.

4. Results and Numerical Examples

The authors validate their theory through two numerical examples:

• Example 1 (Gaussian to GMM): A simple case steering a single Gaussian to a GMM with 8 components. The results illustrate the difference between the SB solution (minimizing deviation from random walk) and the Density Steering solution (minimizing control effort).
• Example 2 (Double Integrator): A more complex control problem involving a double integrator system steering a GMM with 2 components to a GMM with 3 components. The simulation confirms the feasibility of the approach in higher-dimensional state spaces and with dynamic constraints.

5. Significance and Impact

• Bridging Discrete and Continuous: The work provides a rigorous link between discrete-time control algorithms and continuous-time diffusion models, which is crucial for applications in generative AI (where discrete diffusion is popular) and robotics (which often operates in discrete time).
• Practical Implementation: By shifting from memory-based to Markovian policies, the method becomes significantly more practical for autonomous systems and real-time control, where storing latent variables is often infeasible.
• Generative Modeling: The approach offers a faster inference alternative to continuous-time diffusion models by utilizing discrete-time Markov chains with fewer steps, while handling complex, multi-modal distributions (GMMs) better than standard Gaussian assumptions.
• Theoretical Rigor: The paper establishes that the "randomize at every step" strategy is theoretically superior (or at least equal) to "randomize once" for these specific classes of problems, correcting a potential misconception in prior literature regarding the optimality of memory-based approaches.

In summary, the paper presents a computationally tractable, theoretically sound, and practically implementable framework for steering Gaussian Mixture Models in discrete time, offering improvements in both control cost (for SB) and implementation simplicity (Markovian property) over existing methods.