optimiz-rs/examples/notebooks/mean_field_games_tutorial.ipynb

Setup

First, let's import the necessary libraries and set up the environment.

In [ ]:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
import seaborn as sns

# Set plotting style
sns.set_style('whitegrid')
plt.rcParams['figure.figsize'] = (12, 8)
plt.rcParams['font.size'] = 11

print("✓ Libraries loaded")

Example 1: Congestion Game

Consider a simple congestion game where agents want to move from an initial distribution to a target distribution, but experience congestion costs.

Problem Formulation

• Hamiltonian: H(x,p) = \frac{1}{2}|p|^2 (quadratic control cost)
• Running cost: f(x,m) = \lambda m(x,t) (congestion penalty)
• Terminal cost: g(x,m(T)) = \frac{1}{2}(x - x_{\text{target}})^2
• Initial distribution: m_0(x) = \mathcal{N}(0.3, 0.05^2)

Numerical Solution

We solve this using finite difference methods with:

• Spatial domain: \Omega = [0, 1]
• Time horizon: T = 1.0
• Grid: n_x = 100, n_t = 100
• Viscosity: \nu = 0.01

In [ ]:

# Problem parameters
nx = 100  # Spatial grid points
nt = 100  # Time steps
T = 1.0   # Time horizon
nu = 0.01 # Viscosity
lambda_congestion = 0.5  # Congestion penalty
x_target = 0.7  # Target location

# Spatial and temporal grids
x = np.linspace(0, 1, nx)
t = np.linspace(0, T, nt)
dx = x[1] - x[0]
dt = t[1] - t[0]

# Initial distribution: Gaussian centered at 0.3
m0 = np.exp(-((x - 0.3)**2) / (2 * 0.05**2))
m0 /= np.sum(m0) * dx  # Normalize

# Plot initial distribution
plt.figure(figsize=(10, 4))
plt.plot(x, m0, 'b-', linewidth=2, label='Initial distribution $m_0(x)$')
plt.axvline(x_target, color='r', linestyle='--', label=f'Target: $x={x_target}$')
plt.xlabel('Space $x$')
plt.ylabel('Density')
plt.title('Initial Agent Distribution')
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

print(f"Grid: {nx} × {nt}")
print(f"dx = {dx:.4f}, dt = {dt:.4f}")

Fixed-Point Iteration Algorithm

The algorithm proceeds as follows:

1. Initialize: Start with uniform distribution m^{(0)}(x,t) = 1
2. Iterate until convergence:
   • Solve HJB backward: -\partial_t u^{(k)} - \nu \Delta u^{(k)} + \frac{1}{2}|\nabla u^{(k)}|^2 = \lambda m^{(k-1)}
   • Solve FP forward: \partial_t m^{(k)} - \nu \Delta m^{(k)} - \text{div}(m^{(k)} \nabla u^{(k)}) = 0
   • Update: m^{(k)} \leftarrow \alpha m^{(k)} + (1-\alpha) m^{(k-1)} (relaxation)
3. Check convergence: \|m^{(k)} - m^{(k-1)}\|_{L^2} < \epsilon

This is implemented using:

• Upwind finite differences for first derivatives
• Central differences for second derivatives
• Implicit time stepping for stability

In [ ]:

def solve_hjb(m, u_T):
    """Solve HJB equation backward in time"""
    u = np.zeros((nx, nt))
    u[:, -1] = u_T  # Terminal condition
    
    for n in range(nt-2, -1, -1):
        for i in range(1, nx-1):
            # Laplacian (central difference)
            u_xx = (u[i+1, n+1] - 2*u[i, n+1] + u[i-1, n+1]) / dx**2
            
            # Hamiltonian with upwind scheme
            u_x_plus = (u[i+1, n+1] - u[i, n+1]) / dx
            u_x_minus = (u[i, n+1] - u[i-1, n+1]) / dx
            H = 0.5 * min(u_x_plus**2, u_x_minus**2)  # Upwind
            
            # Running cost
            f = lambda_congestion * m[i, n]
            
            # Update (implicit Euler)
            u[i, n] = u[i, n+1] - dt * (nu * u_xx - H + f)
        
        # Boundary conditions (Neumann)
        u[0, n] = u[1, n]
        u[-1, n] = u[-2, n]
    
    return u

def solve_fp(u, m0):
    """Solve Fokker-Planck equation forward in time"""
    m = np.zeros((nx, nt))
    m[:, 0] = m0  # Initial condition
    
    for n in range(nt-1):
        for i in range(1, nx-1):
            # Laplacian
            m_xx = (m[i+1, n] - 2*m[i, n] + m[i-1, n]) / dx**2
            
            # Velocity field
            u_x = (u[i+1, n] - u[i-1, n]) / (2*dx)
            v = u_x  # For quadratic Hamiltonian: H_p = p
            
            # Upwind for advection
            if v > 0:
                flux_diff = v * (m[i, n] - m[i-1, n]) / dx
            else:
                flux_diff = v * (m[i+1, n] - m[i, n]) / dx
            
            # Update (forward Euler)
            m[i, n+1] = m[i, n] + dt * (nu * m_xx - flux_diff)
            m[i, n+1] = max(m[i, n+1], 0)  # Non-negativity
        
        # Boundary conditions
        m[0, n+1] = m[1, n+1]
        m[-1, n+1] = m[-2, n+1]
        
        # Normalize
        m[:, n+1] /= (np.sum(m[:, n+1]) * dx)
    
    return m

print("✓ Solver functions defined")

In [ ]:

# Fixed-point iteration
max_iter = 50
tol = 1e-5
relax = 0.5

# Initialize
m_old = np.ones((nx, nt)) / nx
errors = []

print("Running fixed-point iteration...")
for iter in range(max_iter):
    # Terminal condition for HJB
    u_T = 0.5 * (x - x_target)**2
    
    # Solve HJB backward
    u = solve_hjb(m_old, u_T)
    
    # Solve FP forward
    m_new = solve_fp(u, m0)
    
    # Check convergence
    error = np.sqrt(np.sum((m_new - m_old)**2)) / np.sqrt(np.sum(m_old**2))
    errors.append(error)
    
    if iter % 5 == 0:
        print(f"  Iteration {iter:3d}: error = {error:.6f}")
    
    if error < tol:
        print(f"✓ Converged in {iter+1} iterations")
        break
    
    # Relaxation
    m_old = relax * m_new + (1 - relax) * m_old

# Plot convergence
plt.figure(figsize=(10, 4))
plt.semilogy(errors, 'b-', linewidth=2)
plt.axhline(tol, color='r', linestyle='--', label=f'Tolerance: {tol}')
plt.xlabel('Iteration')
plt.ylabel('Relative L² error')
plt.title('Convergence of Fixed-Point Iteration')
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

Results Visualization

Now let's visualize the solution: value function u(x,t) and distribution m(x,t).

In [ ]:

# Create meshgrid for plotting
X, T = np.meshgrid(x, t)

# Plot distribution evolution
fig = plt.figure(figsize=(16, 6))

# 3D surface plot of distribution
ax1 = fig.add_subplot(121, projection='3d')
surf1 = ax1.plot_surface(X, T, m_new.T, cmap=cm.viridis, alpha=0.8)
ax1.set_xlabel('Space $x$')
ax1.set_ylabel('Time $t$')
ax1.set_zlabel('Density $m(x,t)$')
ax1.set_title('Distribution Evolution')
fig.colorbar(surf1, ax=ax1, shrink=0.5)

# 3D surface plot of value function
ax2 = fig.add_subplot(122, projection='3d')
surf2 = ax2.plot_surface(X, T, u.T, cmap=cm.plasma, alpha=0.8)
ax2.set_xlabel('Space $x$')
ax2.set_ylabel('Time $t$')
ax2.set_zlabel('Value $u(x,t)$')
ax2.set_title('Value Function')
fig.colorbar(surf2, ax=ax2, shrink=0.5)

plt.tight_layout()
plt.show()

In [ ]:

# Temporal snapshots
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
time_indices = [0, nt//2, nt-1]
times = [0.0, T/2, T]

for ax, idx, time_val in zip(axes, time_indices, times):
    ax.plot(x, m_new[:, idx], 'b-', linewidth=2, label='Distribution')
    ax.axvline(x_target, color='r', linestyle='--', alpha=0.5, label='Target')
    ax.set_xlabel('Space $x$')
    ax.set_ylabel('Density')
    ax.set_title(f'$m(x, t={time_val:.1f})$')
    ax.legend()
    ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

print("✓ Solution computed and visualized")

Analysis

From the results, we observe:

1. Agent Migration: Agents move from initial position (0.3) toward target (0.7)
2. Congestion Effect: The distribution spreads out due to congestion penalty \lambda m
3. Nash Equilibrium: The solution represents a Nash equilibrium where no agent can improve their cost by deviating
4. Value Function: Shows the optimal cost-to-go from each position at each time

The numerical method successfully captures:

• Mass conservation: \int_\Omega m(x,t)\,dx = 1 for all t
• Non-negativity: m(x,t) \geq 0
• Convergence to equilibrium within 30-50 iterations

Conclusion

This notebook demonstrated:

• Mathematical formulation of Mean Field Games
• Numerical solution via finite difference methods
• Fixed-point iteration for coupled HJB-FP system
• Visualization and analysis of equilibrium solutions

Further Extensions

The optimizr Rust library provides additional capabilities:

• Primal-dual methods for faster convergence
• Higher dimensions (2D, 3D spatial domains)
• Non-quadratic Hamiltonians
• State constraints and obstacle problems
• Parallel computation for large-scale problems

Related Methods

For mean-field variational inference via optimal transport, see the work by Jiang, Chewi, and Pooladian (2023), which develops polyhedral optimization methods in the Wasserstein space for computing mean-field approximations.

---

Next Steps:

• Example 2: Multi-population games
• Example 3: Mean field type control
• Example 4: Comparison with particle methods