MFAX: JAX-Based MFG Framework

• MFAX is a JAX-based framework for mean-field games that integrates analytic and sample-based methods, enabling high-performance simulation and solver development.
• It unifies white-box access to known transition dynamics with black-box environments, leveraging JAX vectorization and GPU/TPU acceleration for efficient batched mean-field updates.
• The framework offers benchmark environments and reference implementations of RSPG, SPG, and standard RL baselines, supporting common noise, partial observability, and equilibrium evaluation.

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 (Wibault et al., 23 Feb 2026). 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 (Wibault et al., 23 Feb 2026).

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 (Wibault et al., 23 Feb 2026). 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 (Wibault et al., 23 Feb 2026).

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 (Wibault et al., 23 Feb 2026). 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 (Wibault et al., 23 Feb 2026). 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 (Wibault et al., 23 Feb 2026). 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 (Wibault et al., 23 Feb 2026).

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 $z_0\sim p_{z_0},\;z_{t+1}\sim \Xi(\cdot\mid z_t),$0 has entries

$z_0\sim p_{z_0},\;z_{t+1}\sim \Xi(\cdot\mid z_t),$1

then

$z_0\sim p_{z_0},\;z_{t+1}\sim \Xi(\cdot\mid z_t),$2

MFAX’s analytic wrappers implement this push-forward functionally, using JAX vectorized mapping rather than explicit storage of $z_0\sim p_{z_0},\;z_{t+1}\sim \Xi(\cdot\mid z_t),$3 (Wibault et al., 23 Feb 2026).

The expected return conditioned on a mean-field evolution induced by policy $z_0\sim p_{z_0},\;z_{t+1}\sim \Xi(\cdot\mid z_t),$4 is written as

$z_0\sim p_{z_0},\;z_{t+1}\sim \Xi(\cdot\mid z_t),$5

A Mean-Field Nash Equilibrium $z_0\sim p_{z_0},\;z_{t+1}\sim \Xi(\cdot\mid z_t),$6 satisfies

$z_0\sim p_{z_0},\;z_{t+1}\sim \Xi(\cdot\mid z_t),$7

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 (Wibault et al., 23 Feb 2026).

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 (Wibault et al., 23 Feb 2026). On top of this, the framework provides two mean-field wrappers with distinct computational assumptions.

The AnalyticMFWrapper imports $z_0\sim p_{z_0},\;z_{t+1}\sim \Xi(\cdot\mid z_t),$8, $z_0\sim p_{z_0},\;z_{t+1}\sim \Xi(\cdot\mid z_t),$9, $$s_{t+1}\sim\mathcal T(\cdot\mid s_t,a_t,\mu_t,z_t),$$0, and $$s_{t+1}\sim\mathcal T(\cdot\mid s_t,a_t,\mu_t,z_t),$$1, where $$s_{t+1}\sim\mathcal T(\cdot\mid s_t,a_t,\mu_t,z_t),$$2 is the agent’s aggregate observation function, and implements functionalized mean-field push-forward and expectation operators via JAX vectorized mapping (Wibault et al., 23 Feb 2026). The stated purpose is to avoid explicit $$s_{t+1}\sim\mathcal T(\cdot\mid s_t,a_t,\mu_t,z_t),$$3-matrix storage while still exploiting exact transition structure. The SampleBasedMFWrapper instead samples $$s_{t+1}\sim\mathcal T(\cdot\mid s_t,a_t,\mu_t,z_t),$$4 trajectories and approximates $$s_{t+1}\sim\mathcal T(\cdot\mid s_t,a_t,\mu_t,z_t),$$5 through an empirical histogram, allowing black-box environments to be handled in the same code base (Wibault et al., 23 Feb 2026).

Several data structures recur throughout the framework:

• $$s_{t+1}\sim\mathcal T(\cdot\mid s_t,a_t,\mu_t,z_t),$$6 for the mean field.
• $$s_{t+1}\sim\mathcal T(\cdot\mid s_t,a_t,\mu_t,z_t),$$7 for common noise.
• $$s_{t+1}\sim\mathcal T(\cdot\mid s_t,a_t,\mu_t,z_t),$$8 for observations.
• $$s_{t+1}\sim\mathcal T(\cdot\mid s_t,a_t,\mu_t,z_t),$$9 and $\mu_t\in\Delta_{\mathcal S}$0 as recurrent state and time variables for RSPG.
• Policy networks $\mu_t\in\Delta_{\mathcal S}$1, implemented with JAX plus Haiku or Flax, returning logits $\mu_t\in\Delta_{\mathcal S}$2 (Wibault et al., 23 Feb 2026).

The analytic environment loop is specified as follows: reset $\mu_t\in\Delta_{\mathcal S}$3; for $\mu_t\in\Delta_{\mathcal S}$4, evaluate the policy $\mu_t\in\Delta_{\mathcal S}$5, update $\mu_t\in\Delta_{\mathcal S}$6 via mf_update, compute the reward matrix $\mu_t\in\Delta_{\mathcal S}$7, and sample $\mu_t\in\Delta_{\mathcal S}$8 and $\mu_t\in\Delta_{\mathcal S}$9 (Wibault et al., 23 Feb 2026). 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 (Wibault et al., 23 Feb 2026).

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

$r_t \;=\; R\bigl(s_t,a_t,\mu_t,z_t\bigr).$0

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 (Wibault et al., 23 Feb 2026).

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 (Wibault et al., 23 Feb 2026). In the benchmark summary, the baselines named are IPPO, RIPPO, and M-OMD (Wibault et al., 23 Feb 2026).

RSPG is described as the first history-aware HSM for settings involving public information (Wibault et al., 23 Feb 2026). 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, $a_t \sim \pi(\cdot\mid s_t,\mu_t,z_t)$0 environments are rolled out in parallel, tracing $a_t \sim \pi(\cdot\mid s_t,\mu_t,z_t)$1, and values are computed through the backward recursion

$a_t \sim \pi(\cdot\mid s_t,\mu_t,z_t)$2

The objective is

$a_t \sim \pi(\cdot\mid s_t,\mu_t,z_t)$3

and the update is

$a_t \sim \pi(\cdot\mid s_t,\mu_t,z_t)$4

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 (Wibault et al., 23 Feb 2026).

The implementation details emphasize three design principles. First, environments are vectorized over $a_t \sim \pi(\cdot\mid s_t,\mu_t,z_t)$5 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 (Wibault et al., 23 Feb 2026). The training loop is correspondingly compact:

$r_t \;=\; R\bigl(s_t,a_t,\mu_t,z_t\bigr).$1

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 (Wibault et al., 23 Feb 2026).

5. Environment support, observability, and usage patterns

MFAX provides benchmark environments spanning Linear–Quadratic, Beach Bar, and Macroeconomics with heterogeneous agents (Wibault et al., 23 Feb 2026). These environments are paired with scripts, defaults, and an API intended for reproducible end-to-end experiments, including training, exploitability evaluation, and algorithmic comparison (Wibault et al., 23 Feb 2026).

The framework explicitly supports common noise, partial observability, and multiple initial mean-field distributions (Wibault et al., 23 Feb 2026). In the RSPG setting, the recurrent state $a_t \sim \pi(\cdot\mid s_t,\mu_t,z_t)$6, time variable $a_t \sim \pi(\cdot\mid s_t,\mu_t,z_t)$7, common-noise history $a_t \sim \pi(\cdot\mid s_t,\mu_t,z_t)$8, and observation history $a_t \sim \pi(\cdot\mid s_t,\mu_t,z_t)$9 are all part of the modeling interface (Wibault et al., 23 Feb 2026). This support is significant because the motivating paper states that previous HSMs had not been scaled to partially observable settings (Wibault et al., 23 Feb 2026).

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 (Wibault et al., 23 Feb 2026). The provided example includes the following sequence:

$r_t \;=\; R\bigl(s_t,a_t,\mu_t,z_t\bigr).$2

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 (Wibault et al., 23 Feb 2026). The paper also states that RSPG solves, for the first time, a macroeconomics MFG with heterogeneous agents, common noise, and history-aware policies (Wibault et al., 23 Feb 2026).

6. Performance characteristics and empirical findings

The framework reports mean-field update times on an NVIDIA L40S GPU for four systems (Wibault et al., 23 Feb 2026):

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 (Wibault et al., 23 Feb 2026). The same source states that MFAX demonstrates order-of-magnitude speed-ups over prior MFG libraries, specifically OpenSpiel and MFGLib (Wibault et al., 23 Feb 2026).

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 (Wibault et al., 23 Feb 2026). 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 (Wibault et al., 23 Feb 2026).

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 (Wibault et al., 23 Feb 2026). 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 (Wibault et al., 23 Feb 2026). 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 (Wibault et al., 23 Feb 2026).

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 (Wibault et al., 23 Feb 2026).

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 (Wibault et al., 23 Feb 2026).

MFAX is publicly available at the repository specified by the source paper: https://github.com/CWibault/mfax (Wibault et al., 23 Feb 2026).