---
title: 'MFAX: JAX-Based MFG Framework'
url: https://www.emergentmind.com/topics/mfax
type: topic
---

# MFAX: JAX-Based MFG Framework

MFAX is a fully JAX-based framework for specifying, simulating, and solving large-population Mean-Field Games (MFGs), introduced alongside Recurrent Structural Policy Gradient (RSPG) in work on partially observable mean field games with common noise [2602.20141]. Its stated scope is to unify “white-box” access to known individual transition kernels with “black-box” environments, exploit JAX vectorization and GPU/TPU acceleration for batched mean-field updates and expectation computations, provide wrappers for finite-horizon MFGs with common noise and partial observability, supply reference implementations of Hybrid Structural Methods (HSMs) and standard RL baselines, and deliver benchmark environments and end-to-end scripts for training, exploitability measurement, and algorithmic comparison [2602.20141].

## 1. Conceptual scope within mean-field game computation

MFAX is positioned in the algorithmic setting of MFGs in which a large population induces deterministic population dynamics at scale, while uncertainty enters through aggregate shocks or common noise [2602.20141]. In the formulation used with MFAX, an MFG with common noise is defined by the tuple
\[
\bigl(p_{\mu_0},\,p_{z_0},\,\mathcal S,\,\mathcal A,\,\mathcal Z,\,\mathcal T,\,\Xi,\,R,\,\gamma\bigr),
\]
with finite state, action, and common-noise spaces, a common-noise Markov chain, state transitions depending on the current mean field, and rewards depending on the agent state, action, mean field, and common noise [2602.20141].

The framework is designed to support both finite-horizon and infinite-horizon settings, although the environment wrappers are described specifically as supporting finite-horizon MFGs with common noise, partial observability, and multiple initial mean-field distributions [2602.20141]. Its objectives are therefore not confined to environment simulation alone. MFAX also serves as a research platform for solver development, benchmark comparison, and exploitability-based evaluation.

A central distinction encoded in MFAX is between analytic and sample-based treatment of mean-field evolution. This division reflects two complementary computational regimes. In one regime, known transition dynamics permit analytic mean-field updates and exact expectation operators. In the other, empirical histograms from sampled trajectories approximate the evolving mean field for black-box environments [2602.20141]. This suggests that MFAX is intended to bridge structural MFG methods and more generic RL-style experimentation within a single software stack.

## 2. Mathematical model and equilibrium objects

The framework adopts the standard MFG state evolution under common noise:
\[
z_0\sim p_{z_0},\;z_{t+1}\sim \Xi(\cdot\mid z_t),
\]
and
\[
s_{t+1}\sim\mathcal T(\cdot\mid s_t,a_t,\mu_t,z_t),
\]
where \(\mu_t\in\Delta_{\mathcal S}\) is the mean field [2602.20141]. Agents choose actions according to
\[
a_t \sim \pi(\cdot\mid s_t,\mu_t,z_t)
\]
and receive reward
\[
r_t \;=\; R\bigl(s_t,a_t,\mu_t,z_t\bigr).
\]
These definitions make explicit that policy dependence may include the current mean field and the current common-noise state [2602.20141].

For analytic environments, the mean-field evolution is given by the push-forward
\[
\mu_{t+1}(s') \;=\;\int_{s,a}\!\mathcal T\bigl(s'\!\mid s,a,\mu_t,z_t\bigr)\,\pi\bigl(a\!\mid s,\mu_t,z_t\bigr)\,\mu_t(s)\,da\,ds.
\]
In vector-matrix form, if \(\boldsymbol\mu_t\) is the \(|\mathcal S|\)-vector of \(\mu_t\), and \(\mathbf A^\pi_{\mu_t,z_t}\in\mathbb R^{|\mathcal S|\times|\mathcal S|}\) has entries
\[
\bigl[\mathbf A^\pi_{\mu,z}\bigr]_{s,s'}=\mathbb E_{a\sim\pi(\cdot\mid s,\mu,z)}[\mathcal T(s'\mid s,a,\mu,z)],
\]
then
\[
\boldsymbol\mu_{t+1}  = \bigl(\mathbf A^\pi_{\mu_t,z_t}\bigr)^\top\, \boldsymbol\mu_t.
\]
MFAX’s analytic wrappers implement this push-forward functionally, using JAX vectorized mapping rather than explicit storage of \(\mathbf A\) [2602.20141].

The expected return conditioned on a mean-field evolution induced by policy \(\pi'\) is written as
\[
J_{\mathrm{evol}(\pi,\pi')  = \mathbb E\Bigl[\sum_{t=0}^\infty \gamma^t R\bigl(s_t,a_t,\mu'_t,z_t\bigr)\Bigr].
\]
A Mean-Field Nash Equilibrium \(\pi^*\) satisfies
\[
\pi^* \in\arg\max_\pi  J_{\mathrm{evol}\bigl(\pi,\pi^*\bigr).
\]
These definitions situate MFAX squarely within equilibrium computation and approximate equilibrium evaluation in MFGs, rather than only policy optimization in the single-agent RL sense [2602.20141].

## 3. Software architecture and core abstractions

MFAX organizes environments through a layered architecture. At the base is a **BaseEnv** implementing deterministic single-state steps and rewards, together with terminal logic, in either white-box or black-box form [2602.20141]. On top of this, the framework provides two mean-field wrappers with distinct computational assumptions.

The **AnalyticMFWrapper** imports \(\mathcal T\), \(R\), \(\Xi\), and \(\mathcal O\), where \(\mathcal O\) is the agent’s aggregate observation function, and implements functionalized mean-field push-forward and expectation operators via JAX vectorized mapping [2602.20141]. The stated purpose is to avoid explicit \(\mathbf A\)-matrix storage while still exploiting exact transition structure. The **SampleBasedMFWrapper** instead samples \(N\) trajectories and approximates \(\mu_t\) through an empirical histogram, allowing black-box environments to be handled in the same code base [2602.20141].

Several data structures recur throughout the framework:

- \(\boldsymbol\mu_t\in \mathbb R^{|\mathcal S|}\) for the mean field.
- \(\mathbf z_{0:t}\) for common noise.
- \(o_{0:t}\) for observations.
- \(h_t\) and \(d_t\) as recurrent state and time variables for RSPG.
- Policy networks \(\pi_\theta\), implemented with JAX plus Haiku or Flax, returning logits \(\boldsymbol\Pi_t\) [2602.20141].

The analytic environment loop is specified as follows: reset \(\mu_0,z_0,o_0\); for \(t=0\!:\!T-1\), evaluate the policy \(\boldsymbol\Pi_t=\pi_\theta(s,\mu_t,z_t)\), update \(\mu_{t+1}\) via `mf_update`, compute the reward matrix \(\mathbf R_t\), and sample \(z_{t+1}\) and \(o_{t+1}\) [2602.20141]. In implementation terms, MFAX uses `jax.vmap` to vectorize over states, actions, and parallel environments, and `jax.lax.scan` to unroll time in recurrent policies and backward recursions [2602.20141].

A representative analytic stepping pattern is given in the framework’s pseudocode:

```python
def mf_pushforward(mu, z, policy_logits):
  # policy_logits: [S,A] log‐probs
  # call transition_fn to yield A_matrix_fn: (S,) -> (S,A) → S  expectation operator
  return A_matrix_fn(mu, policy_logits)  # returns new mu of shape [S]

def step_env(carry, t):
  mu, z, h, d = carry
  logits, h_new = policy.apply(theta, s_grid, mu, z, h, d)
  mu_next = jax.lax.stop_gradient(mf_pushforward(mu, z, logits))
  r = jax.lax.stop_gradient(reward_fn_grid(s_grid, logits, mu, z))
  z_next = sample_noise(z)
  o = obs_fn_grid(mu, z)
  return (mu_next, z_next, h_new, d+1), (logits, r, o)
```

This architecture makes explicit that MFAX is not a monolithic simulator. It is a framework for composing environment models, mean-field operators, recurrent policies, and solver logic in a functional JAX style [2602.20141].

## 4. Algorithms implemented in MFAX

MFAX includes reference implementations of Hybrid Structural Methods, notably Structural Policy Gradient (SPG) and Recurrent Structural Policy Gradient (RSPG), and also standard RL baselines including Independent PPO [2602.20141]. In the benchmark summary, the baselines named are IPPO, RIPPO, and M-OMD [2602.20141].

RSPG is described as the first history-aware HSM for settings involving public information [2602.20141]. It extends structural policy-gradient ideas to partially observable settings by differentiating through exact per-step returns while sampling only the common noise. In the implementation, \(E\) environments are rolled out in parallel, tracing \(\{\mu_t,z_t,o_{0:t},h_t\}\), and values are computed through the backward recursion
\[
\mathbf v_t  = \bigl(\boldsymbol\Pi_t\odot R_t\bigr)\,\mathbf1 + \gamma\,\mathbf A^\pi_{\mu_t,z_t,o_{0:t}\,\mathbf v_{t+1}.
\]
The objective is
\[
J(\theta) =\mathbb E_{\mu_0,z_0}\bigl[\boldsymbol\mu_0^\top\,\mathbf v_0\bigr],
\]
and the update is
\[
\theta\leftarrow\theta+\alpha\nabla_\theta J,
\]
with gradients flowing through the exact expected-return recursion while being blocked through `mf_update` and the environment reward function by `stop_gradient` in JAX [2602.20141].

The implementation details emphasize three design principles. First, environments are vectorized over \(E\) particles. Second, temporal unrolling and backward Bellman-style recursion are handled with `jax.lax.scan`. Third, non-differentiable or intentionally excluded components—common-noise sampling, mean-field update, and random reward noise—are passed through `stop_gradient` [2602.20141]. The training loop is correspondingly compact:

```python
for _ in range(num_iters):
  _, traj = jax.lax.scan(step_env, init, None, length=T);
  vals = backward_eval(traj);
  grads = jax.grad(lambda θ: expected_return(vals));
  apply `optax.apply_updates`.
```

The framework therefore operationalizes a distinction between exact expectation computation over state-action structure and Monte Carlo treatment of common noise. This is the defining algorithmic feature of the HSM implementations as presented in the source paper [2602.20141].

## 5. Environment support, observability, and usage patterns

MFAX provides benchmark environments spanning Linear–Quadratic, Beach Bar, and Macroeconomics with heterogeneous agents [2602.20141]. These environments are paired with scripts, defaults, and an API intended for reproducible end-to-end experiments, including training, exploitability evaluation, and algorithmic comparison [2602.20141].

The framework explicitly supports common noise, partial observability, and multiple initial mean-field distributions [2602.20141]. In the RSPG setting, the recurrent state \(h_t\), time variable \(d_t\), common-noise history \(\mathbf z_{0:t}\), and observation history \(o_{0:t}\) are all part of the modeling interface [2602.20141]. This support is significant because the motivating paper states that previous HSMs had not been scaled to partially observable settings [2602.20141].

The usage pattern described for analytic and sample-based environments is straightforward. One can instantiate either `AnalyticMeanFieldEnv` or `SampleBasedMeanFieldEnv`, define a recurrent policy such as `RNNPolicy`, choose a solver such as `RSPG`, `SPG`, or `IPPO`, train for multiple epochs, and then evaluate exploitability [2602.20141]. The provided example includes the following sequence:

```python
import mfax
from mfax.envs import AnalyticMeanFieldEnv, SampleBasedMeanFieldEnv
from mfax.algos import RSPG, SPG, IPPO

# 1) Instantiate environment
env = AnalyticMeanFieldEnv("LinearQuadratic", config=...)
# or:
env = SampleBasedMeanFieldEnv("BeachBar", config=...)

# 2) Define policy network
policy = RNNPolicy(hidden_size=64, action_dim=env.action_dim)

# 3) Set up RSPG optimizer
solver = RSPG(policy, env,
             lr=3e-4, gamma=0.99, E=256, T=30)

# 4) Train
for epoch in range(1000):
    loss, metadata = solver.step()
    print(f"Epoch {epoch}, loss {loss}, exploit {metadata['exploit']}")

# 5) Evaluate
exploit = solver.evaluate_exploitability(num_seqs=100)
print("Estimated exploitability:", exploit)
```

A separate macroeconomic example uses `MacroEconEnv`, an analytic mean-field environment with `grid=(100,100)` and `T=128`, optimized by RSPG with `lr=1e-4`, `gamma=0.96`, and `E=128` over `epochs=5000` [2602.20141]. The paper also states that RSPG solves, for the first time, a macroeconomics MFG with heterogeneous agents, common noise, and history-aware policies [2602.20141].

## 6. Performance characteristics and empirical findings

The framework reports mean-field update times on an NVIDIA L40S GPU for four systems [2602.20141]:

| System | Mean-field update time |
|---|---:|
| MFAX (analytic) | 0.000298 s/update |
| MFAX (sample-based) | 0.000435 s/update |
| OpenSpiel (C++) | 0.00544 s/update |
| MFGLib (Python) | 0.358 s/update |

These timings are presented as environment mean-field update time benchmarks rather than full end-to-end solver runtimes [2602.20141]. The same source states that MFAX demonstrates order-of-magnitude speed-ups over prior MFG libraries, specifically OpenSpiel and MFGLib [2602.20141].

For exploitability versus wall-clock on an NVIDIA L40S over 10 seeds, the reported qualitative findings are that RSPG and SPG converge an order of magnitude faster—“minutes vs. hours”—than RL baselines such as IPPO, Recurrent IPPO, and M-OMD, and that RSPG attains lower exploitability, described as closest to zero, than SPG and other baselines across Linear–Quadratic, Beach Bar, and Macroeconomics [2602.20141]. The abstract likewise states that by leveraging known transition dynamics, RSPG achieves state-of-the-art performance as well as an order-of-magnitude faster convergence [2602.20141].

The qualitative benchmark descriptions focus on policy behavior under history dependence. In Beach Bar, RSPG and Recurrent IPPO capture anticipatory exit on bar-closure. In Macroeconomics, history-aware RSPG agents anticipate terminal date by adjusting consumption profile, driving endogenous interest rate dynamics [2602.20141]. These observations indicate that MFAX is intended not only for scalar optimization metrics such as exploitability, but also for inspecting mean-field evolution and behaviorally interpretable responses to public information.

## 7. Research significance and positioning

MFAX occupies a specific niche in the MFG software and methods landscape. It is neither purely an exact solver library nor purely a generic RL benchmark suite. Instead, it combines JAX-accelerated analytic mean-field operators, sample-based approximation pathways, recurrent policy support, and structural policy-gradient algorithms in a unified framework [2602.20141]. This combination is especially relevant for settings with common noise and partial observability, where structural methods had previously not been scaled according to the motivating paper [2602.20141].

Its stated contributions can be summarized in four points drawn directly from the source. MFAX delivers JAX-accelerated environment wrappers for tractable analytic or sample-based mean-field updates in MFGs with partial observability. It provides RSPG and SPG as white-box HSMs together with standardized RL baselines including IPPO, RIPPO, and M-OMD. It includes scripts, defaults, and an API for benchmark reproduction in Linear–Quadratic, Beach Bar, and heterogeneous Macroeconomics MFGs. It demonstrates strong empirical performance in exploitability and qualitative agent behavior, together with substantial speed advantages over prior libraries [2602.20141].

A notable implication is that MFAX formalizes a workflow in which known transition structure can be exploited without abandoning modern accelerator-oriented software practice. A second implication is that recurrent, history-aware policy optimization in MFGs can be embedded directly into the same environment abstractions used for analytic or sample-based mean-field simulation. Both implications are consistent with the paper’s broader claim that MFAX supports the first history-aware HSM in settings involving public information [2602.20141].

MFAX is publicly available at the repository specified by the source paper: `https://github.com/CWibault/mfax` [2602.20141].

Source: https://www.emergentmind.com/topics/mfax