Bridging Mean-Field Games and Normalizing Flows with Trajectory Regularization^††thanks: This work is supported in part by an NSF Career Award (DMS-1752934) and NSF DMS-2134168.

The study of MFGs considers classic $N$-player games at the limit of $N\to\infty$ [45, 33, 32]. Assuming homogeneity of the objectives and anonymity of the players, we set up an MFG with a continuum of non-cooperative rational agents distributed in a spatial domain $\Omega\subseteq\mathbb{R}^{d}$ and a time interval $[0,T]$. For any fixed $t\in[0,T]$, we denote the population density of agents by $p(\cdot,t)\in\mathcal{P}(\Omega)$, where $\mathcal{P}(\Omega)$ is the space of all probability densities over $\Omega$. For an agent starting at ${\bm{x}}_{0}\in\Omega$, their position over time follows a trajectory ${\bm{x}}:[0,T]\to\Omega$ governed by

where ${\bm{v}}:\Omega\times[0,T]\to\Omega$ specifies an agent’s action at a given time. For simplicity, there are no stochastic terms in (1) in this work. Then, each agent’s trajectory is completely determined by ${\bm{v}}({\bm{x}},t)$. To play the game over an interval $[t,T]$, each agent seeks to minimize their objective:

Each term in the objective denotes a particular type of cost. The running cost $L:\Omega\times\Omega\to\mathbb{R}$ is incurred due to each agent’s action. A commonly used example for $L$ is the kinetic energy $L({\bm{x}},{\bm{v}})=\|{\bm{v}}\|^{2}$, which accounts for the total amount of movement along the trajectory. The running cost $I:\Omega\times\mathcal{P}(\Omega)\to\mathbb{R}$ is accumulated through each agent interacting with one another or with the environment. For example, this term can be an entropy that discourages the agents from grouping together, or a penalty for the agents to collide with an obstacle in the space [59]. The terminal cost $M:\Omega\times\mathcal{P}(\Omega)\to\mathbb{R}$ is computed from the final state of the agents. Typically, $M$ measures a discrepancy between the final distribution $p(\cdot,T)$ and a desirable distribution $p_{1}\in\mathcal{P}(\Omega)$.

To solve the MFG, we define the value function $u:\Omega\times[0,T]\to\mathbb{R}$ as

One can show that $u({\bm{x}},t)$ and $p({\bm{x}},t)$ are solutions to the following HJB equation and continuity equations [45, 33, 32]

where $H:\Omega\times\Omega\to\mathbb{R}$ is the Hamiltonian $H({\bm{x}},{\bm{q}})\coloneqq\sup_{\bm{v}}-\langle{\bm{q}},{\bm{v}}\rangle-L({\bm{x}},{\bm{v}})$; $p_{0}\in\mathcal{P}(\Omega)$ is the initial population density at $t=0$; and ${\bm{v}}({\bm{x}},t)$ denotes the optimal strategy for an agent at position ${\bm{x}}$ and time $t$. It can be shown that the optimal strategy satisfies ${\bm{v}}({\bm{x}},t)=-\nabla_{p}H({\bm{x}},\nabla_{{\bm{x}}}u({\bm{x}},t))$. We note that the continuity equations in the last two lines of (4) are a special case of the Fokker-Planck equation with no diffusion terms.

The setup (1)–(2) concerns the strategy of an individual agent. Under suitable assumptions, the work [45] developed a macroscopic formulation of the MFG that models the collective strategy of all agents. Suppose there exist functionals $\mathcal{I},\mathcal{M}:\mathcal{P}(\Omega)\to\mathbb{R}$ such that

where $\frac{\delta}{\delta p}$ is the variational derivative. Then, the functions $p({\bm{x}},t)$ and ${\bm{v}}({\bm{x}},t)$ satisfying (4) coincide with the optimizers of the following variational problem:

This formulation is termed the variational MFG and it serves as the main optimization problem this work solves. It is known that traditional numerical PDE methods that solve (4) are computationally intractable in dimensions higher than three, as the number of grid points grows exponentially. This challenge motivates us to solve the MFG problem with parameterizations of the agent trajectory based on NFs.

NFs are a family of invertible neural networks that find applications in density estimation, data generation, and variational inference [16, 17, 22, 30, 57, 40, 31]. In machine learning, NFs are considered deep generative models, which include variational autoencoders [38] and generative adversarial networks [25] as other well-known examples. The key advantage of NFs over these alternatives is the exact computation of the data likelihood, which is made possible by invertibility.

Mathematically, suppose we have a dataset ${\mathcal{X}}=\{{\bm{x}}_{n}\}_{n=1}^{N}\subseteq\mathbb{R}^{d}$ generated by an underlying data distribution $P_{1}$; that is, ${\bm{x}}_{n}\overset{\mathrm{iid}}{\sim}P_{1}$. We define a flow to be an invertible function $f_{\theta}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ parameterized by $\theta\in\mathbb{R}^{W}$, and a normalizing flow to be the composition of a sequence of flows: $F_{\theta}=f_{\theta_{K}}\circ f_{\theta_{K-1}}\circ...\circ f_{\theta_{2}}\circ f_{\theta_{1}}$ parameterized by $\theta=(\theta_{1},\theta_{2},...,\theta_{K})$. To model the complex data distribution $P_{1}$, the idea is to use an NF to gradually transform a simple base distribution $P_{0}$ to $P_{1}$. Formally, the transformation of the base distribution is conducted through the push-forward operation $F_{\theta*}P_{0}$, where $F_{\theta*}P(A)\coloneqq P(F_{\theta}^{-1}(A))$ for all measurable sets $A\subset\Omega$. The aim is that the transformed distribution resembles the data distribution. Thus, a commonly used loss function for training the NF is the KL divergence:

If the measure $P$ admits density $p$ with respect to the Lebesgue measure, the density $F_{\theta*}p$ associated with the push-forward measure $F_{\theta*}P$ satisfies the change-of-variable formula [41]:

where $\nabla F_{\theta}^{-1}$ is the Jacobian of $F_{\theta}^{-1}$. This allows one to tractably compute the KL term in (6), which is equivalent to the negative log-likelihood of the data up to a constant (independent of $\theta$):

where $g_{\theta_{i}}\coloneqq f_{\theta_{i}}^{-1}$ are the flow inverses, and ${\bm{y}}_{j}\coloneqq g_{\theta_{K-j}}\circ...\circ g_{\theta_{K}}({\bm{x}}),\forall j=1,2,...,K-1,{\bm{y}}_{0}\coloneqq{\bm{x}}$.

We remark that there exists flexibility in choosing the base distribution $P_{0}$, as long as it admits a known density to evaluate on. In practice, a typical choice is the standard normal $\mathrm{N}(0,I)$. In addition, (8) requires the log-determinant of the Jacobians, which take $O(d^{3})$ time to compute in general. Existing NF architectures sidestep this issue by meticulously using functions that follow a (block) triangular structure [41], so that the log-determinants can be computed in $O(d)$ time from the (block) diagonal elements.

One popular family of NF models is the coupling flow [41], which composes individual flows in the following form

where $({\bm{x}}_{1},{\bm{x}}_{2})\in\mathbb{R}^{a}\times\mathbb{R}^{d-a}$, $\theta:\mathbb{R}^{d-a}\to\mathbb{R}^{d}$ is the called the conditioner, and $h:\mathbb{R}^{d}\to\mathbb{R}^{d}$ is the coupling function. Suppose $h$ is invertible, then the inverse $f^{-1}$ is:

In common scenarios, $h$ is defined element-wise, such that invertibility is ensured by strict monotonicity. For example, a simple yet expressive choice is the affine function $h({\bm{x}};\theta)=\theta_{1}{\bm{x}}+\theta_{2},\theta_{1}\neq 0,\theta_{2}\in\mathbb{R}$. This gives rise to the RealNVP flow [17]. Other popular coupling functions include splines, mixture of CDFs, and functions parameterized by neural networks, which are shown to be more expressive than the affine coupling function [22, 50, 21, 30, 35, 31].

Rather than solving for the density $p({\bm{x}},t)$ and the action ${\bm{v}}({\bm{x}},t)$, we reparameterize the problem (5) to directly work with agent trajectories, from which $p({\bm{x}},t)$ and ${\bm{v}}({\bm{x}},t)$ can be derived. Let $P(\cdot,t)$ be the measure that admits $p(\cdot,t)$ as its density for all $t\in[0,T]$. Define the agent trajectory as $F:\mathbb{R}^{d}\times\mathbb{R}\to\mathbb{R}^{d}$, where $F({\bm{x}},t)$ is the position of the agent starting at ${\bm{x}}$ and having traveled for time $t$. It satisfies the following ordinary differential equation:

The evolution of the population density is determined by the movement of agents. Thus, $P(\cdot,t)$ is simply the push-forward of $P_{0}$ under $F(\cdot,t)$; namely, $P({\bm{x}},t)=F(\cdot,t)_{*}P_{0}({\bm{x}})$, whose associated density satisfies

where $\mathrm{d}(F(\cdot,t)_{*}p_{0}({\bm{x}}))$ is a Radon-Nikodym derivative. Note that this definition of $p({\bm{x}},t)$ automatically satisfies the continuity equation as well as the initial condition for $p({\bm{x}},t)$ in (5). Unless specified otherwise, we assume $T=1$ and use the $L_{2}$ transport cost $L({\bm{x}},{\bm{v}}({\bm{x}},t),p({\bm{x}},t))=\lambda_{L}p({\bm{x}},t)\|{\bm{v}}({\bm{x}},t)\|_{2}^{2}$ from now on. Here, $\lambda_{L}\geq 0$ is treated as a model (hyper-)parameter. Furthermore, we can apply a change of variables on (12) to rewrite the integral involving $L({\bm{x}},{\bm{v}},p)$:

Similarly, we can rewrite the interaction cost $\mathcal{I}(p(\cdot,t))$ and the terminal cost $\mathcal{M}(p(\cdot,T))$ as:

Therefore, (5) becomes:

It is worth noting that (14) removes the continuity-equation constraint. The reformulated problem optimizes over agent trajectories $F$, which allows one to parameterize with an NF. In addition, the remaining constraint $F(\cdot,0)=\mathrm{Id}$ will be automatically integrated into the NF, such that it can be trained in an unconstrained manner.

Starting from the base density $p_{0}$, each flow in the NF advances the density one step forward. Hence, the NF model specifies a discretized evolution of the agent trajectories. Formally, we discretize the trajectory $F({\bm{x}},t)$ in time with regular grid points $t_{i}\coloneqq i\cdot\Delta t,\forall i=0,1,...K$. Thus, the grid spacing is $\Delta t\coloneqq\frac{1}{K}.$ Denote $F({\bm{x}},t_{i})\coloneqq F_{i}({\bm{x}})$. By approximating $\partial_{t}F({\bm{x}},t)$ with the forward difference, the transport cost becomes:

Moreover, suppose that the terminal cost is computed as the KL divergence as in standard NFs:

For simplicity of exposition, for now we omit the interaction term (that is, $\mathcal{I}\equiv 0$). We can discretize the interaction cost $\mathcal{I}(F(\cdot,t)_{*}P_{0})$, similarly to (15), given its exact form. Subsequently, we will give numerical examples for the case $\mathcal{I}\not\equiv 0$ (see Sections 5.1.2 and 5.1.3).

For any $F$, the discretized objective value converges to the continuous version at $O(\frac{1}{K})$. However, we show later that the optimal discretized and continous objective values in fact agree for any fixed number of grid points used. Hence, there introduce no additional errors by discretizing the original MFG problem, in the absence of interaction costs.

To solve the discretized MFG, we define the flow maps for $0\leq t_{i}\leq t_{j}\leq 1$ as $\Phi_{t_{i}}^{t_{j}}({\bm{x}}_{0})={\bm{x}}(t_{j})$, where

Here, we leverage the semi-group property $\Phi_{t_{b}}^{t_{c}}\circ\Phi_{t_{a}}^{t_{b}}=\Phi_{t_{a}}^{t_{c}}$ to decompose the agent trajectories into a sequence of flow maps. It follows that $F_{k}=\Phi_{t_{k-1}}^{t_{k}}\circ\Phi_{t_{k-2}}^{t_{k-1}}\circ...\circ\Phi_{t_{0}}^{t_{1}},\forall k=1,2,...K$, and we parameterize each $\Phi_{t_{i}-1}^{t_{i}}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ by a flow model $f_{\theta_{i}}$. Denote $F_{\theta^{k}}\coloneqq f_{\theta_{k}}\circ...\circ f_{\theta_{1}}$, $F_{\theta^{0}}\coloneqq Id$, and $F_{\theta}\coloneqq F_{\theta^{K}}$. The discretized MFG problem with flow-parameterized trajectories is:

where $\lambda_{\mathcal{M}},\lambda_{L}\geq 0$ are weights.

By comparing (18) to the NF optimization problem (8), we see that training a standard NF is equivalent to solving a special MFG with no transport and interaction costs. In the past, many highly flexible flows were developed that are capable of solving this special MFG [31, 22, 17, 39, 30]. As a result, the trajectory-based formulation opens the door to employ expressive NFs to accurately solve MFGs. In addition, note that the density evolution can be directly obtained from the agent trajectories, which are the outputs of the individual flows in the NF model. Thus, this formulation sidesteps the additional effort required to solve the Fokker-Planck equation (4), resulting in a more efficient solution of the population dynamics.

We also note that the computation of (18) involves both a forward (normalizing) and an inverse (generating) pass of the NF: the forward pass computes $\mathcal{D}_{KL}(P_{1}||F_{\theta_{*}}P_{0})$ via (8) and the inverse pass yields the transport cost. For NFs that have comparable run time for passes in both directions (e.g. coupling flows), adding the transport cost only doubles the overall computation time. For NFs where the inverse pass is much slower than the forward pass (or vice versa), such as autoregressive flows [41], the additional computational effort associated with the transport cost may be prohibitive. To remedy this, we propose an alternative, equivalent formulation that incurs less overhead.

The aforementioned method for solving high-dimensional MFGs using NFs leads to the view of NFs as a special type of MFG problems. Comparing the MFG formulation (18) with the standard NF training (8), one can interpret the MFG transport cost as a regularization term in the NF training with strength $\lambda_{L}/\lambda_{\mathcal{M}}$; we call this regularization a trajectory regularization. It is straightforward to implement it on top of any existing discrete NF models (not necessarily restricted to coupling flows). Since the transport cost aggregates squared differences between the input and the output of each flow, the regularization encourages smoothness in the intermediate transforms.

From the perspective of inverse problems, there exist many flows (transport plans) that can transform the base $P_{0}$ to $P_{1}$. Therefore, regularization helps ameliorate the inevitably ill-posed nature of NF training. Through OT theory, the $L_{2}$ regularized evolution of densities is characterized as the geodesic connecting $P_{0}$ and $P_{1}$ in the space of measures equipped with the Wasserstein-2 metric [4]. Since $P_{0}$ is always chosen to be absolutely continuous for NFs, the OT between the two densities will be a unique and bijective mapping known as the Monge map, $T({\bm{x}})$. The same result holds for the discretized MFG problem, shown later in Theorem 4.4. Therefore, the NF training problem is no longer ill-posed. Furthermore, the optimal trajectory is then a linear interpolation between each point ${\bm{x}}$ and its destination $T({\bm{x}})$: $F^{*}({\bm{x}},t)=(1-t){\bm{x}}+tT({\bm{x}})$.

From the perspective of machine learning, the transport cost provides an effective framework to control the flow’s Lipschitz bound, which closely correlates with the robustness and the generalization capability [53, 14]. Note that by enforcing smoothness in the time domain, the trajectory regularization also encourages smoothness in the spatial domain, owing to the triangle inequality:

where $F_{0}=Id,F_{i}=f_{i}\circ...f_{1},i=1,...,K$, and $F_{K}$ is the NF model. By regularizing with the transport cost as in (18), the right-hand side can be controlled.

As observed by the authors of [7], explicit regularization methods such as $l_{2}$ weight decay has little impact on the model’s generalization gap. In contrast, the trajectory regularization is implicit in that it imposes no direct penalty on the parameters. We find it to be much more effective than weight decay in reducing the NF’s Lipschitz bound. We will demonstrate in the numerical experiments that a regularized evolution can reduce overfitting and improve density estimation.

Our analysis focuses on the MFG problem without the interaction cost. We recall the discretized MFG problem with $\mathcal{I}\equiv 0$ here for convenience:

First, we note that in the absence of the interaction cost, the discretized objective converges to the continuous analogue in $O(\frac{1}{K})$ when evaluated on any $F$.

The proof is based on a straightfoward Taylor expansion, similar to that used in finite difference error analysis. In general, one can use a higher order scheme to discretize (14) in time, yielding a higher order convergence rate. For example, the composite Simpson’s rule is used in our numerical experiments.

Next, we consider the optimization and establish two lemmas that characterize the optimal trajectories of both the discretized and the continuous MFG problems.

Lemma 4.2 can be shown by recasting the MFG as a suitable OT problem. Then, the straight trajectories follow from the OT theory. A similar result holds for the discretized MFG problem.

Lemma 4.3 can be shown by recasting the problem in a similar fashion, followed by solving its KKT system. Combining Lemma 4.2 and Lemma 4.3, we can show that the optimal value of the discretized problem agrees with its continuous counterpart.

Theorem 4.4 has two important implications. First, when we use an NF to parameterize and solve the MFG problem with no interaction costs, the optimizer incurs no discretization error. Therefore, it is not necessary to use higher order discretization schemes for the problem. Second, the discretized MFG problem, which corresponds to the training objective (18) modulo parameterization, admits a unique solution. In other words, the regularized NF training is well-posed.

The universality results from [60] indicate that NFs built with affine coupling transforms and sufficiently expressive conditioners are universal approximators for ${\mathcal{D}}^{2}$ locally in $L_{p},p\in[1,\infty)$. Here, ${\mathcal{D}}^{2}=\{f:Dom(f)\to Im(f)\subseteq\mathbb{R}^{d}\}$, where $f$ is $C^{2}$ diffeomorphic, and $Dom(f)\subseteq\mathbb{R}^{d}$ is open and diffeomorphic to $\mathbb{R}^{d}$. As a result, affine coupling flows are well-suited for approximating the optimal MFG trajectories. To be more precise, we consider the following sets introduced in [60].

Note that the affine coupling flows used in our numerical experiments belong to $\mathbf{INN_{{\mathcal{H}}\textbf{-}ACF}}$, where ${\mathcal{H}}$ is the class of ReLU networks. It turns out that the universality of $\mathbf{INN_{{\mathcal{H}}\textbf{-}ACF}}$ provides an ideal characterization for the approximate solutions of (14).

This result can be directly obtained from the conclusion in [60]. We remark that the conditions on ${\mathcal{H}}$ can be satisfied by fully connected networks equipped with the ReLU activation. In addition, we show that the push-forward measures from the approximated optimal trajectories converge to the ground truth in distribution. This ensures the validity of the learned evolution of the agent population.

The proof is based on the Lipschitz-bounded definition of convergence in distribution. See the proof in the appendix.

In the first part, we follow the settings in [59] to construct two MFG problems. They are designed so that the behavior of the MFG solutions is invariant to dimensionality, when projected onto the first two components. This property allows for qualitative evaluation through visualization. As the dimensionality increases from 2 to 100, the problem becomes more challenging in terms of the neural network’s capacity and the computational complexity. We further extend our numerical experiments to multi-group path planning. In these experiments, we use spline flows as it is commonly believed that they are more expressive than coupling flows.