Mean-Field Sampling for Cooperative Multi-Agent Reinforcement Learning

Reinforcement Learning (RL) has become a popular framework to solve sequential decision making problems in unknown environments and has achieved tremendous success in a wide array of domains such as playing the game of Go (Silver et al., 2016), robotic control (Kober et al., 2013), and autonomous driving (Kiran et al., 2022; Lin et al., 2023). A key feature of most real-world systems is their uncertain nature, and thus, RL has emerged as a powerful tool for learning optimal policies for multi-agent systems to operate in unknown environments (Kim & Giannakis, 2017; Zhang et al., 2021; Lin et al., 2024; Anand & Qu, 2024). While early RL works focused on the single-agent setting, recently, multi-agent RL (MARL) has achieved impressive success for many applications, such as coordinating robotic swarms (Preiss et al., 2017), self-driving vehicles (DeWeese & Qu, 2024), real-time bidding (Jin et al., 2018), ride-sharing (Li et al., 2019), and stochastic games (Jin et al., 2020).

Despite growing interest in multi-agent RL (MARL), extending RL to multi-agent settings poses significant computational challenges due to the curse of dimensionality (Sayin et al., 2021): even if the individual agents’ state or action spaces are small, the global state space or action space can take values from a set with size that is exponentially large as a function of the number of agents. For example, even model-free RL algorithms such as temporal difference (TD) learning (Sutton et al., 1999) or tabular $Q$-learning require computing and storing a $Q$-function (Bertsekas & Tsitsiklis, 1996) that is as large as the state-action space, which is exponential in the number of agents. These scalability issues have been observed in the multi-agent reinforcement learning literature in a variety of settings (Blondel & Tsitsiklis, 2000; Papadimitriou & Tsitsiklis, 1999; Littman, 1994). To overcome this challenge, Tan (1997) proposes to perform independent $Q$-learning, where each agent models other agents as part of the environment to decouple decision-making and reduce the state-action space; however, this simplification often fails to capture a fundamental feature of MARL: inter-agent interactions.

MARL is fundamentally difficult because agents in the real-world not only interact with the environment but also with each other (Shapley, 1953). An exciting line of work that addresses this intractability is mean-field MARL (Lasry & Lions, 2007; Yang et al., 2018; Gu et al., 2021, 2022; Hu et al., 2023). The mean-field approach assumes all agents are homogeneous in their state and action spaces, enabling their interactions to be approximated by a two-agent setting: here, each agent interacts with a representative “mean agent” which evolves as the empirical distribution of states of all other agents. Under these assumptions, mean-field MARL allows learning of optimal policies with a sample complexity that is polynomial in the number of agents. However, if the number of homogeneous agents is large, storing a polynomially-large $Q$ table (where the polynomial’s degree depends on the size of the state space of a single agent) may remain infeasible. In this paper, we study the cooperative setting– where agents work collaboratively to maximize a structured global reward–and ask: Can we design an efficient and approximately optimal MARL algorithm for policy-learning in a cooperative multi-agent system as the number of agents increases asymptotically?

Contributions. We answer this question affirmatively. Our key contributions are outlined below.

Subsampling Algorithm. We propose SUBSAMPLE-MFQ to address the challenge of MARL with a large number of local agents. We model the problem as a Markov Decision Process (MDP) with a global agent and $n$ local agents. SUBSAMPLE-MFQ selects $k\leq n$ local agents to learn a deterministic policy $\hat{\pi}^{\mathrm{est}}_{k}$. It does this by applying mean-field value iteration on the $k$-local-agent subsystem to learn $\hat{Q}_{k}^{\mathrm{est}}$, which can be viewed as a smaller $Q$ function. It then deploys a stochastic policy ${\pi}^{\mathrm{est}}_{k}$, where the global agent samples $k$ local agents uniformly at each step and uses $\hat{\pi}^{\mathrm{est}}_{k}$ to determine its action, while each local agent samples $k-1$ other local agents and uses $\hat{\pi}^{\mathrm{est}}_{k}$ to determine its action.

Sample Complexity and Theoretical Guarantee. As the number of local agents increases, the size of $\hat{Q}_{k}$ scales polynomially with $k$, rather than polynomially with $n$ as in mean-field MARL. Analogously, when the size of the local agent’s state space grows, the size of $\hat{Q}_{k}$ scales exponentially with $k$, rather than exponentially with $n$, as in traditional $Q$-learning. The key analytic technique underlying our results is a novel MDP sampling result. Through it, we show that the performance gap between ${{\pi}_{k}^{\mathrm{est}}}$ and the optimal policy ${\pi^{*}}$ is $\tilde{O}(1/\sqrt{k})$. The choice of $k$ reveals a fundamental trade-off between the size of the $Q$-table and the optimality of ${\pi}_{k}^{\mathrm{est}}$. When $k=O(\log n)$, SUBSAMPLE-MFQ is the first centralized MARL algorithm to achieve a polylogarithmic run-time in $n$, representing an exponential speedup over the previously best-known polytime mean-field MARL methods, while maintaining a decaying optimality gap.

Numerical Simulations. We evaluate the effectiveness of SUBSAMPLE-MFQ in a bounding box problem, and a Gaussian squeeze problem. Our experiments reveal a monotonic improvement in the learned policies as $k\to n$, while providing a substantial speedup over mean-field $Q$-learning.

While our results are theoretical in nature, it is our hope that SUBSAMPLE-MFQ will lead to a further exploration of sampling in Markov games, and inspire the development of practical algorithms for multi-agent settings.

In this section, we propose the SUBSAMPLE-MFQ algorithm to overcome the polynomial (in $n$) sample complexity of mean-field value iteration and the exponential (in $n$) sample complexity of traditional $Q$-learning. In the algorithm, the global agent randomly samples a subset of local agents $\Delta\subseteq[n]$ such that $|\Delta|=k$, for $k\leq n$. It ignores all other local agents $[n]\setminus\Delta$, and performs value iteration to learn the $Q$-function $\hat{Q}_{k,m}^{\mathrm{est}}$ and policy $\hat{\pi}_{k,m}^{\mathrm{est}}$ for this surrogate subsystem of $k$ local agents, where $m$ is the number of samples used to update the $Q$-functions’ estimates of the unknown system. When $|\mathcal{Z}_{l}|^{k-1}<k^{|\mathcal{Z}_{l}|}$, the algorithm uses traditional value-iteration (Algorithm 1), and when $|\mathcal{Z}_{l}|^{k-1}>k^{|\mathcal{Z}_{l}|}$, it uses mean-field value iteration (Algorithm 2). We denote the surrogate reward gained by this subsystem at each time step by $r_{\Delta}:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$, where

To convert the optimality of each agent’s action in the $k$ local-agent subsystem to an approximate optimality guarantee on the full $n$-agent system, we use a randomized policy $\tilde{\pi}^{\mathrm{est}}_{k,m}$ (Algorithm 3), where the global agent samples $\Delta\in\mathcal{U}\binom{[n]}{k}$ at each time-step to derive an action $a_{g}\leftarrow\hat{\pi}_{k,m}^{\mathrm{est}}(s_{g},s_{\Delta})$, and where each local $i$ agent samples $k-1$ other local agents $\Delta_{i}$ to derive a action $\hat{\pi}_{k,m}^{\mathrm{est}}(s_{g},s_{i},s_{\Delta_{i}})$. Finally, Theorem 3.4 shows that the policy $\tilde{\pi}_{k,m}^{\mathrm{est}}$ converges to the optimal policy $\pi^{*}$ as $k\to n$. We first present Algorithms 2 and 1 (SUBSAMPLE-MFQ: Learning) and Algorithm 3 (SUBSAMPLE-MFQ: Execution), which we describe below. For this, a crucial characterization is the notion of the empirical distribution function:

Let $\mu_{k}(\mathcal{Z}_{l})\coloneq\left\{\frac{b}{k}|b\in\{0,\dots,k\}\right\}^{|\mathcal{S}_{l}|\times|\mathcal{A}_{l}|}$ be the space of $|\mathcal{S}_{l}|\times|\mathcal{A}_{l}|$-length vectors where each entry is an element of $\{0,\frac{1}{k},\frac{2}{k},\dots,1\}$. Here, $F_{z_{\Delta}}\in\mu_{k}(\mathcal{Z}_{l})$ is the proportion of agents in the $k$-local-agent subsystem at each state.

Algorithms 2 and 1 (Offline learning). Let $m\in\mathbb{N}$ denote the sample size for the learning algorithm with sampling parameter $k\leq n$. When $|\mathcal{Z}_{l}|^{k-1}\leq k^{|\mathcal{Z}_{l}|}$, we iteratively learn the optimal $Q$-function for a subsystem with $k$-local agents denoted by $\hat{Q}_{k,m}^{t}:\mathcal{S}_{g}\times\mathcal{S}_{l}^{k}\times\mathcal{A}_{g}\times\mathcal{A}_{l}^{k}\to\mathbb{R}$, which is initialized to 0. At time $t$, we update $\hat{Q}_{k,m}^{t+1}(s_{g},{s_{\Delta}},a_{g},a_{\Delta})={\tilde{\mathcal{T}}}_{k,m}\hat{Q}_{k,m}^{t}(s_{g},s_{\Delta},a_{g},a_{\Delta})$, where ${\tilde{\mathcal{T}}}_{k,m}$ is the empirically adapted Bellman operator in Equation 8.

When $|\mathcal{Z}_{l}|^{k-1}>k^{|\mathcal{Z}_{l}|}$, we learn the optimal mean-field $Q$-function for a $k$ local agent system, denoted by $\hat{Q}_{k,m}^{t}:\mathcal{S}_{g}\times\mathcal{S}_{1}\times\mu_{k-1}(\mathcal{Z}_{l})\times\mathcal{A}_{g}\times\mathcal{A}_{l}\to\mathbb{R}$, which is initialized to 0. At time $t$, we update $\hat{Q}_{k,m}^{t+1}(s_{g},s_{1},F_{z_{\tilde{\Delta}}},a_{1},a_{g})=\hat{\mathcal{T}}_{k,m}\hat{Q}_{k,m}^{t}(s_{g},s_{1},F_{z_{\tilde{\Delta}}},a_{1},a_{g})$, where $\hat{\mathcal{T}}_{k,m}$ is the empirically adapted mean-field Bellman operator in Equation 9.

Since the system is unknown, $\tilde{\mathcal{T}}_{k,m}$ and $\hat{\mathcal{T}}_{k,m}$ cannot directly perform the Bellman update and instead use $m$ random samples $s_{g}^{j}\sim P_{g}(\cdot|s_{g},a_{g})$ and $s_{i}^{j}\sim P_{l}(\cdot|s_{i},s_{g},a_{i})$ for each $j\in[m]$, $i\in\Delta$ to approximate the system:

$\hat{Q}_{k,m}^{t}$ depends on $s_{\Delta}$ and $a_{\Delta}$ through $F_{z_{\Delta}}$, and $\tilde{\mathcal{T}}_{k,m}$/$\hat{\mathcal{T}}_{k,m}$ are $\gamma$-contractive. So, Algorithms 1 and 2 apply value iteration with their respective Bellman operators $\tilde{\mathcal{T}}$ and $\hat{\mathcal{T}}$ until $\hat{Q}_{k,m}$ converges to a fixed point satisfying $\tilde{\mathcal{T}}_{k,m}\hat{Q}_{k,m}^{\mathrm{est}}=\hat{Q}_{k,m}^{\mathrm{est}}$ and $\hat{\mathcal{T}}_{k,m}\hat{Q}_{k,m}^{\mathrm{est}}=\hat{Q}_{k,m}^{\mathrm{est}}$, yielding equivalent deterministic policies $\hat{\pi}^{\mathrm{est}}_{k,m}(s_{g},{s_{\Delta}})$ and $\hat{\pi}^{\mathrm{est}}_{k,m}(s_{g},s_{1},F_{s_{\tilde{\Delta}}})$:

For $\hat{\pi}_{k,m}^{\mathrm{est}}(s_{g},s_{i},s_{\Delta\setminus i})\!=\!a_{g}^{*},a_{i}^{*},a_{\Delta\setminus i}^{*}$, let $[\hat{\pi}_{k,m}^{\mathrm{est}}(s_{g},s_{\Delta})]_{g}\!=\!a_{g}^{*}$ and $[\hat{\pi}_{k,m}^{\mathrm{est}}(s_{g},s_{i},s_{\Delta\setminus i})]_{l}\!=\!a_{i}^{*}$.

Algorithm 3 (Online implementation). In Algorithm 3, (SUBSAMPLE-MFQ: Execution) the global agent samples local agents $\Delta(t)\sim\mathcal{U}\binom{[n]}{k}$ at each time step to derive action $a_{g}(t)=[\hat{\pi}_{k,m}^{\mathrm{est}}(s_{g},s_{\Delta}(t))]_{g}$, and each local agent $i$ samples other local agents $\Delta_{i}(t)\sim\mathcal{U}{[n]\setminus i\choose k-1}$ to derive action $a_{i}(t)=[\hat{\pi}_{k,m}^{\mathrm{est}}(s_{g},s_{i},s_{\Delta}(t))]_{l}$. The system then incurs a reward $r(s,a)$. This procedure of first sampling agents and then applying $\hat{\pi}^{\mathrm{est}}_{k,m}$ is denoted by a stochastic policy $\tilde{\pi}_{k,m}^{\mathrm{est}}(a|s)$, where $\tilde{\pi}_{k,m}^{\mathrm{est}}(a_{g}|s)$ is the global agent’s action distribution and $\tilde{\pi}_{k,m}^{\mathrm{est}}(a_{l}|s)$ is the local agent’s action distribution:

The agents then transition to their next states.