Discrete-Time Maximum Likelihood Neural Distribution Steering
George Rapakoulias Panagiotis Tsiotras

Laser cooling a centimeter-scale torsion pendulum
    Vivishek Sudhir
    September 9, 2024
=================================================



§ ABSTRACT 
This paper studies the problem of steering the distribution of a discrete-time dynamical system from an initial distribution to a target distribution in finite time.
The formulation is fully nonlinear, allowing the use of general control policies, parametrized by neural networks.
Although similar solutions have been explored in the continuous-time context, extending these techniques to systems with discrete dynamics is not trivial.
The proposed algorithm results in a regularized maximum likelihood optimization problem, which is solved using machine learning techniques. 
After presenting the algorithm, we provide several numerical examples that illustrate the capabilities of the proposed method. 
We start from a simple problem that admits a solution through semidefinite programming, serving as a benchmark for the proposed approach.
Then, we employ the framework in more general problems that cannot be solved using existing techniques, such as problems with non-Gaussian boundary distributions and non-linear dynamics.




§ INTODUCTION


The problem of controlling distributions of dynamical systems has attracted increased attention in the past few years, mainly due to its practical applications in machine learning and generative AI <cit.>.
Compared to standard optimal control problems, where a feedback policy that minimizes a functional in the presence of uncertainty is sought, distribution control aims at directly steering the distribution of the system's state while minimizing a cost function <cit.>.
Apart from generative AI applications, this approach is suitable for controlling systems whose behavior is better captured by a distribution rather than a deterministic state variable.

Such applications include, for example, swarm and multiagent control <cit.>, mean field games <cit.>, opinion dynamics <cit.>, and safe stochastic model predictive control <cit.>, among many others.

This work focuses on the problem of controlling the distribution of a deterministic, discrete-time, control-affine system with nonlinear drift.  
The problem of interest can be cast as the following infinite-dimensional optimization problem

 

    min_ x_k, u_k   J =  ∑_k = 0^N-1[  u_k ^2 + V_k(x_k)], 
        x_k+1 = f_k(x_k) + B_k u_k, 
        x_0 ∼ρ_i, 
        x_N ∼ρ_f,



where ρ_i,  ρ_f are the boundary state distributions, f_k, B_k describe the system's prior dynamics, and V_k(x_k) is a state-dependent cost that penalizes potentially undesirable regions of the state space.
Problem (<ref>) is an Optimal Mass Transport problem with (nonlinear) prior dynamics (OMTwpd).

For continuous linear prior dynamics and V(x_k) ≡ 0, this problem is equivalent to an Optimal Transport (OT) problem in a transformed set of coordinates defined by the linear dynamics <cit.>. 
Furthermore, for a quadratic state cost V(x_k) = x_k Q_k x_k, linear prior dynamics, and for Gaussian initial and terminal distributions, globally optimal solutions exist through the Covariance Steering (CS) framework for both continuous and discrete time settings. 
These can be computed efficiently using semidefinite programming (SDP) <cit.>.

For general, nonlinear systems or systems with non-Gaussian boundary distributions, a globally optimal solution is difficult to obtain. 
Local solutions through linearization have been explored in <cit.>, while solutions utilizing characteristic functions, for non-Gaussian boundary distributions and noise, have been proposed in <cit.>. 
However, both of these methods utilize linear feedback policies and therefore, may be suboptimal.

Furthermore, the authors of <cit.> propose a randomized policy for steering between Gaussian Mixture Models (GMM) with deterministic linear prior dynamics using randomized linear policies.
While the algorithm results in numerically efficient solutions with guaranteed convergence leveraging analytical results from the Covariance Steering theory, the optimality of the policy compared to more general nonlinear policies is not explored in <cit.>. 

Focusing on formulations with nonlinear dynamics, most existing works concern systems in continuous time. 
In the simplest case where ẋ_t = u_t, the problem has been studied in the context of Continuous Normalizing Flows (CNFs) <cit.> due to its applications in generative AI <cit.>.  
Its stochastic counterpart, referred to as the Schrödinger Bridge Problem (SPB), has also been explored in a similar context <cit.>.
Works that account for more complicated prior dynamics usually impose some assumptions on their structure. 
For example, in <cit.> the authors account for dynamics with a nonlinear drift term but otherwise require control in all states by restricting their analysis to cases where B_k = I, and focus on mean field games and generative AI applications.

A more general framework, allowing dynamics with fewer control channels than states is considered in <cit.> where the authors extended the results of <cit.> to feedback linearizable systems. 
Finally, a framework that does not require feedback linearizable dynamics is explored in <cit.> but the deterministic drift term is required to be the gradient of a potential function. 

The discrete-time problem has been studied much less.
The case of degenerate prior dynamics, i.e., x_k+1 = x_k + u_k has been addressed within the framework of Discrete Normalizing Flows (DNFs) <cit.>. 

To the best of the authors' knowledge, the more general problem with nonlinear prior dynamics has not been addressed in the literature.

This problem is of interest for several reasons.
First, it captures the case where the system dynamics are inherently discrete, as in the digital implementation of control algorithms.
Furthermore, most types of neural networks can be analyzed through the lens of discrete dynamical systems, <cit.>, providing a significant drive towards more research in the discrete setting. 
Finally, from a computational point of view, solving the discrete-time steering problem requires storing the state vector at a constant number of intervals, requiring, therefore, a fixed memory budget, compared to continuous formulations which perform a temporal discretization based on the stiffness of the trained dynamics <cit.>.  

To this end, in this work, we study Problem (<ref>) using tools from machine learning and the DNF literature, specifically combining 
control-theoretic ideas and tools from <cit.>. 
To bring Problem (<ref>) to the framework of Normalizing Flows, we first relax the terminal distribution constraint (<ref>) to a Kullback–Leibler (KL) divergence soft constraint, giving rise to the problem

 

    min_ x_k, u_k   J = ∑_k = 0^N-1[  u_k ^2 + V_k(x_k) ] + λ (ρ_N ρ_f), 
        x_k+1 = f_k(x_k) + B_k u_k, 
        x_0 ∼ρ_i,



where λ > 0. 
Henceforth, this is the general problem formulation we will use in this paper.




§ NOTATION

We use lowercase letters to denote vectors and vector random variables and capital letters to denote matrices. 
Given a function F: ^n→^m, its Jacobian is denoted by ∇ F: ^n →^n × m. Distributions in ^n are denoted by ρ and probability density functions (PDFs) by p. Given a random variable x ∼ρ_x and a transformation y = F(x), the pushforward of x is denoted by ρ_y = F_#ρ_x. 



§ PRELIMINARIES




 §.§ Normalizing flows
 

Let ρ_1, ρ_2 be two distributions with probability density functions p_1(x), p_2(x) respectively. Their KL divergence is given by

 
    
    (ρ_1  ρ_2)     = ∫_-∞^∞ p_1(x) log p_1(x)/ p_2(x)dx 
        = _x ∼ρ_1[ logp_1(x)/p_2(x) ] 
        = _x ∼ρ_1[ log p_1(x) ] - _x ∼ρ_1[ log p_2(x)  ] .



Equation (<ref>) suggests that calculating the KL divergence requires the analytic expressions of the PDFs of the two functions ρ_1 and ρ_2.

In the context of Problem (<ref>), ρ_1 would be a target distribution, which we know explicitly, while ρ_2 would correspond to the distribution of the state at some time step of interest. 
The calculation of ρ_2 would therefore require propagating the initial state distribution through the nonlinear dynamics, which is generally intractable.
To overcome this issue, one can use the change of variables formula connecting the PDFs of two random variables that are linked through a diffeomorphic (invertible and differentiable) transformation. 
This is summarized in the following lemma:

 <cit.> (Change of Variables) 
Let x be an n-dimensional random variable with known PDF, denoted p_x(x), and let F: ^n →^n be a diffeomorphism. 
The PDF of y = F(x), denoted p_y(y), can be calculated using the formula

    p_y(y) = p_x (F^-1(y) )  ∇ F^-1(y).








Transformations that correspond to flow maps of continuous-time systems are invertible, as far as unique solutions exist for the differential equation describing their dynamics. 
In CNF problems without prior dynamics, this condition is satisfied because neural networks with finite weights and Lipschitz nonlinearities result in Lipshitz ODE dynamics, and the uniqueness of solutions is guaranteed through Picard's existence theorem  <cit.>.
In the discrete-time case, however, the invertibility of the network needs to be addressed explicitly. 
In this paper, we make use of the following lemma:  

  <cit.> (Flow Invertibility) 
Consider the discrete-time nonlinear system described by 

    x_k+1 = x_k + f_k(x_k),    k = 0, 1, …, N-1,

and the transformation x_N = F(x_0). The transformation F is invertible if,
for all k = 0, 1, …, N-1, the mappings f_k are contractive.


For a different technique that requires modeling the individual state transition functions as gradients of a convex potential, we refer the reader to <cit.>. 
Finally, one result that facilitates the derivation of the proposed algorithm is the equivalence of the KL-divergence before and after a diffeomorphic transformation.

  <cit.> (Forward and Backward KL divergence) 
    Let F:^n→^n be a diffeomorphism and let ρ_1, ρ_2 be two distributions. Then 
    
    (ρ_1 ρ_2) =  (F_#ρ_1  F_#ρ_2 ).

 



 §.§ Exact Steering Using Semidefinite Programming
  

For the special case of a system with linear dynamics of the form

    x_k+1 = A x_k + B u_k,

and for Gaussian initial and terminal state distributions, the optimal policy for Problem (<ref>) is parametrized by an affine feedback controller of the form <cit.>

    π_k(x_k) = K_k (x_k - μ_k) + v_k,

where μ_k = [x_k], while its solution, i.e., the calculation of {K_k, v_k } for k = 0, 1, … N-1, can be attained efficiently through semidefinite programming.

The soft-constrained version (<ref>)
has only been studied for a Wasserstein-2 soft constraint penalty function in <cit.>.
The KL-divergence soft-constrained version can be solved similarly.
Although this paper focuses on the nonlinear version of the problem, the case of linear prior dynamics with Gaussian boundary conditions will be used as a benchmark, to validate the accuracy of the proposed algorithm. 
To this end, we briefly present how one can calculate the optimal solution to Problem (<ref>) with a controller of the form (<ref>) using semidefinite programming. 
Although we do not explicitly prove that this family of policies is globally optimal for the problem in this paper, this has been shown to be the case for the hard constrained Problem (<ref>) 
in <cit.>.



 §.§ KL-divergence Covariance Steering
 

When the dynamics of the system are linear and a control policy of the form (<ref>) is used, the first two moments of the state can be calculated explicitly at any time instant k in the steering horizon.
The corresponding equations are

 

    μ_k+1 = Aμ_k + B v_k, 
       Σ_k+1 = (A + B K_k) Σ_k (A + B K_k),



where Σ_k = (x_k). The KL-divergence between the terminal distributions ρ_N = 𝒩(μ_N, Σ_N) and ρ_f = 𝒩(μ_f, Σ_f) can be calculated via  <cit.>


    (ρ_N ρ_f) =      ( (Σ_f^-1Σ_N) + (μ_f - μ_N)Σ_f^-1 (μ_f - μ_N)
        + log(Σ_f) - log (Σ_N)  - n ),



where n is the dimension of the state vector. 
Using equations (<ref>), the change of variables K_k = Σ_k^-1 U_k, and Y_k = U_k Σ_k^-1 U_k in 
Problem (<ref>) yields

 

    min  J = ∑_k = 0^N-1 (Y_k) + v_k^2 + λ(ρ_Nρ_f),  
         U_k Σ_k^-1 U_k = Y_k , 
       Σ_k+1 = A Σ_k A + B U_k + U_k B  + B Y_k B, 
       μ_k+1 = Aμ_k + B v_k, 
       Σ_0 = Σ_i, 
       μ_0 = μ_i, 
       Σ_N = Σ_f, 
       μ_N = μ_f.



Relaxing (<ref>) to the semidefinite inequality U_k Σ_k^-1 U_k ≼ Y_k turns Problem (<ref>) into a semidefinite program. 
This relaxation has been proven to be lossless in <cit.>. 
Finally, we note that the term -log(Σ_f) in the KL-divergence is convex with respect to Σ_f and can be added to the cost function using appropriate slack variables accompanied by an LMI constraint <cit.>. 



§ MAXIMUM LIKELIHOOD DISTRIBUTION STEERING


This section contains the main results of the paper. 
To this end, consider Problem (<ref>) with a policy parametrized by a neural network, i.e., u_k = π_k(x_k ; θ) where θ corresponds to the trainable policy parameters.

To optimize (<ref>), tractable expressions for the cost function (<ref>)





must be developed. 
The first step is arguably the calculation of the KL divergence. 
Calculating it directly would involve the explicit calculation of the probability density of the state at the end of the steering horizon, p_N, which is challenging due to the nonlinearities of the system.
Instead of calculating p_N directly, let F(x_0) = Φ_N-1∘Φ_N-2∘…∘Φ_0 (x_0) denote the transformation linking the initial and terminal states under the discrete dynamic model (<ref>), where 

    Φ_k(x_k) = f_k(x_k) + B_k π_k(x_k),
 
is the closed-loop state transition function at time step k.
Under certain conditions on the control policy that will be specified later in the section, this transformation is diffeomorphic.

Therefore, its inverse x_0 = F^-1(x_N) satisfies the conditions of Lemma <ref>.
Applying this result to the KL divergence yields 

    (ρ_N ρ_f) =  (ρ_i  F^-1_#ρ_f).

Further expanding the second term using the definition of the KL divergence, yields 

    (ρ_i  F^-1_#ρ_f)) =  _x ∼ρ_i [log p_i(x)] - _x ∼ρ_i [log p_0(x)],

where p_i(x) is the PDF of ρ_i and p_0(x) is the PDF of the distribution F^-1_#ρ_f, that is, the density of a random variable sampled from ρ_f and pushed through the inverse transformation F^-1.
Notice that the term _x ∼ρ_i [p_i(x)] does not depend on the control policy parameters, and can therefore be omitted from the cost function of (<ref>). 

The calculation of log p_0(x) can be facilitated through Lemma <ref>. 
Specifically, one can link the density p_0 with the density of p_f using 

    log p_0(x) = log p_f (F(x)) + log∇ F(x).

The second term can also be efficiently calculated using the chain rule as follows 

    log∇ F(x) = ∑_k=0^N-1log∇Φ_k (x_k).


In our implementation, the Jacobian of the state transition functions ∇Φ_k(x_k) were calculated using automatic differentiation. 



Finally, we discuss conditions for π_k(x_k; θ) that preserve the invertibility of the discrete-time dynamics.
Without loss of generality, we assume that the state transition functions (<ref>) are of the form 


    x_k+1 = x_k + ϕ_k(x_k) + B_k π_k(x_k ; θ).
 

 
    Let the system dynamics be described by (<ref>), and let L_ϕ_k, L_π_k be the Lipschitz constants of ϕ_k, π_k, respectively, and let σ_B_k be the spectral norm of the matrix B_k. 
    Then, if L_π_k < (1 -L_ϕ_k)/σ_B_k, the state transition function defined in (<ref>) is a diffeomorphism.  


    Based on Lemma <ref>, it suffices to show that ϕ(x_k) + B_k π_k(x_k) is a contraction.
    To upper bound its Lipschitz constant note that ∇( ϕ_k + B_k π_k(x_k) ) _2 ≤∇ϕ_k _2 +  B_k ∇π_k _2 ≤ L_ϕ_k + σ_B_k L_π_k, 
    due to the subadditivity and submultiplicativity of the spectral norm <cit.>. 
    
    
    
    
    
    
    
    Constraining this upper bound yields the desired result. 



    In the case where f_k, B_k in (<ref>) are discretized versions of the continuous dynamics ẋ_t = f_t(x_t) + B_t u_t, then the first order approximation of the terms in (<ref>) are ϕ_k(x_k) = Δ T f_t(x_k) and B_k = Δ T B_t, where Δ T is the discritzation step size. 
    Therefore, for Lipschitz continuous-time dynamics, L_ϕ_k and σ_B_k can be made sufficiently small by reducing the discretization step Δ T.


Note that training Neural Networks with bounded Lipschitz constants can be achieved using spectral normalization <cit.>. 
In this work, we use π_k = α L_π_kπ̂_k where L_π_k = (1 -L_ϕ_k)/σ_B_k,  α∈ (0 , 1) and π̂_k is a Multilayer Perceptron (MLP) with spectral normalization in all of its weights, having therefore a Lipschitz constant of 1.


After calculating the loss function, optimization is carried out using standard gradient-based optimizers. 
For our implementation, we used the AdamW <cit.> scheme, implemented in pytorch <cit.>.































§ NUMERICAL EXAMPLES




In the first numerical example, we study the problem of driving a double integrator system from an initial to a terminal Gaussian distribution. 
Since an exact solution can be obtained for this problem using the results from Section <ref>, we use this example as a benchmark.

To this end, consider the discrete-time deterministic dynamics of the form (<ref>) with


    A = [     I_2 Δ T I_2;     0_2     I_2 ],    B = [     0_2; Δ T I_2 ],   Δ T = 0.1,


a horizon of N=30 time steps, V_k(x_k) ≡ 0 and λ = 60, which is equivalent with normalizing u_k ^2 with 1/(N m) where m is the number of input channels. We use this value for λ for all the subsequent examples. 
The boundary distributions are x_0 ∼(μ_i, Σ_f) and x_N ∼𝒩(μ_f, Σ_f) with parameters μ_i = [0,  0,   5,    8],   Σ_i = (1, 1, 0.2, 0.2), μ_f = [0,  0,   10,    0],   Σ_f = 0.4 I_4.
For this example, the Lipschitz constants of the prior dynamics are L_ϕ = σ_B = Δ T. 
To this end, we set L_π = 9, α = 0.9 and π_k = α L_ππ̂_k.

Each policy π̂_k(·) is modeled using a fully connected MLP with spectral normalization, and five layers with {4, 64, 64, 64, 64, 2} neurons per layer.
The convergence plot, along with the optimal solution calculated using the SDP technique described in Section <ref> can be viewed in Figure <ref>. 










In the second numerical example, we use double integrator dynamics, a horizon of N=40, but this time, we opt to steer from a Gaussian Mixture Model with 8 modes to the Normal distribution in the presence of obstacles.
The obstacles are modeled using appropriate potential fields with Gaussian kernels <cit.> of the form

    V_k(x) = λ_obsexp( -(x - x_0)^2/r_obs^2).

The policy at each time step has the same structure as in the first example but with 128 neurons in the hidden layers. 
The results are illustrated in Figure <ref>.

Another, more complicated example, which capitalizes on the fact that only a batch of samples is required for the initial distribution rather than its PDF, is depicted in Figure <ref>, where the initial distribution is arbitrary and the terminal is the normal distribution. 
The prior dynamics and policy parametrization are identical to Example 2.
We note that the inverse problem, i.e., steering from a Gaussian distribution to an arbitrary distribution for which only samples are available, would be significantly harder and cannot currently be solved with the proposed approach, since explicit information about the PDF of the terminal distribution is required for the maximum likelihood training.
We leave the investigation of this case as part of future work.

Finally, we test the proposed method in the 2D nonlinear model 

    
    x_k+1 = x_k + 0.1 √(1 + y_k^2) + u_k, 
        y_k+1 = y_k + 0.1 x_k,


for N=40.
By bringing these equations in the form of (<ref>) with  L_B = σ_B = 0.1, 
we may use the same policy parametrization as in the first example. 
The results are illustrated in Figure <ref>.

Since this is a two-dimensional example, we can overlay the samples with the contour lines of the PDF at each time step to validate the precision of the solution.


Table <ref> demonstrates a quantitative evaluation of the algorithm's performance. 
The first column corresponds to the experiment number. 
To measure the distance between the distribution of the state at the final time step of the steering horizon and the target distribution, we calculate the 2-Wasserstein distance using discrete Optimal Transport <cit.> and report its value in the second column.
We use the 2-Wasserstein distance because it can be computed exactly on empirical distributions that are available only through samples and accurately reflects the actual distance between the continuous distributions given enough samples.
In the third column, we report the minimum value of the log-determinant of the Jacobian of the optimal map linking the initial and final state in order to validate the invertibility of the computed map and finally provide the total training time in minutes in the last column.
Training was performed on an Nvidia RTX-3070 GPU. 







§ CONCLUSIONS


This paper presents a method for solving the distribution steering problem for discrete-time nonlinear systems by formulating it as a regularized maximum likelihood optimization problem.
The control policies are parametrized using neural networks with appropriate Lipschitz constraints to ensure the invertibility of the discrete-time dynamics.

A general cost function is considered, allowing state-dependent terms to model obstacles in the state space using potential fields. 
In parallel, a KL-divergence soft constraint version of the Covariance Steering problem is developed as a benchmark to compare with the proposed nonlinear maximum likelihood methods. 
Finally, four comprehensive numerical examples are presented and analyzed with respect to how closely they achieve the target distribution, as well as in terms of run time.
For the linear dynamics with Gaussian boundary distributions, the solution is also compared against the globally optimal solution calculated as the solution of a semidefinite program. 


§ ACKNOWLEDGMENT


The authors would like to sincerely thank Dr. Ali Reza Pedram for his comments in an initial version of the manuscript and for multiple fruitful discussions. 
Support for this work has been provided by 
ONR award N00014-18-1-2828 and NASA ULI award #80NSSC20M0163.
This article solely reflects the opinions and conclusions of its authors and not of any NASA entity.



ieeetr



