# Deep Learning for Mean Field Optimal Transport 00footnotetext:
Acknowledgements:  This project was realized during CEMRACS 2022. The authors
would like to thank the CIRM for welcoming the CEMRACS 2022, the organizers of
the CEMRACS 2022 for the opportunity to work on this project as well as their
respective institutions. They are also grateful to the NYU-ECNU Institute of
Mathematical Sciences at NYU Shanghai, the CIMPA fellowships program, and the
ENS Rennes for their support. This work was supported in part through the
NYUSH IT High Performance Computing resources, services, and staff expertise,
as well as the ASCC Toubkal cluster resources.

Sebastian Baudelet, Brieuc Frénais, Mathieu Laurière, Amal Machtalay, Yuchen
Zhu Université Côte d’Azur, 28 Avenue de Valrose, 06103 Nice, France,
<EMAIL_ADDRESS>UMR 7501, Université de Strasbourg, 7
Rue René Descartes, 67000 Strasbourg, France,
<EMAIL_ADDRESS>Institute of Mathematical Sciences at
NYU Shanghai, 3663 Zhongshan Road North, Shanghai, 200062, China,
<EMAIL_ADDRESS>VI Polytechnic University. Lot 660, Hay
Moulay Rachid Ben Guerir, 43150, Morocco<EMAIL_ADDRESS>University, 433 Temple Street, New Haven, 06511, USA<EMAIL_ADDRESS>
###### Abstract

Mean field control (MFC) problems have been introduced to study social optima
in very large populations of strategic agents. The main idea is to consider an
infinite population and to simplify the analysis by using a mean field
approximation. These problems can also be viewed as optimal control problems
for McKean-Vlasov dynamics. They have found applications in a wide range of
fields, from economics and finance to social sciences and engineering.
Usually, the goal for the agents is to minimize a total cost which consists in
the integral of a running cost plus a terminal cost. In this work, we consider
MFC problems in which there is no terminal cost but, instead, the terminal
distribution is prescribed. We call such problems mean field optimal transport
problems since they can be viewed as a generalization of classical optimal
transport problems when mean field interactions occur in the dynamics or the
running cost function. We propose three numerical methods based on neural
networks. The first one is based on directly learning an optimal control. The
second one amounts to solve a forward-backward PDE system characterizing the
solution. The third one relies on a primal-dual approach. We illustrate these
methods with numerical experiments conducted on two families of examples.

##  1 Introduction

Mean field games (MFGs) have been introduced by Lasry and Lions [41, 42, 43]
and Caines, Huang and Malhamé [40, 39] to approximate Nash equilibria in games
with a very large number of players. At a high level, the main idea is to use
a mean field approximation to represent the state of the population, and then
to focus on the interactions between a single representative player and the
distribution of the states of the other players. Mean field control (MFC) [14]
relies on a similar approximation but aims at representing situations in which
a large number of agents cooperate to minimize a common social cost. The
problem can be interpreted as an optimal control problem for a McKean-Vlasov
(MKV) stochastic differential equation (SDE) or an optimal control for a
Kolmogorov-Fokker-Planck (KFP) partial differential equation (PDE). In the
past decade, the analysis of both MFGs and MFC problems has been extensively
developed, see e.g. [14] for an introduction to this topic, and [21] for a
probabilistic viewpoint.

In the most common setup, the players try to minimize a total cost which is
composed of a running cost integrated over time and a terminal cost. These
costs generally account for the efforts made to control the dynamics as well
as the preferences for some states over others. Another class of models has
been introduced, in which there is no terminal cost and instead the terminal
distribution of the population is imposed as a constraint. Nash equilibria
have been studied under the name of planning problem for mean field game. This
class of problems has been analyzed mostly using PDE-based techniques [1, 50,
51, 48, 34, 16]. In the special case of linear dynamics and a quadratic
running cost in the control, the problem is related to the Schrödinger bridge
problem, and equilibrium conditions can be phrased in terms of ordinary
differential equations (ODEs) [25, 27, 26, 28, 29].

This research direction is tightly connected to optimal transport (OT).
Benamou and Brenier proposed in [10] a fluid mechanics framework for the
$\textstyle L^{2}$ Monge-Kantorovich mass transport problem. Many works built
on this approach to relate optimal transport and optimal control problems for
continuity equations. Of particular interest for MFGs is the work [19], which
clarified the link between geodesics for a class of distances between
probability measures and a PDE system similar to the one arising in MFGs. For
more background on OT, we refer the interested reader to the monographs [58,
59, 55, 49, 8]. However, the solutions of MFGs correspond to Nash equilibria,
and hence, in general, MFGs do not admit a variational structure. Furthermore,
in many applications, it is not immediately clear to us why selfish players
caring only about their individual costs would manage to agree and reach a
target terminal distribution. Imposing a fixed terminal distribution seems
more natural in the MFC setting, where the agents behave in a cooperative way
to minimize the social cost. In the present work, we focus on such MFC with
planning problems, in which a mean field of agents try to collectively
minimize a social cost while ensuring that a fixed distribution is attained at
the terminal time.

Since the work of Benamou and Brenier [10], several numerical methods have
been investigated for similar problems, including MFGs with planning. Achdou
et al. have proposed in [1] a method based on finite differences and Newton
method to solve the PDE system of MFG with planning. Benamou and Carlier have
used in [11, 13] an Augmented Lagrangian method approach with the alternating
direction method of multipliers to solve OT and MFG (without planning).
Similar methods have been used in [9, 5] to solve MFGs and MFC problems (still
without planning). Benamou et al. proposed in [12] a method to solve MFG with
planning through entropy regularization and Sinkhorn algorithm.

Recently, several deep learning methods have been proposed to solve high-
dimensional optimal control problems and PDEs, such as the DeepBSDE method
[36, 37, 35], the Deep Galerkin Method [57] and physics-informed neural
networks [52]. Some of these methods have been extended to MFGs and MFC
problems. In particular, [7, 22] proposed deep learning methods to solve the
PDE systems arising in mean field problems, [23, 32, 31] introduced deep
learning methods for differential MFC problems. Ruthotto et al. introduced a
deep learning method for variational MFG with degenerate diffusion in [54].
Lin et al. introduced a deep learning method in [45] that utilizes the primal-
dual relationship of variational MFG. Cao et al. noticed a connection between
MFGs, generative adversarial networks and OT in [18]. We refer the interested
reader to e.g. [24, 33, 38] for recent surveys on this topic. The work most
related to ours is the work of Liu et al. in [46], where they considered the
planning problems in a class of MFGs based on a generalized version of the
Schrödinger bridge problem and proposed a neural network-based numerical
method to solved it.

The main goal of this paper is to propose numerical methods based on deep
learning to solve MFC problems with planning constraint, that we will call
mean field optimal transport problems. To the best of our knowledge, the
theory remains to be investigated in detail, and this is beyond the scope of
the present work. Here, we proceed formally when needed, and we focus on the
numerical aspects using machine learning tools. The rest of the paper is
organized as follows. In Section 2, we introduce the problem and discuss
several examples. In Section 3, we describe three numerical methods, each
based on a different approach for the problem. In Section 4, we present
numerical results on several benchmark problems.

##  2 Definition of the problem

Before presenting the mean field optimal transport problem, let us first
recall the definition of a typical mean field control problem. Let $\textstyle
T$ be a time horizon. Let $\textstyle\mathcal{Q}={{\mathbb{R}}}^{d}$ and
$\textstyle\mathcal{Q}_{T}=[0,T]\times\mathcal{Q}$ denote the space domain and
the time-space domain. Denote by $\textstyle{\mathcal{P}}_{2}(\mathcal{Q})$
the set of square-integrable probability measures on $\textstyle\mathcal{Q}$.
Let $\textstyle
f:\mathcal{Q}\times{\mathcal{P}}_{2}(\mathcal{Q})\times{{\mathbb{R}}}^{k}\to{{\mathbb{R}}}$
be a running cost function, $\textstyle
g:\mathcal{Q}\times{\mathcal{P}}_{2}(\mathcal{Q})\to{{\mathbb{R}}}$ be a
terminal cost function, $\textstyle
b:\mathcal{Q}\times{\mathcal{P}}_{2}(\mathcal{Q})\times{{\mathbb{R}}}^{k}\to{{\mathbb{R}}}^{d}$
be a drift function and $\textstyle\sigma\in{{\mathbb{R}}}$ be a non-negative
constant diffusion coefficient. In a classical MFC problem with given initial
distribution $\textstyle\rho_{0}$ in
$\textstyle{\mathcal{P}}_{2}(\mathcal{Q})$, the goal is to find a feedback
control $\textstyle v^{*}:\mathcal{Q}_{T}\to{{\mathbb{R}}}^{k}$ minimizing:

$\textstyle\displaystyle
J^{MFC}:v\mapsto{{\mathbb{E}}}\left[\int_{0}^{T}f(X_{t}^{v},\mu^{v}(t),v(t,X_{t}^{v}))\mathrm{d}t+g(X_{T}^{v},\mu^{v}(T))\right]$
(1)

where $\textstyle\mu^{v}(t)$ is the distribution of $\textstyle X_{t}^{v}$,
under the constraint that the process $\textstyle X^{v}=(X_{t}^{v})_{t\geq 0}$
solves the SDE

$\textstyle\displaystyle\begin{cases}X_{0}^{v}\sim\rho_{0}\\\
\mathrm{d}X_{t}^{v}=b(X_{t}^{v},\mu^{v}(t),v(t,X_{t}^{v}))\mathrm{d}t+\sigma\mathrm{d}W_{t},\qquad
t\geq 0,\end{cases}$ (2)

where $\textstyle W$ is a standard $\textstyle d$-dimensional Brownian motion.
It would also be interesting to consider open-loop controls, but since we are
motivated by numerical applications, we restrict our attention to feedback
controls. The cost (1) can be interpreted either as the expected cost for a
single representative player, or as the average cost for the whole population,
which we refer to as the social cost.

In this work, we are interested in a modified version of the above problem,
where instead of having a terminal cost, a terminal distribution is imposed.
This type of problem encompasses optimal transport as a special case, but it
may incorporate mean field interactions in the drift and the running cost. For
this reason, we will refer to this class of problems as mean field optimal
transport (MFOT for short).111By analogy with MFG of planning type, we could
also call such problems “MFC of planning type”. But referring to “optimal
transport” seems clearer so we will stick to the MFOT terminology. Given two
distributions $\textstyle\rho_{0}$ and
$\textstyle\rho_{T}\in{\mathcal{P}}_{2}(\mathcal{Q})$, the goal is to find a
feedback control $\textstyle v^{*}:\mathcal{Q}_{T}\to{{\mathbb{R}}}^{k}$
minimizing

$\textstyle\displaystyle
J^{MFOT}:v\mapsto{{\mathbb{E}}}\left[\int_{0}^{T}f(X_{t}^{v},\mu^{v}(t),v(t,X_{t}^{v}))\mathrm{d}t\right],$
(3)

where $\textstyle\mu^{v}(t)$ is the distribution of $\textstyle X_{t}^{v}$,
under the constraint that the process $\textstyle X^{v}=(X_{t}^{v})_{t\geq 0}$
solves the SDE

$\textstyle\displaystyle\begin{cases}X_{0}^{v}\sim\rho_{0},\qquad
X_{T}^{v}\sim\rho_{T}\\\
dX_{t}^{v}=b(X_{t}^{v},\mu^{v}(t),v(t,X_{t}^{v}))\mathrm{d}t+\sigma\mathrm{d}W_{t},\qquad
t\geq 0.\end{cases}$ (4)

We stress that the terminal constraint implicitly restricts the class of
admissible controls since we are interested in minimizing only over controls
$\textstyle v$ that make $\textstyle X_{T}^{v}$ have distribution
$\textstyle\rho_{T}$.

We now present a few useful examples, some of which will be revisited in the
numerical experiments (see Section 4).

######  Example 1 (Optimal transport).

When $\textstyle b(x,\mu,a)=a$, $\textstyle f(x,\mu,a)=\frac{1}{2}a^{\top}a$
and $\textstyle\sigma=0$, the MFOT problem reduces to a _standard OT_ problem.
See e.g. [10].

######  Example 2 (Linear-quadratic).

Take $\textstyle b(x,\mu,a)=Ax+\bar{A}\bar{\mu}+Ba$, $\textstyle
f(x,\mu,a)=x^{\top}Qx+\bar{\mu}^{\top}\bar{Q}\bar{\mu}+a^{\top}Ra$, and
$\textstyle g(x,\mu)=x^{\top}Q_{T}x+\bar{\mu}^{\top}\bar{Q}_{T}\bar{\mu}$,
where $\textstyle\bar{\mu}=\int\xi\mu(\mathrm{d}\xi)$, where $\textstyle
A,\bar{A},B,Q,\bar{Q},R,Q_{T}$ and $\textstyle\bar{Q}_{T}$ are matrices of
suitable sizes. In this setting, the MFC problem has an explicit solution, up
to solving a forward-backward system of ODEs. Furthermore, if the initial
distribution is Gaussian, then the optimal flow of distribution remains
Gaussian. See e.g. [14, Chapter 6]. To the best of our knowledge, in the MFOT
setting, a similar result is available in the literature only when $\textstyle
Q=\bar{Q}=0$, which corresponds to the Schrödinger bridge problem. See [29,
Section 7.1].

######  Example 3 (Crowd motion with congestion).

Take $\textstyle b(x,\mu,a)=a$, $\textstyle
f(x,\mu,a)=(c+\rho\star\mu(x))^{\gamma}|a|^{2}+\ell(x,\mu(x))$, where
$\textstyle c\geq 0$ is a constant, $\textstyle\rho$ is a regularizing kernel
and $\textstyle\star$ denotes the convolution. For $\textstyle\gamma=0$, the
model is linear-quadratic in the control. If $\textstyle\gamma>0$, the cost of
moving increases with the density surrounding the agent, which represents the
fact that the “energy” spent to move is higher in regions with higher density.
This models a _congestion effect_. The last term in $\textstyle f$ can be used
to represent _crowd aversion_ if $\textstyle\ell$ is increasing with respect
to $\textstyle\mu(x)$, and it can be used to represent spatial preferences by
taking for instance $\textstyle\ell(x,\mu(x))=|x_{*}-x|^{2}$, where
$\textstyle x_{*}$ is a preferred position. The terminal cost $\textstyle g$
can also be used to represent crowd aversion or spatial preferences. See e.g.
[3, 4] for more details on the analysis of the MFC PDE system for this class
of models and [5] for numerical aspects. When $\textstyle\ell=0$ and
$\textstyle\sigma=0$, the corresponding MFOT problem has been studied e.g. in
[19]. Similar models have also been studied in the context of MFGs, see e.g.
[6, 2].

##  3 Numerical methods

In this section, we introduce three different numerical methods to solve MFOT.
Section 3.1 introduces a direct approach to solve a MFC problem that
approximates the MFOT problem. Section 3.2 discusses the Deep Galerkin Method
(DGM) to solve the underlying PDE system that characterizes the optimal
solution to MFOT, which is composed of a coupled Hamilton-Jacobi-Bellman
equation and a Kolomogrov-Fokker-Planck equation. Section 3.3 introduces the
DeepADMM algorithm that solves a variational reformulation of the MFOT problem
based on an augmented Lagrangian approach.

###  3.1 Direct approach for the optimal control formulation

We first introduce the direct approach, which does not require any derivation
of optimality conditions. In order to make the problem numerically tractable,
we make approximations on several levels. Motivated by the deep learning
method for MFC problems proposed in [22] (see also the first algorithm in
[24]), we first approximate the MFOT problem (3) by an MFC problem in which a
terminal penalty is incurred based on the distance between the terminal
distribution and the target distribution. We can then apply the algorithm of
[22], which trains a neural network to learn the optimal control of the MFC
problem. This method itself relies on three approximations.

####  3.1.1 Problem Approximation

Instead of directly tackling the MFOT problem (3), we first consider the
following MFC problem as an approximation of the original problem: Find a
feedback control $\textstyle v^{*}:\mathcal{Q}_{T}\to{{\mathbb{R}}}^{k}$
minimizing (1) under the constraint (2) when the terminal cost is:

$g(x,\mu)=G(\mathcal{W}_{2}(\mu,\rho_{T})),\qquad\mu\in{\mathcal{P}}_{2}({\mathcal{Q}}).$
(5)

where $\textstyle G:{{\mathbb{R}}}_{+}\to{{\mathbb{R}}}_{+}$ is an increasing
function and $\textstyle\mathcal{W}_{2}$ denotes the Wasserstein distance on
$\textstyle{\mathcal{P}}_{2}(\mathcal{Q})$. A typical example that we will use
in the experiments is a linear function. The purpose of introducing
$\textstyle G(\mathcal{W}_{2}(\mu,\rho_{T}))$ is to add a penalty that
enforces the planning constraint for the terminal distribution. Here we focus
on the Wasserstein distance because of its connection with optimal transport,
see e.g. [10, 55], although other similarity measures could be used. In our
numerical experiments, we will take an increasing linear function for
$\textstyle G$.

Then, we use the following approximations:

  * •

Since it is not possible to optimize overall feedback controls, we restrict
the space of controls to the space of neural networks with a given
architecture. We will denote by $\textstyle v_{\theta}$ a representative
neural network of this class with parameter $\textstyle\theta$. The problem
becomes a finite-dimensional optimization problem, in which the goal is to
find a value for the parameter $\textstyle\theta$ that minimizes the loss
$\textstyle J(v_{\theta})$, i.e., the total cost of the MFC problem when using
control $\textstyle v_{\theta}$.

  * •

Since it is not possible to represent the mean field state
$\textstyle\mu^{v_{\theta}}(t)$ or to compute its evolution exactly, we
approximate it by the empirical distribution
$\textstyle\bar{\mu}^{N,v_{\theta}}(t)=\frac{1}{N}\sum_{i=1}^{N}\delta_{X^{i,v_{\theta}}_{t}}$,
where each $\textstyle X^{i,v_{\theta}}$ is a solution of,

$\textstyle\displaystyle\begin{cases}X_{0}^{i,v_{\theta}}\sim\rho_{0}\quad\hbox{
i.i.d.}\\\
dX_{t}^{i,v_{\theta}}=b(X_{t}^{i,v_{\theta}},\bar{\mu}^{N,v_{\theta}}(t),v_{\theta}(t,X_{t}^{i,v_{\theta}}))\mathrm{d}t+\sigma\mathrm{d}W_{t}^{i},\qquad
t\geq 0,\end{cases}$ (6)

where $\textstyle(W^{i})_{i=1,\dots,N}$ is a family of $\textstyle N$
independent $\textstyle d$-dimensional Brownian motions, which represent
idiosyncratic noises affecting each particle independently. All the SDEs are
based on the same control function $\textstyle v_{\theta}$.

  * •

Last, in order to be able to compute these dynamics using Monte Carlo
simulations, we discretize the time variable $\textstyle t$. Letting
$\textstyle N_{T}$ be a number of regular time steps of length
$\textstyle\Delta t=T/N_{T}$, we replace the interval $\textstyle[0,T]$ by the
time steps $\textstyle\\{t_{0}=0,t_{1}=\Delta t,\dots,t_{N_{T}}=N_{T}\Delta
t\\}$. The time steps are $\textstyle t_{n}=n\Delta t$, $\textstyle
n=0,\dots,N_{T}$. We then approximate the SDE system (6) using an Euler-
Maruyama scheme. The family of trajectories
$\textstyle((X^{i,v_{\theta}}_{t})_{t\in[0,T]})_{i=1,\dots,N}$ is approximated
by the family of sequences
$\textstyle((X^{i,v_{\theta},N_{T}}_{t_{n}})_{n=0,\dots,N_{T}})_{i=1,\dots,N}$
satisfying:

$\begin{cases}X^{i,v_{\theta},N_{T}}_{0}\sim\rho_{0}\quad\hbox{ i.i.d.}\\\
X^{i,v_{\theta},N_{T}}_{t_{n+1}}=X^{i,v_{\theta},N_{T}}_{t_{n}}+b(X^{i,v_{\theta},N_{T}}_{t_{n}},\bar{\mu}^{N,v_{\theta},N_{T}}_{t_{n}},v_{\theta}(t_{n},X^{i,v_{\theta},N_{T}}_{t_{n}}))\Delta
t+\sigma\Delta W^{i}_{n},\end{cases}$ (7)

where
$\textstyle\bar{\mu}^{N,v_{\theta},N_{T}}_{t_{n}}=\frac{1}{N}\sum_{i=1}^{N}\delta_{X^{i,v_{\theta},N_{T}}_{t_{n}}}$
is the empirical distribution associated with the samples $\textstyle
X^{i,\theta,N_{T}}_{t_{n}}$. Here, $\textstyle(\Delta
W^{i}_{n})_{i=1,\dots,N,n=0,\dots,N_{T}-1}$ are independent Gaussian random
variables with variance $\textstyle\Delta t$.

To summarize, the new problem is to find $\textstyle\theta^{*}$ minimizing:

$J^{N,N_{T}}(\theta)=\mathbb{E}\left[\dfrac{1}{N}\sum_{i=1}^{N}\sum_{n=0}^{N_{T}-1}f(X^{i,\theta,N_{T}}_{t_{n}},\bar{\mu}^{N,\theta,N_{T}}_{t_{n}},v_{\theta}(t_{n},X^{i,\theta,N_{T}}_{t_{n}}))\Delta
t+g(X^{N,\theta,N_{T}}_{T},\bar{\mu}^{N,\theta,N_{T}}_{T})\right]$

subject to the dynamics (7). The full analysis of this problem and its
rigorous connection with the original MFOT problem (3) is beyond the scope of
this paper and is left for future work. We expect the control $\textstyle
v_{\theta^{*}}$, with the parameter value that is optimal for the above
problem, to be approximately optimal for (3), under suitable assumptions on
$\textstyle b$ and $\textstyle f$. In particular, $\textstyle b$ and
$\textstyle f$ should probably depend smoothly on the distribution so that
they can be evaluated in a meaningful way at the empirical distribution
$\textstyle\bar{\mu}^{N,\theta,N_{T}}_{t_{n}}$.

####  3.1.2 Description of the algorithm

Optimization method.  To find an approximate minimizer, we use stochastic
gradient descent (SGD) or one of its variants. At iteration $\textstyle k$, we
have a parameter $\textstyle\theta_{k}$ that we wish to update. We sample the
initial positions $\textstyle(X^{i,v_{\theta_{k}},N_{T}}_{0})_{i=1,\dots,N}$
and the Brownian motion increments $\textstyle(\Delta
W^{i}_{n})_{i=1,\dots,N,n=0,\dots,N_{T}-1}$. We then compute the empirical
cost for this realization of the $\textstyle N$-particle population, and use
its gradient with respect to $\textstyle\theta_{k}$ to update the parameter.
In other words, we apply SGD to the following loss function:

$\textstyle\mathcal{L}(\theta)=J^{N,N_{T}}(\theta)=\mathbb{E}_{S}[\mathcal{L}(\theta;S)],$

with:

$\textstyle\mathcal{L}(\theta;S)=\dfrac{1}{N}\sum_{i=1}^{N}\sum_{n=0}^{N_{T}-1}f(X^{i,\theta,N_{T}}_{t_{n}},\bar{\mu}^{N,\theta,N_{T}}_{t_{n}},v_{\theta}(t_{n},X^{i,\theta,N_{T}}_{t_{n}}))\Delta
t+g(X^{N,\theta,N_{T}}_{T},\bar{\mu}^{N,\theta,N_{T}}_{T})$

where $\textstyle S=\left((X^{i,v_{\theta},N_{T}}_{0})_{i=1,\dots,N},(\Delta
W^{i}_{n})_{i=1,\dots,N,n=0,\dots,N_{T}-1}\right)$ denotes one random sample.

Computation of the Wasserstein distance.  As shown in (5), the new problem we
considered involves a Wasserstein distance between two continuous
distributions, namely, the mean field distribution at terminal time
$\textstyle\mu^{v}_{T}$ and the target distribution $\textstyle\rho_{T}$. This
is in general hard to compute. However, in our implementation, the mean field
distribution is approximated by an empirical distribution obtained by Monte
Carlo simulations, as is explained above. We then sample the same number of
points from the target distribution and compute the Wasserstein distance
between the two empirical distributions. This is done in the following way.
Let $\textstyle X$ and $\textstyle Y$ be two sets of $\textstyle N$ points
each sampled from distributions $\textstyle\mu,\nu$, $\textstyle M_{p}$ the
distance matrix, $\textstyle(M_{p})_{ij}=|X_{i}-Y_{j}|^{p}$, and the following
set:

$\textstyle\displaystyle U_{N}=\left\\{A\in\mathbb{R}^{N\times
N}\Big{|}\sum\limits_{j=1}^{N}A_{ij}=\sum\limits_{i=1}^{N}A_{ij}=\dfrac{1}{N}\right\\}.$
(8)

Then

$\textstyle\displaystyle\Bigl{(}\mathcal{W}_{p}\left(\mu,\nu\right)\Bigr{)}^{p}=\lim\limits_{N\to\infty}\min\limits_{T\in
U_{N}}\left\langle T,M_{p}\right\rangle.$

In order to efficiently compute the Wasserstein distance, we follow the
algorithm proposed by Cuturi in [30]. We consider an extra entropy
regularization of the following form. Let $\textstyle\alpha>0$, we want to
find $\textstyle T^{*}_{\alpha}$, which is the solution to the following
program:

$\textstyle\displaystyle\min\limits_{T\in U_{N}}\left\langle
T,M_{p}\right\rangle-\alpha\left\langle T\log(T),1\right\rangle.$

Optimality conditions and Sinkhorn-Knopp [56] theorem give us the existence
and uniqueness of the solution, as well as a unique decomposition of
$\textstyle T_{\alpha}^{*}$ using two vectors $\textstyle u$ and $\textstyle
v$ such that:

$\textstyle\displaystyle
T^{*}_{\alpha}=\operatorname{Diag}(u)\exp\left(-\dfrac{M_{p}}{\alpha}\right)\operatorname{Diag}(v).$

We can then compute $\textstyle u$ and $\textstyle v$ with Sinkhorn-Knopp
algorithm. Further explanations on this algorithm can be found in [30]. This
method allows for fast computations and is easy to export to greater
dimensions, at the cost of adding a layer of approximation due to the extra
parameter $\textstyle\alpha$. It can be noticed that as $\textstyle\alpha$
tends to zero, the regularized solution tends to the solution of discrete
optimal transport. In practice, reducing $\textstyle\alpha$ to zero increases
the number of iterations required for Sinkhorn algorithm to converge. However,
in our numerical experiments we usually obtain good results with a small but
non-zero $\textstyle\alpha$.

######  Remark 1.

Notice that using our approach, we have one empirical distribution and one
continuous distribution. Indeed, we have the empirical distribution obtained
by Monte Carlo simulation and the target distribution $\textstyle\rho_{T}$,
which is generally given by a closed-form formula for its density. We could
thus try to use the designated methods, such as Semi-discrete Optimal
transport [47]. While being more accurate, these methods do not scale well in
higher dimensions compared to Sinkhorn’s alternative.

Terminal penalty.  In our implementation, we take $\textstyle G$ as a linear
function $\textstyle G(r)=C_{W}r$, where $\textstyle C_{W}$ is a positive
constant that weighs the importance of the terminal penalty in comparison with
the running cost. This leads to a trade-off between minimizing the running
cost and satisfying the terminal constraint. We noticed that when $\textstyle
C_{W}$ is too small, the algorithm minimizes the running cost without much
consideration for the terminal condition and hence the terminal distribution
is far from the target distribution. Therefore, the penalization has to be a
significant component of the total loss if we want the terminal planning
constraint to be approximately satisfied with good accuracy.

###  3.2 Deep Galerkin Method for the PDE system

We now turn our attention to a method based on solving a forward-backward PDE
system that characterizes the solution. We first discuss the PDE system and
then use a deep learning method to solve this system.

####  3.2.1 PDE system for MFOT

As recalled above, in a standard MFC, the whole population uses a given
feedback control $\textstyle v$. Assuming that the distribution
$\textstyle\mu^{v}_{t}={\mathcal{L}}(X^{v}_{t})$ of a representative agent
with dynamics (2) admits a smooth enough density $\textstyle m_{t}$, the
latter satisfies the Kolmogorov-Fokker-Planck (KFP) PDE:

$\begin{cases}&\frac{\partial m}{\partial t}(t,x)-\nu\Delta
m(t,x)+\operatorname{div}\bigl{(}m(t,x)b(x,m(t,\cdot),v(t,x))\bigr{)}=0\qquad
t\in(0,T],x\in\mathcal{Q}\\\ &m(0,x)=m_{0}(x),\qquad
x\in\mathcal{Q},\end{cases}$

where $\textstyle m_{0}$ is the density of the initial distribution
$\textstyle\rho_{0}$ and $\textstyle\nu=\frac{\sigma^{2}}{2}$. The MFC problem
(1) can then be viewed as an optimal control problem driven by the above KFP
PDE. Under suitable conditions, the optimal control can be characterized
through an adjoint PDE, which can be derived for instance via calculus of
variations. See e.g. [14, Chapter 4] for more details.

Let $\textstyle H:\mathcal{Q}\times
L^{2}(\mathcal{Q})\times{{\mathbb{R}}}^{d}\to{{\mathbb{R}}}$ be the
Hamiltonian of the control problem faced by an infinitesimal agent in the
first point above, which is defined by:

$\textstyle\displaystyle H:(x,m,p)\mapsto
H(x,m,p)=\max_{v\in{{\mathbb{R}}}^{k}}\\{-L(x,m,v,p)\\},$ (9)

where $\textstyle m$ denotes the density of $\textstyle\mu$, $\textstyle
L:\mathcal{Q}\times
L^{2}(\mathcal{Q})\times{{\mathbb{R}}}^{k}\times{{\mathbb{R}}}^{d}\to{{\mathbb{R}}}$
is the Lagrangian, defined by:

$\textstyle\displaystyle L:(x,m,v,p)\mapsto L(x,m,v,p)=f(x,m,v)+\langle
b(x,m,v),p\rangle.$ (10)

A necessary condition for the existence of an optimal control $\textstyle
v^{*}$ is that:

$\textstyle
v^{*}(t,x)=\operatorname*{\arg\\!\max}_{v\in{{\mathbb{R}}}^{k}}\big{\\{}-L(x,m(t,\cdot),v,\nabla
u(t,x))\big{\\}},$

where $\textstyle(u,m)$ solve the following system of partial differential
equations:

$\textstyle\displaystyle 0$
$\textstyle\displaystyle=\displaystyle-\frac{\partial u}{\partial
t}(t,x)-\nu\Delta u(t,x)+H(x,m(t,\cdot),\nabla u(t,x))$
$\textstyle\displaystyle\qquad+\int_{\mathcal{Q}}\frac{\partial H}{\partial
m}(\zeta,m(t,\cdot),\nabla u(t,\zeta))(x)m(t,\zeta)\mathrm{d}\zeta,$
$\textstyle\displaystyle\hbox{ in }(0,T]\times\mathcal{Q},$ (11a)
$\textstyle\displaystyle 0$
$\textstyle\displaystyle=\displaystyle\frac{\partial m}{\partial
t}(t,x)-\nu\Delta
m(t,x)-\operatorname{div}\Bigl{(}m(t,x)\partial_{p}H(x,m(t,\cdot),\nabla
u(t,x))\Bigr{)},$ $\textstyle\displaystyle\hbox{ in }[0,T)\times\mathcal{Q},$
(11b) $\textstyle\displaystyle
u(T,x)=g(x,m(T,\cdot))+\int_{\mathcal{Q}}\frac{\partial g}{\partial
m}(\zeta,m(T,\cdot))(x)m(T,\zeta)\mathrm{d}\zeta,$
$\textstyle\displaystyle\hbox{ in }\mathcal{Q},$ (11c)
$\textstyle\displaystyle m(0,x)=m_{0}(x),$ $\textstyle\displaystyle\hbox{ in
}\mathcal{Q}.$ (11d)

The partial derivatives with respect to $\textstyle m$ appear in the backward
PDE due to the fact that the population distribution changes when the control
changes. These partial derivatives with respect to $\textstyle m$ should be
understood in the following sense: if
$\textstyle\varphi:L^{2}({{\mathbb{R}}}^{d})\to{{\mathbb{R}}}$ is
differentiable,

$\textstyle\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\varphi(m+\varepsilon\tilde{m})(x)_{\big{|}\varepsilon=0}=\int_{{{\mathbb{R}}}^{d}}\frac{\partial\varphi}{\partial
m}(m)(\zeta)\tilde{m}(\zeta)\mathrm{d}\zeta.$

We refer to e.g. [14, Chapter 4] for more details and for the derivation using
calculus of variations, which clarifies why the partial derivatives with
respect to $\textstyle m$ appear. If the cost functions and the drift function
depend on the density only locally (i.e., only on the density at the current
position of the agent), $\textstyle\frac{\partial}{\partial m}$ becomes a
derivative in the usual sense.

In this PDE system, $\textstyle m$ plays the role of the MFC problem’s state.
The forward equation is a Kolmogorov-Fokker-Planck (KFP) equation which
describes the evolution of the mean field distribution. The other unknown
function, $\textstyle u$, plays the role of an adjoint state. Although the
backward PDE has the form of a Hamilton-Jacobi-Bellman (HJB) equation,
$\textstyle u$ cannot, in general, be interpreted as the value function
associated to problem (1) because the value function depends on the population
distribution, see e.g. [44, 15]. We refer the interested reader to e.g. [14,
Chapters 3 and 4] for the comparison with the MFG PDE system, in which the
terms involving a derivative with respect to $\textstyle m$ are absent, and
$\textstyle u$ can be interpreted as the value function of an infinitesimal
player.

Now, for the MFOT problem, we can proceed formally in a similar way. We derive
an analogous PDE system, except that the terminal condition for $\textstyle u$
disappears, and a terminal condition for $\textstyle m$ is added to the
system. More precisely, we (formally) obtain the following PDE system:

$\textstyle\displaystyle 0$
$\textstyle\displaystyle=\displaystyle-\frac{\partial u}{\partial
t}(t,x)-\nu\Delta u(t,x)+H(x,m(t,\cdot),\nabla u(t,x))$
$\textstyle\displaystyle\qquad+\int_{\mathcal{Q}}\frac{\partial H}{\partial
m}(\zeta,m(t,\cdot),\nabla u(t,\zeta))(x)m(t,\zeta)\mathrm{d}\zeta,$
$\textstyle\displaystyle\hbox{ in }(0,T]\times\mathcal{Q},$ (12a)
$\textstyle\displaystyle 0$
$\textstyle\displaystyle=\displaystyle\frac{\partial m}{\partial
t}(t,x)-\nu\Delta
m(t,x)-\operatorname{div}\Bigl{(}m(t,x)\partial_{p}H(x,m(t,\cdot),\nabla
u(t,x))\Bigr{)},$ $\textstyle\displaystyle\hbox{ in }[0,T)\times\mathcal{Q},$
(12b) $\textstyle\displaystyle m(0,x)=m_{0}(x),\qquad m(T,x)=m_{T}(x)$
$\textstyle\displaystyle\hbox{ in }\mathcal{Q}.$ (12c)

where $\textstyle m_{0}$ and $\textstyle m_{T}$ are respectively the densities
of $\textstyle\rho_{0}$ and $\textstyle\rho_{T}$. To the best of our
knowledge, this PDE system has not been derived nor analyzed in a general
setting. Notice that even the existence of a solution is a non trivial
question due to the fact that there is both an initial and a terminal
constraint on the density. However, this system has been analyzed in special
cases corresponding to optimal transport [10, 19] or to MFGs with planning [1,
50, 51, 34]. In the numerical examples of Section 4, we will mostly focus on
cases that have been previously studied, such as standard optimal transport or
MFOT with congestion effects captured by the running cost.

####  3.2.2 Description of the algorithm

To solve the PDE system (12), we follow the idea of the Deep Galerkin Method
(DGM) introduced by Sigignano and Spiliopoulos [57] and adapted to the MFG and
PDE systems in [7, 23, 24]. The main motivation underlying this approach is to
learn the PDE solutions using parameterized functions. This avoids computing
the functions on a mesh, which is not feasible in high dimensions. In the DGM,
we replace the function(s) solving the PDE(s) with a neural network(s), which
are trained to minimize the PDE residual(s) as well as the boundary
condition(s).

To be specific, in our setting, we replace the functions $\textstyle m$ and
$\textstyle u$ with neural networks, denoted by $\textstyle m_{\theta}$ and
$\textstyle u_{\omega}$ and parameterized by $\textstyle\theta$ and
$\textstyle\omega$ respectively. When the state $\textstyle x$ is in high
dimension, i.e., $\textstyle d$ is large, we expect $\textstyle m_{\theta}$
and $\textstyle u_{\omega}$ to provide good approximations of $\textstyle m$
and $\textstyle u$ using much fewer parameters than the number of points in a
grid. Furthermore, for the numerical implementation, we restrict our attention
to a compact domain $\textstyle\tilde{\mathcal{Q}}$. We denote
$\textstyle\tilde{\mathcal{Q}}_{T}=[0,T]\times\tilde{\mathcal{Q}}$. We expect
the density to have a negligible mass outside a compact set so that by solving
the PDE system on a large enough compact set, we obtain a good approximation
of the solution, at least in the region where the density is significantly
positive. We then define the loss function:

${\mathcal{L}}(\theta,\omega)=\mathcal{L}^{({\mathrm{KFP}})}(m_{\theta},u_{\omega})+\mathcal{L}^{({\mathrm{HJB}})}(m_{\theta},u_{\omega}),$

where, for any
$\textstyle(m,u)\in{\mathcal{C}}^{1,2}(\tilde{\mathcal{Q}}_{T})\times{\mathcal{C}}^{1,2}(\tilde{\mathcal{Q}}_{T})$,
the two losses are as:

$\textstyle\displaystyle\mathcal{L}^{({\mathrm{KFP}})}(m,u)$
$\textstyle\displaystyle=C^{({\mathrm{KFP}})}\left\|\displaystyle\frac{\partial
m}{\partial t}-\nu\Delta m-\operatorname{div}\Bigl{(}m\partial_{p}H(m,\nabla
u)\Bigr{)}\right\|_{L^{2}(\tilde{\mathcal{Q}}_{T})}^{2}$
$\textstyle\displaystyle\qquad+C_{0}^{({\mathrm{KFP}})}\left\|m(0,\cdot)-m_{0}\right\|_{L^{2}(\tilde{\mathcal{Q}})}^{2}+C_{T}^{({\mathrm{KFP}})}\left\|m(T,\cdot)-m_{T}\right\|_{L^{2}(\tilde{\mathcal{Q}})}^{2},$
(13)

and

$\textstyle\displaystyle\mathcal{L}^{({\mathrm{HJB}})}(m,u)$
$\textstyle\displaystyle=C^{({\mathrm{HJB}})}\Big{\|}\displaystyle-\frac{\partial
u}{\partial t}-\nu\Delta u+H(m,\nabla
u)+\int_{\tilde{\mathcal{Q}}}\frac{\partial H}{\partial
m}(\zeta,m(t,\cdot),\nabla
u(t,\zeta))(\cdot)m(t,\zeta)\mathrm{d}\zeta\Big{\|}_{L^{2}(\tilde{\mathcal{Q}}_{T})}^{2}.$

Here, $\textstyle
C^{({\mathrm{KFP}})},C_{0}^{({\mathrm{KFP}})},C_{T}^{({\mathrm{KFP}})},C^{({\mathrm{HJB}})}$
are positive weights that give more or less importance to each component. If
the space domain is bounded, we must include more penalty terms. Note that any
smooth enough solution $\textstyle(m,u)$ to the PDE system (12) makes
$\textstyle\mathcal{L}^{({\mathrm{KFP}})}$ and
$\textstyle\mathcal{L}^{({\mathrm{HJB}})}$ vanish. The goal is to find two
neural networks which approximately minimize these losses.

Since it is not possible to compute exactly the above residuals, we
approximate the $\textstyle L^{2}$ norms using Monte Carlo samples. For
example, we rewrite:

$\textstyle\displaystyle\left\|\displaystyle\frac{\partial m}{\partial
t}-\nu\Delta
m-\operatorname{div}\Bigl{(}m\partial_{p}H(m,\operatorname{\nabla}u)\Bigr{)}\right\|_{L^{2}(\tilde{\mathcal{Q}}_{T})}^{2}$
$\textstyle\displaystyle=C(\tilde{\mathcal{Q}}_{T})\cdot\mathbb{E}_{\tau,\xi}\left[\left|\displaystyle\frac{\partial
m}{\partial t}(\tau,\xi)-\nu\Delta
m(\tau,\xi)-\operatorname{div}\Bigl{(}m(\tau,\xi)\partial_{p}H(m(\tau),\nabla
u(\tau,\xi))\Bigr{)}\right|^{2}\right],$

where $\textstyle(\tau,\xi)$ follows a uniform distribution over
$\textstyle\tilde{\mathcal{Q}}_{T}$, and $\textstyle
C(\tilde{\mathcal{Q}}_{T})$ is a normalizing constant that depends on the
domain. Likewise, for the other norms, it would also be possible to use
different norms and different distributions to sample $\textstyle(\tau,\xi)$.
But for the sake of simplicity, we will stick to this setting for the present
work. We obtain the following probabilistic formulation of the loss function
$\textstyle{\mathcal{L}}$:

${\mathcal{L}}(\theta,\omega)=\mathbb{E}_{S}\left[{\mathcal{L}}(\theta,\omega;S)\right],\qquad{\mathcal{L}}(\theta,\omega;S)=\mathcal{L}^{({\mathrm{KFP}})}(m_{\theta},u_{\omega};S)+\mathcal{L}^{({\mathrm{HJB}})}(m_{\theta},u_{\omega};S),$

where $\textstyle
S=(\tau,\xi,\xi_{0},\xi_{T})\in[0,T]\times\tilde{\mathcal{Q}}\times\tilde{\mathcal{Q}}\times\tilde{\mathcal{Q}}$
denotes one sample, and for any
$\textstyle(m,u)\in{\mathcal{C}}^{1,2}(\mathcal{Q}_{T})\times{\mathcal{C}}^{1,2}(\mathcal{Q}_{T})$,
the two losses at $\textstyle S$ are as:

$\textstyle\displaystyle\mathcal{L}^{({\mathrm{KFP}})}(m,u;S)$
$\textstyle\displaystyle=C^{({\mathrm{KFP}})}\left|\displaystyle\frac{\partial
m}{\partial t}(\tau,\xi)-\nu\Delta
m(\tau,\xi)-\operatorname{div}\Bigl{(}m(\tau,\xi)\partial_{p}H(m(\tau),\nabla
u(\tau,\xi))\Bigr{)}\right|^{2}$
$\textstyle\displaystyle\qquad+C_{0}^{({\mathrm{KFP}})}|m(0,\xi_{0})-m_{0}(\xi_{0})|^{2}+C_{T}^{({\mathrm{KFP}})}|m(T,\xi_{T})-m_{T}(\xi_{T})|^{2},$

and

$\textstyle\displaystyle\mathcal{L}^{({\mathrm{HJB}})}(m,u;S)$
$\textstyle\displaystyle=C^{({\mathrm{HJB}})}\Big{|}\displaystyle-\frac{\partial
u}{\partial t}(\tau,\xi)-\nu\Delta u(\tau,\xi)+H(\xi,m(\tau),\nabla
u(\tau,\xi))$
$\textstyle\displaystyle\qquad\qquad\qquad+\int_{\mathcal{Q}}\frac{\partial
H}{\partial m}(\zeta,m(t,\cdot),\nabla
u(t,\zeta))(\xi)m(t,\zeta)\mathrm{d}y\Big{|}^{2}.$

Finally, to optimize over $\textstyle(\theta,\omega)$, we use SGD (or one of
its variants) on the loss $\textstyle{\mathcal{L}}$. In practice, we use a
mini-batch of samples at each iteration, which amounts to approximate the
expectation by an empirical average over several samples.

###  3.3 Augmented Lagrangian Method with Deep Learning

In this subsection, we present an approach based on a primal-dual formulation
of the MFOT problem. We then introduce a deep learning adaptation of the
alternating direction method of multipliers. We focus on the case when the
interactions are local, and the drift is the control.

####  3.3.1 Primal and dual problems

Under suitable assumptions, the MFOT problem admits a variational formulation,
which can be tackled using a direct optimization approach. As in the previous
subsection, we assume that $\textstyle\rho_{0}$ and $\textstyle\rho_{T}$ have
respectively density $\textstyle m_{0}$ and $\textstyle m_{T}$.

We focus on a model with local interactions, meaning that an agent at state
$\textstyle x$ interacts with the density of the population at $\textstyle x$.
To alleviate the presentation, we will use the same notations for the costs
and the drift functions, but now their second input is a real number
$\textstyle m$ instead of an element
$\textstyle\mu\in{\mathcal{P}}_{2}(\mathcal{Q})$. So we have $\textstyle
f:\mathcal{Q}\times{{\mathbb{R}}}\times{{\mathbb{R}}}^{k}\to{{\mathbb{R}}}$,
$\textstyle g:\mathcal{Q}\times{{\mathbb{R}}}\to{{\mathbb{R}}}$, and
$\textstyle
b:\mathcal{Q}\times{{\mathbb{R}}}\times{{\mathbb{R}}}^{k}\to{{\mathbb{R}}}^{d}$.
We also modify accordingly the definition of the Hamiltonian $\textstyle H$ in
(9) and the Lagrangian $\textstyle L$ in (10) in subsection 3.2. We further
assume that $\textstyle f(x,m,v)$ is convex in $\textstyle v$ for every
$\textstyle(x,m)$, and $\textstyle mf(x,m,v)$ is convex in $\textstyle m$ for
every $\textstyle(x,v)$. For simplicity, we consider that $\textstyle
b(x,m,v)=v$, i.e., the drift is the control. We remark that the setting here
is not restrictive and satisfied by a large class of problems.

Primal problem.  The MFOT problem (3) introduced in Section 2 is formally
equivalent to the following PDE-constrained optimization problem:

$\textstyle\displaystyle\inf_{v:\mathcal{Q}_{T}\rightarrow{{\mathbb{R}}}^{k}}\,\int_{\mathcal{Q}_{T}}f\bigl{(}x,m(t,x),v(t,x)\bigr{)}m(x,t)\mathrm{d}x\,\mathrm{d}t$
subject to $\textstyle\displaystyle\frac{\partial m}{\partial
t}(t,x)-\nu\Delta
m(t,x)+\operatorname{div}\bigl{(}m(t,x)v(t,x)\bigr{)}=0\qquad
t\in(0,T],x\in\mathcal{Q}$ $\textstyle\displaystyle m(0,x)=m_{0}(x),\qquad
m(T,x)=m_{T}(x)$ (14)

The PDE constraint is the KFP equation corresponding to the stochastic
dynamics in (4). Note that the formulation in terms of $\textstyle(m,v)$,
while intuitive, is not convex in general. For this reason, we consider an
equivalent formulation in terms of $\textstyle(m,z)=(m,mv)$. We define:

$\textstyle\displaystyle\tilde{f}(x,m,z)=\begin{cases}mf\left(x,m,\frac{z}{m}\right)&\text{
if }m>0\\\ 0&\text{ if }(m,z)=(0,0)\\\ +\infty&\text{ otherwise }\end{cases}$
(15)

Note that $\textstyle(m,z)\mapsto\tilde{f}(x,m,z)$ is LSC on
$\textstyle{{\mathbb{R}}}\times{{\mathbb{R}}}^{k}$. Under suitable conditions,
it can be proved that $\textstyle(m,z)\mapsto\tilde{f}(x,m,z)$ is convex on
$\textstyle{{\mathbb{R}}}\times{{\mathbb{R}}}^{d}$. We also define the space
$\textstyle\mathbf{K}$,

$\textstyle\displaystyle\mathbf{K}=\Bigl{\\{}(m,z)\,\Big{|}\,\frac{\partial
m}{\partial t}(t,x)-\nu\Delta
m(t,x)+\operatorname{div}z(t,x)=0,m(0,x)=m_{0}(x),m(T,x)=m_{T}(x),m\geq
0\Bigr{\\}}$ (16)

With all these definitions, we are ready to present the primal problem:

$\textstyle\displaystyle\inf_{(m,z)\in\mathbf{K}}{\mathcal{B}}(m,z)=\inf_{(m,z)\in\mathbf{K}}\int_{\mathcal{Q}_{T}}\tilde{f}\bigl{(}x,m(t,x),z(t,x)\bigr{)}\mathrm{d}x\mathrm{d}t$
(17)

Assuming that problem (17) has a unique optimal solution
$\textstyle(m^{*},z^{*})$ and that problem (3.3.1) has a unique optimal
control $\textstyle v^{*}$, then the following connection holds: $\textstyle
v^{*}(t,x)=z^{*}(t,x)/m^{*}(t,x)$ if $\textstyle m^{*}(t,x)>0$, $\textstyle
v^{*}(t,x)=0$ if $\textstyle m^{*}(t,x)=0$.

Dual problem.  We now introduce a dual optimization problem. We define the
following functionals:

$\textstyle\displaystyle{\mathcal{A}}(u)=\inf_{m\geq
0}\int_{\mathcal{Q}_{T}}m(t,x)\Bigl{(}\frac{\partial u}{\partial
t}(t,x)+\nu\Delta u(t,x)-H\bigl{(}x,m(t,x),\nabla
u(t,x)\bigr{)}\Bigr{)}\mathrm{d}x\,\mathrm{d}t$ (18)
$\textstyle\displaystyle{\mathcal{F}}(u)=\int_{\mathcal{Q}}\left(m_{T}(x)u(T,x)-m_{0}(x)u(0,x)\right)\mathrm{d}x$
(19)
$\textstyle\displaystyle{\mathcal{G}}(\mathfrak{a},\mathfrak{b})=-\inf_{m\geq
0}\int_{\mathcal{Q}_{T}}m(t,x)\Bigl{(}\mathfrak{a}(t,x)-H\bigl{(}x,m(t,x),\mathfrak{b}(t,x)\bigr{)}\Bigr{)}\mathrm{d}x\,\mathrm{d}t.$
(20)

Note that if we define the linear differential operator $\textstyle\Lambda
u=\bigl{(}\frac{\partial u}{\partial t}+\nu\Delta u,\nabla u\bigr{)}$, then
$\textstyle{\mathcal{A}}(u)={\mathcal{G}}(\Lambda u)$. Consider the following
problem:

$\textstyle\displaystyle\inf_{u}{\mathcal{F}}(u)+{\mathcal{G}}(\Lambda u).$
(21)

Based on Fenchel-Rockafellar duality theorem (see Section 31, Theorem 31.1 in
[53]), we expect problems (17) and (21) to be in duality, meaning:

$\textstyle\displaystyle\inf_{(m,z)\in\mathbf{K}}{\mathcal{B}}(m,z)=\sup_{u}{\mathcal{A}}(u)=-\inf_{u}{\mathcal{F}}(u)+{\mathcal{G}}(\Lambda
u)$ (22)

Note that this primal-dual relationship also plays an important role in
demonstrating the uniqueness and existence of solutions to MFG and MFC PDE
systems, see e.g. [43, 20, 4]. Here, we expect a similar result to hold for
MFOT under suitable conditions. The rigorous definition of the two problems
and the analysis of this duality relationship is left for future work. For
now, we proceed formally.

We can at least formally establish a connection between the primal problem,
the dual problem, and the optimal control in the following way. Let
$\textstyle u^{*}$ be the optimal solution to the dual (21) and let
$\textstyle(m^{*},z^{*})$ be the optimal solution to the primal problem (17).
Then the optimal control for the original problem (3.3.1) is given by:

$v^{*}(t,x)=\partial_{p}H\bigl{(}x,m^{*}(t,x),\nabla u^{*}(t,x)\bigr{)}.$

We notice that $\textstyle(u^{*},m^{*})$ forms a solution to the MFOT PDE
system (12). This fact suggests that we can work on the dual problem (21)
directly to solve the MFOT problem. Under suitable assumptions, it can be
shown that the dual problem (21) is a strongly convex, unconstrained
optimization problem over function space, which motivates the use of classic
algorithms in convex optimization. However, the presence of the infinite
dimensional linear operator $\textstyle\Lambda$ makes the problem hard to
solve efficiently in general. Fortunately, the structure of the objective as a
sum of two convex functionals makes the problem amenable to algorithms based
on splitting schemes, such as the Alternating Direction Method of Multipliers
(ADMM) [17].

Data: Initial Guess $\textstyle\bigl{(}u^{(0)},q^{(0)},\lambda^{(0)}\bigr{)}$;
number of iterations $\textstyle N$; hyperparameter $\textstyle r>0$

Result: Function $\textstyle\bigl{(}u^{(N)},q^{(N)},\lambda^{(N)}\bigr{)}$
that are close to the saddle point of $\textstyle{\mathcal{L}}_{r}$ defined in
(24)

1 begin

2 for  _$\textstyle k=1,\cdots,N,$_ do

3 $\textstyle
u^{(k)}=\underset{u:\mathcal{Q}_{T}\rightarrow{{\mathbb{R}}}}{\mathop{\arg\min}}\;\mathcal{F}(u)-\langle\lambda^{(k-1)},\Lambda
u\rangle+\frac{r}{2}\bigl{\|}\Lambda u-q^{(k-1)}\bigr{\|}^{2}$

4 $\textstyle
q^{(k)}=\underset{q:\mathcal{Q}_{T}\rightarrow{{\mathbb{R}}}^{k+1}}{\mathop{\arg\min}}\;\mathcal{G}(q)+\langle\lambda^{(k-1)},q\rangle+\frac{r}{2}\bigl{\|}\Lambda
u^{(k)}-q\bigr{\|}^{2}$

5 $\textstyle\lambda^{(k)}=\lambda^{(k-1)}-r\bigl{(}\Lambda
u^{(k)}-q^{(k)}\bigr{)}$

6

7

Algorithm 1 Vanilla ADMM for MFOT

####  3.3.2 Description of the algorithm

Introducing a new variable $\textstyle q$ that will play the role of
$\textstyle\Lambda u$, we can rewrite problem (21) as the following
constrained optimization program:

$\textstyle\displaystyle\inf_{u,q:\,q=\Lambda
u}{\mathcal{F}}(u)+{\mathcal{G}}(q).$ (23)

The goal is now to find a saddle point of the associated Lagrangian. In fact,
for numerical purposes, we will consider an augmented Lagrangian, defined as
follows.

Let $\textstyle r>0$ be a constant and introduce
$\textstyle\lambda:\mathcal{Q}_{T}\rightarrow\mathbb{R}^{k+1}$, the Lagrangian
multiplier associated with the constraint $\textstyle q=\Lambda u$. Let
$\textstyle\langle\cdot,\cdot\rangle$ denote the inner product on $\textstyle
L^{2}(\mathcal{Q}_{T})$. We introduce the augmented Lagrangian:

$\textstyle\displaystyle\mathcal{L}_{r}(u,q,\lambda)=\mathcal{F}(u)+\mathcal{G}(q)-\langle\lambda,\Lambda
u-q\rangle+\frac{r}{2}\|\Lambda u-q\|^{2}$ (24)

Now, the original MFOT problem is reduced to finding a saddle point of
$\textstyle{\mathcal{L}}_{r}$. Here, we state the original ADMM method in
Algorithm 1 that finds the saddle point via an alternating optimization
procedure.

This general procedure can be implemented, for example, when the functions
$\textstyle(u,q,\lambda)$ are approximated by their values on a finite-
difference grid. Such a procedure has been used for MFG and MFC problems,
using finite elements [11, 9] or finite differences [5]. Furthermore, [11]
proved the convergence of this method under suitable conditions. However, as
already mentioned, approximating functions by their values on a mesh is not
feasible in high dimensions. We thus propose a different implementation of the
ADMM based on neural network approximations.

In Algorithm 1, the objectives in the steps are given by functionals to be
minimized over functional spaces, which is not tractable in general. We
restrict our attention to spaces of parameterized functions that can be
expressed as neural networks, denoted by
$\textstyle\bigl{(}u_{\theta},q_{\omega},\lambda_{\psi}\bigr{)}$ with
parameter $\textstyle\theta,\omega,\psi$ respectively. We then follow the
strategy introduced with the DGM [57] and already used in Section 3.2 to
create computable loss functions that are stochastic approximations of the
functionals.

Recall that the truncated space domain $\textstyle\tilde{\mathcal{Q}}$ and the
associated time-space domain $\textstyle\tilde{\mathcal{Q}}$. Let $\textstyle
X\sim\mathcal{U}(\tilde{\mathcal{Q}}_{T})$, $\textstyle
Y\sim\mathcal{U}(\tilde{\mathcal{Q}})$ be two random variables with uniform
distribution in the time-space domain and the space domain respectively. Let
$\textstyle\rho_{X}$, $\textstyle\rho_{Y}$ be the value of the uniform density
on $\textstyle\tilde{\mathcal{Q}}_{T}$ and $\textstyle\tilde{\mathcal{Q}}$
respectively. Here, we overload the notation
$\textstyle\langle\cdot,\cdot\rangle$ and $\textstyle\|\cdot\|$ to represent
the Euclidean inner product and norm on both $\textstyle
L_{2}(\tilde{\mathcal{Q}}_{T})$ and $\textstyle{{\mathbb{R}}}^{d+1}$:

$\textstyle{\mathcal{L}}^{(u)}(\theta;\omega,\psi)={\mathcal{L}}_{1}(u_{\theta},q_{\omega},\lambda_{\psi}),\quad{\mathcal{L}}^{(q)}(\omega;\theta,\psi)={\mathcal{L}}_{2}(u_{\theta},q_{\omega},\lambda_{\psi}),\quad{\mathcal{L}}^{(\lambda)}(\psi;\theta,\omega,\psi_{old})={\mathcal{L}}_{3}(u_{\theta},q_{\omega},\lambda_{\psi_{old}},\lambda_{\psi}),$

where

$\textstyle\displaystyle{\mathcal{L}}_{1}(u,q,\lambda)$
$\textstyle\displaystyle=\frac{1}{\rho_{Y}}\mathbb{E}_{Y}\left[u(T,Y)m_{T}(Y)-u(0,Y)m_{0}(Y)\right]+\frac{1}{\rho_{X}}\mathbb{E}_{X}\left[\frac{r}{2}\|\Lambda
u(X)-q(X)\|^{2}-\langle\Lambda u(X),\lambda(X)\rangle\right]$ (25)
$\textstyle\displaystyle{\mathcal{L}}_{2}(u,q,\lambda)$
$\textstyle\displaystyle=\frac{1}{\rho_{X}}{{\mathbb{E}}}_{X}\left[{\mathcal{G}}(q(X))+\langle\lambda(X),q(X)\rangle+\frac{r}{2}\|\Lambda
u(X)-q(X)\|^{2}\right]$
$\textstyle\displaystyle{\mathcal{L}}_{3}(u,q,\lambda_{old},\lambda)$
$\textstyle\displaystyle=\frac{1}{\rho_{X}}{{\mathbb{E}}}_{X}\Bigl{[}\|\lambda_{old}(X)-r\left(\Lambda
u(X)-q(X)\right)-\lambda(X)\|^{2}\Bigr{]}.$ (26)

Here, the subscript $\textstyle old$ is used to refer to the previous
iteration: the loss for $\textstyle\lambda$ involves the previous estimate
$\textstyle\lambda_{old}$. When using a neural network, it amounts to using
the previous neural network parameters $\textstyle\psi_{old}$. The loss
function aims at mimicking the effect of the direct update in the third step
of standard ADMM (Algorithm 1) when $\textstyle\lambda$ is approximated by a
neural network.

The algorithm DeepADMM is presented in Algorithm 2.

Data: Initial parameter $\textstyle\theta^{(0)},\omega^{(0)},\psi^{(0)}$;
number of ADMM iterations $\textstyle K$; SGD parameters

Result: Final parameter $\textstyle\theta^{(K)},\omega^{(K)},\psi^{(K)}$

1 begin

2 for  _$\textstyle k=1,\cdots,K$_ do

3 Compute $\textstyle\theta^{(k)}$ using SGD to (approximately) minimize the
loss $\textstyle{\mathcal{L}}^{(u)}(\cdot;\omega^{(k-1)},\psi^{(k-1)})$

4 Compute $\textstyle\omega^{(k)}$ using SGD to (approximately) minimize the
loss $\textstyle{\mathcal{L}}^{(q)}(\cdot;\theta^{(k)},\psi^{(k-1)})$

5 Compute $\textstyle\theta^{(k)}$ using SGD to (approximately) minimize the
loss
$\textstyle{\mathcal{L}}^{(\lambda)}(\cdot;\theta^{(k)},\omega^{(k)},\psi^{(k-1)})$

6

7

Algorithm 2 DeepADMM for MFOT

We have several remarks regarding DeepADMM and the augmented Lagrangian
formulation in order for readers to better understand this approach. First,
compared with Algorithm 1, the updates in Algorithm 2 for functions
$\textstyle u$ and $\textstyle q$ are quite straightforward to understand.
Instead of searching optimizer over function space, we reduce the problem to a
finite dimension through stochastic approximation of the objective and search
in the parameter space instead. The computed stochastic gradient can be
considered as an unbiased estimator of the population gradient with respect to
the functional, and the variance of this stochastic gradient decreases as the
batch size increases.

In Appendix B, we discuss the computation of $\textstyle{\mathcal{G}}$ for
several typical models.

##  4 Numerical experiments

In this section, we present numerical experiments obtained with the three
methods discussed in the previous section. For brevity, we refer to the three
methods respectively introduced in sections 3.1, 3.2 and 3.3 as Method 1,
Method 2 and Method 3 (and M1, M2, and M3 for short in the plots).

We first consider two test cases for which we have explicit solutions (up to
solving ODE systems) and can thus be used to benchmark our algorithms in any
dimension. We then consider two test cases that can be viewed as modifications
of standard OT with crowd aversion or congestion effects.

###  4.1 Case 1: Linear Quadratic Problem

The first class of models that we consider has a linear-quadratic structure,
which falls in the setting discussed in Example 2.

####  4.1.1 Description of the problem

In this model, we take:

$\textstyle\displaystyle b(x,\mu,a)=Ax+Ba,\qquad
f(x,\mu,a)=a^{\top}Ra,\qquad\rho_{0}={\mathcal{N}}(\bar{x}_{0};\Sigma_{0}),\qquad\rho_{T}={\mathcal{N}}(\bar{x}_{T};\Sigma_{T}),$

where $\textstyle A,B,R,\Sigma_{0},\Sigma_{T}$ are (constant) matrices of
suitable sizes. The vectors $\textstyle\bar{x}_{0}$ and
$\textstyle\bar{x}_{T}$ correspond to the initial and terminal means. We will
consider two settings. In order to have a benchmark solution, we will take
$\textstyle\sigma=B$. This enables us to use the solution provided by [29,
Section 7.1], which boils down to solving a system of ODEs. For the sake of
completeness, we provide the details in Appendix A.

####  4.1.2 Evaluation Metrics

In this model, since we have access to the optimal solution, we can evaluate
the learnt solutions given by the three methods we proposed with respect to
the ground-truth solution. We denote by $\textstyle v^{*}$ the optimal control
and $\textstyle\hat{v}$ a learnt control. As explained below in detail, we use
the following metrics: the total cost (namely $\textstyle J^{MFOT}(\hat{v})$,
with $\textstyle J^{MFOT}$ introduced in (3)), the relative error between the
achieved cost, and the optimal cost (namely $\textstyle J^{MFOT}(\hat{v})$ and
$\textstyle J^{MFOT}(v^{*})$), the deviation from the terminal distribution
(i.e., the Wasserstein distance between the achieved terminal distribution and
the target terminal distribution, $\textstyle\rho_{T}$), and the weighted
$\textstyle L^{2}$ error between the learnt control $\textstyle\hat{v}$ and
the optimal control $\textstyle v^{*}$, weighted by the population
distribution.

Computation of the control. The control is parameterized in different ways
across different methods. For Method 1,
$\textstyle\hat{v}(t,x)=v_{\theta}(t,x)$. For Method 2 and Method 3,
$\textstyle\hat{v}(t,x)=-\frac{1}{2}BR^{-1}\nabla\hat{u}_{\theta}(t,x)$, where
$\textstyle\hat{u}_{\theta}$ is the neural network that approximates the dual
variable, solution to the HJB equation.

Total cost. Recall the definition of the objective $\textstyle J^{MFOT}$
defined in (1). Let $\textstyle v$ be a control. In the present Linear-
Quadratic case, we have that,

$\textstyle\displaystyle
J^{MFOT}(v)=\int_{0}^{T}\int_{{\mathcal{Q}}}m^{v}(t,x)f(x,m^{v},v)\mathrm{d}x\,\mathrm{d}t=\int_{0}^{T}\int_{{\mathcal{Q}}}m^{v}(t,x)v(t,x)^{\top}Rv(t,x)\mathrm{d}x\,\mathrm{d}t,$

where $\textstyle m^{v}$ is the density of mean field distribution driven by
$\textstyle v$, which satisfies the KFP PDE (12). In order to evaluate
$\textstyle J^{MFOT}(v)$, we use Monte Carlo simulations. We discretize the
time variable $\textstyle t$. Let $\textstyle N_{T}$ be a number of time steps
of length $\textstyle\Delta t=T/N_{T}$. We consider a equi-distanced time
discretization with time-steps $\textstyle\\{t_{0}=0,t_{1}=\Delta
t,\dots,t_{N_{T}}=N_{T}\Delta t\\}$. Again, we simulate solutions to the
underlying SDE using an Euler-Maruyama scheme similar to the one used in (7).
We simulate a family of $\textstyle N$ sequences
$\textstyle((X_{t_{n}}^{i,v})_{n=0,\dots,N_{T}})_{i=1,\cdots,N})$ using the
following update

$\begin{cases}X^{i}_{0}\sim\rho_{0}\quad\hbox{ i.i.d.}\\\
X^{i,v}_{t_{n+1}}=X^{i,v}_{t_{n}}+(AX^{i,v}_{t_{n}}+Bv(t_{n},X^{i,v}_{t_{n}}))\Delta
t+\sigma\sqrt{\Delta t}\Delta W^{i}_{n},\end{cases}$ (27)

where $\textstyle\\{\Delta W^{i}_{n}\\}$ are i.i.d standard Gaussian random
variables in $\textstyle\mathbb{R}^{d}$. With these sampled sequences, we
compute the objective as

$\textstyle\displaystyle
J^{MFOT}(v)=\dfrac{1}{N}\sum_{i=1}^{N}\sum_{n=0}^{N_{T}-1}v(t_{n},X_{t_{n}}^{i,v})^{\top}Rv(t_{n},X_{t_{n}}^{i,v})\Delta
t.$

Relative Error. The relative error between $\textstyle J^{MFOT}(\hat{v})$ and
$\textstyle J^{MFOT}(v^{*})$ is defined as

$\textstyle\frac{|J^{MFOT}(\hat{v})-J^{MFOT}(v^{*})|}{|J^{MFOT}(v^{*})|}.$

The major reason to consider relative error instead of absolute error is
because the scale of the running cost varies greatly across different
problems. Also, we want to stress that even though $\textstyle v^{*}$ is the
analytical optimal solution, it may happen that $\textstyle
J^{MFOT}(\hat{v})<J^{MFOT}(v^{*})$ if $\textstyle\hat{v}$ does not satisfy
exactly the constraint (in contrast with $\textstyle v^{*}$).

Expected $\textstyle L^{2}$ error for control. The expected $\textstyle L^{2}$
error between the learnt control $\textstyle\hat{v}$ and the ground-truth
control $\textstyle v^{*}$ is defined as,

$\textstyle\displaystyle
d_{L^{2}}(\hat{v},v^{*})=\int_{0}^{T}\int_{\mathcal{Q}}m^{*}(t,x)\|\hat{v}(t,x)-v^{*}(t,x)\|^{2}\mathrm{d}x\,\mathrm{d}t$

where $\textstyle m^{*}$ is the density of the optimal mean-field associated
with $\textstyle v^{*}$. For the LQ problem, the optimal mean field
$\textstyle m^{*}$ is Gaussian for any $\textstyle t\in[0,T]$, with mean
$\textstyle\mu_{t}$ and variance $\textstyle\Sigma_{t}$ given by analytical
formulas in Appendix A. We can thus evaluate the $\textstyle L^{2}$ error
again with Monte Carlo samples for each time step. As above, we discretize the
time variable $\textstyle t$ with $\textstyle N_{T}+1$ points and for
$\textstyle t_{n}=n\Delta t,n=0,\dots,N_{T}$, we generate i.i.d. samples
$\textstyle(X^{i,v}_{t_{n}})_{i=1,\dots,N}\sim\mathcal{N}(\mu_{t_{n}},\Sigma_{t_{n}})$.
We then estimate the $\textstyle L^{2}$ error as:

$\textstyle\displaystyle
d_{L^{2}}(\hat{v},v^{*})\approx\dfrac{1}{N}\sum_{i=1}^{N}\sum_{n=0}^{N_{T}-1}\|\hat{v}(t_{n},X_{t_{n}}^{i,v})-v^{*}(t_{n},X_{t_{n}}^{i,v})\|^{2}\Delta
t.$

Deviation of distribution. The deviation of the mean field
$\textstyle\hat{\rho}_{T}$ from the terminal target distribution
$\textstyle\rho_{T}$ is quantified by two different metrics: Wasserstein-2
distance $\textstyle{\mathcal{W}}_{2}(\hat{\rho}_{T},\rho_{T})$ and
$\textstyle L^{2}$ distance $\textstyle d_{L^{2}}(\hat{m}_{T},m_{T})$. Here,
$\textstyle\hat{\rho}_{T}$ is the measure of the mean field distribution
driven by the learnt control $\textstyle\hat{v}$ at time $\textstyle T$.
$\textstyle\hat{m}_{T}$ and $\textstyle m_{T}$ are the density of
$\textstyle\hat{\rho}_{T}$ and $\textstyle\rho_{T}$ respectively.

  * •

Wasserstein-2 distance. We adopt a similar method to compute the Wasserstein-2
distance as the one discussed in Section 3.1. We simulate $\textstyle N$
particles following the dynamics (27), and obtain a collection of $\textstyle
N$ samples $\textstyle(X_{T}^{i,\hat{v}})_{i=1,\dots,N}$, which forms an
empirical distribution approximating $\textstyle\hat{\rho}_{T}$. We also
generate $\textstyle N$ samples directly from the target distribution
$\textstyle\rho_{T}$, denoted by $\textstyle(Y^{i}_{T})_{i=1,\dots,N}$. Then,
we define the distance matrix $\textstyle{\mathcal{M}}$ by
$\textstyle{\mathcal{M}}_{ij}=|X_{T}^{i,\hat{v}}-Y_{T}^{j}|^{2}$, and we
recall that the set $\textstyle U_{N}$ is defined by (8). We approximate the
Wasserstein-2 distance between the two empirical distributions formed by
$\textstyle X_{T}^{i,\hat{v}}$ and $\textstyle Y_{T}^{i}$ through the
following linear program:

$\textstyle\displaystyle\mathcal{W}_{2}\left(\hat{\rho}_{T},\rho_{T}\right)\approx\Bigl{(}\min\limits_{{\mathcal{A}}\in
U_{N}}\langle{\mathcal{A}},M\rangle\Bigr{)}^{1/2}.$

  * •

$\textstyle L^{2}$ distance. We will also use as a metric the $\textstyle
L^{2}$ distance between $\textstyle\hat{m}_{T}$ and $\textstyle m_{T}$ on the
truncated domain $\textstyle\tilde{\mathcal{Q}}$. To evaluate the integral, in
the absence of analytical formula for $\textstyle m_{T}$, we again use a Monte
Carlo approach. We uniformly sample $\textstyle N$ points in
$\textstyle\tilde{\mathcal{Q}}$ denoted by $\textstyle(X_{j})_{j=1,\dots,N}$.
Let $\textstyle C(\tilde{\mathcal{Q}})$ denotes the inverse of the value of
the uniform density. Then we can approximate the $\textstyle L^{2}$ distance
by:

$\textstyle\displaystyle\frac{C(\tilde{\mathcal{Q}})}{N}\sum_{i=1}^{N}\|\hat{m}_{T}(X_{i})-m_{T}(X_{i})\|^{2}.$

####  4.1.3 Numerical results

In the numerical tests, we take the values given in Table 1 for the parameters
of the model, with the time horizon $\textstyle T=1.0$.

Test | $\textstyle d$ | $\textstyle A$ | $\textstyle B$ | $\textstyle\sigma$ | $\textstyle R$ | $\textstyle\bar{x}_{0}$ | $\textstyle\Sigma_{0}$ | $\textstyle\bar{x}_{T}$ | $\textstyle\Sigma_{T}$
---|---|---|---|---|---|---|---|---|---
| LQ Test 1
---
$\textstyle 1$ | $\textstyle 1$ | $\textstyle 1$ | $\textstyle 1$ | $\textstyle\frac{1}{2}$ | $\textstyle 0.0$ | $\textstyle 1$ | $\textstyle 2.0$ | $\textstyle 0.5$
| LQ Test 2
---
$\textstyle 2$ | $\textstyle I_{d}$ | $\textstyle I_{d}$ | $\textstyle I_{d}$ | $\textstyle\frac{1}{2}I_{d}$ | $\textstyle[0.0,0.0]$ | $\textstyle I_{d}$ | $\textstyle[2.0,2.0]$ | $\textstyle\frac{1}{2}I_{d}$
Table 1: Parameters for the two linear-quadratic test cases

For LQ test 1, in dimension 1, Figure 1 displays the evolution of the density.
For methods 1, 2, and 3, we obtain the control learnt using the neural
networks and simulate $\textstyle N$ trajectories by the Monte Carlo method
following the dynamics (27), with $\textstyle v$ replaced by the learnt
control. We then estimate the mean field distribution using kernel density
estimation (KDE). We see that the distributions obtained with the three
methods match well the ground-truth one obtained with ODEs. The distributions
move towards the right and concentrate around the final mean. Figure 2 shows
the evolution of the control. We see that the three methods provide good
approximations of the true optimal control, at least in the region where the
density is high. In regions where the density is very low, the control is not
well approximated, but this is not an issue as far as the optimal behavior of
the population is concerned. The first part of Table 2 shows the results
obtained for the metrics introduced above. We see that each of the three
methods achieves a smaller total cost than the true optimal control. This
would not be possible for controls satisfying perfectly the terminal
constraint, but it is possible here due to the fact that the methods satisfy
only approximately the planning constraint. The optimal control is well
approximated, as shown by the $\textstyle L^{2}$ distance to the true optimal
control. Furthermore, we see that Methods 2 and 3 have a higher Wasserstein-2
distance between the terminal distribution and the target distribution, but
the $\textstyle L^{2}$ distance is much lower.

As for LQ test 2, in dimension 2, Figure 3 displays the evolution of the
density for each of the methods. The densities move from the bottom left
corner to the top right corner. Furthermore, the terminal distribution is more
concentrated because the terminal variance is smaller than the initial
variance. Figure 4 shows the evolution of the first dimension of the control
(the second dimension is similar, so we omit it for brevity). The ground-truth
control is linear in space for each time step. We see that the three methods
manage to learn approximately linear controls, at least in the region where
the density is significantly positive. Table 3 shows the results obtained for
the metrics introduced above. We see that here again, each of the three
methods achieves a smaller total cost than the true optimal control due to the
fact that the terminal constraint is not perfectly satisfied. The optimal
control is well approximated, and the terminal distribution is matched with
good accuracy.

Figure 1: Evolution of the density in the LQ Test case 1. Each plot corresponds to one time step and displays the densities as functions of the space variable, $\textstyle x$. The densities are: The density obtained by applying the control learnt by each of the three deep learning methods as well as the ground-truth density given by the ODE method. Figure 2: Evolution of the control in the LQ Test case 1. Each plot corresponds to one time step and displays the controls as functions of the space variable, $\textstyle x$. The controls are: The control learnt by each of the three deep learning methods as well as the ground-truth control given by the ODE method. Test case | Method | Total Cost | Relative Error | $\textstyle d_{L^{2}}(\hat{v},v^{*})$ | $\textstyle\mathcal{W}_{2}(\hat{\rho}_{T},\rho_{T})$ | $\textstyle d_{L^{2}}(\hat{m}_{T},m_{T})$
---|---|---|---|---|---|---
Linear Quadratic LQ Test 1 |  ODE $\textstyle(v^{*})$ | $\textstyle 2.126$ |  - |  - |  - |  -
M1 | $\textstyle 2.099$ | $\textstyle 1.24\%$ | $\textstyle 0.021$ | $\textstyle 0.002$ | $\textstyle 0.006$
M2 | $\textstyle 2.096$ | $\textstyle 1.41\%$ | $\textstyle 0.003$ | $\textstyle 0.043$ | $\textstyle 0.00004$
M3 | $\textstyle 2.077$ | $\textstyle 2.29\%$ | $\textstyle 0.011$ | $\textstyle 0.031$ | $\textstyle 0.001$
Table 2: Comparison of three different methods v.s. the analytical solution on the LQ test case 1. The evaluation metrics are described in Section 4.1.2. Figure 3: Evolution of the density in the Linear Quadratic Test case 2. Each column corresponds to one time step, and each row corresponds to one of the methods. Each plot displays the density as a function of the space variable, i.e., $\textstyle(m(t,x)_{1})_{x\in[-4,6]^{2}}$. The first row corresponds to the solution obtained by the ground-truth ODE method. The second, third and fourth rows correspond respectively to methods 1, 2 and 3. Figure 4: Evolution of the control in the Linear Quadratic Test case 2. Each column corresponds to one time step, and each row corresponds to one of the methods. Each plot displays the first dimension of the control as a function of the space variable, i.e., $\textstyle(v(t,x)_{1})_{x\in[-4,6]^{2}}$. The first row corresponds to the solution obtained by the ground-truth ODE method. The second, third and fourth rows correspond respectively to methods 1, 2 and 3. Test case | Method | Total Cost | Relative Error | $\textstyle d_{L^{2}}(\hat{v},v^{*})$ | $\textstyle\mathcal{W}_{2}(\hat{\rho}_{T},\rho_{T})$ | $\textstyle d_{L^{2}}(\hat{m}_{T},m_{T})$
---|---|---|---|---|---|---
Linear Quadratic LQ Test 2 |  ODE $\textstyle(v^{*})$ | $\textstyle 4.175$ |  - |  - |  - |  -
M1 | $\textstyle 4.117$ | $\textstyle 1.39\%$ | $\textstyle 0.074$ | $\textstyle 0.043$ | $\textstyle 0.007$
M2 | $\textstyle 3.935$ | $\textstyle 5.63\%$ | $\textstyle 0.043$ | $\textstyle 0.403$ | $\textstyle 0.00005$
M3 | $\textstyle 4.054$ | $\textstyle 2.89\%$ | $\textstyle 0.131$ | $\textstyle 0.561$ | $\textstyle 0.015$
Table 3: Comparison of three different methods v.s. the analytical solution on
the LQ test case 2. The evaluation metrics are described in Section 4.1.2.

###  4.2 Case 2: Transport with congestion effects

The second class of models that we consider is inspired by crowd motion and
falls in the setting discussed in Example 3.

####  4.2.1 Description of the problem

In this model, intuitively, the cost is higher when moving through a crowded
region, i.e., where the density is high. Specifically, we take:

$\textstyle\displaystyle b(x,\mu,a)=a,\qquad
f(x,\mu,a)=R|\ell(x,\mu)|^{\gamma}|a|^{2},\qquad\rho_{0}={\mathcal{N}}(\bar{x}_{0};\Sigma_{0}),\qquad\rho_{T}={\mathcal{N}}(\bar{x}_{T};\Sigma_{T}).$

For the function $\textstyle\ell$, we take two different models. We consider
the following non-local dependence:

$\textstyle\ell(x,\mu)=c+\rho_{\epsilon}\star\mu(x),$

where $\textstyle c>0$ is a constant, $\textstyle\rho_{\epsilon}$ is a
Gaussian kernel and $\textstyle\star$ denotes the convolution. We use Method 1
to solve the MFOT problem with this function $\textstyle\ell$. Since it is
based on Monte Carlo simulations of trajectories, it is straightforward to
compute a convolution with the empirical distribution at a given time step.

We also consider a variation with a local dependence.

$\textstyle\ell(x,\mu)=c+m(x)$

where $\textstyle c>0$ is a constant and $\textstyle m$ denotes the density of
$\textstyle\mu$. For this type of model, Methods 2 and 3 are better suited
since, in these methods, we directly have access to the approximate density in
the form of a neural network.

####  4.2.2 Numerical results

We focus on one test case called "Congestion" below in dimension $\textstyle
d=1$. In this model, $\textstyle\gamma=1$. For the sake of comparison, we also
consider the corresponding model with the same choice of parameters except
that $\textstyle\gamma=0$, i.e., there are no congestion effects in the
running cost. The values that we take in the numerical tests are given in
Table 4 below.

Test case | $\textstyle d$ | $\textstyle\gamma$ | $\textstyle c$ | $\textstyle\sigma$ | $\textstyle R$ | $\textstyle\bar{x}_{0}$ | $\textstyle\Sigma_{0}$ | $\textstyle\bar{x}_{T}$ | $\textstyle\Sigma_{T}$
---|---|---|---|---|---|---|---|---|---
Case 1, No congestion | $\textstyle 1$ | $\textstyle 0$ | $\textstyle 0.1$ | $\textstyle 0.1$ | $\textstyle 0.5$ | $\textstyle 0$ | $\textstyle 0.04$ | $\textstyle 2$ | $\textstyle 0.04$
Case 2, Congestion | $\textstyle 1$ | $\textstyle 1$ | $\textstyle 0.1$ | $\textstyle 0.1$ | $\textstyle 0.5$ | $\textstyle 0$ | $\textstyle 0.04$ | $\textstyle 2$ | $\textstyle 0.04$
Case 3, Congestion | $\textstyle 5$ | $\textstyle 1$ | $\textstyle 1$ | $\textstyle 1$ | $\textstyle 0.5$ | $\textstyle 0$ | $\textstyle 0.1$ | $\textstyle 2$ | $\textstyle 0.1$
Table 4: Parameters for the test case with congestion and the benchmark model
without congestion effects

In Figure 5, we present the evolution of the density under the control learnt
by each of the three methods for congestion cases 1 and 2. Each row
corresponds to one method. We see that, in the case where $\textstyle\gamma=0$
(no congestion effect), the mass is transported directly towards the terminal
distribution without much change in its shape. In contrast, in the case with
$\textstyle\gamma=1$, the mass spreads in space and one part starts moving
towards the target mean $\textstyle\bar{x}_{T}=2$ whereas another part stays
behind and catches up at later time steps. This is consistent with the idea
that moving in congested regions is more expensive, so some agents would agree
to wait until the density decreases before moving forward.

Finally, in Figure 6, we present the evolution of the density under the
control learnt by each of the three methods for congestion case 3, which is in
dimension 5. Each row corresponds to one method. To visualize density
evolution in dimension 5, we plot the marginal distribution of the mean field
distribution on the first and second dimensions. We see that, similarly to the
congestion case 2, the mass spreads in space and gradually moves towards the
target mean. Compared with congestion case 2, the difference in the moving
pattern and extent of spreading is due to the difference of parameters in
Table 4. With a larger value $\textstyle c$, the behavior of the density would
be closer to a direct transport to the terminal distribution without changes
in the shape of the distribution.

Figure 5: Visualization of the mean field density $\textstyle\hat{m}(t,x)$ in
Congestion test Case 1,2 Figure 6: Visualization of the mean field density
$\textstyle\hat{m}(t,x)$ in Congestion test Case 3

###  4.3 Remarks on the choice of hyperparameters

Each method has several hyperparameters, including the architecture of the
neural networks. We provide below some remarks about the choice of
hyperparameters in our implementation.

Method 1.  In our implementation, we choose $\textstyle G(r)=C_{W}r$ where
$\textstyle C_{W}$ is a hyperparameter that we adjust dynamically. We increase
the constant $\textstyle C_{W}$ when we expect a higher running cost (for
instance, in a higher dimension) in order to give enough importance to the
penalty. The coefficient $\textstyle\alpha$ of regularization for the
computation of the Wasserstein distance is also a hyperparameter that we
adjust dynamically using the following heuristics. We start with a given value
for $\textstyle\alpha$ and, when the estimated Wasserstein distance is small
enough, we reduce the value of $\textstyle\alpha$. The idea is that, as long
as the terminal distribution does not match well enough the target
distribution, we need a high level of regularization in order to estimate
efficiently the Wasserstein distance between them. As the two distributions
get closer, we can decrease the degree of regularization in order to have a
more accurate estimation of the Wasserstein distance. The way we adjust
$\textstyle C_{W}$ also depends on the dimension of the state variable. There
is also a computational time aspect to take into account: as
$\textstyle\alpha$ becomes smaller, the computations take more time (see 3.1.2
for more details). For the neural network, we take a feedforward fully
connected neural network with 6 layers of 60 neurons each. The other
hyperparameters are the number of particles $\textstyle N$ and the number of
time steps $\textstyle N_{T}$. We take $\textstyle N=300$ and $\textstyle
N_{T}=20$.

Method 2.  In the second method, no time or space discretization is needed,
and the density is directly approximated by a neural network, so we do not
need to use a finite number of particles. However, we need to choose the
values of the weights $\textstyle C_{0}^{({\mathrm{KFP}})},$ $\textstyle
C_{T}^{({\mathrm{KFP}})},$ $\textstyle C^{({\mathrm{KFP}})},$ and $\textstyle
C^{({\mathrm{HJB}})}$ in the loss function. We used $\textstyle
C_{0}^{({\mathrm{KFP}})}=20,$ $\textstyle C_{T}^{({\mathrm{KFP}})}=50,$
$\textstyle C^{({\mathrm{KFP}})}=20,$ $\textstyle C^{({\mathrm{HJB}})}=1$. As
for the neural network, we used the architecture proposed in the DGM article
[57], with 2 layers and a width equal to 40. During the training, at each
iteration of SGD, we use a minibatch of 500 points in time and space, and 500
points in space for the initial and terminal conditions.

Method 3.  The main hyperparameter in this method is $\textstyle r$, which is
used in the definition of the augmented Lagrangian (24). For the experiments,
we select $\textstyle r=0.1$. Even though, in theory, the convergence of ADMM
is independent of the choice of $\textstyle r$, in practice, we often find
that a large $\textstyle r$ value could potentially increase numerical
instability and lead the algorithm to diverge. Similarly, a small $\textstyle
r$ value could slow down the convergence. As for the neural networks, we use
the following architectures. For both $\textstyle u_{\theta}$, $\textstyle
q_{\omega}$, and $\textstyle\lambda_{\psi}$, in general, we use a fully
connected neural network with residual connections, sigmoid activation
function, and appropriate output dimension. We use $\textstyle 6$ layers and
$\textstyle 100$ neurons per layer. For LQ test cases, we further consider an
extra quadratic correction in addition to the neural networks: the output of
$\textstyle u_{\theta}$ is the sum of neural network output and a quadratic
function with trainable weights. To effectively model the mean field density,
a sigmoid activation function is applied to the first dimension of the output
of the neural network $\textstyle\lambda_{\psi}$, and then the result is
multiplied by a constant $\textstyle C$. In this way, the first dimension of
$\textstyle\lambda_{\psi}$ takes values in $\textstyle(0,C)$. In the
experiments, we take $\textstyle C=1$ for the LQ test cases and $\textstyle
C=5$ for the congestion test cases. During training, at each iteration of SGD,
we use a minibatch of 512 points in time and space, and 512 points in space
for the initial and terminal conditions.

##  5 Conclusion and future directions

In this work, we have proposed three numerical methods based on deep learning
for mean field optimal transport problems. The three methods can tackle a
larger class of problems than deep learning methods proposed previously, which
were mostly focusing on the Schrödinger bridge problem or MFGs with a specific
structure. The first method replaces the terminal constraint with a penalty
and then directly learns the optimal control using Monte Carlo trajectories.
The second method solves a PDE system which is obtained as the optimality
conditions for the MFOT problem. The third method relies on an augmented
Lagrangian approach for the variational formulation of the problem. The
numerical results show that the three methods match the analytical solution on
an LQ problem, and that they are able to handle non-trivial mean field
interactions modeling congestion effects.

From here, we can envision several research directions. First of all, the
theoretical analysis of the MFOT problem remains to be tackled. For example,
the existence and uniqueness of the solution to the PDE system have been
proved only in relatively specific cases, see e.g. [1, 19, 50, 34]. It would
be interesting to extend the analysis to more general forms of dynamics and
cost functions. From the numerical point of view, it would be interesting to
scale-up the methods proposed in this work to a higher dimension, and to
explore other deep learning methods. The numerical analysis and the
convergence proof of the proposed methods also remain to be investigated in
future work.

## References

  * [1] Yves Achdou, Fabio Camilli, and Italo Capuzzo-Dolcetta, _Mean field games: numerical methods for the planning problem_ , SIAM Journal on Control and Optimization 50 (2012), no. 1, 77–109.
  * [2] Yves Achdou and Jean-Michel Lasry, _Mean field games for modeling crowd motion_ , Contributions to partial differential equations and applications (2019), 17–42.
  * [3] Yves Achdou and Mathieu Laurière, _On the system of partial differential equations arising in mean field type control_ , Discrete and Continuous Dynamical Systems 35 (2015), no. 9, 3879–3900.
  * [4] , _Mean field type control with congestion_ , Applied Mathematics & Optimization 73 (2016), no. 3, 393–418.
  * [5] , _Mean field type control with congestion (II): An augmented Lagrangian method_ , Applied Mathematics & Optimization 74 (2016), no. 3, 535–578.
  * [6] Yves Achdou and Alessio Porretta, _Mean field games with congestion_ , Annales de l’Institut Henri Poincaré C, Analyse non linéaire, vol. 35, Elsevier, 2018, pp. 443–480.
  * [7] Ali Al-Aradi, Adolfo Correia, Danilo Naiff, Gabriel Jardim, and Yuri Saporito, _Solving nonlinear and high-dimensional partial differential equations via deep learning_ , arXiv preprint arXiv:1811.08782 (2018).
  * [8] Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré, _Gradient flows: in metric spaces and in the space of probability measures_ , Springer Science & Business Media, 2005.
  * [9] Roman Andreev, _Preconditioning the augmented Lagrangian method for instationary mean field games with diffusion_ , SIAM Journal on Scientific Computing 39 (2017), no. 6, A2763–A2783.
  * [10] Jean-David Benamou and Yann Brenier, _A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem_ , Numerische Mathematik 84 (2000), no. 3, 375–393.
  * [11] Jean-David Benamou and Guillaume Carlier, _Augmented Lagrangian methods for transport optimization, mean field games and degenerate elliptic equations_ , Journal of Optimization Theory and Applications 167 (2015), no. 1, 1–26.
  * [12] Jean-David Benamou, Guillaume Carlier, Simone Di Marino, and Luca Nenna, _An entropy minimization approach to second-order variational mean-field games_ , Mathematical Models and Methods in Applied Sciences 29 (2019), no. 08, 1553–1583.
  * [13] Jean-David Benamou, Guillaume Carlier, and Maxime Laborde, _An augmented Lagrangian approach to Wasserstein gradient flows and applications_ , ESAIM: Proceedings and surveys 54 (2016), 1–17.
  * [14] Alain Bensoussan, Jens Frehse, Phillip Yam, et al., _Mean field games and mean field type control theory_ , vol. 101, Springer, 2013.
  * [15] Alain Bensoussan, Jens Frehse, and Sheung Chi Phillip Yam, _The master equation in mean field theory_ , Journal de Mathématiques Pures et Appliquées 103 (2015), no. 6, 1441–1474.
  * [16] Charles Bertucci, Jean-Michel Lasry, and Pierre-Louis Lions, _Master equation for the finite state space planning problem_ , Archive for Rational Mechanics and Analysis 242 (2021), no. 1, 327–342.
  * [17] Stephen Boyd, Neal Parikh, Eric Chu, Borja Peleato, Jonathan Eckstein, et al., _Distributed optimization and statistical learning via the alternating direction method of multipliers_ , Foundations and Trends® in Machine learning 3 (2011), no. 1, 1–122.
  * [18] Haoyang Cao, Xin Guo, and Mathieu Laurière, _Connecting GANs, MFGs, and OT_ , arXiv preprint arXiv:2002.04112 (2020).
  * [19] Pierre Cardaliaguet, Guillaume Carlier, and Bruno Nazaret, _Geodesics for a class of distances in the space of probability measures_ , Calculus of Variations and Partial Differential Equations 48 (2013), no. 3, 395–420.
  * [20] Pierre Cardaliaguet and P Jameson Graber, _Mean field games systems of first order_ , ESAIM: Control, Optimisation and Calculus of Variations 21 (2015), no. 3, 690–722.
  * [21] René Carmona, François Delarue, et al., _Probabilistic theory of mean field games with applications I-II_ , Springer, 2018.
  * [22] René Carmona and Mathieu Laurière, _Convergence analysis of machine learning algorithms for the numerical solution of mean field control and games: II: The finite horizon case_ , arXiv preprint arXiv:1908.01613 (2019).
  * [23] , _Convergence analysis of machine learning algorithms for the numerical solution of mean field control and games I: the ergodic case_ , SIAM Journal on Numerical Analysis 59 (2021), no. 3, 1455–1485.
  * [24] René Carmona and Mathieu Laurière, _Deep learning for mean field games and mean field control with applications to finance_ , 2021.
  * [25] Yongxin Chen, Tryphon T Georgiou, and Michele Pavon, _Optimal steering of a linear stochastic system to a final probability distribution, part I_ , IEEE Transactions on Automatic Control 61 (2015), no. 5, 1158–1169.
  * [26] , _Optimal steering of a linear stochastic system to a final probability distribution, part II_ , IEEE Transactions on Automatic Control 61 (2015), no. 5, 1170–1180.
  * [27] , _On the relation between optimal transport and schrödinger bridges: A stochastic control viewpoint_ , Journal of Optimization Theory and Applications 169 (2016), no. 2, 671–691.
  * [28] , _Optimal steering of a linear stochastic system to a final probability distribution—part III_ , IEEE Transactions on Automatic Control 63 (2018), no. 9, 3112–3118.
  * [29] , _Steering the distribution of agents in mean-field games system_ , Journal of Optimization Theory and Applications 179 (2018), no. 1, 332–357.
  * [30] Marco Cuturi, _Sinkhorn distances: Lightspeed computation of optimal transport_ , Advances in neural information processing systems 26 (2013).
  * [31] Jean-Pierre Fouque and Zhaoyu Zhang, _Deep learning methods for mean field control problems with delay_ , Frontiers in Applied Mathematics and Statistics 6 (2020), 11.
  * [32] Maximilien Germain, Joseph Mikael, and Xavier Warin, _Numerical resolution of McKean-Vlasov FBSDEs using neural networks_ , Methodology and Computing in Applied Probability (2022), 1–30.
  * [33] Maximilien Germain, Huyên Pham, and Xavier Warin, _Neural networks-based algorithms for stochastic control and PDEs in finance_ , arXiv preprint arXiv:2101.08068 (2021).
  * [34] P Jameson Graber, Alpár R Mészáros, Francisco J Silva, and Daniela Tonon, _The planning problem in mean field games as regularized mass transport_ , Calculus of Variations and Partial Differential Equations 58 (2019), no. 3, 1–28.
  * [35] Jiequn Han, Ruimeng Hu, and Jihao Long, _Learning high-dimensional McKean-Vlasov forward-backward stochastic differential equations with general distribution dependence_ , arXiv preprint arXiv:2204.11924 (2022).
  * [36] Jiequn Han, Arnulf Jentzen, and Weinan E, _Solving high-dimensional partial differential equations using deep learning_ , Proceedings of the National Academy of Sciences 115 (2018), no. 34, 8505–8510.
  * [37] Jiequn Han, Arnulf Jentzen, et al., _Deep learning-based numerical methods for high-dimensional parabolic partial differential equations and backward stochastic differential equations_ , Communications in mathematics and statistics 5 (2017), no. 4, 349–380.
  * [38] Ruimeng Hu and Mathieu Laurière, _Recent developments in machine learning methods for stochastic control and games_ , Preprint. SSRN:4096569 (2022).
  * [39] Minyi Huang, Peter E Caines, and Roland P Malhamé, _Large-population cost-coupled lqg problems with nonuniform agents: individual-mass behavior and decentralized $\textstyle\varepsilon$-Nash equilibria_, IEEE transactions on automatic control 52 (2007), no. 9, 1560–1571.
  * [40] Minyi Huang, Roland P Malhamé, and Peter E Caines, _Large population stochastic dynamic games: closed-loop McKean-Vlasov systems and the Nash certainty equivalence principle_ , Communications in Information & Systems 6 (2006), no. 3, 221–252.
  * [41] Jean-Michel Lasry and Pierre-Louis Lions, _Jeux à champ moyen. I–le cas stationnaire_ , Comptes Rendus Mathématique 343 (2006), no. 9, 619–625.
  * [42] , _Jeux à champ moyen. II–Horizon fini et contrôle optimal_ , Comptes Rendus Mathématique 343 (2006), no. 10, 679–684.
  * [43] , _Mean field games_ , Japanese journal of mathematics 2 (2007), no. 1, 229–260.
  * [44] Mathieu Laurière and Olivier Pironneau, _Dynamic programming for mean-field type control_ , Comptes Rendus Mathematique 352 (2014), no. 9, 707–713.
  * [45] Alex Tong Lin, Samy Wu Fung, Wuchen Li, Levon Nurbekyan, and Stanley J Osher, _Apac-net: Alternating the population and agent control via two neural networks to solve high-dimensional stochastic mean field games_ , arXiv preprint arXiv:2002.10113 (2020).
  * [46] Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos A Theodorou, _Deep generalized Schrodinger bridge_ , arXiv preprint arXiv:2209.09893 (2022).
  * [47] Quentin Merigot and Boris Thibert, _Optimal transport: discretization and algorithms_ , Handbook of Numerical Analysis, vol. 22, Elsevier, 2021, pp. 133–212.
  * [48] Carlo Orrieri, Alessio Porretta, and Giuseppe Savaré, _A variational approach to the mean field planning problem_ , Journal of Functional Analysis 277 (2019), no. 6, 1868–1957.
  * [49] Gabriel Peyré, Marco Cuturi, et al., _Computational optimal transport: With applications to data science_ , Foundations and Trends® in Machine Learning 11 (2019), no. 5-6, 355–607.
  * [50] Alessio Porretta, _On the planning problem for a class of mean field games_ , Comptes Rendus Mathematique 351 (2013), no. 11-12, 457–462.
  * [51] , _On the planning problem for the mean field games system_ , Dynamic Games and Applications 4 (2014), no. 2, 231–256.
  * [52] Maziar Raissi, Paris Perdikaris, and George E Karniadakis, _Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations_ , Journal of Computational physics 378 (2019), 686–707.
  * [53] R Tyrrell Rockafellar, _Convex analysis_ , vol. 18, Princeton university press, 1970.
  * [54] Lars Ruthotto, Stanley J Osher, Wuchen Li, Levon Nurbekyan, and Samy Wu Fung, _A machine learning framework for solving high-dimensional mean field game and mean field control problems_ , Proceedings of the National Academy of Sciences 117 (2020), no. 17, 9183–9193.
  * [55] Filippo Santambrogio, _Optimal transport for applied mathematicians_ , Birkäuser, NY 55 (2015), no. 58-63, 94.
  * [56] Richard Sinkhorn and Paul Knopp, _Concerning nonnegative matrices and doubly stochastic matrices_ , Pacific Journal of Mathematics 21 (1967), no. 2, 343–348.
  * [57] Justin Sirignano and Konstantinos Spiliopoulos, _DGM: A deep learning algorithm for solving partial differential equations_ , Journal of computational physics 375 (2018), 1339–1364.
  * [58] Cédric Villani, _Optimal transport: old and new_ , vol. 338, Springer, 2009\.
  * [59] , _Topics in optimal transportation_ , vol. 58, American Mathematical Soc., 2021.

##  Appendix A Solution for LQ problem

In this section, we provide an explicit solution for the LQ setting considered
in Section 4.1. We summarize the analytical solution derived in [29, Section
7.1] and we add the analytical formula for the variance, which is useful to
fully describe the evolution of the population distribution. We consider the
following setting, which is actually slightly more general than the one used
in Section 4.1. For any probability distribution $\textstyle\mu$ admitting a
first moment, we use the notation $\textstyle\bar{\mu}=\int x\mu(dx)$, to be
understood coordinate-wise. In this model, we take:

$\textstyle\displaystyle
b(x,\mu,a)=Ax+\bar{A}\bar{\mu}+Ba,\qquad\sigma=B,\qquad
f(x,\mu,a)=\frac{1}{2}a^{\top}a,\qquad\rho_{0}={\mathcal{N}}(\bar{x}_{0};\Sigma_{0}),\qquad\rho_{T}={\mathcal{N}}(\bar{x}_{T};\Sigma_{T}).$

We denote by $\textstyle v^{*}$ the optimal control and by
$\textstyle(\mu_{t})_{t\in[0,T]}$ the optimal flow of distributions. Since the
dynamics are linear and the initial and terminal distributions are Gaussians,
it can be shown that the distribution remains Gaussian at every $\textstyle
t$. At each $\textstyle t$, we denote by $\textstyle\bar{\mu}_{t}$ its mean
and by $\textstyle\Sigma_{t}$ its covariance matrix. In the case without mean
field interactions, the optimal control is linear in space, and the
coefficient depends on the solution of a Ricatti equation in time. The mean
field interaction being uniform in space and the fact that our problem is
linear explains why we can set

$\textstyle v^{*}(t,x)=-B^{\top}\Pi_{t}x+B^{\top}n_{t}$

and solve for $\textstyle(\Pi_{t})_{t\in[0,T]}$ and
$\textstyle(n_{t})_{t\in[0,T]}$, as time-dependent functions only. Note that
$\textstyle\bar{\mu}$ is now solution of the following ODE depending on
$\textstyle n_{t}$:

$\dot{\bar{\mu}}_{t}=(A+\bar{A}-BB^{\top}\Pi_{t})\bar{\mu}_{t}+BB^{\top}n_{t},\quad\bar{\mu}_{0}=\bar{x}_{0}.$
(28)

Furthermore, if we apply Itô’s formula to $\textstyle
V_{t}=\mathbb{E}[X_{t}X_{t}^{\top}]=\Sigma_{t}+\bar{\mu}_{t}\bar{\mu}_{t}^{\top}$,
we also obtain a matrix ODE depending on $\textstyle(n_{t})_{t\in[0,T]}$ for
$\textstyle(V_{t})_{t\in[0,T]}$:

$\textstyle\displaystyle\begin{cases}\dot{V}_{t}=V_{t}(A^{\top}-\Pi_{t}^{\top}BB^{\top})+(A-BB^{\top}\Pi_{t})V_{t}+\bar{\mu}_{t}n_{t}^{\top}BB^{\top}+BB^{\top}n_{t}\bar{\mu}_{t}^{\top}+\bar{\mu}_{t}\bar{\mu}_{t}^{\top}\bar{A}^{\top}+\bar{A}\bar{\mu}_{t}\bar{\mu}_{t}^{\top}+BB^{\top}\\\
V_{0}=\Sigma_{0}+\bar{\mu}_{0}\bar{\mu}_{0}^{\top}.\end{cases}$ (29)

We will also need the state transition matrices associated to $\textstyle A$
and $\textstyle A+\bar{A}$: we define $\textstyle\Phi$ as the solution of

$\frac{\partial\Phi(t,s)}{\partial t}=A\Phi(t,s),\quad\Phi(s,s)=I,\quad 0\leq
s\leq t\leq T.$

Since $\textstyle A$ is constant, we have $\textstyle\Phi(t,s)=e^{(t-s)A}$.
Similarly, we define $\textstyle\bar{\Phi}(t,s)=e^{(t-s)(A+\bar{A})}$. We also
set

$\textstyle\displaystyle\begin{cases}M(t,s)=\displaystyle\int_{s}^{t}\Phi(t,\tau)BB^{\top}\Phi(t,\tau)^{\top}\mathrm{d}\tau,\quad
0\leq s\leq t\leq T,\\\
\bar{M}(t,s)=\displaystyle\int_{s}^{t}\bar{\Phi}(t,\tau)BB^{\top}\bar{\Phi}(t,\tau)^{\top}d\tau,\quad
0\leq s\leq t\leq T,\end{cases}$ (30)

and denote for brevity $\textstyle\Phi_{T,0}=\Phi(T,0),$
$\textstyle\bar{\Phi}_{T,0}=\bar{\Phi}(T,0),$ $\textstyle M_{T,0}=M(T,0),$ and
$\textstyle\bar{M}_{T,0}=\bar{M}(T,0)$. Notice that if $\textstyle BB^{\top}$
commutes with $\textstyle\Phi$ and $\textstyle\bar{\Phi}$, and if the matrices
$\textstyle A+A^{\top}$ and $\textstyle A+\bar{A}+A^{\top}+\bar{A}^{\top}$ are
non-singular, then we have the following closed form expressions for
$\textstyle M(t,s)$ and $\textstyle\bar{M}(t,s)$, with $\textstyle 0\leq s\leq
t\leq T$:

$\textstyle\displaystyle
M(t,s)=BB^{\top}\int_{s}^{t}e^{A+A^{\top}}(t-\tau)\mathrm{d}\tau=BB^{\top}\bigl{[}e^{(A+A^{\top})t}-e^{(A+A^{\top})s}\bigr{]}(A+A^{\top})^{-1}$
$\textstyle\displaystyle\bar{M}(t,s)=BB^{\top}\int_{s}^{t}e^{A+A^{\top}+\bar{A}+\bar{A}^{\top}}(t-\tau)\mathrm{d}\tau=BB^{\top}\bigl{[}e^{(A+A^{\top}+\bar{A}+\bar{A}^{\top})t}-e^{(A+A^{\top}+\bar{A}+\bar{A}^{\top})s}\bigr{]}(A+A^{\top}+\bar{A}+\bar{A}^{\top})^{-1}.$

Now, $\textstyle(\Pi_{t})_{t\in[0,T]}$ is the solution to the following
Riccati ODE, which is independent from the other variables:

$\textstyle\displaystyle\begin{cases}\dot{\Pi}_{t}=-A^{\top}\Pi_{t}-\Pi_{t}A+\Pi_{t}BB^{\top}\Pi_{t}\\\
\Pi_{0}=\Sigma_{0}^{-1/2}\left[\frac{I}{2}+\Sigma_{0}^{1/2}\Phi_{T,0}^{\top}M_{T,0}^{-1}\Phi_{T,0}\Sigma_{0}^{1/2}-\left(\frac{I}{4}+\Sigma_{0}^{1/2}\Phi_{T,0}^{\top}M_{T,0}^{-1}\Sigma_{T}M_{T,0}^{-1}\Phi_{T,0}\Sigma_{0}^{1/2}\right)^{1/2}\right]\Sigma_{0}^{-1/2}.\end{cases}$
(31)

Under some conditions, Riccati equations admit explicit solutions in dimension
$\textstyle 1$, see e.g. page 110 in [21]. More generally, we can solve (31)
using a forward time-marching method.

We can then show that

$n_{t}=\Pi_{t}\bar{\Phi}(t,T)\bar{M}(T,t)\bar{M}_{T,0}^{-1}\bar{\Phi}_{T,0}z_{0}+\Pi_{t}\bar{M}(T,0)\bar{\Phi}(T,t)^{\top}\bar{M}_{T,0}^{-1}z_{T}+\bar{\Phi}(T,t)^{\top}\bar{M}_{T,0}^{-1}(z_{T}-\bar{\Phi}_{T,0}z_{0}),$
(32)

where $\textstyle(z_{t})_{t\in[0,T]}$ is defined by

$z_{t}=\bar{\Phi}(T,t)^{\top}\bar{M}_{T,0}^{-1}(\bar{x}_{T}-\bar{\Phi}_{T,0}\bar{x}_{0}).$
(33)

Now, up to the computation of $\textstyle\Pi$, we have an explicit formula for
$\textstyle n$ and we can obtain the mean $\textstyle\bar{\mu}_{t}$ from (28)
and then the covariance matrix $\textstyle\Sigma_{t}$ from (29). The optimal
mean field distribution at time $\textstyle t$ is the Gaussian distribution
with mean $\textstyle\bar{\mu}_{t}$ and variance $\textstyle\Sigma_{t}$.

##  Appendix B Computation of $\textstyle{\mathcal{G}}$

In this section, we discuss some of the issues arising when computing
$\textstyle{\mathcal{G}}$ in practice, as well as our solutions. As the reader
may notice, the losses defined in (25) and (26) depend on the exact form of
functionals $\textstyle{\mathcal{F}}$ and $\textstyle{\mathcal{G}}$. As
$\textstyle{\mathcal{F}}$ is already defined in an explicit, easy-to-compute
form, we are left with the problem of figuring out a good approach to compute
$\textstyle{\mathcal{G}}$. Recalling the definition of
$\textstyle{\mathcal{G}}$, we have the following observation,

$\textstyle\displaystyle{\mathcal{G}}(\mathfrak{a},\mathfrak{b})$
$\textstyle\displaystyle=-\inf_{m\geq
0}\int_{\mathcal{Q}_{T}}m(t,x)\Bigl{(}\mathfrak{a}(t,x)-H\bigl{(}x,m(t,x),\mathfrak{b}(t,x)\bigr{)}\Bigr{)}\mathrm{d}x\,\mathrm{d}t$
$\textstyle\displaystyle=-\int_{\mathcal{Q}_{T}}\inf_{m\geq
0}\Bigl{[}m\bigl{(}\mathfrak{a}(t,x)-H\bigl{(}x,m,\mathfrak{b}(t,x)\bigr{)}\bigr{)}\Bigr{]}\mathrm{d}x\,\mathrm{d}t$
(34)
$\textstyle\displaystyle=-\int_{\mathcal{Q}_{T}}{\mathcal{K}}\bigl{(}\mathfrak{a}(t,x),\mathfrak{b}(t,x)\bigr{)}\mathrm{d}x\,\mathrm{d}t,$
(35)

where we denoted $\textstyle{\mathcal{K}}\bigl{(}a,b\bigr{)}=\inf_{m\geq
0}\Bigl{[}m\bigl{(}\mathfrak{a}-H\bigl{(}x,m,\mathfrak{b}\bigr{)}\bigr{)}\Bigr{]}$.

Therefore, the form of $\textstyle{\mathcal{G}}$ depends on
$\textstyle{\mathcal{K}}$. In general, we do not know any closed form of
$\textstyle{\mathcal{K}}$ in terms of $\textstyle\mathfrak{a}$ and
$\textstyle\mathfrak{b}$. However, for several Hamiltonian functions
$\textstyle H(x,m,p)$ of interest, we can derive such closed-form solution.
Here we demonstrate some of the calculations to deliver a general idea.

######  Example 4 (Mean-field Aversion).

Consider the running cost $\textstyle f(m,v)=\frac{1}{4}\|v\|^{2}+m$ and drift
$\textstyle b(x,m,v)=v$. The corresponding Hamiltonian is $\textstyle
H(x,m,p)$ = $\textstyle\|p\|^{2}-m$. Then

$\textstyle\displaystyle{\mathcal{K}}(\mathfrak{a},\mathfrak{b})=$
$\textstyle\displaystyle\inf_{m\geq
0}\;m^{2}+(\mathfrak{a}-\|\mathfrak{b}\|^{2})m=\begin{cases}-\frac{1}{4}\bigl{(}\mathfrak{a}-\|\mathfrak{b}\|^{2}\bigr{)}^{2}\quad&\hbox{
if }\mathfrak{a}-\|\mathfrak{b}\|^{2}\geq 0\\\ 0\quad&\hbox{ otherwise.
}\end{cases}$ (36)

######  Example 5 (Mean-field Maximum Entropy).

Consider the running cost $\textstyle f(m,v)=\frac{1}{4}\|v\|^{2}+\log(m)$ and
drift $\textstyle b(x,m,v)=v$, the corresponding Hamiltonian is $\textstyle
H(x,m,p)=\frac{1}{2}\|p\|^{2}-\log(m)$, then

$\textstyle\displaystyle{\mathcal{K}}(\mathfrak{a},\mathfrak{b})=$
$\textstyle\displaystyle\inf_{m\geq 0}\;m\log
m+(\mathfrak{a}-\frac{1}{2}\|\mathfrak{b}\|^{2})m=-\exp\bigl{(}\frac{1}{2}\|\mathfrak{b}\|^{2}-\mathfrak{a}-1\bigr{)}.$
(37)

######  Example 6 (Continuous Optimal Transport).

Consider the running cost $\textstyle f(m,v)=\frac{1}{2}\|v\|^{2}$ and drift
$\textstyle b(x,m,v)=v$, the corresponding Hamiltonian is $\textstyle
H(x,m,p)=\frac{1}{2}\|p\|^{2}$, then

$\textstyle\displaystyle{\mathcal{K}}(\mathfrak{a},\mathfrak{b})=$
$\textstyle\displaystyle\inf_{m\geq
0}\;m(\mathfrak{a}-\frac{1}{2}\|\mathfrak{b}\|^{2})=\begin{cases}0\quad\quad&\hbox{
if }\mathfrak{a}-\frac{1}{2}\|\mathfrak{b}\|^{2}\geq 0\\\ -\infty\quad&\hbox{
otherwise.}\end{cases}$ (38)

######  Example 7 (Mean-field Congestion).

Consider the running cost $\textstyle f(m,v)=\frac{1}{4}m\|v\|^{2}$ and drift
$\textstyle b(x,m,v)=v$, the corresponding Hamiltonian is $\textstyle
H(x,m,p)=\frac{\|p\|^{2}}{m}$, then

$\textstyle\displaystyle{\mathcal{K}}(\mathfrak{a},\mathfrak{b})=$
$\textstyle\displaystyle\inf_{m\geq
0}\;m(\mathfrak{a}-\frac{\|\mathfrak{b}\|^{2}}{m})=\begin{cases}-\|\mathfrak{b}\|^{2}\quad\quad&\hbox{
if }\mathfrak{a}\geq 0\\\ -\infty\quad&\hbox{ otherwise. }\end{cases}$ (39)

From these examples, it can be seen that $\textstyle{\mathcal{K}}$ has a
closed form for many smooth Hamiltonian. However, the existence of closed form
expressions of $\textstyle{\mathcal{K}}$ alone is not enough for making the
training loss tractable. In the example of (38) and (39),
$\textstyle{\mathcal{K}}$ takes values $\textstyle-\infty$, which makes
$\textstyle{\mathcal{G}}$ singular and computationally intractable at some
points. Moreover, gradient-based training cannot be carried out successfully
in the presence of infinite values as well. The presence of
$\textstyle-\infty$ in $\textstyle{\mathcal{K}}$ is due to the degeneracy of
$\textstyle H$ in terms of order in $\textstyle m$. For cases with singular
$\textstyle{\mathcal{K}}$ and $\textstyle{\mathcal{G}}$, we need an additional
trick to tackle this issue.

Recall that in Algorithm 1, the update for function $\textstyle q$ is given
by,

$\textstyle\displaystyle q^{(k)}=$
$\textstyle\displaystyle\underset{q:\mathcal{Q}_{T}\rightarrow{{\mathbb{R}}}^{d+1}}{\mathop{\arg\min}}\;\mathcal{G}(q)+\langle\lambda^{(k-1)},q\rangle+\frac{r}{2}\bigl{\|}\Lambda
u^{(k)}-q\bigr{\|}^{2}$ $\textstyle\displaystyle=$
$\textstyle\displaystyle\underset{q:\mathcal{Q}_{T}\rightarrow{{\mathbb{R}}}^{d+1}}{\mathop{\arg\min}}\;\int_{\mathcal{Q}_{T}}\left(-{\mathcal{K}}(q(t,x))+\langle\lambda^{(k-1)}(t,x),q(t,x)\rangle+\frac{r}{2}\bigl{\|}\Lambda
u^{(k)}(t,x)-q(t,x)\bigr{\|}^{2}\right)\;\mathrm{d}x\;\mathrm{d}t$ (40)

Since we don’t have any additional constraint on the value of $\textstyle q$,
the function $\textstyle q^{(k)}$ that minimizes the integral in (B) should
minimize the integrand point wisely. Therefore, it holds that:

$\textstyle\displaystyle q^{(k)}(t,x)=\;$
$\textstyle\displaystyle\underset{q\in{{\mathbb{R}}}^{d+1}}{\mathop{\arg\min}}\quad\left(-{\mathcal{K}}(q)+\langle\lambda^{(k-1)}(t,x),q\rangle+\frac{r}{2}\bigl{\|}\Lambda
u^{(k)}(t,x)-q\bigr{\|}^{2}\right)$ $\textstyle\displaystyle=\;$
$\textstyle\displaystyle\underset{q\in{{\mathbb{R}}}^{d+1}}{\mathop{\arg\min}}\sup_{m\geq
0}\;m\bigl{(}H(x,m,q_{2})-q_{1}\bigr{)}+\langle\lambda^{(k-1)}(t,x),q\rangle+\frac{r}{2}\bigl{\|}\Lambda
u^{(k)}(t,x)-q\bigr{\|}^{2},$ (41)

where we denote $\textstyle q=(q_{1},q_{2})$, with $\textstyle
q_{1}\in{{\mathbb{R}}}$ and $\textstyle q_{2}\in{{\mathbb{R}}}^{d}$. For fixed
$\textstyle\lambda^{(k-1)}$ and $\textstyle u^{(k)}$, let us define
$\textstyle{\mathcal{L}}(q,m)=m\bigl{(}H(x,m,q_{2})-q_{1}\bigr{)}+\langle\lambda^{(k-1)}(t,x),q\rangle+\frac{r}{2}\bigl{\|}\Lambda
u^{(k)}(t,x)-q\bigr{\|}^{2}$. Under certain conditions, the following minimax
equality holds for the right-hand side of (B),

$\textstyle\displaystyle\inf_{q\in{{\mathbb{R}}}^{d+1}}\sup_{m\geq
0}{\mathcal{L}}(q,m)=\sup_{m\geq
0}\inf_{q\in{{\mathbb{R}}}^{d+1}}{\mathcal{L}}(q,m)$ (42)

Therefore, we can exploit (42) to address the issues of singular
$\textstyle{\mathcal{K}}$. We notice that for the cases of congestion and
general linear quadratic, $\textstyle{\mathcal{L}}(q,m)$ is quadratic in
$\textstyle q$ for fixed $\textstyle m$. This means that we can solve
explicitly in a closed form for $\textstyle\inf_{q}{\mathcal{L}}(q,m)$ for
fixed $\textstyle m$, then we solve for the maximization problem over
$\textstyle m\geq 0$. In this way, we can effectively avoid the issues
generated by infinity values in $\textstyle{\mathcal{K}}$.

Moreover, using this trick, we directly obtain the value for $\textstyle
q^{(k)}$ at any $\textstyle(t,x)$ based on the value of $\textstyle
u^{(k)}(t,x)$ and $\textstyle\lambda^{(k-1)}(t,x)$. Therefore, we can skip the
neural network training in Algorithm 2 for function $\textstyle q$ in each
DeepADMM iteration. Instead, we compute the values for $\textstyle
q^{(k)}(t,x)$ directly following the above procedure whenever the values are
needed to compute $\textstyle{\mathcal{L}}^{(u)}$ and
$\textstyle{\mathcal{L}}^{(\lambda)}$.

##  Appendix C Some loss plots

Contrary to the LQ case, we do not have any benchmark solution to compare our
numerical results to for the congestion case. We already explained why the
methods gave consistent results, namely that we see the density spreading
before reforming the terminal distribution, in contrast with the LQ case where
the shape of the population remains the same along the trajectory (see 4.2.2).
In this section, we provide some plots showing the evolution of the different
losses through the training of the neural networks.

Figure 7: Evolution of the loss in the congestion case in dimension 5 for
method 1.

Figure 7 shows the losses for Method 1. The total loss is the sum of the
“running cost” loss and the “terminal penalization” loss multiplied by
$\textstyle C_{W}$. In this test case, we took $\textstyle C_{W}=10$, meaning
that the penalization is ten times the Wasserstein distance between the
effective terminal distribution of agents and the desired one. We observe that
for the first iterations, the distance is high and the running cost almost
zero, which comes from the fact that the first try is not to move. Then the
algorithm makes the penalization decrease by paying a trade-off in the form of
the running cost, and seems to reach a plateau.

Figure 8 shows the losses for Method 2. The “HJB loss” is the squared
$\textstyle L^{2}$ residual for the HJB equation. The “KFP loss” is the
squared $\textstyle L^{2}$ residual for the KFP equation. The “initial BC”
loss and the “terminal BC” loss respectively correspond to the squared
$\textstyle L^{2}$ error on the initial and terminal distributions. The “total
loss” is the sum of the other losses, up to multiplicative weights.

Figure 9 shows, for Method 3, the squared $\textstyle L^{2}$ residuals for the
HJB and KFP equations, as well as the squared $\textstyle L^{2}$ loss for the
initial and terminal conditions. Note that these losses are not directly
minimized during the algorithm of Method 3, but they are minimized as a by-
product of the iterations.

Figure 8: Evolution of the loss in the congestion case in dimension 5 for
method 2.  Figure 9: Evolution of the loss in the congestion case in dimension
5 for method 3.