# High-dimensional Optimal Density Control with Wasserstein Metric Matching

Shaojun Ma  Mengxue Hou  Xiaojing Ye  Haomin Zhou This work was supported in
part by NSF under grants DMS-1925263 and DMS-2152960, and ONR
N00014-21-1-2891.S. Ma and H. Zhou are with the School of Mathematics, Georgia
Tech, Atlanta, GA 30332, USA<EMAIL_ADDRESS>Hou is with the
College of Engineering, Purdue University, West Lafayette, IN 47906, USA
<EMAIL_ADDRESS>Ye is with the Department of Mathematics and Statistics,
Georgia State University, Atlanta, GA 30303, USA<EMAIL_ADDRESS>
###### Abstract

We present a novel computational framework for density control in high-
dimensional state spaces. The considered dynamical system consists of a large
number of indistinguishable agents whose behaviors can be collectively modeled
as a time-evolving probability distribution. The goal is to steer the agents
from an initial distribution to reach (or approximate) a given target
distribution within a fixed time horizon at minimum cost. To tackle this
problem, we propose to model the drift as a nonlinear reduced-order model,
such as a deep network, and enforce the matching to the target distribution at
terminal time either strictly or approximately using the Wasserstein metric.
The resulting saddle-point problem can be solved by an effective numerical
algorithm that leverages the excellent representation power of deep networks
and fast automatic differentiation for this challenging high-dimensional
control problem. A variety of numerical experiments were conducted to
demonstrate the performance of our method.

##  I INTRODUCTION

Recent years have witnessed an emerging research interest in optimal control
problems with large collective of agents, such as drones, robots, and vehicles
[4, 11]. Such problems appear in many important real-world applications,
including surveillance over a large region, time-sensitive search-and-rescue,
infrastructure inspection and maintenance, among many others. Meanwhile, the
cost to manufacture robots, including the builtin electrical components and
batteries, has reduced significantly during the past decades. The individual
robots used in these swarm control applications are often of small size and
thus have limited computation and communication power. However, using such
swarms of robots is particularly necessary and cost-effective to achieve goals
in the aforementioned applications. Optimal control of swarm robots is to
design efficient algorithms that can steer them to achieve specific tasks
achieve specified high-level tasks in a way that accommodates the high
dimensionality of the system [36].

###  I-A Related Work

Optimal control of large collective of agents has been commonly considered in
the context of density control based on mean-field models [12, 42, 27]. The
idea of mean-field models is to characterize the collective behaviors of the
agents in the swarm using a time-evolving probability density function defined
on the specified state space. Mean-field models allow for both deterministic
and even stochastic behaviors if the randomness of individual agents can be
appropriately described by stochastic processes. The dynamics of the density
function is governed by the associated continuity equation in the
deterministic case and Fokker-Planck equation in the stochastic case, which
are evolution partial differential equations (PDEs). In these models, the
probability density characterizing the swarm behavior does not depend on the
number of agents, but only on the state space of individual agents.
Nevertheless, direct computation of these PDEs can still be computationally
prohibitive due to the high dimension of the state space. Thus, most xisting
works on density controls approximate the density functions based on spatial
discretization [1] or using kernel density estimation [2, 14].

While the agents and the population dynamics are steered by the control, an
optimal density control problem aims at finding the controls that achieve the
specified task with minimum control cost. The work [14] designed velocity
field to make swarms density converges to the desired/commanded density
distribution. In [40], the authors presented a density control by improving
Smoothed Particle Hydrodynamic (SPH) [33] method with collision free
condition, where one defines the density of a robot as the weighted sum of
distances to its nearby robots within a certain range. A data-driven approach
was proposed in [9] to learn the control by estimating the Koopman operators.
In [39], the authors provided a control strategy to steer the state from an
initial distribution to a terminal one with specified mean and covariance,
which is also called nonlinear covariance control. For other more general
works we refer to [7, 8, 21].

A specific type of density control problem, known as optimal transport, has
received significant attentions in recent years. Several works have combined
Wasserstein metric and optimal control to develop relevant control algorithms.
For instance, [5] develops Wasserstein proximal algorithms to solve density
control problem, where a nonlinear drift term is considered. In [3], the
authors discussed the optimality conditions for optimal control problems in
Wasserstein spaces. In [38], a Wasserstein based robust control was developed
to resolve the issue that uncertain variables is unavailable. In [22], the
authors proposed a novel model-predictive control (MPC) method for limiting
the safety risk when the distribution of the obstacles is set to be within an
ambiguity set which is measured by the Wasserstein metric. In [23], terminal
cost as Wasserstein distance and finding the parameters of normalizing flow
[10] by minimizing the terminal cost have been considered. In [13], the
authors presented a primal-dual formulation of optimal control problem with
existence proof and efficient numerical method, similar method can also be
found in [17] where the problem is solved by a Mean Field Game (MFG) approach.
In [16], the authors calculated the individual robot trajectories by
alternating two gradient flows that involve an attractive potential, a
repelling function, and a process of intermittent diffusion, in which way a
large group of robots are employed to accomplish the task of shape formation.

Optimal control over discrete state spaces can be cast as Markov chain models
as shown in [1]. For optimal control with continuous state spaces, the density
dynamics are determined by the corresponding continuity equation or Fokker-
Planck equation. For instance, [34, 35] provided both theoretical analysis and
numerical methods of the optimal control to achieve target equilibrium
density. In [6], the dynamics of the control was discretized in time and a
multi-marginal optimal transport scheme was used to solve the problem. In
[24], the authors developed a method to solve the optimal transport based
control based on Voronoi Tessellation coverage. A distributed algorithm for
density estimation to reduce the computational cost is proposed in [41].
Multi-agent optimal control has been widely applied in the domain of robotics
in recent years. For a general introduction we refer to [15]. In particular,
[43] presents the work of using reinforcement learning to control multi-robot
system. Futhermore, [29, 37, 32] carefully designed an objective function to
minimize the control energy, $L_{2}$ based difference between the initial and
target positions, meanwhile keeping collision-free conditions. Parameterizing
controls using neural networks was also considered in [32].

###  I-B Main contributions of this work

In this work, we propose a novel method to compute the optimal control $u$
which steers an initial density as close as possible to a target density at a
prescribed terminal time, while minimizing the control cost. The major
advantage of the proposed method is that it can be readily applied to the
cases with high-dimensional state spaces, where computation of the associated
continuity equation or Fokker-Planck equations is prohibitive using existing
approaches. In our approach, this issue is effectively tackled by using a
large number of synthesized agents whose trajectories can be easily evaluated
using decoupled ODEs or SDEs. Therefore, our approach does not require any
spatial discretization using finite difference and element methods or any
basis representations, and the expected cost function values can be accurately
approximated using empirical expectations, which allows our method to scale to
much higher dimensions that are traditionally considered computationally
infeasible. Moreover, the proposed method can readily adopt parallel
computation because of the decoupling, and hence the computation can be done
very fast by leveraging the power of standard deep learning computation
framework. We also provide convergence analysis, including the rate of
convergence, of the numerical algorithm for solving the resulting saddle-point
problem. We validate our approach in numerical simulations that include
complex scenarios in various tests with dimension up to 100.

###  I-C Paper organization

The remainder of the paper is organized as follows. In Section II, we present
the proposed framework to solve the optimal density control with Wasserstein
metric matching to any target distribution at a prescribed terminal time. The
framework renders a saddle-point problem of the control and the dual variable
of the Wasserstein distance. We also provide convergence analysis of the
proposed method including the convergence rate. In Section III, we conduct
several numerical experiments on synthetic and real data to demonstrate the
performance of the proposed algorithm. Section IV concludes this work.

###  I-D Notations

Throughout this paper, we use $|\cdot|$ to denote the absolute value of a
scalar and Euclidean norm of a vector. Furthermore, $\|\cdot\|$ denotes the
matrix norm induced by vector Euclidean norm. We also use $\|\cdot\|_{\cdot}$
to denote function norms where the space is specified in the subscript
$\cdot$. The spatial domain is denoted by $\Omega$, which is assumed to be a
closed and quasiconvex subset of $\mathbb{R}^{d}$. Integrals over $\Omega$ or
$\Omega\times\Omega$ are written without subscript for notation simplicity
unless otherwise specified. We use $\mu$ to denote standard Lebesgue measure
on $\mathbb{R}^{d}$. We use $\nabla$ to denote the gradient with respect to
the state variable $x$. Gradient and partial derivatives with respect to any
other variables will be indicated by their subscripts. We denote $\mathbb{N}$
the set of positive integers, and $[n]\mathrel{\mathop{\mathchar
58\relax}}=\mathinner{\left\\{1,\dots,n\right\\}}$ for notation simplicity.

##  II PROPOSED METHOD

In this section, we develop a general computational framework for solving
density control over high-dimensional state spaces. To maximize applicability
of our framework, we use the following prototype mean-field model of density
control with any user-defined (running) cost functional $c$:

$\min_{u}\
\mathbb{E}\mathinner{\Bigl{[}\int_{0}^{T}c(x_{t},u(x_{t}))\,dt\Bigr{]}}+\gamma
D(\rho_{T},\nu)$ (1)

where $x_{t}\in\Omega$ is the dynamic following

$\begin{cases}dx_{t}=u(x_{t})dt+\sigma(x_{t})dW_{t},\\\
x_{0}\sim\rho_{0},\end{cases}$ (2)

which represents any indistinguishable agent in the swarm to be controlled. In
(1), $\rho_{0}$ is a given initial population density, $\nu$ is a given target
population density, $D$ is a metric that measures the distance between two
probability densities, $\gamma>0$ is a user-chosen weight parameter, and
$\rho(t,\cdot)$ is the density of $x_{t}$ for any $t$ with short hands
$\rho(0,\cdot)=\rho_{0}(\cdot)$ and $\rho(T,\cdot)=\rho_{T}(\cdot)$. The time
evolution of $\rho(t,\cdot)$ is governed by the Fokker-Planck equation

$\partial_{t}\rho=-\nabla\cdot(\rho
u)+\frac{1}{2}\langle\nabla^{2},\sigma\sigma^{\top}\rho\rangle,$ (3)

where $\langle\nabla^{2},A(x)\rangle\mathrel{\mathop{\mathchar
58\relax}}=\sum_{i,j}\partial_{x_{i}x_{j}}^{2}a_{ij}(x)$ for any matrix-valued
function $A(x)=[a_{ij}(x)]\in\mathbb{R}^{d\times d}$. The goal of (1) is to
find the optimal control $u\mathrel{\mathop{\mathchar
58\relax}}\Omega\to\mathbb{R}^{n}$ that steers the population characterized by
$x_{t}$ with initial distribution $\rho_{0}$ to approximate the target
distribution $\nu$ at terminal time $T$ with minimal overall cost. The
parameter $\gamma$ weighs the matching of terminal density to $\nu$ against
the overall cost. If an exact matching to $\nu$ at $T$ is desired, then
$\gamma$ can be cast as the Lagrangian multiplier and solved jointly with $u$.
For simplicity, we will only consider the case with soft-penalty on the
matching as given in (1) in the present work. In what follows, we will discuss
the choice of cost functional $c$ and density distance $D$ to instantiate (1).

#### Cost functional

The density control problem (1) allows for a general class of cost functional
$c\mathrel{\mathop{\mathchar
58\relax}}\Omega\times\mathbb{R}^{n}\to\mathbb{R}$. Typical optimal control
problems with minimal energy use $c(x,u)=\frac{1}{2}|u(x)|^{2}$. If there is a
component $u_{0}(x_{t})$ due to environmental force in $u$, then the cost
functional can be set to $c(x,u)=\frac{1}{2}|u(x)-u_{0}(x)|^{2}$. In many
real-world applications, the swarm with density $\rho(t,\cdot)$ is realized by
a number of robotic agents. In such cases, it is necessary that the agents do
not collide with each other throughout the control process. To this end, one
can impose additional penalty when any two robots become too close. For
example, at any time $t$, such penalty function can be set to (the average
should be over all agents different from the input variable $x$ but we use the
following for notation simplicity):

$\frac{1}{N}\sum_{i=1}^{N}V(x,x_{t}^{(i)})\quad\mathinner{\left(\approx\int
V(x,y)\rho(t,y)\,dy\right)},$ (4)

where the approximation in the parenthesis holds as $N\to\infty$, and $V$ is
an interactive potential defined as $V(x,y)=c\ln|x-y|$ or
$V(x,y)=c|x-y|^{-\alpha}$ for some constants $c,\alpha>0$, etc. The purpose of
$V$ is to incur larger cost when two robots become closer. A cost functional
including such interactive potential can effectively prevent robot collisions.

#### Distance $D$ between densities

We use the 1-Wasserstein distance, also known as the Earth Mover’s Distance
(EMD), as the distance $D$ to measure the difference between the terminal
density $\rho_{T}$ and target density $\nu$. For any pair of densities
$(\varrho,\nu)$, their 1-Wasserstein distance is defined by

$D(\varrho,\nu)=\inf\mathinner{\Bigl{\\{}\int|x-y|\,d\pi(x,y)\mathrel{\mathop{\mathchar
58\relax}}\ \pi\in\Pi(\varrho,\nu)\Bigr{\\}}},$ (5)

where $\Pi(\rho,\nu)$ denotes the set of joint distributions on
$\Omega\times\Omega$ whose marginal densities are exactly $\varrho$ and $\nu$.
The 1-Wasserstein distance defined in (5) quantifies the minimum total cost to
transfer probability density $\varrho$ to $\nu$ (and vice versa), and its
optimal solution $\pi^{*}$ describes how the transfer should be made. In real-
world applications, the 1-Wasserstein distance is an appropriate metric for
the density control problem (1)—its value indicates how much additional
efforts, measured using the standard Euclidean distance, are needed to move
the robots to reach the target distribution in case the matching is not exact
at terminal time $T$. Moreover, 1-Wasserstein distance allows two densities to
have different supports, which could be problematic for other density distance
measure such as Kullback-Leibler (KL) divergence.

Despite of the various advantages of Wasserstein distance, its computation can
be challenging due to the optimization process required in (5). However, for
1-Wasserstein distance, one can derive its dual form given by

$D(\rho,\nu)=\sup\\{\mathbb{E}_{\rho}[\phi]-\mathbb{E}_{\nu}[\phi]\mathrel{\mathop{\mathchar
58\relax}}\ \mathrm{Lip}(\phi)\leq 1\\},$ (6)

where $\phi$ is the dual variable and $\mathrm{Lip}(\phi)$ denotes the
Lipschitz constant of $\phi$. As $\Omega$ is quasi-convex, we can relax the
constraint $\mathrm{Lip}(\phi)\leq 1$ to $\|\nabla\phi(x)\|\leq 1$ for all
$x\in\Omega$ a.e., where $\nabla\phi(x)$ is the weak derivative of $\phi$ with
respect to $x$, and $\|\nabla\phi(x)\|$ is the induced 2-norm of
$\nabla\phi(x)$ and thus its maximal singular value. This allows us to use
nonlinear reduced-order models, such as deep neural networks, to parameterize
$\phi$ and apply various techniques such as spectral normalization to enforce
the constraint $\|\nabla\phi(x)\|\leq 1$.

#### Complete model and approximations using deep networks

The optimization problem for solving the control problem (1) is given by

$\inf_{u}\sup_{\phi}\ \mathcal{E}(u,\phi)$ (7)

with constraint $\|\nabla\phi(x)\|\leq 1$ to hold for $x\in\Omega$ a.e., and
the objective functional $\mathcal{E}$ is given by

$\mathcal{E}(u,\phi)\mathrel{\mathop{\mathchar
58\relax}}=\mathbb{E}\mathinner{\Bigl{[}\int_{0}^{T}c(x_{t},u(x_{t}))\,dt\Bigr{]}}+\gamma\mathbb{E}_{\rho_{T}}[\phi]-\gamma\mathbb{E}_{\nu}[\phi].$
(8)

As both $u$ and $\phi$ are in infinite-dimensional function spaces, we need a
finite-dimensional representation of them for numerical computation. This can
be achieved by approximating them using linear or nonlinear reduced-order
models. In particular, we will use deep neural networks to parameterize $u$
and $\phi$ as $u_{\theta}$ and $\phi_{\eta}$ respectively. Here
$\theta,\eta\in\mathbb{R}^{m}$ (we assume both are $m$-dimensional for
notation simplicity here, but they can be different in practice) denote their
trainable network parameters, such as their weights and biases. The
approximation power of deep neural networks has been justified by the
universal approximation theorem, which has been extended to various cases with
different choices of activation functions and function approximation norm. In
particular, the following approximation result ensures that $u$ can be
approximated by a standard feed-forward network with rectified power unit
function (RePU) $a(x)=\max(0,x)^{2}$ as the activation for sufficient
smoothness requirement.

######  Lemma 1 (Theorem 4.9 [20]).

Let $d\in\mathbb{N}$, $B>0$, and $\|u\|_{W^{1,\infty}(\Omega)}\leq B$. For any
$\epsilon\in(0,1/2)$, there exists a feed-forward neural network $u_{\theta}$
with properly chosen network parameter $\theta$ and RePU activation such that

$\|u_{\theta}-u\|_{W^{1,\infty}(\Omega)}\leq\epsilon.$

Lemma 1 implies that a sufficiently large deep network is capable to represent
any function in a specified $W^{1,\infty}$ ball with any prescribed accuracy
$\epsilon>0$. Despite of the often pessimetic bound on the network complexity
in theoretical analysis, neural networks have been verified as very powerful
function approximators in numerous empirical studies in recent years. Hence
neural networks have become a special type of nonlinear reduced order models
to approximate general functions in infinite-dimensional spaces using
parameterized functions with finite-dimensional parameters, and they have
proven to be much more effective than traditional linear basis models.

In addition to using finite-dimensional representations $u_{\theta}$ and
$\phi_{\eta}$, we also need to approximate the integrals and expectations in
(8) for numerical computations. With high dimension $d$, the Fokker-Planck
equation (3) (or the continuity equation with $\sigma=0$) cannot be computed
using any classical numerical methods (e.g., finite difference and finite
element methods) due to the curse of dimensionality. In this work, we use the
particle dynamics (2) if the robots are given and fixed, or simulate a large
number of particles following (2) otherwise, to avoid direct computation of
$\rho(t,\dot{)}$. More precisely, we approximate the saddle-point problem (7)
with deep neural network representations $u_{\theta}$ and $\phi_{\eta}$ and
particle dynamics as follows

$\min_{\theta}\max_{\eta}E(\theta,\eta)$ (9)

where the objective function given by

$E(\theta,\eta)\mathrel{\mathop{\mathchar
58\relax}}=\frac{1}{N}\sum_{i=1}^{N}\big{(}c(x_{t}^{(i)},u_{\theta}(x_{t}^{(i)}))+\gamma\phi_{\eta}(x^{(i)}_{T})\big{)}-\gamma\mathbb{E}_{\nu}[\phi_{\eta}].$

The last term in $E$ can also be replaced with empirical expectation
$(\gamma/M)\cdot\sum_{j=1}^{M}\phi_{\eta}(z^{j})$ for $M$ i.i.d. samples
$\\{z^{(j)}\mathrel{\mathop{\mathchar 58\relax}}j\in[M]\\}$ drawn from $\nu$.
For simplicity, we also assume that $\theta,\eta\in\Theta$ and $\Theta$ is a
compact set in $\mathbb{R}^{d}$. Then the optimal solution $\theta$ to (9) is
the parameter of $u_{\theta}$ approximating the optimal control $u$ in (1).

######  Remark 1.

If the target distribution $\nu$ are given as $M$ fixed discrete points
$Z\mathrel{\mathop{\mathchar 58\relax}}=\\{z^{(j)}\mathrel{\mathop{\mathchar
58\relax}}j\in[M]\\}$, then we can also set $D$ to the Chamfer’s distance
between the two point clouds $X_{T}\mathrel{\mathop{\mathchar
58\relax}}=\\{x^{(i)}\mathrel{\mathop{\mathchar 58\relax}}i\in[N]\\}$ and $Z$:

$\displaystyle\sum_{i=1}^{N}\min_{1\leq j\leq
M}|z_{T}^{(i)}-z^{(j)}|^{2}+\sum_{j=1}^{M}\min_{1\leq i\leq
N}|x_{T}^{(i)}-z^{(j)}|^{2}.$ (10)

The Chamfer’s distance between $X_{T}$ and $Z$ given in (10) is a promising
alternative to measure the difference between these two clouds when $N$ and
$M$ are moderately small (e.g., dozens). We will show in our experiment that
the proposed method can also be applied using this Chamfer distance.

#### Numerical computations

An important feature of the saddle-point formulation (9) for solving the
optimal control problem (1) is that we can leverage parallel computation to
effectively handle the large number of simulated trajectories. More precisely,
for any fixed $\theta$, we can use the Euler-Maruyama (or the Euler method if
the diffusion $\sigma=0$) to generate trajectories

$x_{t+h}^{(i)}=x_{t}^{(i)}+hu(x_{t}^{(i)})+\sqrt{h}\sigma(x_{t}^{(i)})\delta_{t}^{(i)}$
(11)

for discrete time points $t=0,\dots,\lfloor T/h\rfloor-1$ where $h>0$ is the
time step size and $\delta_{t}^{(i)}$ are i.i.d. standard normal random
variables for all $t$ and $i$. Therefore, the objective function
$E(\theta,\eta)$ given in (9) can be computed explicitly. The we can use
automatic differentiation to obtain the gradients of $E$ with respect to
$\theta$ and $\eta$, which is crucial to solving the saddle-point problem (9).
In particular, we consider the objective function

$L(\theta)\mathrel{\mathop{\mathchar
58\relax}}=\max_{\eta\in\Theta}E(\theta,\eta).$ (12)

The optimal solution $\theta$ to (9) is thus the minimizer of $L(\theta)$,
whose full gradient with respect to $\theta$ can be immediately computed by
solving the maximization problem (12), as shown in the following lemma.

######  Lemma 2 (Gradient of $L$).

Let $L$ be defined in (12). Then for any fixed $\theta$, the gradient
$\nabla_{\theta}L(\theta)$ at $\theta$ is given by

$\nabla_{\theta}L(\theta)=\partial_{\theta}E(\theta,\eta(\theta)),$ (13)

where
$\eta(\theta)\in\operatorname*{arg\,max}_{\eta}\\{E(\theta,\eta)\mathrel{\mathop{\mathchar
58\relax}}\|\nabla\phi_{\eta}\|\leq 1\\}$.

###### Proof.

The full gradient $\nabla_{\theta}L(\theta)$ is given by

$\displaystyle\nabla_{\theta}L(\theta)=\partial_{\theta}E(\theta,\eta(\theta))+\partial_{\eta}E(\theta,\eta(\theta))\nabla_{\theta}\eta(\theta).$
(14)

For any fixed $\theta$, consider the problem

$\max_{\eta}\ \\{E(\theta,\eta)\mathrel{\mathop{\mathchar
58\relax}}\|\nabla\phi(x)\|\leq 1,\ \forall\,x,\ a.e.\\}$ (15)

The Lagrangian associated with (15) is given by

$\displaystyle
E(\theta,\eta)+\frac{1}{2}\int\lambda_{\theta}(x)(\|\nabla\phi_{\eta}(x)\|^{2}-1)\,dx,$
(16)

and the Karush-Kuhn-Tucker (KKT) system for the necessary optimality condition
of $\eta(\theta)$ is given by

$\displaystyle\partial_{\eta_{j}}E(\theta,\eta(\theta))+H_{j}(\theta,\eta(\theta))=0,\
\forall\,j\in[m],$ (17a) $\displaystyle\lambda_{\theta}(x)\geq 0,\
\|\nabla\phi_{\eta(\theta)}(x)\|\leq 1,\ \forall\,x,\ a.e.,$ (17b)
$\displaystyle\textstyle\int\lambda_{\theta}(x)(|\nabla\phi_{\eta(\theta)}(x)|^{2}-1)\,dx=0.$
(17c)

where the function $H_{j}$ is defined by

$\displaystyle
H_{j}(\theta,\eta(\theta))=\int\lambda_{\theta}(x){z_{\theta}^{x}}^{\top}\nabla\phi_{\eta(\theta)}(x)^{\top}\partial_{\eta_{j}}\nabla\phi_{\eta(\theta)}(x)z_{\theta}^{x}\,dx,$

and the vector $z_{\theta,x}\in\mathbb{R}^{d}$ is defined by

$z_{\theta,x}=\operatorname*{arg\,max}\\{|\nabla\phi_{\eta(\theta)}(x)z|\mathrel{\mathop{\mathchar
58\relax}}|z|\leq 1\\}.$

Now we partition the state space $\Omega$ into the following disjoint
subdomains:

$\displaystyle\Omega_{0}\mathrel{\mathop{\mathchar 58\relax}}=$
$\displaystyle\\{x\mathrel{\mathop{\mathchar
58\relax}}\lambda_{\theta}(x)<0\mbox{ or
}\|\nabla\phi_{\eta(\theta)}(x)\|>1\\},$ (18)
$\displaystyle\Omega_{1}\mathrel{\mathop{\mathchar 58\relax}}=$
$\displaystyle\\{x\mathrel{\mathop{\mathchar
58\relax}}\lambda_{\theta}(x)>0\mbox{ and
}\|\nabla\phi_{\eta(\theta)}(x)\|<1\\},$ (19)
$\displaystyle\Omega_{2}\mathrel{\mathop{\mathchar 58\relax}}=$
$\displaystyle\\{x\mathrel{\mathop{\mathchar
58\relax}}\lambda_{\theta}(x)=0\\},$ (20)
$\displaystyle\Omega_{3}\mathrel{\mathop{\mathchar 58\relax}}=$
$\displaystyle\\{x\mathrel{\mathop{\mathchar
58\relax}}\lambda_{\theta}(x)>0\mbox{ and
}\|\nabla\phi_{\eta(\theta)}(x)\|=1\\}.$ (21)

Due to (17b), we know $\mu(\Omega_{0})=0$. Furthermore, we can see that

$\lambda_{\theta}(x)(\|\nabla\phi_{\eta(\theta)}(x)\|^{2}-1)\leq 0$ (22)

for all $x\in\Omega\setminus\Omega_{0}$. In particular, the equality in (22)
holds for all $x$ in $\Omega_{2}$ and $\Omega_{3}$.

Now we claim that $\mu(\Omega_{1})=0$. If not, we define for any
$k,\ell\in\mathbb{N}$ the set

$\Omega_{1}^{k,\ell}\mathrel{\mathop{\mathchar
58\relax}}=\mathinner{\Bigl{\\{}x\mathrel{\mathop{\mathchar
58\relax}}\lambda_{\theta}(x)>\frac{1}{k}\Bigr{\\}}}\cap\mathinner{\Bigl{\\{}x\mathrel{\mathop{\mathchar
58\relax}}\|\nabla\phi_{\eta(\theta)}(x)\|^{2}<1-\frac{1}{\ell}\Bigr{\\}}}.$

Then it is easy to show that

$\Omega_{1}=\cup_{k=1}^{\infty}\cup_{\ell=1}^{\infty}\Omega_{1}^{k,\ell},$
(23)

which is a countable union of sets. Hence, there must exist
$k^{\prime},\ell^{\prime}\in\mathbb{N}$ such that
$\mu(\Omega_{1}^{k^{\prime},\ell^{\prime}})>0$. Combining this fact and (22),
we obtain

$\displaystyle\textstyle\int\lambda_{\theta}(x)(\|\nabla\phi_{\eta(\theta)}(x)\|^{2}-1)\,dx$
$\displaystyle=$
$\displaystyle\textstyle\int_{\Omega_{1}}\lambda_{\theta}(x)(\|\nabla\phi_{\eta(\theta)}(x)\|^{2}-1)\,dx$
$\displaystyle\leq$
$\displaystyle\textstyle\int_{\Omega_{1}^{k^{\prime},\ell^{\prime}}}\lambda_{\theta}(x)(\|\nabla\phi_{\eta(\theta)}(x)\|^{2}-1)\,dx$
$\displaystyle\leq$
$\displaystyle-\textstyle\frac{1}{k^{\prime}\ell^{\prime}}<0,$

which contradicts to (17c). Therefore $\mu(\Omega_{1})=0$.

Since $\|\nabla\phi_{\eta(\theta)}(x)\|^{2}=1$ for all $x\in\Omega_{3}$, the
following equality holds over $\Omega_{3}$:

$\displaystyle\partial_{\eta_{j}}\|\nabla\phi_{\eta(\theta)}(x)\|^{2}=2z_{\theta}^{x}{}^{\top}\nabla\phi_{\eta(\theta)}(x)^{\top}\partial_{\eta_{j}}\nabla\phi_{\eta(\theta)}(x)z_{\theta}^{x}=0.$

Together with the definition of $\Omega_{3}$ and
$\mu(\Omega_{0})=\mu(\Omega_{1})=0$, we see that
$H_{j}(\theta,\eta(\theta))=0$. Combining this and (17a), we have
$\partial_{\eta}E(\theta,\eta(\theta))=0$. Therefore by (14) we have
$\nabla_{\theta}L(\theta)=\partial_{\theta}E(\theta,\eta(\theta))$, which
completes the proof. ∎

Lemma 2 suggests a practical implementation to compute
$\nabla_{\theta}L(\theta)$: calculate the partial derivative of
$E(\theta,\eta)$ with respect to $\theta$ with $\eta$ held fixed, then
substitute $\eta$ by the maximizer $\eta(\theta)$ of the problem (12). Since
(12) is constrained, the gradient mapping, defined by
$\mathcal{G}(\theta)\mathrel{\mathop{\mathchar
58\relax}}=\tau^{-1}[\theta-\Pi(\theta-\tau\nabla_{\theta}L(\theta))]$, is
used as the convergence criterion of $\theta$ [18, 28, 25]. Note that the
definition of gradient mapping takes the normalization of step size $\tau$
into consideration. Moreover, the gradient mapping reduces to the gradient
$\nabla_{\theta}L(\theta)$ for unconstrained case. In the following theorem,
we show that for any $\varepsilon>0$ the standard stochastic gradient descent
scheme based on $\nabla_{\theta}L(\theta)$ (i.e., stochastic gradient is
unbiased and has bound variance) is guaranteed to push the gradient mapping
$\mathcal{G}(\theta)$ to $0$ with $\varepsilon$ accuracy within
$O(\varepsilon^{-1})$ iterations in the ergodic sense. For simplicity, we
assume both $\theta$ and $\eta$ are of the same dimension $m$ and lie in the
same closed Euclidean ball in $\mathbb{R}^{m}$. The results can be easily
generalized to the case where they have different dimensions $m$ and
$m^{\prime}$ and are in some convex compact subsets of $\mathbb{R}^{m}$ and
$\mathbb{R}^{m^{\prime}}$ respectively.

######  Theorem 1.

Suppose the parameters $\theta$ and $\eta$ are bounded in a ball centered at
origin with radius $R>0$, namely, we have $\Theta\mathrel{\mathop{\mathchar
58\relax}}=\\{\theta\mathrel{\mathop{\mathchar 58\relax}}|\theta|\leq R\\}$.
For any $\varepsilon>0$, let $\\{\theta_{j}\\}$ be a sequence of the network
parameters of $u_{\theta}$ generated by the stochastic gradient descent
algorithm, where $\nabla_{\theta}L(\theta)$ is approximated by the empirical
expectation using samples as in (9). If the sample complexity is
$N=O(\varepsilon^{-1})$ in each iteration, then $min_{1\leq j\leq
J}\mathbb{E}[|\mathcal{G}(\theta_{j})|^{2}]\leq\varepsilon$ after
$J=O(\varepsilon^{-1})$ iterations.

###### Proof.

Since $u_{\theta}$ and $\phi_{\eta}$ are parameterized using neural networks
with finitely many neurons, and the parameters are bounded, we know by Lemma 2
that $\nabla L(\theta)$ is globally $M$-Lipschitz continuous on $\Theta$ for
some $M>0$. Following the procedure of the projected stochastic gradient
descent (SGD) scheme, started from any initial $\theta_{0}$, the sequence
$\\{\theta_{j}\\}$ generated by the scheme is

$\displaystyle\theta_{j+1}=\Pi(\theta_{j}-\tau
G_{j})=\operatorname*{arg\,min}_{\theta\in\Theta}(G_{j}^{\top}\theta+\frac{1}{2\tau}|\theta-\theta_{j}|^{2}),$
(24)

where $G_{j}$ denotes the stochastic gradient of $L(\theta)$ at $\theta_{j}$
at the $j$th iteration using a mini-batch of samples, i.e., the average of the
gradient over a mini-batch of the samples to approximate the full average (9),
$\Pi$ denotes the projection of $\theta$ to $\Theta$, and $\tau$ is the
learning rate. Let $g_{j}\mathrel{\mathop{\mathchar
58\relax}}=\nabla_{\theta}L(\theta_{j})$ denote the full (but unknown)
gradient of $L$ at $\theta_{j}$, then we define the auxiliary sequence
$\\{\bar{\theta}_{j}\\}$ generated by $g_{j}$ as follows:

$\displaystyle\bar{\theta}_{j+1}=\Pi(\theta_{j}-\tau
g_{j})=\operatorname*{arg\,min}_{\theta\in\Theta}(g_{j}^{\top}\theta+\frac{1}{2\tau}|\theta-\theta_{j}|^{2}).$
(25)

Note that the sequence $\\{\bar{\theta}\\}_{j}$ is only for proof but not
necessarily computed in practice (in fact $g_{j}$ can be the actual gradient
of the exact integral and expectations in (8) and the proof below still
follows). To prove the convergence of the projected SGD iterations, first we
utilize the property of $M$-Lipschitz continuity of $\nabla_{\theta}L(\theta)$
which implies that

$\displaystyle L(\theta_{j+1})\leq
L(\theta_{j})+g_{j}^{\top}e_{j}+\frac{M}{2}|e_{j}|^{2},$ (26)

and

$\displaystyle-L(\bar{\theta}_{j+1})\leq-L(\theta_{j})-g_{j}^{\top}\bar{e}_{j}+\frac{M}{2}|\bar{e}_{j}|^{2},$
(27)

where we denote $e_{j}\mathrel{\mathop{\mathchar
58\relax}}=\theta_{j+1}-\theta_{j}$, $\bar{e}_{j}\mathrel{\mathop{\mathchar
58\relax}}=\bar{\theta}_{j+1}-\theta_{j}$ for all $i$. Due to the optimality
of $\theta_{j+1}$ in (24), we have

$\displaystyle 0$
$\displaystyle\leq(G_{j}+\frac{\theta_{j+1}-\theta_{j}}{\tau})^{\top}(\bar{\theta}_{j+1}-\theta_{j+1})$
$\displaystyle=(G_{j}+\frac{e_{j}}{\tau})^{\top}(\bar{e}_{j}-e_{j}).$ (28)

Combining (26), (27) and (II) yields

$\displaystyle
L(\theta_{j+1})-L(\bar{\theta}_{j+1})\leq(g_{j}-G_{j})^{\top}(e_{j}-\bar{e}_{j})$
(29)
$\displaystyle\qquad\qquad+\frac{e_{j}^{\top}(\bar{e}_{j}-e_{j})}{\tau}+\frac{M}{2}|e_{j}|^{2}+\frac{M}{2}|\bar{e}_{j}|^{2}.$

Again due to the $M$-Lipschitiz continuity of $\nabla_{\theta}L(\theta)$, we
have that

$\displaystyle L(\bar{\theta}_{j+1})\leq
L(\theta_{j})+g_{j}^{\top}\bar{e}_{j}+\frac{M}{2}|\bar{e}_{j}|^{2}.$ (30)

By the optimality of $\hat{\theta}_{j+1}$ in (25) with $g_{j}$, we also have

$\displaystyle 0$
$\displaystyle\leq(g_{j}+\frac{\bar{\theta}_{j+1}-\theta_{j}}{\tau})^{\top}(\theta_{j}-\bar{\theta}_{j+1})=-(g_{j}+\frac{\bar{e}_{j}}{\tau})^{\top}\bar{e}_{j}.$
(31)

Combining (30) and (31), we obtain

$\displaystyle
L(\bar{\theta}_{j+1})-L(\theta_{j})\leq-(\frac{1}{\tau}-\frac{M}{2})|\bar{e}_{j}|^{2}.$
(32)

We add (29) and (32) which yields

$\displaystyle
L(\theta_{j+1})-L(\theta_{j})\leq(g_{j}-G_{j})^{\top}(e_{j}-\bar{e}_{j})$ (33)
$\displaystyle\qquad\qquad\quad+\frac{e_{j}^{\top}(\bar{e}_{j}-e_{j})}{\tau}+\frac{M}{2}|e_{j}|^{2}-(\frac{1}{\tau}-M)|\bar{e}_{j}|^{2}.$

Now due to the Cauchy-Schwarz inequality and the definitions of $\theta_{j+1}$
and $\bar{\theta}_{j+1}$ in (24) and (25), we show that

$\displaystyle(g_{j}-G_{j})^{\top}(e_{j}-\bar{e}_{j})$
$\displaystyle=(g_{j}-G_{j})^{\top}(\theta_{j+1}-\bar{\theta}_{j+1})$
$\displaystyle=(g_{j}-G_{j})^{\top}(\Pi(\theta_{j}-\tau
G_{j})-\Pi(\theta_{j}-\tau g_{j}))$ (34)
$\displaystyle\leq|g_{j}-G_{j}||\Pi(\theta_{j}-\tau G_{j})-\Pi(\theta_{j}-\tau
g_{j})|$ $\displaystyle\leq\tau|g_{j}-G_{j}|^{2},$

where the last inequality is due to the fact that the projection $\Pi$ onto
the convex set $\Theta$ is non-expansive. Moreover we notice that

$\displaystyle\frac{e_{j}^{\top}(\bar{e}_{j}-e_{j})}{\tau}=\frac{1}{2\tau}(|\bar{e}_{j}|^{2}-|e_{j}|^{2}-|e_{j}-\bar{e}_{j}|^{2}).$
(35)

Substituting (34) and (35) into (33) we have

$\displaystyle
L(\theta_{j+1})-L(\theta_{j})\leq\tau|g_{j}-G_{j}|^{2}-(\frac{1}{2\tau}-M)|\bar{e}_{j}|^{2}$
(36)
$\displaystyle\quad-(\frac{1}{2\tau}-\frac{M}{2})|e_{j}|^{2}-\frac{1}{2\tau}|e_{j}-\bar{e}_{j}|^{2}.$

We denote $\sigma^{2}$ the variance of the stochastic gradient $G_{j}$, i.e.,
$\mathbb{E}[|G_{j}-g_{j}|^{2}]\leq\varepsilon$. Note that by the sampling
complexity of Monte-Carlo integration there is $\sigma^{2}=CN^{-1}$ for some
$C>0$ dependent on the maximum value of $L(\theta)$ (which is bounded due to
the continuity of $L$ and boundedness of $\Theta$) but not on $N$. Taking
expectation on both sides of (36) and reordering lead to

$\displaystyle(\frac{1}{2\tau}-M)\mathbb{E}[|\bar{e}_{j}|^{2}]\leq\mathbb{E}[L(\theta_{j})]-\mathbb{E}[L(\theta_{j+1})]$
(37)
$\displaystyle\qquad\qquad\qquad+\tau\sigma^{2}-(\frac{1}{2\tau}-\frac{M}{2})\mathbb{E}[|e_{j}|^{2}],$

Recall the definition of the gradient mapping, we have

$\displaystyle\mathcal{G}(\theta_{j})$
$\displaystyle=\tau^{-1}[\theta_{j}-\Pi(\theta_{j}-\tau g_{j})]$ (38)
$\displaystyle=\tau^{-1}(\theta_{j}-\bar{\theta}_{j+1})=-\tau^{-1}\bar{e}_{j}.$

Combining (37) and (38), taking sum over $j=1,2,...,J$, dividing both sides by
$(\frac{1}{2\tau}-M)\tau J$, setting $\tau=\frac{1}{4M}$, and noticing
$\sigma^{2}=C/N$, we obtain

$\displaystyle\min_{1\leq j\leq J}\mathbb{E}$
$\displaystyle[|\mathcal{G}(\theta_{j})|^{2}]\leq\frac{1}{J}\sum_{j=1}^{J}\mathbb{E}[\mathcal{G}(\theta_{j})|^{2}]$
(39)
$\displaystyle\leq\frac{16M(\mathbb{E}[L(\theta_{1})]-\mathbb{E}[L(\theta_{J+1})])}{J}+\frac{4C}{N}.$

Notice that $\mathbb{E}[L(\theta_{J+1})]\geq L^{*}\mathrel{\mathop{\mathchar
58\relax}}=\min_{\theta\in\Theta}L(\theta)$. By choosing
$N=J=[16M(\mathbb{E}[L(\theta_{1})-L^{*})+4C]\varepsilon^{-1}=O(\varepsilon^{-1})$,
we obtain that $\min_{1\leq j\leq
J}\mathbb{E}[|\mathcal{G}(\theta_{j})|^{2}]\leq\varepsilon$, as claimed. ∎

In practice, we found that using a small number of iterations to solve the
inner maximization problem with any fixed $\theta$ appears to perform the
best. The pseudo code to solve the saddle-point problem (9) is given in
Algorithm 1. The symbol $\hat{\nabla}$ and $\hat{\partial}$ stand for
stochastic gradient and partial derivative using mini-batch samples of
$X_{t}$. The selection of parameters will be given in Section III.

Algorithm 1 Optimal Density Control (ODC)

1: Input: Samples $X_{0}=\\{x_{0}^{(i)}\mathrel{\mathop{\mathchar
58\relax}}i\in[N]\\}$ following initial distribution $\rho_{0}$, samples
following target distribution $\nu$, time horizon $T>0$, step size $h$,
maximum outer and inner iteration numbers $K_{1}$ and $K_{2}$.

2: Initialize: Neural networks $u_{\theta},\phi_{\eta}$.

3: for $k_{out}=1,\dots,K_{1}$ do

4: Generate trajectories $X_{t}=\\{x_{t}^{(i)}\mathrel{\mathop{\mathchar
58\relax}}i\in[N]\\}$ using (11) with $0\leq t\leq T$.

5: Compute running cost $E(\theta,\eta)$ in (9).

6: for $k_{inn}=1,\dots,K_{2}$ do

7: Update
$\eta\leftarrow\eta+\tau\hat{\nabla}\\{\mathbb{E}_{\rho_{X_{T}}}[\phi_{\eta}]-\mathbb{E}_{\nu}[\phi_{\eta}]\\}$.

8: end for

9: Update $\theta\leftarrow\theta-\tau\hat{\partial_{\theta}}E(\theta,\eta)$.

10: end for

11: Output: $u_{\theta}$.

##  III EXPERIMENTS

In this section we validate Algorithm 1 (ODC) on both synthetic and realistic
data sets. For specific examples, we also test Algorithm 1 using Chamfer
distance 10, refered to as ODC-Chamfer. Note that it is easier to implement
ODC-Chamfer since the Chamfer distance has a closed-form expression and hence
the inner iteration 1 of ODC can be eliminated.

###  III-A Synthetic Data

We first evaluate our model on several synthetic data sets, termed by
Synthetic-1, Synthetic-2 and Synthetic-3. The state spaces of Synthetic-1 and
Synthetic-2 are 3 and 100 dimensional, while Synthetic-3 is is 2 dimensional
respectively. We set the control $u_{\theta}=\nabla\psi_{\theta}$, where
$\psi_{\theta}$ is a fully connected neural network with 3 hidden layers and
36 nodes per layer. We also set the dual function $\phi_{\eta}$ as a fully-
connected neural network with 6 hidden layers and 256 nodes per layer. We use
tanh as the activation function for all layers. We implement and test
Algorithm 1 using PyTorch and train the networks using the builtin optimizer
Adam [26] with learning rate $10^{-4}$. This optimizer appears to have
improved efficiency compared to SGD empirically. Furthermore, we use spectral
normalization to enforce $\lVert\nabla\phi_{\eta}\rVert\leq 1$ [31]. We
initialize all network parameters with Xavier initialization [19] and train
$u_{\theta}$ according to Algorithm 1 ODC (or ODC-Chamfer if the number of
agents is small and the target locations are fixed). When ODC is used, we set
the size of simulated agents at each time point as $N=2,000$, and use $1,500$
points as the training set and the other $500$ data points for testing. The
step size $h$ and time horizon $T$ are set to $0.1$ and $1$ respectively, we
choose outer iteration number as $10^{4}$ and inner iteration number for
approximating Wasserstein distance as $6$, for all synthetic experiments.

Synthetic-1: We first apply Algorithm 1 ODC with noise term in a state space
of dimension $d=3$ for two instances (two different pairs of initial-target
densities), denoted by Syn-1-a and Syn-1-b, respectively. The initial and
target distributions are set to mixtures of Gaussians. We set the perturbation
noise $\sigma=0.01$. The controlled density and target density are
respectively shown as the green and blue point clouds in Figure 1, which shows
that the controlled density in green correctly moved from the initial at $t=0$
to the blue target density at $t=1$ in the presence of perturbations in both
Syn-1-a and Syn-1-b.

(a) Syn-1-a:t=0

(b) Syn-1-a:t=1

(c) Syn-1-b:t=0

(d) Syn-1-b:t=1

Figure 1: Results of ODC on Synthetic-1 datasets for two instances. The green
point clouds indicate the controlled density, and the blue cloud points
indicate the target density. (a) and (b) show the densities at $t=0$ and $t=1$
respectively for the first instance. (c) and (d) show for the second instance.

Synthetic-2: We test the control using a large number of agents in a state
space of dimension $d=100$. We simulate 2 pairs of initial-target densities
and name them as Syn-2-a and -b. For Syn-2-a and Syn-2-b, we plot the
controlled density in green point clouds and target in blue as above, and show
the projection of them onto the first three coordinates for $t=0$ and $t=1$ in
Figure 2 (a) and (b) respectively.

(a) Syn-2-a:t=0

(b) Syn-2-a:t=1

(c) Syn-2-b:t=0

(d) Syn-2-b:t=1

Figure 2: Results of ODC on Synthetic-2 datasets for four instances of
dimension 100 as projections onto the first three dimensions (top row) and
first two dimensions (bottom row). Note that we used ResNet and Normalizing
Flow to parameterize the control $u$ in Syn-2-a and Syn-2-b respectively in
the bottom row.

(a) Syn-3:t=0

(b) Syn-3:t=10

(c) Syn-3:t=20

(d) Closest Dist

Figure 3: Results of ODC on Synthetic-3 to avoid collison on pairwise agent
distance. The black point clouds indicate the controlled density, the blue
cloud points indicate the target density. The agents move towards target
distribution from (a) to (c) and in (d) shows the closest pairwise distance
among agents at different time steps.

Synthetic-3: We also test Algorithm 1 (ODC) with an interaction potential with
$V(x,y)=|x-y|^{-2}$ in (4) to avoid agent collision on a 2D example. The
agents are indicated by the black points and the target distribution is
indicated by the blue points in Figure 3. The left three panels in Figure 3
show them at initial time $t=0$, middle time $10$, and terminal time $20$. The
rightmost figure show the minimum pairwise distance of the controlled agents
stays around 1.4 until near the terminal time. This minimum distance is
consistently above the minimum pairwise distance 1 of the target sample
points, showing that the learned control tends to avoid collisions between the
agents during the movement.

(a) Real: t=0

(b) Real: t=0.5

(c) Real: t=1

Figure 4: Results on real data. The controlled density (red) and the target
density (blue) are shown at normalized times (a) $t=0$, (b) $t=0.5$ and (c)
$t=1$.

###  III-B Real Data

In this experiment we aim to control a group of Autonomous Underwater Vehicles
(AUVs) to reach a target distribution. This experiment is motivated by the
16-month AUV ocean deployment (Processes driving Exchange At Cape Hatteras,
PEACH) near Cape Hatteras, North Carolina, a highly dynamic region
characterized by confluent western boundary currents and convergence in the
adjacent shelf and slope waters. Due to the AUVs limited forward speed, the
motion of AUV is highly influenced by the ocean flow field $v$. We define the
cost functional as
$r(x,u)=\frac{1}{2}\mathbb{E}[\int_{0}^{T}|u(x_{t})-v(x_{t})|^{2}dt]$. The
input flow map in this simulation is given by a 1-km horizontal resolution
version of the Navy Coastal Ocean Model [30] made available by J. Book and J.
Osborne (Naval Research Laboratory, Stennis Space Center). In the AUV
experiment, multiple vehicles were deployed repeatedly in the same domain, to
sample the variability in the position of the Hatteras Front. Since the AUVs
were deployed at similar locations throughout the $16$ months experiment, we
collect the starting and goal positions from all experiments as the initial
and target distribution of the system. We aim to move the AUVs to the surround
of target location on distribution level. Since the data size is small ($10$
samples only), we augment the data by generating Gaussian samples around the
given locations, thus the final distributions are mixture of Gaussians. The we
apply Algorithm 1 (ODC) to the data and show the controlled density (red) and
target density (blue) at normalized time $t=0,\ 0.5,1.0$ in Figure 4. The
network $u_{\theta}$ is parameterized as a ResNet with three layers and each
layer has 36 nodes.As we can observe, the initial density gradually moves to
the target density under the guidance of the learned control $u_{\theta}$.

##  IV CONCLUSION

In this paper, we develop a novel computational framework for density control.
We formulate the control problem with running cost and penalize inaccurate
matching to target density using Wasserstein metric. The control problem is
reformulated as a saddle-point optimization problem of the control and the
dual function of the Wasserstein distance. Both of the control and the dual
function are parameterized as neural networks which can handle high-
dimensional problems without any spatial discretization or basis
representations. We validated the promising performance of this method on
several synthetic and real data sets.

## References

  * [1] B. Açikmeşe and D. S. Bayard.  A markov chain approach to probabilistic swarm guidance.  In 2012 American Control Conference (ACC), pages 6300–6307. IEEE, 2012.
  * [2] S. Bandyopadhyay, S.-J. Chung, and F. Y. Hadaegh.  Probabilistic and distributed control of a large-scale swarm of autonomous agents.  IEEE Transactions on Robotics, 33(5):1103–1123, 2017.
  * [3] B. Bonnet and H. Frankowska.  Necessary optimality conditions for optimal control problems in wasserstein spaces.  Applied Mathematics & Optimization, 84(2):1281–1330, 2021.
  * [4] M. Brambilla, E. Ferrante, M. Birattari, and M. Dorigo.  Swarm robotics: a review from the swarm engineering perspective.  Swarm Intelligence, 7:1–41, 2013.
  * [5] K. Caluya and A. Halder.  Wasserstein proximal algorithms for the schr$\backslash$”$\\{$o$\\}$ dinger bridge problem: Density control with nonlinear drift.  IEEE Transactions on Automatic Control, 2021.
  * [6] Y. Chen.  Density control of interacting agent systems.  arXiv preprint arXiv:2108.07342, 2021.
  * [7] Y. Chen, T. Georgiou, and M. Pavon.  On the relation between optimal transport and schrödinger bridges: A stochastic control viewpoint.  In Journal of Optimization Theory and Applications, 2016.
  * [8] Y. Chen, T. Georgiou, and M. Pavon.  Optimal transport in systems and control.  Annual Review of Control, Robotics, and Autonomous Systems, 4, 2020\.
  * [9] H. Choi, U. Vaidya, and Y. Chen.  A convex data-driven approach for nonlinear control synthesis.  Mathematics, 9(19):2445, 2021.
  * [10] L. Dinh, D. Krueger, and Y. Bengio.  Nice: Non-linear independent components estimation.  arXiv preprint arXiv:1410.8516, 2014.
  * [11] M. Dorigo, G. Theraulaz, and V. Trianni.  Swarm robotics: Past, present, and future [point of view].  Proceedings of the IEEE, 109(7):1152–1165, 2021.
  * [12] K. Elamvazhuthi and S. Berman.  Mean-field models in swarm robotics: A survey.  Bioinspiration & Biomimetics, 15(1):015001, 2019.
  * [13] K. Elamvazhuthi, S. Liu, W. Li, and S. Osher.  Dynamical optimal transport of nonlinear control-affine systems.
  * [14] U. Eren and B. Açıkmeşe.  Velocity field generation for density control of swarms using heat equation and smoothing kernels.  IFAC-PapersOnLine, 50(1):9405–9411, 2017.
  * [15] Z. Feng, G. Hu, Y. Sun, and J. Soon.  An overview of collaborative robotic manipulation in multi-robot systems.  Annual Reviews in Control, 49:113–127, 2020.
  * [16] C. Frederick, M. Egerstedt, and H. Zhou.  Collective motion planning for a group of robots using intermittent diffusion.  Journal of Scientific Computing, 90(1):1–20, 2022.
  * [17] H. Gao, W. Lee, W. Li, Z. Han, S. Osher, and H. V. Poor.  Energy-efficient velocity control for massive numbers of rotary-wing uavs: A mean field game approach.  In GLOBECOM 2020-2020 IEEE Global Communications Conference, pages 1–6. IEEE, 2020.
  * [18] S. Ghadimi, G. Lan, and H. Zhang.  Mini-batch stochastic approximation methods for nonconvex stochastic composite optimization.  Math. Program., 155(1-2, Ser. A):267–305, 2016.
  * [19] X. Glorot and Y. Bengio.  Understanding the difficulty of training deep feedforward neural networks.  In International Conference on Artificial Intelligence and Statistics, 2010.
  * [20] I. Gühring and M. Raslan.  Approximation rates for neural networks with encodable weights in smoothness spaces.  Neural Networks, 134:107–130, 2021.
  * [21] I. Haasler, Y. Chen, and J. Karlsson.  Optimal steering of ensembles with origin-destination constraints.  IEEE Control Systems Letters, 5(3):881–886, 2020.
  * [22] A. Hakobyan and I. Yang.  Wasserstein distributionally robust motion control for collision avoidance using conditional value-at-risk.  IEEE Transactions on Robotics, 2021.
  * [23] K. Hoshino.  Finite-horizon control of nonlinear discrete-time systems with terminal cost of wasserstein distance.  In 2020 59th IEEE Conference on Decision and Control (CDC), pages 4268–4274. IEEE, 2020.
  * [24] D. Inoue, Y. Ito, and H. Yoshida.  Optimal transport-based coverage control for swarm robot systems: Generalization of the voronoi tessellation-based method.  In 2021 American Control Conference (ACC), pages 3032–3037. IEEE, 2021.
  * [25] S. J Reddi, S. Sra, B. Poczos, and A. J. Smola.  Proximal stochastic methods for nonsmooth nonconvex finite-sum optimization.  Advances in neural information processing systems, 29, 2016.
  * [26] D. P. Kingma and J. Ba.  Adam: A method for stochastic optimization, 2014.
  * [27] J. Lasry and P. Lions.  Mean field games.  In Japanese journal of mathematics, 2007.
  * [28] Z. Li and J. Li.  A simple proximal stochastic gradient method for nonsmooth nonconvex optimization.  Advances in neural information processing systems, 31, 2018.
  * [29] C. E. Luis, M. Vukosavljev, and A. P. Schoellig.  Online trajectory generation with distributed model predictive control for multi-robot motion planning.  IEEE Robotics and Automation Letters, 5(2):604–611, 2020.
  * [30] P. J. Martin.  Description of the navy coastal ocean model version 1.0.  Technical Report NRL/FR/7322–00-9962, Naval Research Lab, 2000.
  * [31] T. Miyato, T. Kataoka, M. Koyama, and Y. Yoshida.  Spectral normalization for generative adversarial networks.  arXiv preprint arXiv:1802.05957, 2018.
  * [32] D. Onken, L. Nurbekyan, X. Li, S. W. Fung, S. Osher, and L. Ruthotto.  A neural network approach for high-dimensional optimal control.  arXiv preprint arXiv:2104.03270, 2021.
  * [33] L. C. Pimenta, N. Michael, R. C. Mesquita, G. A. Pereira, and V. Kumar.  Control of swarms based on hydrodynamic models.  In 2008 IEEE international conference on robotics and automation, pages 1948–1953. IEEE, 2008.
  * [34] S. Roy, M. Annunziato, A. Borzì, and C. Klingenberg.  A fokker–planck approach to control collective motion.  Computational Optimization and Applications, 69:423–459, 2018.
  * [35] C. Sinigaglia, A. Manzoni, and F. Braghin.  Density control of large-scale particles swarm through pde-constrained optimization.  IEEE Transactions on Robotics, 38(6):3530–3549, 2022.
  * [36] C. Sinigaglia, A. Manzoni, F. Braghin, and S. Berman.  Robust optimal density control of robotic swarms.  arXiv preprint arXiv:2205.12592, 2022.
  * [37] L. Wang, A. D. Ames, and M. Egerstedt.  Safety barrier certificates for collisions-free multirobot systems.  IEEE Transactions on Robotics, 33(3):661–674, 2017.
  * [38] I. Yang.  Wasserstein distributionally robust stochastic control: A data-driven approach.  IEEE Transactions on Automatic Control, 66(8):3863–3870, 2020.
  * [39] Z. Yi, Z. Cao, E. Theodorou, and Y. Chen.  Nonlinear covariance control via differential dynamic programming.  In 2020 American Control Conference (ACC), pages 3571–3576. IEEE, 2020.
  * [40] S. Zhao, S. Ramakrishnan, and M. Kumar.  Density-based control of multiple robots.  In Proceedings of the 2011 American control conference, pages 481–486. IEEE, 2011.
  * [41] T. Zheng, Q. Han, and H. Lin.  Distributed mean-field density estimation for large-scale systems.  IEEE Transactions on Automatic Control, 67(10):5218–5229, 2021\.
  * [42] T. Zheng, Q. Han, and H. Lin.  Backstepping mean-field density control for large-scale heterogeneous nonlinear stochastic systems.  In 2022 American Control Conference (ACC), pages 4832–4837. IEEE, 2022.
  * [43] P. Zhu, W. Dai, W. Yao, J. Ma, Z. Zeng, and H. Lu.  Multi-robot flocking control based on deep reinforcement learning.  IEEE Access, 8:150397–150406, 2020.