PaperlayerView exact revision on arXiv

Source-linked AI summary

Towards Scaling Reinforcement Learning to Massive Populations: Learning Mean-Field Representations

Aditya Makkar, Benjamin Unger, Jeongyeol Kwon, Mathieu Laurière, Eugene Vinitsky, Yonathan Efroni

arXiv:2609.02928v1cs.MAcs.LGstat.ML

TL;DR

Large state-action spaces make full population distributions difficult to learn even in mean-field RL, motivating a lower-dimensional representation. The paper develops an offline model-based framework to learn that representation and reports improved reward prediction and equilibrium quality in a one-step routing benchmark, with open limitations in computation and coverage theory.

  • Problem

    In large state-action spaces, representing the full mean-field distribution can be intractable even when it captures population effects relevant to rewards and transitions.

    Source: §1 · pp. 1–2

  • Method

    The paper formalizes mean-field rewards and transitions through an unknown low-dimensional population statistic and learns a model-based offline policy using logged data.

    Source: §1.1 · p. 2, §3.3 · p. 7

  • Results

    In the one-step routing benchmark, learned mean-field representations improve reward prediction and planned-policy equilibrium quality, outperform a structurally incomplete raw mean, approach representation oracles as resources grow, and outperform a parameter-matched full-population-law model.

    Source: §5.4 · p. 13, §7 · p. 15

  • Takeaways & Limitations

    Under matched parameter and optimization budgets, learning a population representation is more effective in the routing benchmark than modeling the complete population law directly.

    Source: §5.4 · p. 13

    The algorithm is statistically motivated but computationally inefficient, and the weakest coverage conditions for offline equilibrium learning remain unknown.

    Source: §7 · p. 15

Abstract

from arXiv · show

Modern multi-agent systems are increasingly deployed at scale over large populations of agents in settings such as ad-auctions, traffic routing, and recommendation systems. The dominant approach in such settings is to optimize each agent's policy independently, treating the other agents as part of a fixed single-agent environment rather than modeling the population dynamics. In many large-population systems, the dynamics depend on an aggregate summary of the population rather than the identity of any individual. Mean-field RL exploits such structure, providing a principled framework that models each agent's environment as an explicit function of the population distribution. However, in large state-action spaces or high-dimensional control problems, modeling the population distribution is itself intractable. How can we design a scalable framework for high-dimensional control problems with large populations? This work explores this question from the perspective of representation learning. We introduce a mean-field RL framework in which the rewards and transition dynamics depend on the population only through an unknown low-dimensional aggregate statistic. We then study this framework in the offline setting and design a provable approach that learns a near-optimal policy by learning a low-dimensional representation. Motivated by real-life supply-chain optimization problems, we design a one-step routing game to test the hypothesis that learning a low-dimensional population representation improves reward prediction and Nash gap estimation relative to baselines that don't exploit this structure. We show that under a fixed neural-network parameter count and optimization budget, learning a low-dimensional population representation improves reward prediction and the equilibrium quality of the resulting policies.

1 Introduction

Large-population mean-field models avoid the exponential joint state space, but the full population distribution can still be intractable in high-dimensional state-action spaces. This work learns an unknown low-dimensional population representation offline and develops a model-based approach with a finite-sample Nash-gap guarantee.

1 Introduction

  • The key challenge is that even mean-field distributions can be too large to model, although rewards and transitions may depend on them only through lower-dimensional statistics.The paper asks how to exploit this structure for learning that scales to high-dimensional decision problems with many agents.Source: §1 · pp. 1–2
  • The routing illustration reduces 400 state-action combinations to a 56-dimensional representation of the network edges used by each route.The feature map records which directed edges a route uses.Source: Figure 1 · p. 2

1.1 Contributions

  • The framework makes rewards and transitions depend on the population through an unknown low-dimensional statistic that must be learned from representative-agent trajectories and finite population samples.Unlike ordinary state or state-action representations, this statistic summarizes the surrounding population and changes with its policy.Source: §1.1 · pp. 2–3
  • The model-based offline algorithm learns the representation, reward, and transition models and proves a finite-sample Nash-gap guarantee under adaptive coverage.The guarantee concerns the offline setting.Source: §1.1 · p. 3
  • In a one-step congestion-routing game, the learned representation outperforms baselines that underuse population structure and approaches an oracle with privileged mean-field access as data grow.The comparison uses a fixed neural-network parameter count and optimization budget.Source: §1.1 · p. 3

2 Mean-Field Games with Low-Dimensional Representations

The section formulates mean-field games using a low-dimensional population statistic in place of the full population law. It defines representative-agent values and equilibrium, including why values may depend on future representations and when an equilibrium exists.

2 Mean-Field Games with Low-Dimensional Representations

  • A mean-field game models a representative agent interacting with the population through its law, but that full distribution may remain too large when the state-action space is large.The section presents low-dimensional aggregate statistics as an additional structure that rewards and transitions can use instead.Source: §2 · p. 3

2.1 Mean-Field Game as a Representative-Agent Control Problem

  • Against a fixed population flow, a unilateral deviation leaves that flow unchanged, so the representative agent faces a single-agent control problem.The population policy must be optimal against the flow it induces for equilibrium.Source: §2.1 · p. 4

2.2 Low-Dimensional Mean-field representations

  • The framework assumes rewards and transitions use the population law only through a representation of dimension d, ideally much smaller than |X||A|−1.The full-law mean-field game is recovered by taking the representation dimension to be |X||A|.Source: §2.2 · p. 4

2.3 Value Functions Depend on Future Representations

  • An action value cannot in general be determined by the current representation alone, because different population policies can induce different future representation paths.The value functions therefore depend on the future representation suffix.Source: §2.3 · p. 5

2.4 Equilibrium and Well-Posedness

  • A Nash equilibrium is a policy that is a best response to itself, and one exists under finite state and action spaces, bounded features, and continuous reward and transition maps.The section also notes that a similar approach could address social optimum, or mean-field control.Source: §2.4 · p. 5

3 Offline Mean-Field RL Framework

The offline framework assumes an unknown mean-field representation belongs to an accessible function class and uses logged representative-agent trajectories alongside population samples. It fits candidate models and plans over a confidence set using an interval Nash-gap objective.

3 Offline Mean-Field RL Framework

  • The offline framework treats the population representation as unknown but belonging to a function class accessible to the learner.The section introduces a model class and empirical criterion for learning the model.Source: §3 · p. 6

3.1 Candidate mean-field models

  • The candidate class includes a representation, reward model, and transition model, with assumptions including realizability and continuity.Realizability requires the true model to belong to the candidate class.Source: §3.1 · p. 6

3.2 Offline data

  • Each of N independent logged episodes pairs a representative-agent trajectory with K independent population snapshots over time.Behavior population policies are drawn from a distribution over policies, and the snapshots support empirical representation estimates.Source: §3.2 · pp. 6–7

3.3 Learning algorithm

  • For each candidate model, the algorithm estimates representations, fits reward and transition components, and retains near-minimizers in a confidence set before planning with an interval Nash-gap objective.It evaluates deviations optimistically and self-play pessimistically; if the true model is in the confidence set, the objective upper bounds the true Nash gap. Support stability also lets the realizable model assign positive probability to transitions possible at the exact representation.Source: §3.3 · p. 7

4 Provable Representation Learning in Mean-Field RL

The offline procedure learns population representations and reward and transition models, then plans over a confidence set using optimistic deviation values and pessimistic self-play values. Its Nash-gap guarantee relies on transferring logged-data errors to target evaluations and requires coverage of self-play and a good approximate response, not every unilateral deviation.

4 Provable Representation Learning in Mean-Field RL

  • For each model, the procedure fits a representation, reward, and transition components, retains near-minimizers in a confidence set, and plans over plausible models.It scores candidate population policies using optimistic deviation values and a pessimistic self-play value.Source: §4 · p. 7

4.1 Model confidence from data

  • With probability at least 1 −δ, the finite-class confidence set contains the true model and is contained in ΘsN,K.The likelihood-based argument contributes an N −1 term, while evaluating models at a finite-population representation estimate contributes a K−1 term.Source: §4.1 · p. 8

4.2 Value transfer through Bellman residuals

  • For θ ∈Θs, the transfer radius Rs bounds the value error |Jθ(π; η) −J(π; η)| and supports the planner’s one-sided interval errors.The radius uses the Bellman-transfer coefficient κ to transfer logged residual control to target evaluations.Source: §4 · p. 7, §4.2 · p. 9

4.3 Nash-gap guarantee

  • The adaptive Nash-gap guarantee requires data to evaluate a low-gap policy’s self-play value and at least one good approximate response, not every unilateral deviation.The planner combines optimistic deviation values with pessimistic self-play values when scoring policies.Source: §4 · p. 7, §4.3 · p. 10
  • Finite K creates representation-context plug-in error even with infinitely many episodes, while the confidence and Nash-gap bounds have no explicit dependence on state-space cardinality.The N −1/2 term reflects logged representative trajectories; the K−1/2 term reflects finite population samples.Source: §4.3 · p. 10

5 Experiments

The experiments test learned population representations in a one-step congestion-routing game, comparing reward prediction and the true Nash gaps of policies planned from fitted reward models. Across the high-resource setting, the learned mean field improves on incomplete population encodings and a parameter-matched full-law model, while approaching representation oracles.

5.1 The routing game environment

  • The routing benchmark is a one-step congestion game: agents choose routes, and their rewards depend on edge congestion from the population flow.The full population law has 400 state-action coordinates, while the true representation records 56 edge loads.Source: §5 · p. 10, §5.1 · p. 11

5.2 The reward models

  • The learned model encodes each population sample before averaging, while the raw mean can miss joint frequencies needed to recover outer-edge loads.The six reward models share the focal encoder, reward head, optimizer, data, and checkpoint-selection rule.Source: §5.2 · p. 12

5.3 Evaluation metrics

  • For each fitted reward model, candidate policies are selected by fitted-model Nash residual, then evaluated with the true reward; reward error is measured over population-weighted predictions.The reported primary curves average five paired training seeds and show 95% Student-t half-widths.Source: §5.3 · p. 13

5.4 Results

  • As K grows, the models separate by recoverable population information, and the finite-K oracle approaches the infinite-population oracle.At K = 1, encode-then-average has no operation-order advantage, as expected, and the data show this sanity check.Source: §5.4 · p. 13
  • At (N, K) = (200000, 1024), infinite-population and finite-K oracles ≈ learned mean field < monolithic raw mean < single-agent on both population RMSE and Nash gap.The full-population-law baseline also performs worse on average than the learned mean-field model under the matched parameter and optimization budget.Source: §5.4 · p. 13, Figure 3 · p. 14

6 Related Work

The related work spans mean-field RL, offline RL and game learning, coverage and distribution shift, and representation learning. This paper distinguishes its setting by learning an unknown population representation from fixed offline data and handling shifts in both population flow and unilateral deviations.

Mean-field RL

  • Prior mean-field RL work studies equilibrium computation and guarantees under prescribed embeddings or model classes; this paper learns the low-dimensional representation from offline data.The cited offline mean-field precedents do not focus on mean-field representation learning.Source: §6 · pp. 14–15

Offline reinforcement learning

  • Offline Markov-game methods use pessimism and coverage notions for policy values and responses, whereas this setting adds population-flow shift to the coverage problem.Related single-agent offline RL methods and unilateral-coverage work are reviewed as context for these distinctions.Source: §6 · pp. 14–15

Coverage and distribution shift in offline mean-field learning

  • Offline game learning must control candidate-policy values and unilateral deviations, while population-policy changes also shift the mean-field flow and representation path.The paper uses a Bellman-transfer condition for a low-gap population policy and one approximate response.Source: §6 · p. 15

Representation learning

  • Unlike standard low-rank RL representations, the mean-field representation here depends on the population policy being evaluated; deviations are assessed with that policy-induced sequence fixed.The paper situates its model-learning component closest to representation learning for low-rank MDPs.Source: §6 · p. 15

7 Conclusion and Future Work

The paper formalizes offline mean-field RL with unknown low-dimensional population representations, proposes a model-based algorithm, and establishes a finite-sample Nash-gap guarantee under stated conditions. Its routing benchmark supports the usefulness of learned representations, while the conclusion identifies computational efficiency and other extensions as open problems.

7 Conclusion and Future Work

  • In the one-step routing benchmark, learned representations improve reward prediction and policy equilibrium quality, approach representation oracles as resources grow, and outperform a parameter-matched full-law model in the high-resource regime.The learned model also performs better than the structurally incomplete raw mean.Source: §7 · p. 15
  • The finite-sample guarantee assumes realizability and adaptive Bellman-transfer coverage for the offline equilibrium-learning framework.The paper formalizes the representation-learning problem and proposes a model-based algorithm.Source: §7 · p. 15
  • The proposed algorithm is statistically motivated but computationally inefficient; scalable algorithms with provably comparable guarantees remain open.The authors also identify lower-bound analysis and an online extension with new exploration criteria as open directions.Source: §7 · p. 15

Notation Table

The appendix develops technical support for offline equilibrium learning: it distinguishes representation paths, establishes equilibrium existence, and builds statistical control of fitted models. Its arguments include support stability and the effect of finite-population representation estimates.

Notation Table

  • The notation distinguishes true and candidate models, representative-agent and population policies, lifted measures, empirical risks, confidence sets, and Bellman-transfer quantities.The appendix also defines state and state-action flows and candidate representations used in planning.Source: Notation Table · p. 19

A.1.1 One-step mean-field representations are not Bellman-sufficient

  • A two-stage construction shows that matching current state, action, and representation inputs need not make mean-field representations Bellman-sufficient.The example yields different action values despite coincident current inputs.Source: §A.1.1 · pp. 19–20

A.1.2 Existence of Nash equilibria

  • A fixed point of the representation best-response correspondence yields a population policy with zero Nash gap.The proof constructs the correspondence from optimal occupancy flows under a frozen representation path.Source: §A.1.2 · pp. 21–22

A.1.3 Support and distance facts

  • Support stability establishes that transition support does not vary with the representation within the specified feature ball.The appendix also records distance bounds for probability distributions.Source: §A.1.3 · p. 22

A.2 Exact statistical confidence event

  • The statistical argument controls prediction error at sampled representations and then incorporates finite-population representation error.The empirical criterion combines negative log likelihood for transitions with squared reward loss.Source: §A.2 · p. 23, §A.2.2 · p. 24

A.2.1 Episode quantities

  • The episode-level analysis defines likelihood differences and prediction losses for transitions and rewards, using the observed transition probabilities’ positive lower bound.The proof treats the likelihood decomposition in extended reals when a candidate assigns zero probability to an observed transition.Source: §A.2.1 · p. 23

A.2.2 Likelihood and prediction loss

  • A likelihood-to-Hellinger bound supports simultaneous prediction control over the candidate model class with high probability.The proof applies this bound across episode-stage pairs and adds reward errors.Source: §A.2.2 · pp. 23–24

A.2.3 The realizable model

  • The realizable-model analysis bounds the true model’s score with high probability, and the finite-population analysis accounts for prediction error from snapshot-based representations.The appendix separately identifies finite empirical risk for the true model.Source: §A.2.3 · p. 24, §A.2.4 · p. 25

A.2.5 Confidence theorem

The appendix develops confidence-set guarantees and interval-based control of Nash gaps, connecting Bellman residuals to value-transfer bounds. It also distinguishes two Bellman-coverage measures by how they support interval proofs.

A.2.5 Confidence theorem

  • Under finite Θ and Assumptions 1–3, Theorem 3 states a confidence event with probability at least 1−δ.The supplied statement does not expose the bound inside that event.Source: §A.2.5 · p. 25

A.3 Interval gaps and deterministic oracle inequalities

  • When the true model lies in the candidate set, the interval objective upper-bounds the true Nash gap for every policy.Source: §A.3 · p. 26

A.4 Bellman residuals and value-transfer radii

  • Bellman-residual identities and prediction-loss control underpin one-sided interval bounds through Bellman-transfer radii.Source: §A.4 · pp. 27–28

A.5 Proof of the adaptive Nash-gap theorem and rate corollary

  • The adaptive Nash-gap proof combines the deterministic oracle inequality with Bellman-transfer bounds for the confidence set.Source: §A.5 · p. 29

A.6 Scalar and L2 Bellman coverage

  • C measures full target L2 Bellman-residual energy, whereas κ measures the signed residual integral tied to value error and is sharper for the interval proof.Source: §A.6 · p. 30

B Routing Experiment Details

The empirical study isolates whether a task-relevant population representation can be learned from finite offline snapshots, using a one-step routing setting. Its contexts and population snapshots are generated to examine generalization across population laws and resource sizes.

B.1 Scope and relation to the theory

  • The empirical study uses H=1, deterministic reward labels, continuous neural reward models, and planning in one fitted reward game; it does not implement the confidence-set algorithm or validate the finite-sample bound.Source: §B.1 · p. 30

B.2 Environment and offline context distribution

  • The routing population is represented by edge loads across three network layers, which sum to one within each layer for every population law.Source: §B.2 · p. 30
  • Product-form origin–destination demand makes the raw population mean identify the complete demand law through its origin and destination marginals.Source: §B.2 · p. 31
  • The study varies population size N from 1,000 to 200,000 and snapshot size K from 1 to 1,024, using nested prefixes of shared data streams.Source: §B.2 · p. 31

B.3 Reward-model architectures and fitting

The experiment compares reward-model architectures under matched parameter counts and a common fitting procedure. Checkpoint selection uses validation MSE rather than final population RMSE or true Nash gap.

B.3 Reward-model architectures and fitting

  • The models share a reward head, while the full-law branch uses a 400-entry histogram and the single-agent network is widened to keep parameter counts comparable.Source: §B.3 · p. 31
  • The full-law and learned mean-field models have 37,297 and 37,209 parameters, respectively, a 0.237% difference.The full-law branch takes a 400-entry empirical histogram as input.Source: §B.3 · p. 31
  • Checkpoint selection minimizes validation MSE, while final population RMSE and true Nash gap are unavailable to the trainer and selector.Source: §B.3 · p. 31
  • Every newly fitted checkpoint reached its minimum validation MSE at the final scheduled evaluation, giving these fits a common fixed update budget.Source: §B.3 · p. 31

B.4 Exact reward evaluation

The evaluation computes reward predictions over focal categories and population inputs, then plans in each fitted reward game using a common candidate portfolio. Candidate selection uses fitted-model residuals, with true rewards reserved for final policy evaluation.

B.4 Exact reward evaluation

  • The evaluation enumerates all 400 focal categories and computes the learned model’s expected population encoding and the full-law model’s exact 400-entry input.The learned and full-law models are not given the edge-load oracle unless their definitions include it.Source: §B.4 · p. 32
  • The outer expectation over contexts uses eight independently Owen-scrambled Sobol blocks shared across models and seeds.Source: §B.4 · p. 32

B.5 Fitted-game planning and exact controls

  • Each fitted reward model uses the same 12-candidate optimization portfolio, with candidates stopping when fitted-model residual reaches 10^-5 or at 5,000 iterations.Source: §B.5 · p. 32
  • The selected policy minimizes fitted-model residual, and its true-reward evaluation occurs only after selection.Source: §B.5 · p. 32
  • The exact-reward control applies the same planner to the analytic reward to isolate optimization error from reward-model error.Source: §B.5 · p. 32

B.6 Statistical protocol and reproducibility

The protocol computes uncertainty across five training seeds after averaging each seed’s metric over evaluation context samples. Reported ± values are 95% Student-t half-widths over paired seeds, with QMC integration uncertainty assessed separately.

B.6 Statistical protocol and reproducibility

  • Each seed’s metric is averaged over evaluation context samples before uncertainty is computed across seeds.The protocol defines five seeds, indexed 0 through 4.Source: §B.6 · pp. 32–33
  • Within each seed, models share contexts, focal samples, ordered population streams, nested resource prefixes, and fixed quasi–Monte Carlo evaluation context samples.Source: §B.6 · p. 32
  • Reported ± values are 95% Student-t half-widths over five paired training seeds, not integration error bars.QMC integration uncertainty is assessed separately from variation across training seeds.Source: §B.6 · p. 33
Loading 2609.02928v1…