Robust $Q$-learning for mean-field control under Wasserstein uncertainty in common noise

Abstract: In this article, we present a robust $Q$-learning algorithm for discrete-time mean-field control problems under Wasserstein uncertainty in the common noise law. The algorithm combines a quantization-and-projection scheme with a Wasserstein dual reformulation on the common-noise space. We establish its convergence together with finite-time iteration bounds for both synchronous and asynchronous learning schemes. Numerical experiments on systemic risk and epidemic models compare the asynchronous implementation with an idealized Bellman iteration, illustrate the robustness-performance tradeoff under common-noise misspecification, and report the observed convergence behavior of the asynchronous $Q$-learning algorithm.

• The paper develops a robust Q-learning algorithm that optimizes mean-field control under Wasserstein noise uncertainty using a min-max Bellman framework.
• It employs quantization-projection and duality techniques to transform the high-dimensional robust expectation into a tractable one-dimensional optimization.
• Empirical evaluations in systemic risk and epidemic models demonstrate improved performance under adverse noise conditions compared to non-robust policies.

Robust Q-Learning for Mean-Field Control with Wasserstein Common Noise Uncertainty

Introduction and Problem Formulation

This paper develops a robust $Q$-learning algorithm for mean-field control (MFC) problems with common noise, where uncertainty in the law of the common noise is modeled via a Wasserstein ambiguity set. The setting is a discrete-time cooperative agent population, where the dynamics and rewards are influenced by both individual (idiosyncratic) and shared (common) noise sources, with mean-field interaction arising through population distributions. Unlike conventional MFC works, the law of the common noise is not known precisely; instead, only a reference distribution and Wasserstein radius are given, necessitating a distributionally robust control formulation.

The authors formulate a robust MFC objective as a min-max problem, seeking policies that maximize the worst-case expected cumulative reward across all laws within a Wasserstein ball around the reference common-noise distribution. Key technical challenges arise from (i) the infinite-dimensionality of the lifted state-control-measure space even when the original state and action spaces are finite, and (ii) the need for robustification under common noise law misspecification.

Robust Bellman Operator and Q-Function Characterization

The main theoretical construct is a robust dynamic programming principle on the space of probability measures, yielding a Bellman–Isaacs operator that computes, for each lifted state-policy pair $(\mu, \pi)$, the maximum over policies and minimum over admissible common noise laws. This operator takes the form

$$Q^*(\mu, \pi) = \overline r(\mu, \pi) + \beta \inf_{p \in \mathcal{B}^0_{m,q}(\hat{p}_{\varepsilon^0})} \int_{E^0} \sup_{\pi' \in \Pi} Q^*(\overline{\operatorname{F}}(\mu, \pi, e^0), \pi') p(de^0)$$

where $\mathcal{B}^0_{m,q}$ is the Wasserstein ball, $\overline r$ is the lifted reward, and $\overline{\operatorname{F}}$ the lifted transition operator. When the uncertainty radius is zero, this reduces to the non-robust MFC Bellman operator. Existence and uniqueness of a fixed point $Q^*$ is established under standard Lipschitz and boundedness conditions.

To make $Q$-learning feasible, the authors introduce (a) a quantization-projection technique that yields finite grids for measure-valued state/policy spaces, and (b) a dual formulation for the minimization over Wasserstein balls, which enables reducing the robust expectation to an optimization over a one-dimensional dual variable using the cost transform.

Robust $Q$-Learning Algorithm and Convergence

Building on these constructs, a tabular robust $Q$-learning scheme is formulated for the quantized grid, capturing both synchronous and asynchronous learning. The Bellman target in each update is efficiently computed using the dual Wasserstein representation, without explicit access to the true worst-case common noise law. The core stochastic iteration can be expressed as

$(\mu, \pi)$0

where $(\mu, \pi)$1 is the target value function (max over next policies), $(\mu, \pi)$2 is the cost-transformed function, and $(\mu, \pi)$3 are sampled from the reference common noise law.

Rigorous convergence guarantees are established:

• Almost sure convergence to the optimal $(\mu, \pi)$4 (up to discretization error) holds for both synchronous and asynchronous variants under standard diminishing-stepsize conditions.
• Non-asymptotic finite-sample complexity bounds are given, showing that $(\mu, \pi)$5 iterations suffice to achieve $(\mu, \pi)$6 error at high probability.
• Explicit high-probability bounds and stochastic approximation rates leveraging Azuma–Hoeffding-type concentration for martingale difference updates are provided, matching known tight rates for $(\mu, \pi)$7-learning.

The analysis demonstrates that the quantization error can be uniformly controlled and separated from the sample complexity error, so algorithmic performance can be predicted precisely.

Numerical Study: Robustness and Convergence

The methodology is empirically validated in three classes of mean-field control scenarios:

• Systemic risk in financial networks: Population-level capital buffer control under aggregate shock and contagion, under common-noise misspecification.
• Epidemic control in SIS/SEIR models: Social planner optimizes intervention via distancing policies subject to adverse transmission regime shifts.

In all cases, the asynchronous robust $(\mu, \pi)$8-learning algorithm is benchmarked against an idealized finite-grid Bellman iteration.

Convergence diagnostics

The convergence rate of the asynchronous $(\mu, \pi)$9-learning is numerically validated in various robustness regimes. The median and quantiles of the error $$Q^*(\mu, \pi) = \overline r(\mu, \pi) + \beta \inf_{p \in \mathcal{B}^0_{m,q}(\hat{p}_{\varepsilon^0})} \int_{E^0} \sup_{\pi' \in \Pi} Q^*(\overline{\operatorname{F}}(\mu, \pi, e^0), \pi') p(de^0)$$0 between the learned $$Q^*(\mu, \pi) = \overline r(\mu, \pi) + \beta \inf_{p \in \mathcal{B}^0_{m,q}(\hat{p}_{\varepsilon^0})} \int_{E^0} \sup_{\pi' \in \Pi} Q^*(\overline{\operatorname{F}}(\mu, \pi, e^0), \pi') p(de^0)$$1-function and fixed-point reference are shown to follow the predicted $$Q^*(\mu, \pi) = \overline r(\mu, \pi) + \beta \inf_{p \in \mathcal{B}^0_{m,q}(\hat{p}_{\varepsilon^0})} \int_{E^0} \sup_{\pi' \in \Pi} Q^*(\overline{\operatorname{F}}(\mu, \pi, e^0), \pi') p(de^0)$$2 rate over multiple independent seeds.

Figure 1: Error $$Q^*(\mu, \pi) = \overline r(\mu, \pi) + \beta \inf_{p \in \mathcal{B}^0_{m,q}(\hat{p}_{\varepsilon^0})} \int_{E^0} \sup_{\pi' \in \Pi} Q^*(\overline{\operatorname{F}}(\mu, \pi, e^0), \pi') p(de^0)$$3 for asynchronous Q-learning versus the idealized fixed-point, for selected Wasserstein uncertainty radii.

Systemic risk and epidemic results

The robust $$Q^*(\mu, \pi) = \overline r(\mu, \pi) + \beta \inf_{p \in \mathcal{B}^0_{m,q}(\hat{p}_{\varepsilon^0})} \int_{E^0} \sup_{\pi' \in \Pi} Q^*(\overline{\operatorname{F}}(\mu, \pi, e^0), \pi') p(de^0)$$4-learning agent demonstrably outperforms the non-robust agent under common-noise misspecification (i.e., when the law of the shocks shifts away from the training distribution), as evidenced by an increase in reward under the adverse law for moderate robustness radii, with only mild performance degradation under the nominal law.

Figure 2: Systemic risk scenario—mean reward as a function of Wasserstein robustness radius and law misspecification parameter. Moderate robustness yields significantly improved adverse-law performance.

Figure 3: SIS epidemic control—mean reward profiles closely match the idealized Bellman benchmarks; robust policies achieve improved worst-case outcomes under misspecified transmission regimes.

Figure 4: SEIR (susceptible-exposed-infectious-recovered) epidemic model—robust Q-learning remains close to optimal even for long horizons $$Q^*(\mu, \pi) = \overline r(\mu, \pi) + \beta \inf_{p \in \mathcal{B}^0_{m,q}(\hat{p}_{\varepsilon^0})} \int_{E^0} \sup_{\pi' \in \Pi} Q^*(\overline{\operatorname{F}}(\mu, \pi, e^0), \pi') p(de^0)$$5.

Theoretical and Practical Implications

The robust $$Q^*(\mu, \pi) = \overline r(\mu, \pi) + \beta \inf_{p \in \mathcal{B}^0_{m,q}(\hat{p}_{\varepsilon^0})} \int_{E^0} \sup_{\pi' \in \Pi} Q^*(\overline{\operatorname{F}}(\mu, \pi, e^0), \pi') p(de^0)$$6-learning framework developed here provides a practical paradigm for learning in MFC regimes subject to population-wide regime shifts or risk factors, with explicit robustness against environment/model uncertainty. Incorporation of Wasserstein ambiguity is computationally tractable via duality, and convergence and sample complexity can be rigorously quantified.

Practically, this enables robust control and reinforcement learning in large-population systems where aggregate shocks are difficult to model precisely—from finance to epidemiology and networked systems.

Theoretically, the results bridge the gap between distributionally robust RL in finite-state MDPs and infinite-dimensional mean-field stochastic control, extending rigorous learning guarantees to robust mean-field settings with common noise. The decoupling of quantization/projection error and stochastic error gives clarity on how algorithmic and modeling choices trade off.

Future Directions

Extensions to high-dimensional continuous state-action spaces via function approximation or neural networks would further broaden the impact, as would generalizations to non-cooperative mean-field games, dynamic uncertainty radius adaptation, or online estimation of the Wasserstein uncertainty set in nonstationary environments. Further algorithmic improvements may be achieved with variance reduction, sample-efficient exploration, or batch/offline data utilization.

Conclusion

This work introduces a robust $$Q^*(\mu, \pi) = \overline r(\mu, \pi) + \beta \inf_{p \in \mathcal{B}^0_{m,q}(\hat{p}_{\varepsilon^0})} \int_{E^0} \sup_{\pi' \in \Pi} Q^*(\overline{\operatorname{F}}(\mu, \pi, e^0), \pi') p(de^0)$$7-learning scheme for mean-field control with Wasserstein uncertainty in common noise, achieving provable convergence, finite-sample guarantees, and demonstrated empirical robustness in several high-impact domains. The approach synthesizes dynamic programming on lifted measure spaces with tractable robust reinforcement learning, offering both a blueprint and solid theoretical foundation for future research into robust large-population control under uncertainty (2606.20356).