NF-MKV Net: A Constraint-Preserving Neural Network Approach to Solving Mean-Field Games Equilibrium

We now formalize the MFGs problem without considering common noise. For this purpose, we start with a complete filtered probability space $(\Omega,\mathcal{F},\mathbb{F}=(\mathcal{F}_{t})_{0\leq t\leq T},\mathbb{P}))$ the filtration $\mathbb{F}$ supporting a $d-$dimensional Wiener process $\mathbf{W}=(W_{t})_{0\leq t\leq T}$ with respect to $\mathbb{F}$ and an initial condition $\xi\in L^{2}(\Omega,\mathcal{F}_{0},\mathbb{P};\mathbb{R}^{d})$. This MFGs problem can be described as:

(i) For each fixed deterministic flow $\bm{\mu}=(\mu_{t})_{0\leq t\leq T}$ of probability measures on $\mathbb{R}^{d}$, solve the standard stochastic control problem:

subject to

(ii) Find a flow $\bm{\mu}=(\mu_{t})_{0\leq t\leq T}$such that $\mathcal{L}(\hat{X}_{T}^{\bm{\mu}})=\mu_{t}$ for all $t\in[0,T]$, if $\hat{X}^{\bm{\mu}}$ is a solution of the above optimal control problem.

We can see that the first step provides the best response of a given player interacting with the statistical distribution of the states of the other players if this statistical distribution is assumed to be given by $\mu_{t}$. In contrast, the second step solves a specific fixed point problem in the spirit of the search for fixed points of the best response function.

Usually, the solution of MFGs is transformed into a set of coupled partial differential equations, namely the HJB-FPK equations, which respectively describe the evolution of the value function of the representative element and the density evolution of the group, as shown below:

in which, $u:\mathbb{R}^{n}\times[0,T]\rightarrow\mathbb{R}$ is the value function to guide the agents make decisions; $H:\mathbb{R}^{n}\times\mathbb{R}^{n}\rightarrow\mathbb{R}$ is the Hamiltonian, which describes the physics energy of the system; $\mu(\cdot,t)\in\mathcal{L}(\mathbb{R}^{n})$is the distribution of agents at time t, $f:\mathbb{R}^{n}\times\mathcal{L}(\mathbb{R}^{n})\rightarrow\mathbb{R}^{n}$ denotes the loss during process; and $g:\mathbb{R}^{n}\times\mathcal{L}(\mathbb{R}^{n})\rightarrow\mathbb{R}^{n}$ is the terminal condition, guiding the agents to the final distribution.

Let assumption MFGs Solvability HJB (as shown in Appendix A.1) be in force. Then, for any initial condition $\xi\in L^{2}(\Omega,\mathcal{F}_{0},\mathbb{P};\mathbb{R}^{d})$, the McKean-Vlasov FBSDEs:

for $t\in[0,T]$, with $Y_{T}=g(X_{T},\mathcal{L}(X_{T}))$ as terminal condition, is solvable. Moreover, the flow $(\mathcal{L}(X_{T}))_{0\leq t\leq T}$ given by the marginal distributions of the forward component of any solution is an equilibrium of the MFGs problem associated with the stochastic control problem Eq.(1).

We consider the optimal control step of the formulation of the MFGs problems described earlier. Probabilists have a two-pronged approach to these optimal control problems. We consider that the input $\bm{\mu}=(\mu_{t})_{0\leq t\leq T}$ is deterministic and fixed to search the optimal reaction decision. Then, after the result of fixed decision, the optimal flows of probability measure can be solved. Alternately seeking for the optimal control, the MFGs equilibrium can be finally derived.

NNormalizing Flows (NF), introduced by Tabak & Vanden-Eijnden (2010), enable exact computation of data likelihood through a sequence of invertible mappings. A key feature of NF is its use of arbitrary bijective functions, achieved through stacked reversible transformations. The flow model $R$ consists of a sequence of reversible flows, expressed as $R(x)=r_{1}\circ r_{2}\circ\cdots\circ r_{L}(x)$, where each $r_{i}$ has a tractable inverse and Jacobian determinant.

Our algorithm leverages the volume-preserving property of NF, aligning with the consistency of density flow in MFGs during evolution. This principle is essential for constructing the density flow in the MFGs model.

The connection between MFGs and NF provides inherent advantages. For example, in MFGs, the initial distribution is often represented in a simple analytical form. This parallels NF’s approach, where a simple initial distribution transforms into a more complex one for density estimation. Additionally, one advantage of NF over other generative models is its preservation of total density during transformation, consistent with the MFGs requirement $\int\mu dx=1$. A challenge, however, is that MFGs, unlike Optimal Transport (OT), lacks both initial and terminal density distributions. In MFGs, only an initial distribution exists, and the terminal condition is governed by the terminal value function $g(x,\mu(\cdot,T))$. This complicates framing the problem as a complete density evolution problem.

To address this, we idealize the MFGs model with the assumption that the terminal value function corresponds to an explicitly solvable optimal terminal density. For example, in trajectory planning problems, the terminal value function is often related to the destination, such as $g(x,\mu)=\int_{\mathbb{R}^{d}}-e^{-||x-x_{T}||^{2}}d\mu(x)$. In such cases, we assume the optimal terminal distribution is $\mu(x_{T})=1$, enabling the MFGs problem to be framed with initial and terminal densities. We will show that this assumption is reasonable.

We consider a general class of MFGs problem associated with the stochastic control problem (1), the relative FBSDEs can be represented as in (4), with initial condition $\xi\in L^{2}(\Omega,\mathcal{F}_{0},\mathbb{P};\mathbb{R}^{d})$ and terminal condition $Y_{T}=g(X_{T},\mathcal{L}(X_{T}))$.

Then the solution of Eq.(3) satisfies the following FBSDE:

We apply a temporal discretization to Eq.(4). Given a partition of the time interval$[0,T]:0=t_{0}<t_{1}<\cdots<t_{N}=T$, we consider the simple Euler scheme for$n=1,\cdots,N-1$

and

where

The key in modeling the above FBSDEs is to approximate the value function $x\mapsto u(0,x)$ at $t=t_{0}$, which is $u(0,x)\approx u(0,x|\theta_{0})$; while approximating the function $x\mapsto[\partial_{x}u(x,t)]^{T}\sigma$ at each time step $t=t_{n}$ through a multi-layer feedforward neural network

which represents the adjoint variable $Z_{t}$ of the value function of the element, which is expressed as the product of the gradient and the random variable in the MFGs optimization problem.

Subsequently, all value functions are connected by summing Eq.(7) over $t$. The network uses the generated density flows $\bm{\mu}$ and $W_{t_{n}}$ as inputs and produces the final output $\hat{u}$, approximating $u(x,T)=g(x,\mu(\cdot,T))$. This approximation defines the expected loss function by comparing the maximum likelihoods of the two functions to minimize the difference for $\{x_{i}\}_{i=1}^{N}\sim z=\hat{\mu}_{T}$:

In summary, we reformulate MFGs as MKV FBSDEs (Eq.(5)) and discretize time to establish the relationship between value functions at each time $t$ (Eq.(7)), connecting them via the adjoint variable $Z_{t}$. Next, we parameterize $Z_{t}$ and $u_{0}$, using these relationships to link $u_{T}$ with the terminal condition $g(x,\mu(\cdot,T))$ for maximum likelihood estimation (Eq.(10)) to minimize the final MFGs loss.

Typically, NF methods prioritize density estimation results. In contrast, our approach emphasizes the NF density evolution process, constraining each layer to align with the density evolution in MFGs.

In an NF model, If $r_{i}$ are differentiable and reversible functions, we usually express it like:

To simplify the explanation, we consider a one-dimensional model. During training, $r_{i}(x)$ is represented as a neural network, $r_{i}(x;\phi)$. Multiple $r_{i}$ are combined to obtain the desired function $f$. Typically, training minimizes the negative log-likelihood loss between the final estimated density and the dataset, expressed as

in which there is no need to consider the loss of each $r_{i}$ in the process.

Discretizing the NF construction process reveals that each function corresponds to Euler time discretization. Thus, each sub-function $r_{i}$ transforms the group density $\mu_{i}$ in MFGs into $\mu_{i+1}$. This series of reversible flows represents the time evolution process, which results in that

Additionally, as each sub-function is implemented as a neural network, the loss at each time step can be used to constrain and optimize the sub-functions. The density $\mu_{t}$ can be expressed as $\mu_{t_{n}}=\mathbf{r}_{1,2,\cdots,n}\circ\mu_{0}$, where $\mathbf{r}_{1,2,\cdots,n}=r_{t_{n}}\circ r_{t_{n-1}}\circ\cdots\circ r_{t_{1}}$.

Our way of modeling the limited is to approximate the crowd density of each layer

at each time step $t=t_{n}$ through a NF model consisting of multiple layers of Masked Autoregressive Flow (MAF) and Permute parameterized by $\phi$.

To train the NF, the first step is computing the process loss $l_{\text{HJB}}$. At each time step $t=t_{n}$, the MFGs system satisfies the HJB equation. Section 3.1 describes the value function and its gradient. Thus, samples ${x_{i}}_{i=1}^{M}$ from $\mu_{t}$ at each time step can be used in the HJB equation to compute the loss,

where

Additionally, the NF method must match the terminal density condition, so the terminal loss $l_{\text{T}}$ is also included in the loss calculation.If the terminal condition $g$ is explicitly defined, the corresponding optimal density $\hat{\mu}_{T}(\phi)=\mathbf{f}(x;\Phi)\circ\mu_{0}$ can serve as the target distribution for NF. The negative log-likelihood between $\hat{\mu}T(\phi)$ and $\mu_{T}$ generated by NF is used to compute the terminal loss $l\text{T}$. For ${x_{i}}_{i=1}^{N}\sim\hat{\mu}_{T}(\Phi)$:

In summary, NF, as a generative model, can construct intermediate function compositions and distributions without direct data use, relying solely on the initial distribution $\mu_{0}$ and the terminal distribution $\mu_{T}$, while preserving density consistency. We first compute the terminal density distribution $\mu_{T}$ explicitly and construct an NF to transition from $\mu_{0}$ to $\mu_{T}$. The losses $l_{\text{HJB}}$ (Eq. 15) and $l_{\textit{T}}$ (Eq. 17) constrain NF evolution, ensuring the flow density aligns with the control objectives.

Two processes can be coupled and trained alternately. As NF is a generative model, it can first generate a set of flow density evolution functions along with the corresponding density distributions at each time step. This generated set of density distributions is fixed as the marginal distribution to optimize the value function and gradient under the MKV FBSDE framework. Once the optimal value function for this marginal distribution flow is obtained, it is fixed to update each $r_{n}$ and its corresponding $\mu_{t_{n}}$ in the NF evolution process. This continues until the optimal density distribution flow under the current value function is achieved. This iterative coupled training continues until convergence. Algorithm (1) presents the pseudo-code of the model.

This section presents three examples of solving MFGs using NF-MKV Net, demonstrating its applicability to traffic flow problems, low- and high-dimensional crowd motion problems, and scenarios with obstacles.

To verify volumetric invariance and time continuity, we compare NF-MKV Net with existing MFGs solving methods, including the distribution-based RL-PIDL method proposed by Chen et al. (2023) and the high-dimensional neural network-based APAC-Net by Lin et al. (2021).

Distribution volumetric-invariance. We implement an approximated integral over the dynamics region, widely used in density estimation such as O’Brien et al. (2016). By generating a grid over a specified area, the numerical integration of a specified probability distribution over that area is computed, and the return value should be close to 1. This method verifies the validity of the distribution. Since the approximated integral is a grid-based method, it can only be used in low-dimensional problems. So, when we approximate integral in $d=50$ crowd motion, we use the same process as when showing the high-dimensional trajectories from the experiment above, that is, a projection-like method that accumulates the density distribution function on the other components over the first two components and estimates the density in a 2-dimensional region.

Agents states time-continuity. The Wasserstein distance is generally chosen as the metric for the difference between two density distributions. Laurière et al. (2022) has used the method to assess the difference between distributions. Even without an explicit probability density function, the Wasserstein distance ($W$-dis) can be computed by optimization methods as long as it can be sampled from both distributions.

We set the comparison experiments in traffic flow ($d=1$) and crowd motion ($d=2$ and $50$). The comparison results are shown in Tab.(1) and Fig.(6).

The integral error of the NF-MKV Net method over the dynamic region is below $10^{-2}$ and close to the standard value of $1$, significantly outperforming the other method, which has an error exceeding $0.1$ in low-dimensional problems. In high-dimensional scenarios, the error of the other method is over ten times larger than that of NF-MKV Net, highlighting the distribution volumetric invariance of NF-MKV Net. The NF-MKV Net has the smallest average Wasserstein distance between adjacent time steps among the three methods. This demonstrates smoother evolution and superior agent state time-continuity.

In summary, NF-MKV Net excels in distribution volumetric invariance and agent state time-continuity, making it suitable for solving MFGs problems involving density distributions.