Beyond Conservatism: Diffusion Policies in
Offline Multi-agent Reinforcement Learning

Offline reinforcement learning (RL), commonly referred to as batch RL, aims to learn efficient policies exclusively from previously gathered data without interacting with the environment Lange et al. (2012); Levine et al. (2020). Since the agent has to sample the data from a fixed dataset, naive offline RL approaches fail to learn policies for out-of-distribution actions or states Wu et al. (2019); Kumar et al. (2019), and the obtained Q-value estimation for these actions will be inaccurate with unpredictable consequences. Recent progress in tackling the problem focuses on conservatism by introducing regularization terms in the training of actors and critics Fujimoto et al. (2019); Kumar et al. (2020a); Fujimoto and Gu (2021); Kostrikov et al. (2021a); Lee et al. (2022). Conservatism-based offline RL algorithms have also achieved significant progress in difficult offline multi-agent reinforcement learning settings (MARL) Jiang and Lu (2021); Yang et al. (2021); Pan et al. (2022).

Despite the potential benefits, existing methods have limitations in several aspects. Firstly, the design of the policy network and the corresponding regularizer limits the expressiveness and diversity due to conservatism. Consequently, the resulting policy may be suboptimal and difficult to represent complex strategies Kumar et al. (2019); Wang et al. (2022). Secondly, in multi-agent scenarios, the conservatism-based method is prone to getting trapped in poor local optima. This occurs when each agent is incentivized to maximize its own reward without efficient cooperation with other agents in existing algorithms Yang et al. (2021); Pan et al. (2022).

To demonstrate this phenomenon, we conduct experiment on a simple MARL scenario consisting of $3$ agents and $6$ landmarks (Fig. 1a), to highlight the importance of policy expressiveness and diversity in MARL. In this scenario, the agents are asked to cover $3$ landmarks and are rewarded based on their proximity to the nearest landmark while being penalized for collisions. We first train the agents with $6$ target landmarks and then randomly dismiss $3$ of them in evaluation. Our experiments demonstrate that existing methods (MA-CQL and OMAR detailed in Section 3.2 Pan et al. (2022)), which constrain policies through regularization, limit the expressiveness of each agent and hinder the ability of the agents to cooperate with diversity. As a result, only limited solutions are found. Therefore, in order to design robust algorithms with good generalization capabilities, it is crucial to develop methods beyond conservatism for better performance and more efficient cooperation among agents.

To boost the policy expressiveness and diversity, we propose a novel algorithm based on diffusion for the offline multi-agent setting, called Diffusion Offline Multi-Agent Model (DOM2). Diffusion model has shown significant success in generating data with high quality and diversity Ho et al. (2020); Song et al. (2020a); Croitoru et al. (2023) and our goal is to leverage this advantage to promote expressiveness and diversity in policy learning of RL. Specifically, the policy for each agent is built using the accelerated DPM-solver to sample actions Lu et al. (2022). In order to train an appropriate policy that performs well, we propose a trajectory-based data-augmentation method to facilitate policy training by efficient data sampling. These techniques enable the policy to generate solutions with high quality and diversity and overcome the limitations of conservatism-based approaches. In the $3$-agent example, we show that DOM2 can find a more diverse set of solutions with high performance and generalization (Fig. 1b), compared to conservatism-based methods such as MA-CQL and OMAR Pan et al. (2022). Our contributions are summarized as follows.

• •

  We propose a novel Diffusion Offline Multi-Agent Model (DOM2) algorithm to address the limitations of conservatism-based methods. DOM2 consists of three critical components: diffusion-based policy with an accelerated solver, appropriate policy regularizer, and a trajectory-based data augmentation method for enhancing learning.

• •

  We conduct extensive numerical experiments on Multi-agent Particles Environments (MPE) and Multi-agent MuJoCo HalfCheetah environments. Our results show that DOM2 achieves significantly better performance improvement over state-of-the-art methods.

• •

  We show that our diffusion-based method DOM2 possesses much better generalization abilities and outperforms existing methods in shifted environments (trained in standard environments). Moreover, DOM2 is ultra-data-efficient, and achieves SOTA performance with $20\text{+}$ times less data.

Offline RL and MARL: Distribution shift is a key obstacle in offline RL and multiple methods have been proposed to tackle the problem based on conservatism to constrain the policy or Q-value by regularizers Wu et al. (2019); Kumar et al. (2019); Fujimoto et al. (2019); Kumar et al. (2020a). Policy regularization ensures the policy to be close to the behavior policy via a policy regularizer, e.g., BRAC Wu et al. (2019), BEAR Kumar et al. (2019), BCQ Fujimoto et al. (2019), TD3+BC Fujimoto and Gu (2021), implicit update Peng et al. (2019); Siegel et al. (2020); Nair et al. (2020) and importance sampling Kostrikov et al. (2021a); Swaminathan and Joachims (2015); Liu et al. (2019); Nachum et al. (2019). Critic regularization instead constrains the Q-values for stability, e.g., CQL Kumar et al. (2020a), IQL Kostrikov et al. (2021b), and TD3-CVAE Rezaeifar et al. (2022). On the other hand, Multi-Agent Reinforcement Learning (MARL) has made significant process under the centralized training with decentralized execution (CTDE) paradigm Oliehoek et al. (2008); Matignon et al. (2012), such as MADDPG Lowe et al. (2017), MATD3 Ackermann et al. (2019), IPPO de Witt et al. (2020), MAPPO Yu et al. (2021), VDN Sunehag et al. (2017) and QMIX Rashid et al. (2018) in decentralized critic and centralized critic setting. The offline MARL problem has also attracted attention and conservatism-based methods have been developed, e.g., MA-BCQ Jiang and Lu (2021), MA-ICQ Yang et al. (2021), MA-CQL and OMAR Pan et al. (2022).

Diffusion Models: Diffusion model Ho et al. (2020); Song et al. (2020a); Sohl-Dickstein et al. (2015); Song and Ermon (2019); Song et al. (2020b), a specific type of generative model, has shown significant success in various applications, especially in generating images from text descriptions Nichol et al. (2021); Ramesh et al. (2022); Saharia et al. (2022). Recent works have focused on the foundation of diffusion models, e.g., the statistical theory Chen et al. (2023), and the accelerating method for sampling Lu et al. (2022); Bao et al. (2022). Generative model has been applied to policy modeling, including conditional VAE Kingma and Welling (2013), diffusers Janner et al. (2022); Ajay et al. (2022) and diffusion-based policy Wang et al. (2022); Chen et al. (2022); Lu et al. (2023) in the single-agent setting. Our method successfully introduces the diffusion model with the accelerated solver to offline multi-agent settings.

In this section, we introduce the offline multi-agent reinforcement learning problem and provide preliminaries for the diffusion probabilistic model as the background for our proposed algorithm.

Offline Multi-Agent Reinforcement Learning. A fully cooperative multi-agent task can be modeled as a decentralized partially observable Markov decision process (Dec-POMDP) Oliehoek and Amato (2016) with $n$ agents consisting of a tuple $G=\langle\mathcal{I},\mathcal{S},\mathcal{O},\mathcal{A},\Pi,\mathcal{P},\mathcal{R},n,\gamma\rangle$. Here $\mathcal{I}$ is the set of agents, $\mathcal{S}$ is the global state space, $\mathcal{O}=(\mathcal{O}_{1},...,\mathcal{O}_{n})$ is the set of observations with $\mathcal{O}_{n}$ being the set of observation for agent $n$. $\mathcal{A}=(\mathcal{A}_{1},...,\mathcal{A}_{n})$ is the set of actions for the agents ($\mathcal{A}_{n}$ is the set of actions for agent $n$), $\Pi=(\Pi_{1},...,\Pi_{n})$ is the set of policies, and $\mathcal{P}$ is the function class of the transition probability $\mathcal{S}\times\mathcal{A}\times\mathcal{S}^{\prime}\rightarrow[0,1]$. At each time step $t$, each agent chooses an action $a_{j}^{t}\in\mathcal{A}_{j}$ based on the policy $\pi_{j}\in\Pi_{j}$ and historical observation $o_{j}^{t-1}\in\mathcal{O}_{j}$. The next state is determined by the transition probability $P\in\mathcal{P}$. Each agent then receives a reward $r_{j}^{t}\in\mathcal{R}:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ and a private observation $o_{j}^{t}\in\mathcal{O}_{i}$. The goal of the agents is to find the optimal policies $\bm{\pi}=(\pi_{1},...,\pi_{n})$ such that each agent can maximize the discounted return: $\mathbb{E}[\sum_{t=0}^{\infty}\gamma^{t}r_{j}^{t}]$ (the joint discounted return is $\mathbb{E}[\sum_{j=1}^{n}\sum_{t=0}^{\infty}\gamma^{t}r_{j}^{t}]$), where $\gamma$ is the discount factor. Offline reinforcement learning requires that the data to train the agents is sampled from a given dataset $\mathcal{D}$ generated from some potentially unknown behavior policy $\bm{\pi}_{\bm{\beta}}$ (which can be arbitrary). This means that the procedure for training agents is separated from the interaction with environments.

Conservative Q-Learning. For training the critic in offline RL, the conservative Q-Learning (CQL) method Kumar et al. (2020a) is to train the Q-value function $Q_{{\bm{\phi}}}(\bm{o},\bm{a})$ parameterized by $\bm{\phi}$, by minimizing the temporal difference (TD) loss plus the conservative regularizer. Specifically, the objective to optimize the Q-value for each agent $j$ is given by:

The first term is the TD error to minimize the Bellman operator with the double Q-learning trick Fujimoto et al. (2019); Hasselt (2010); Lillicrap et al. (2015), where $y_{j}=r_{j}+\gamma\min_{k=1,2}\overline{Q}_{\bm{\phi}_{j}}^{k}(\bm{o}^{\prime}_{j},\overline{\pi}_{j}(\bm{o}^{\prime}_{j}))$, $\overline{Q}_{\bm{\phi}_{j}},\overline{\pi}_{j}$ denotes the target network and $\bm{o}^{\prime}_{j}$ is the next observation for agent $j$ after taking action $\bm{a}_{j}$. The second term is a conservative regularizer, where $\tilde{\bm{a}}_{j}$ is a random action uniformly sampled in the action space and $\zeta$ is a hyperparameter to balance two terms. The regularizer is to address the extrapolation error by encouraging large Q-values for state-action pairs in the dataset and penalizing low Q-values of state-action pairs.

Diffusion Probabilistic Model. We present a high-level introduction to the Diffusion Probabilistic Model (DPM) Sohl-Dickstein et al. (2015); Song et al. (2020a); Ho et al. (2020) (detailed introduction is in Appendix A). DPM is a deep generative model that learns the unknown data distribution $\bm{x}_{0}\sim q_{0}(\bm{x}_{0})$ from the dataset. The process of data generation is modeled by a predefined forward noising process characterized by a stochastic differential equation (SDE) ${\rm{d}}\bm{x}_{t}=f(t)\bm{x}_{t}{\rm{d}}t+g(t){\rm{d}}\bm{w}_{t}$ (Eq. (5) in Song et al. (2020a)) and a trainable reverse denoising process characterized by the SDE ${\rm{d}}\bm{x}_{t}=[f(t)\bm{x}_{t}-g^{2}(t)\nabla_{\bm{x}_{t}}\log q_{t}(\bm{x}_{t})]{\rm{d}}t+g(t){\rm{d}}\overline{\bm{w}}_{t}$ (Eq. (6) in Song et al. (2020a)) shown in Fig. 2. Here $\bm{w}_{t},\overline{\bm{w}}_{t}$ are standard Brownian motions, $f(t),g(t)$ are pre-defined functions such that $q_{0t}(\bm{x}_{t}|\bm{x}_{0})=\mathcal{N}(\bm{x}_{t};\alpha_{t}\bm{x}_{0},\sigma^{2}_{t}\bm{I})$ for some constant $\alpha_{t},\sigma_{t}>0$ and $q_{T}(\bm{x}_{T})\approx\mathcal{N}(\bm{x}_{T};\bm{0},\tilde{{\sigma}}^{2}\bm{I})$ is almost a Gaussian distribution for constant $\tilde{{\sigma}}>0$. However, there exists an unknown term $-\sigma_{t}\nabla_{\bm{x}_{t}}\log q_{t}(\bm{x}_{t})$, which is called the score function. In order to generate data close to the distribution $q_{0}(\bm{x}_{0})$ by the reverse SDE, DPM defines a score-based model $\bm{\epsilon}_{\bm{\theta}}(\bm{x}_{t},t)$ to learn the score function and optimize parameter $\bm{\theta}$ such that $\bm{\theta}^{*}={\arg\min}_{\bm{\theta}}\mathbb{E}_{\bm{x}_{0}\sim q_{0}(\bm{x}_{0}),\bm{\epsilon}\sim\mathcal{N}(\bm{0},\bm{I}),t\sim\mathcal{U}(0,T)}[\|\epsilon-\bm{\epsilon}_{\bm{\theta}}(\alpha_{t}\bm{x}_{0}+\sigma_{t}\bm{\epsilon},t)\|^{2}_{2}]$ ($\mathcal{U}(0,T)$ is the uniform distribution in $[0,T]$, same later). With the learned score function, we can sample data by discretizing the reverse SDE. To enable faster sampling, DPM-solver Lu et al. (2022) provides an efficiently faster sampling method and the first-order iterative equation (Eq. (3.7) in Lu et al. (2022)) to denoise is given by $\bm{x}_{t_{i}}=\frac{\alpha_{t_{i}}}{\alpha_{t_{i-1}}}\bm{x}_{t_{i-1}}-\sigma_{t_{i}}(\frac{\alpha_{t_{i}}\sigma_{t_{i-1}}}{\alpha_{t_{i-1}}\sigma_{t_{i}}}-1)\bm{\epsilon}_{\bm{\theta}}(\bm{x}_{t_{i-1}},t_{i-1})$.

In Fig. 2, we highlight a crucial message that we can efficiently incorporate the procedure of data generation into offline MARL as the action generator. Intuitively, we can utilize the fixed dataset to learn an action generator by noising the sampled actions in the dataset, and then denoising it inversely. The procedure assembles data generation in the diffusion model. However, it is important to note that there is a critical difference between the objectives of diffusion and RL. Specifically, in diffusion, the goal is to generate data with a distribution close to the distribution of the training dataset, whereas in offline MARL, one hopes to find actions (policies) that maximize the joint discounted return. This difference influences the design of the action generator. Properly handling it is the key in our design, which will be detailed below in Section 4.

In this section, we present the DOM2 algorithm, which is shown in Fig. 3. In the following, we first discuss how we generate the actions with diffusion in Section 4.1. Next, we show how to design appropriate objective functions in policy learning in Section 4.2. We then present the data augmentation method in Section 4.3. Finally, we present the whole procedure of DOM2 in Section 4.4.

We evaluate our method in different multi-agent environments and datasets. We focus on three primary metrics, performance (how is DOM2 compared to other SOTA baselines), generalization (can DOM2 generalize well if the environment configurations change), and data efficiency (is our algorithm applicable with small datasets).

We propose DOM2, a novel offline MARL algorithm, which contains three key components, i.e., a diffusion mechanism for enhancing policy expressiveness and diversity, an appropriate regularizer, and a data-augmentation method. Through extensive experiments on multi-agent particle and multi-agent MuJoCo environments, we show that DOM2 significantly outperforms state-of-the-art benchmarks. Moreover, DOM2 possesses superior generalization capability and ultra-high data efficiency, i.e., achieving the same performance as benchmarks with $20+$ times less data.