Efficient Multi-agent Offline Coordination via Diffusion-based Trajectory Stitching
Authors: Lei Yuan
Yuqi Bian
Lihe Li
Ziqian Zhang
Cong Guan
Yang Yu
Venue: ICLR 2025 Poster (2025)
Keywords: Multi-agent Reinforcement Learning
Offline MARL
Diffusion based Reinforcement Learning
Trajectory Stitching
Primary area: reinforcement learning
TL;DR: A method develops diffusion-based trajectory stitching for efficient multi-agent offline reinforcement learning

Abstract: Learning from offline data without interacting with the environment is a promising way to fully leverage the intelligent decision-making capabilities of multi-agent reinforcement learning (MARL). Previous approaches have primarily focused on developing learning techniques, such as conservative methods tailored to MARL using limited offline data. However, these methods often overlook the temporal relationships across different timesteps and spatial relationships between teammates, resulting in low learning efficiency in imbalanced data scenarios. To comprehensively explore the data structure of MARL and enhance learning efficiency, we propose Multi-Agent offline coordination via Diffusion-based Trajectory Stitching (MADiTS), a novel diffusion-based data augmentation pipeline that systematically generates trajectories by stitching high-quality coordination segments together. MADiTS first generates trajectory segments using a trained diffusion model, followed by applying a bidirectional dynamics constraint to ensure that the trajectories align with environmental dynamics. Additionally, we develop an offline credit assignment technique to identify and optimize the behavior of underperforming agents in the generated segments. This iterative procedure continues until a satisfactory augmented episode trajectory is generated within the predefined limit or is discarded otherwise. Empirical results on imbalanced datasets of multiple benchmarks demonstrate that MADiTS significantly improves MARL performance.

Published as a conference paper at ICLR 2025
EFFICIENT MULTI-AGENT OFFLINE COORDINATION VIA DIFFUSION-BASED TRAJECTORY STITCHING
Lei Yuan 1,2,3 Yuqi Bian Lihe Li! Ziqian Zhang! Cong Guan 1.2
Yang Yu 12.3.
1 National Key Laboratory for Novel Software Technology, Nanjing University
2 School of Artificial Intelligence, Nanjing University Polixir Technologies {yuanl, bi bianyg, 1ilh, zhangzq, guanc lamda n juu. edu. cn, yuyēnju.edu.
ABSTRACT
Learning from offline data without interacting with the environment is promis- ing way to fully leverage the intelligent decision-making capabilities of multi- agent reinforcement learning (MARL). Previous approaches have primarily fo- cused on developing learning techniques, such as conservative methods tailored MARL using limited offline data. However, these methods often overlook the temporal relationships across different timesteps and spatial relationships be- tween teammates, resulting in low learning efficiency in imbalanced data sce- narios. To comprehensively explore the data structure of MARL and enhance learning efficiency, we propose Multi-Agent offline coordination via Diffusion- based Trajectory Stitching (MADiTS), novel diffusion-based data augmentation pipeline that systematically generates trajectories by stitching high-quality coor- dination segments together. MADITS first generates trajectory segments using a trained diffusion model, followed by applying a bidirectional dynamics constraint ensure that the trajectories align with environmental dynamics. Additionally, we develop an offline credit assignment technique to identify and optimize the behavior of underperforming agents in the generated segments. This iterative procedure continues until a satisfactory augmented episode trajectory is generated within the predefined limit is discarded otherwise. Empirical results on imbal- anced datasets of multiple benchmarks demonstrate that MADiTS significantly improves MARL performance.
INTRODUCTION
Multi-agent reinforcement learning (MARL) is a key technology for addressing complex decision- making problems that involve multiple interacting agents (Gronauer Diepold, 2022). It has demonstrated significant potential in areas such as active voltage control (Wang et al. (021)) large language model (LLM) applications (Sun et al. 2024), and embodied agents (Liu et al. 2024b). simulated Despite the remarkable progress made in MARL, most of its successes are confined environments, where agents can interact with the environment unlimitedly to gather vast amounts of data for policy improvement. However, this is generally impractical in real-world applications such as autonomous driving and financial transactions, where trial and error can be both costly and risky. This challenge has accelerated research in offline reinforcement learning (RL) (Levine al. 2020), which focuses on learning from a fixed dataset without interacting with the environment. Numerous efficient methods have been successfully developed (Prudencio al., 2023). showing promising potential in real-world applications such as industrial process control (Deng et al. 2023) recommender systems (Chen et al. 2024), and legged robot navigation (Weerakoon et al. 2024).
One of the key challenges in offline RL is the issue of distribution shift (Lambert et al., 2022).
a phenomenon where unseen state-action pairs are incorrectly estimated. Early works have intro- duced techniques such as policy constraints (Ran et al. (023), value function regularization (Mao ctt al. 2024), uncertainty estimation (Beeson & Montanal 2024), and world model learning (Luo al. 2024p for data augmentation or direct planning. Recent studies have also leveraged Trans-
formers (Chen et al. 202I) and diffusion models (Janner et al. 2022) to enhance offline RL ef-
"Corresponding author.

---

Published as a conference paper at ICLR 2025
ficiency from various perspectives (Yang et al. 2023 In addition, while single-agent offline RL has seen rapid growth, most works in offline MARL have primarily focused on adapting successful methods from the single-agent to the multi-agent setting. These approaches have aimed to alle- viate extrapolation errors in agent interactions (Yang et al., 2021c), design knowledge distillation mechanisms to bridge the CTDE (Centralized training decentralized execution) gap during policy deployment (Tseng et al.) 2022), discover coordination skills from multi-task data (Zhang et al.) 2023), and implement efficient policy adaptation strategies (Wu al. 2024), showing significant progress in multiple scenarios (Formanek et al. 2023 Guan et al. 20044).
Despite recent advances, current offline RL methods are highly correlated with the quality of datasets (Schweighofer et al. [021]. which may suffer from severe performance degradation (Hong jt al. 2023) under imbalanced data. Various methods have been developed to address these issues, such as policy regularization (Liu et al.) 2024a), data sharing (Yu et al. (02T]. efficient data sampling (Hong al. 2023), and trajectory stitching (Hepburn & Montanal 2022). However, they largely focus utilizing the provided data, while others aim to learn models from offline data and select actions through planning (Rosete-Beas et al., 2023), or perform data augmentation (He 2023) to enhance policy learning, showing widespread improvement in sample efficiency. Nevertheless, these methods primarily address single-agent tasks. Multi-agent tasks are significantly more complex due intricate interactions among agents (Albrecht et al. 2024). which give rise both temporal relationships across different timesteps and spatial relationships between teammates. Furthermore, in open-environment settings (Yuan et al., 2023), the behavior policy used to collect data might be nonidealized, resulting in severely imbalanced offline data and low learning efficiency. This raises an important question: Can we augment suboptimal datasets by stitching together high-quality trajectory segments to overcome the temporal and spatial imbalances inherent in multi agent datasets?
To comprehensively explore the data structure of multi-agent reinforcement learning (MARL) and improve learning efficiency, we propose Multi-Agent offline coordination via Diffusion-based Trajectory Stitching (MADiTS). to explicitly address imbalances in both temporal and spatial dimensions. Initially, for temporal imbalance, a diffusion model is trained to capture the distribution of trajectory segments and is conditioned on high returns to generate trajectory segments for stitching by head-to-tail concatenation. A bidirectional environmental dynamics constraint is then applied to ensure that only trajectories consistent with the environment's dynamics are selected during the stitching process. Next, for spatial imbalance, an integrated gradient-based method is developed to identify the agents responsible for the suboptimal performance. A partial noising approach is subsequently used to optimize the behaviors of these underperforming agents, leveraging diffusion models to stitch together trajectory segments from various agents. This iterative process continues until a satisfactory augmented trajectory achieved within the predefined limit or is discarded otherwise, enabling policy optimization with both the generated and original data. Experiments on the imbalanced datasets of on multiple benchmarks including MPE (Lowe et al. (001)). SMAC Samvelyan et al.] 2019), SMACv2 (Ellis et al.) 2023). and MAMuJoCo (Peng et al.) 2021 demonstrate that MADITS significantly enhances the performance of Behavior Cloning (BC) and other offline MARL algorithms relative to the original datasets, underscoring MADiTS' effectiveness in tackling the challenges posed by imbalanced data.
novel diffusion-based data augmentation pipeline designed
RELATED WORK
Multi-agent reinforcement learning (MARL) Albrechh et al., 2024) has garnered significant atten- tion recently(Du Ding 2021 achieving remarkable success across various complex domains such as active voltage control (Wang et al. 202T) and dynamic algorithm configuration (Xue et al.) 2.22]. A wide range of MARL solutions have been proposed, including value-based approaches like VDN (Sunehag et a 2018) and QMIX (Rashid et al. 2018), as well as policy gradient methods such as MADDPG (Lowe et al. 2017) and MAPPO (Yu et al. 2022), alongside newer variants like Transformer-baset approaches (Wen et al.) 2022). However, mainstream MARL methods rely on continuous interaction with the environment to collect data for policy optimization, this challenge has accelerated interest in offline MARL, which focuses learning policies from pre- collected data. For example, ICQ (Yang et al., 2021c) tackles the extrapolation error by restricting trust to offline data, while MABCQ (Jiang& & Lu 2021) introduces fully decentralized offline
2

---

Published as a conference paper at ICLR 2025
MARL framework and uses techniques such as value deviation and transfer normalization for ef- ficient learning. OMIGA (Wang et al. 2023) bridges multi-agent value decomposition with policy learning by transforming global-level value regularization into implicit local value regularization. CFCQL (Shao et al.] 2023) applies conservative regularization to each agent in a counterfactual manner, then combines them linearly for overall conservative value estimation.
While the aforementioned methods alleviate some challenges, offline RL still faces significant hur- dles due to the limited quality and diversity of pre-collected datasets, leading to low learning ef- ficiency. To tackle this issue, various techniques have been proposed (Yul 2018 Prudencio et al.) 2023), such as data sharing Yu et al. 2021), data augmentation (Yu et al. 2020), knowledge trans- fer (Bose et al. (2024), and leveraging external knowledge like large language models (Shi jet al. 2023). SIT (Tian et al. 2023) further explores offline MARL by explicitly accounting for the diversity of multi-agent trajectories. Among these methods, trajectory stitching has emerged as a promising method of offline trajectory augmentation, synthesizing optimal or near-optimal
trajectories from suboptimal ones. For example, MBTS (Hepburn & Montanal 2022 learns state transition model and value function to stitch together high-quality segments from different trajectories, generating optimal trajectories. BATS (Char et al. 2022) uses an environment model for planning and adds state transitions to fill in missing trajectory segments within offline datasets. More recently, DiffStitch (Li et al., 2024a) leverages the diffusion model (Ho et al. 2020) for trajectory stitching, showing high efficiency in single-agent RI.. However, these trajectory stitching methods struggle to effectively learn complex interactions and cooperative behaviors from offline multi-agent trajectory data, as they fail to account for the temporal and spatial imbalances inherent in such data. More about related work could be found in Appendix
3 BACKGROUND
This paper considers fully cooperative multi-agent task, which is modeled as a Decentralized Partially-Observe Observable Markov Decision Process (Dec-POMDP) (Olichoek et al. 2016). denoted by a tuple MM- = SN A, P. O, r, 7). Here, N = {1..... ,n] is the set of n agents, S is the global state space, A = A x
A is the joint action space of the agents, where A' is the action space of agent i. P x Axx S + [0, is the state transition function. Ω is the observation space, 0 S x N + Ω is the observation function. All agents share the same global reward function r S x A + R. p:S factor. In each episode, so is initially sampled form p(s). At each timestep, agent i receives the observation o' = (s,,)) and outputs an action a² E A'. The joint action = (a' ..... a") leads to the next state s' ~ P( a) and a global reward r(s, a). For the optimization, all agents aim to learn the joint policy
1] is the initial state distribution. 7 E [0, 1] is the discount
(M,, , Tn), that maximizes the expected discounted return E 7 R(s,, R a.)] In our offline setting, agents cannot explore the environment and are trained with static dataset D {7, M where 7 is the joint trajectory composed of transitions (st, Oz,, ar, rt)
J ((T) ((o)))_ff-
This paper utilizes Denoising Diffusion Probabilistic Models (DDPMs) (Ho et al., 2020) for trajec- tory generation, which are a class of generative models that allows for sampling from a distribution via iteratively reversing a forward noising process. DDPMs consist of a forward noising process and
reverse denoising process. Given that Xo q(Xo). forward process gradually adds noise to the = N(xx; v1
B, 1, I). can be further derived that p(x; Xo) A (xx; V Vaggo, (1 ã% I), where 0g (1 B&), a II Og.. reverse denoising process, constructed as po (xx-1 |xx) = N (x,, k), Ex), then optimized to maximize the evidence lower bound on negative log likelihood defined as re(xa.x)). Here po(Xu:k) N (xx; ; 0, I) I (x
data with a pre-defined variance schedule { f0 1K 1== ((xx|xx>>
Instead of directly training (i by optimizing the evidence lower bound, Ho et al. (2020) proposes a simplified surrogate loss: (ovvv + vI k)),k))1
[In 9(3) [
Laenoise (f)
c l
(1)
where (x k) directly estimates the noise added to produce X In the reverse diffusion process, the samples can be generated by following the recursion: X;-1 yar (xx VI-a.
1-ay cg co(xx, k)) +
VBez, with z sampled from \/ I). Furthermore, researchers have proposed two types of methods to guide the conditional generation of samples. Classifier guided methods (Song et 202T Dhani
wal & Nichol (021) train additional classifier and add guidance to predicted noise. Classifier-free
3

---

Published as a conference paper at ICLR 2025
(a) Trajectory generation
(b) Trajectory inspection
Sample
Individual trajectory optimization
(I/(),)) -
Dynamics Constraint
Denoise
Highneturn
Agent performance identification
Diffusion model
Sample init states
Filter & Stitch
Offline MARL
MADiTS
return
High
Imbalanced dataset
Augmented dataset
(c) Data augmentation
Figure 1: The overall framework of MADITS.
methods (Ho & Salimans 202T) modify the original training setup to allow both conditional and un-
conditional model. The perturbed noise is corrected as e = co (x, y, w(c(xx, k)) cx(x, k:)) where y is the attribute of label, and w is guidance scale.
4 METHOD
In this section, we will introduce the design of Multi-agent offline coordination via Diffusion-based Trajectory Stitching (MADiTS), novel data augmentation algorithm that systematically generates high-quality coordination from imbalanced dataset (see Figure [ We first solve the modeling for diffusion in multi-agent scenarios while guaranteeing the bidirectional dynamics consistency in Section 4.1 Next, we propose to overcome the imbalance caused by underperforming agents via integrated gradient in Section 4.2 Finally, an overall pipeline of trajectory stitching for offline MARL within imbalanced dataset is presented in Section 4.3)
4.1 TRAJECTORY GENERATION WITH BIDIRECTIONAL DYNAMICS CONSISTENCY
In imbalanced dataset where high-rewarding and low-rewarding states both exist, it is promising to stitch them together by head-to-tail concatenation so that agents lear to transit to high- rewarding states and achieves the shared goals ultimately. Inspired by the success made by diffusion models in single-agent offline RL (He al. 2024 Li al.] (024a). capture the multimodal distribution the imbalanced trajectories via diffusion model Go:
max Eo, ort.... -1), 00. #+1]|loog Go Ort1,..... 0r.,,..,,,,..1111..... O.-H-I
where is the input trajectory length, and Rg Er=t -+ turn, which can be optimized via the simplified surrogate objective Lieeoise (f) defined in Equation
1011-11''-tp "t' the discounted cumulative re-
We choose diffuse over joint observations because actions are more varied, and are less smooth (Ajay et al., (022). Moreover, it is often multi-discrete in MARL scenarios, making harder directly modeling via diffusion model. After training, the trajectory segments are generated following the denoising process:
O₁, dct1, ÓL-2. Ô = Gg Zt- Z442.-.. 24+H-1(R))
where Z = (4.. 22 #p), VI' E
Here, Oş denotes the joint observation at time step and ô, denotes the predicted joint observation time step !l after denoising. 24. refers the sampled Gaussian noise for all i {1, ....,
4

---

Published as a conference paper at ICLR 2025
R3 is a fixed expected return conditional input. Furthermore, infer the action a, from the generated observation trajectory, we introduce an inverse dynamics model f (Or. 0c+1) = dt.
Despite the powerful generalization ability of diffusion models, the generated trajectories might violate environment dynamies due to excessively prioritizing high returns. Such inconsistency will accumulate over time, and in turn reduces the reliability of trajectory segments, To mitigate the effects of dynamics inconsistency, we propose bidirectional dynamics constraint mechanism to identify generated observations that violate dynamics. Specifically, we instantiate another forward dynamics model ffual (or, ae) to predict the next joint observation. Both the inverse and forward dynamics model f and ford are implemented with three-layer MLPs (multi-layer perceptrons). trained using transitions sampled from the offline dataset:
(4) Once the joint observation trajectory is generated by Go conditioned on the target return, we traverse each observation pair (ōt, 0 generated by Go, and infer the joint action a, = 0(+1). Then predict the legal next-step observation o.+1 fo āt). We discard the subsequent segment after d if ||ō++h ōcthal ll exceeds certain threshold 6 which means that the dynamics consistency is severely violated. In practice, we keep the initial joint observation unchanged and restore generated trajectories with bidirectional dynamics consistency (or, åz, O,.t,. 0z++A-1 1 as candidates for data augmentation. Here 2 < h S H due to the discarded inconsistent segments, the details of trajectory generation will be showed in Section 4.3
Laymamics (#. 1)) ((o......)--)lll.. fm (or, oc+1) a. [1ft a.) o++1 l12)
4.2 BEHAVIOR CORRECTION OF UNDERPERFORMING INDIVIDUALS
Although the diffusion model with bidirectional dynamics consistency can help generate high- rewarding trajectories that accord with environmental dynamics, the unique spatial imbalance in multi-agent system still makes it difficult to generate trajectory segments with high cooperativeness. Such imbalance is caused due to the underperforming individuals in the system, and it is difficult for the diffusion model to correct their behaviors under such large joint trajectory space.
However, the search space can be largely narrowed down if we have access to the identity of under- performing individuals by fixing the joint trajectory of good ones. This is difficult in a cooperative MARL setting as all the agents share the same reward function, known as credit assignment problem. Inspired by integrated gradient (IG) (Sundararajan et al.] 2017). neural network interpretability method, we can quantify each agent's contribution at each timestep so as to identify underperforning individuals. Given the function F: R2d I., IG computes the attribution of each feature by:
(x) (x; - -x() f a==
OF(w' + Fa(x x²)) da,
(5)
Ox:
where x' is the baseline point, and aF(#) is the gradient of F at with respect to the i-th feature. By 0z.,
introducing a path function 7(a) = x' + ax - x') which specifies path from to baseline point r', we can define the path integrated gradients (PathIG) as:
aF(((a)) (m((a) do.
PathIG,(1;1)
(6)
0ni(a) Ox:
We first show the team return can be decomposed into the individual contributions of each agents by specifying F of PathIG according to trainable reward function gomdl St x A + R. Given a joint trajectory segment of length T = (ot; at, 10th--1 0(th-1) 1), x4 = (o., at), then:
R(z,, T() r r(
PathIGy (*6))))
1=1 ,EX,
- (xx)) is the return-to-oo, X, represents the set of observation-action features for agent i, 12th-1 is chosen as the baseline point, and 7/. is the simple path function starting from 2, 214h-1- The detailed proof of Equation 7is is provided in Appendix C and the intuition behind Equation 7 is that the return-to-ga can be decomposed into the sum of integrated gradients of each agent's respective features. Thus, we can approximate the contribution of agent
where R0(2,, TT) = -+h
5

---

Published as a conference paper at ICLR 2025
i at timestep t as Ejex, PathIG 7n) provided with the generated joint trajectory T.. We first train a team reward prediction model g implemented with a three-layer MLP via minimizing:
Lreward (w) = R (o. g (or, ar) - rel12)
To comprehensively evaluate each agent's contribution in the generated trajectory segment (or, â..... 01+h-1, ac+k-1) so that we could find out the underperforming individuals, we first sort the contribution value an ascent order and derive the ranking rank;. Here, ..., is a permutation of (1,....) for any timestep I. Thus, the average ranking of each agent's contribution in the joint trajectory segment can be calculated rank - E rank For agents whose average ranking value is lower than certain threshold 4rank, which means that agent underperforms in the joint trajectory due low contribution to the overall team performance, they will be attributed to underperforming individuals. Finally, to correct the behaviors of these underperforming individuals within a smaller trajectory space and improve the cooperativeness, we conduct resampling process by fixing other agents' generated trajectories:
Or, Oc+1. 0 0 = Go(or, Z1+1, 2142.- 144h-1|8a)
where Z = ($), bye ... 27),, Vt E ft + 1, , 1++ 17,
s rank 2 Grank
and Zy.
Og.. otherwise.
Here, Oy denotes the joint observation time step 1 and oy denotes the predicted joint observation at time step l' after denoising.
4.3 DATA AUGMENTATION PROCESS
We here provide the overall description of the procedure of MADITS, Given the imbalanced dataset D, we first split the trajectories into segments with length H to formulate D Additional tech- niques including return-based filtering and circular shift are applied to facilitate the stitching of high-quality cross-agent individual trajectories (see Appendix [D. Afterwards,
train diffusion defined in Equations 2 and g respectively. Details could be found in Appendix
fowd, and the reward model g via objectives models Gg. inverse and forward dynamics model far f
To generate a high-rewarding trajectory segment that satisfy both environmental dynamics con- sisteney and cooperativeness, we first sample an joint observation o, from dataset D. Then, conditioning the fixed o, and high return, the diffusion model generates a joint observation segment (o,, .. --). We will discard the subsequent part after O:+h-1 if violates the environ- mental dynamics consistency based on the learned inverse and forward dynamics models. To avoid the waste caused by frequent discarding, we will try to re-generate the subsequent segment until the number of failures exceeds threshold. For each generated trajectory segment that accords with the environmental dynamics, we identify underperforming individuals by calculating the path integrated gradient via the learned reward model, and then correct their behaviors through re-generating the cor- responding part in the trajectory segment. It should be noticed that the bidirectional environmental dynamics consistency mechanism will be applied ensure the legacy re-generated segment. This iterative procedure of generation and inspection continues until a pre-defined number of trajectories are generated formulate Daug- Finally, we augment the original dataset with generated trajectories D' 'D U 'Dugg Since our method MADITS is algorithm-agnostic, any offline MARL algorithms can be applied and benefit from the augmented dataset. Specifically, we implement our method to the behavior cloning (BC) (Song et al. 2018), OMIGA (Wang et al. (023) and CFCQL (Shao al. 2023). and the detailed pseudocode could be found in Appendix|F
5
EXPERIMENT
5.1 EXPERIMENT SETTINGS
In this section, we introduce (1) the environments to conduct experiments, (2) the mechanism for data collection, (3) the agnostic offline MARL algorithms for evaluation, and the baselines compared with our method, provide clear understanding of the experimental setup.

---

Published as a conference paper at ICLR 2025
Agent
X
Landmark
Prey
Leader agent
Forest
Food
()) Cooperative Novigation (CN)
(hj Fredator Pray (PP)
10 World
(4) Sturcraft Mubi-gger Challange (SMAC)
Figure 2: Illustrations of the experimental environments in work.
Environments. We evaluate our method on two widely-used multi aulti--gent benchmarks that require the Multi-Agent Particle Environment (MPE) (Lowe al.) (017), the StarCraft Multi-Agent Challenges (SMAC) (Samvelyan et al., 2019). SMACv2 (Ellis et al., (023). and MA- MuloCo (Peng et al. 202T). Some of the environments where particle agents can move, observe each other, and interact with fixed landmarks, including three tasks: Cooperative Navigation (CN), Predator-Prey (PP), and World. In CN, agents must cover different fixed landmarks without colliding with each other. In PP, predator agents collaborate to chase faster, pre-trained prey agent. The World task introduces more elements including food, forests, and leader predator based PP. SMAC consists of suite of StarCraft II battle scenarios where agents must cooperate to defeat enemy forces controlled by the game's built-in AI. We conduct experiments maps 3m, 2s3z, 2m-vs_lz, and 12m. SMACv2 is new version of the benchmark SMAC where scenarios are procedurally generated and add more challenging tasks. We conduct experiments terran 5 vs 5 and zerg 5 vs 5. MAMuJoCo is cooperative multi-agent robotic control where multiple agents within a single robot have to solve a task cooperatively. We 4-agent ant (4ant) configuration in experiments.
agent cooperation:
shown in Figure 2 MPE is 2D game
Datasets. We construct datasets with varying degrees of imbalance both temporal and spatial dimensions evaluate our method. First, behavior policies trained using MATD3 (Ackermann et al. 2019) on MPE and MAMuJoCo and using QMIX (Rashid et al. 2018 on SMAC and SMACv2, for collecting trajectory data for offline training. Next, during the data collection process, we apply perturbation on the policies to introduce imbalance. For the temporal dimension, at each timestep, there small probability Pimb that all agents perform random actions and continue for fimb steps. For the spatial dimension, perturbations are introduced by randomly selecting certain agents to perform random actions for the entire episode. Two levels of spatial perturbation are designed: moderate (exp-m), where only one agent performs random actions throughout the episode, and severe (exp-s), where the number random agents is sampled from {I,... 1). We collect 40k, 20k, 10k, 2k trajectories in MPE, SMAC, SMACv2, MAMuloCo, respectively.
Evaluation algorithms and Baselines. Based on the original dataset, we can apply MADITS to derive augmented dataset, and use any agnostic offline MARL algorithms to learn the policies. For given offline MARL algorithm, the quality of the learned policies can serve as a metric of the augmented dataset's iss quality. We employ the classic algorithm, behavior cloning (BC), which simply imitates the behavior policy and offers an intuitive reflection of dataset quality, and two state-of-the- art algorithms, OMIGA (Wang al.) 2023) and Counterfactual Conservative Q-Learning (CFCQL) (Shao et al. 2023). Additionally, to evaluate the effectiveness of MADITS, we first compare it with a simple baseline OOiginal",, which directly utilizes the original perturbed datasets for policy learning. Then, we include MA-MBTS, by extending the well-developed data augmentation method, MBTS (Hepburn & Montana b022). into multi-agent settings. MBTS is model-based trajectory stitching method that generates new actions to connect high-quality segments from different trajectories. In MA-MBTS, individual observations and actions are replaced by joint ones to perform stitching. MADifI et al. 2023a) is diffusion model-based offline MARL algorithm designed to predict future joint observations for decision-makinn augmentation. We also include another strong baseline where offline MARL algorithms directly learns from balanced datasets, denoted as "Balanced", which are collected by the same behavior policies without adding any perturbation. Details of these algorithms can be found in Appendix G
We directly apply stitching to MADiff for data

---

Published as a conference paper at ICLR 2025
Table 1: Evaluation results on multi-agent imbalanced datasets. The mean and standard error are computed based on the normalized average return average battle won rate of the evaluation algorithms trained on the datasets, with 5 different random seeds. We bold the highest scores on exp- m and exp-s datasets, respectively. The results of exp-m and exp-s on 2m_vs. an.vs_lz 1z are the same since levels of spatial perturbation for environments of 2 agents are the same. Results of augmentation on balanced datasets can be found in Appendix H
Bolanced
Original
MA-MBTS
MADIIT
MADNTS
Algs
BCC
45:22 d 8.00
17:37 6: 3,00
9:03 1: 5-00
40:34 1: 6,89
35.77 1: 6.47
40164 5: 575
50.24 6.82
43.44 :: 6.11
57.82 4.45
OMIGA
55,53 10.08
à: 57.07
55.90
09.17
12.96
79.01
19.25
67.08
7.27
23,00 69.00
22.91
44.94
CRCQ
54.07 + 10.10
-15.40 44.7%
-66.46 40.38
7.60 15 96
-17
25.15
10.0% 17:97
21:88 + 10.71
99,57 16.14
25,00
20.94
BC
52.77 6.15
45,40 11 4,59
40:21 =: 3,90
46.77 4:
48.06 4: 3.34
49:55 6.50
49.93 4: 3.74
54.85 4: 4.23
53.50 :: 4.28
OMIGA
58.34
37.00 25.34
26.73 1: 40.67
30.93 18.02
60.88 4:
47.48 20:97
58.44
4.76
63.02 3.40
63.71 à: 5.67
CRCQL
51.02 6.76
4.66
18:83 6.20
11.74
20.17
45.50 6.76
7.02
3.74
32.25 10.96
BC
54000 0: 5.18
47.29 8: 3:00
48.30 5:5
52,02 #: 7:32
09.24 8: 5.12
50.79 2,02
51.55 8: 4.90
$4.25 4.35
52.84 :: 5.29
World
OMIGA
5.94
54.33
52.92 5.64
43.09 23.45
17.34
49.55 24.99
48.26 24.75
57.35
5.89
58.59 9.32
CIOQ
Average
53,17
27.01
18.35
34,02
28.43
40.14
36,88
45,53
43,50
BC
0.16 0:
0.03 0.06
0.03 à: 0.06
0.32 6 0.23
0.32 0.23
0:30
0.30 0.18
0.36 :: 0.20
0.36 :: 0.20
OMIGA
0.99-+00.15
0.59 = 0.28
0.59 + 0.28
0.68 6 0.30
0.65 0.30
0.97 + 0.04
0.97 +
0.98 11 0.02
0,96 0.02
CRCQL
0.72
0.08
0.89 0.00
0.92
BC
1,00 0.00
0142 + 0.08
0.38 0.20
0.26 0.33
0.28 0.19
0.44 0.1
0.36 6 0.15
0,51 0.19
0.43 11 0.18
3m
OMIGA
0.97 9:
0.90 0.08
0.90 s: 0.03
0:98 6: 0.00
0,94 0.04
1/90 :: 0J00
0.55 à: 0.3
1.00 :: 0.00
0.96 :: 0.03
CRCQL
0.95 4 0.05
0.90 =: 0.06
0.80 =: 0.05
0.87 =2 0.05
0.77 0.18
0.89 :: 0.05
0,81 4 0.09
0.94 1: 0.07
0.89 1: 0.06
BC
0.92 0.05
0.68
0.20
0.67 0.00
0.75
0.67
0.69
OMIGA
0.97 à: 0.03
0.86 0.16
0,67 =: 0.05
0.81 d: 0.05
0.68 0.13
1.00 =: 0.00
0.71 à: 0.15
1,00 0: 0.00
0.76 1: 0.19
CRCQL
0.96 à: 0.02
0.73 0.15
0.66 :: 0.14
0,76 A: 0.13
0.67 0.13
0.80 =:: 0.09
0.63 4: 0.15
0.92 :: 0.02
0.68 :: 0.36
Average
5.2 PERFORMANCE COMPARISON
In this section, we analyze MADITS's s effectiveness on augmenting various kinds of offline MARL datasets, and eventually enhancing policy learning.
Specifically, we apply MADITS and the baseline method to augment the given pre-collected datasets. Offline algoritt are then deployed on the augmented datasets and extra balanced datasets to learn multi-agent policies. The performance of various learned policies are presented in Table [ where we record the normalized average return and battle won rate for MPE and SMAC tasks, respectively. From the results, we can first observe that Original suffers significant performance drop compared with Balanced across all datasets and offline MARL algorithms, showing the necessity for data augmentation techniques in such imbalanced situations. However, simply utilizing such techniques designed for single-agent settings, MA-MBTS shows only marginal improvement compared with Original. It indicates that multi-agent coordination scenarios require specifically designed techniques to handle both temporal and spatial imbalance. MADiff shows obvious improvement compared to the Original, showcasing the necessity of using diffusion to model data distributions. But it does not account for dynamics consistency or data imbalance, which results in a noticeable gap in performance compared to MADITS on imbalanced datasets. On the contrary. our method MADITS achieves the best performance on all imbalanced datasets and offline MARL algorithms, and a comparable average results compared with Balanced. Notably, in some environments such as PP and 2m vs 1z, The performance of MADITS even exceeds policies learned from balanced datasets, highlighting that MADITS successfully leverages the high-quality segments within the collected trajectories to further improve the quality of the datasets.
To thoroughly evaluate the general effectiveness of MADITS, we utilize the MPE balanced datasets constructed Pan et al. (2022). which are collected using behavior policies of varying qualities, and evaluate the above data augmentation techniques. Similarly, MA-MBTS demonstrates only slight improvement over Original, while MADITS consistently outperforms both. These results prove the versatility and generality of MADITS on various offline settings. Detailed configurations and results can be found in Appendix H.2
8

---

Published as a conference paper at ICLR 2025
Underperforming
tes ageit
Inconsistency
Landmark
(a)
Figure 3: Visualization of different trajectories in CN. Lighter-colored trajectories are the ones extracted from observations of other agents. (a) An original trajectory in imbalanced dataset (b) trajectory generated by w/o DC+PN. (c) A trajectory generated by MADiTS. (d) Mean ADE of entity trajectories from different agents' observations, and the average number of underperforming agents in datasets augmented by w/o DC+PN and MADITS.
5.3 ILLUSTRATIVE EXAMPLES OF STITCHED TRAJECTORIES
To further investigate how MADITS addresses temporal and spatial imbalances and improves the quality of imbalanced datasets, we visualize trajectories generated by MADiTS and its variant without dynamics constraint (DC) and partial noising (PN), denoted as w/o DC+PN. Specifically, we first render a pre-collected imbalanced trajectory in the original CN exp-m dataset in Figure 3ja), where the agent colored in red fails to cover a landmark. Given the initial observations of this trajectory, w/o DC+PN generates landmarks, as shown Figure Bb). Nevertheless, on one hand, the new trajectory suffers from dynamics inconsistency, as the individual trajectory of the agent colored in dark green differs from the one observed by the teammates (light green). On the other hand, agents still have a certain distance from the landmarks, indicating the generated trajectory is suboptimal. After applying both DC and PN, our method successfully generates an optimal trajectory with consistent dynamics, as shown in Figure e), which will be added to the dataset enhance policy learning.
improved trajectory where the red agent moves closer to the
We also statistically analyze 1000 randomly sampled episodes from the exp-m dataset. Figure B displays the mean ADE (average displacement error), which measures the average deviation of entities from the observations of all agents, and the average number of underperforming agents per episode. It evident that MADiTS reduces the mean ADE by almost half compared with w/o DC+PN, indicating that it generates a new trajectory that adheres more to environmental dynamics. Additionally, the original imbalanced dataset, the number of random agents is set to 1. In the stitching-ge stitching-generated generated dataset, we consider the number of underperforming agents in trajectory as the number of uncovered landmarks at the end of the episode. This value drops from 0.391 of w/o DC+PN to 0.094 of our method, demonstrating that PN significantly improves the quality of generated trajectories. How each module MADITS influences the coordination ability of the learned policies could be found in the next section.
5.4 ABLATION AND SENSITIVITY STUDIES
In this section, first investigate the effectiveness of each module in MADITS. We conduct abla- tion studies by applying BC on the exp-m dataset from the MPE benchmark after data augmentation. Specifically, we test 3 variants: (1) w/o DC: removes the bidirectional environmental dynamics constraint mechanism, treating all trajectory segments generated by the diffusion model as compliant with environmental dynamics, (2) w/o PN: removes partial noising, not addressing any spatial imbalance between agents, and (3) w/o DC+PN: removes both of the aforementioned modules. As shown in Figure[(a) we can find that the average returns of all the variants significantly exceed that of the original dataset, confirming the necessity of data augmentation for imbalanced datasets. When the dynamics constraint is removed, w/o DC suffers from performance degradation in all MPE tasks, indicating that generated trajectories with dynamics inconsistency learning. Furthermore, the variant w/o PN also demonstrates a decrease in performance, showing that iterative individual trajectory optimization effectively improves data quality. Removing both
be harmful to policy

---

Published as a conference paper at ICLR 2025
50-
3pread
Wors
m-mm--bffore
60
exp-i--ate fee exp-s-belore
Etiloda Rutumn
80-
Epiande
30-
30
Methods
Perturatton durstion
(a)
(b)
Figure 4: (a) Ablation results exp-m datasets in MPE. (b) The impact of reconstruction threshold Orecon- (c) results on CN in various perturbation settings. All the results are computed by the average return of BC policies trained datasets with different random seeds.
modules, w/o DC+PN achieves worse performance in most tasks, emphasizing their indispensable roles in MADiTS. The full method consistently outperforms other variants in all tasks, illustrat- ing that the designed modules within MADiTS can complement each other and contribute to the effective data augmentation.
Next, MADITS includes multiple hyperparameters, conduct experiments to examine their sensitivity. One of the most important hyperparameters the reconstruction threshold Gecon which controls the strictness of the dynamies constraint. If it's too small, stitching complete trajectory would take an unacceptable amount of time. On the other hand, if too large, overly lenient constraint will fail to filter out low-quality trajectories. By grid search, choice in CN as shown in Figure 4bb) More results on other hyperparameters, like augmentation ratio raug and regeneration limit Lima. could be found in Appendix H.3
find that órecon 0.01 is the best
5.5 EVALUATION IN ADDITIONAL PERTURBATION SETTINGS
Finally, we further explore MADiTS's robustness against various perturbation settings. Specifically, we test the data augmentation effects under varying severity of imbalance in the CN tasks, by setting different duration of temporal imbalance, denoted fmmb {2, 6, 8). Using the average return of the BC policy in the online environment as the quality metric, we assess the changes in dataset quality before and after stitching, as shown in Figure [4c] The results demonstrate that MADITS consistently enhances data quality, even with more severe imbalanced datasets, proving its robustness and generalization. In addition to fully trained expert policies, we also collect datasets by applying perturbations to medium-perforee ment could be found in Appendix H.4
policies, and similar results about quality improve-
CLOSING REMARKS
In this work, we introduce MADITS, a novel diffusion-based data augmentation pipeline that sig- nificantly enhances the performance of offline MARL algorithms, particularly on datasets with temporal and spatial imbalances. Empirical evaluations on imbalanced datasets across multiple benchmarks highlight the effectiveness of MADITS improving dataset quality. The method operates under the assumption that the quality of an agent's 's behavior over time is reflected in its contribution team returns, with high returns serving as the criterion for the diffusion model to capture its distribution. However, this assumption may not hold in environments with extremely sparse rewards. Additionally, in extreme scenarios where the diffusion model lacks exposure to high-quality cooperative segments during training, generating Out-Of-Distribund becomes a challenging problem in data generation. Incorporating external knowledge, such as Large Language Models (LLMs) (Sun et al. 2024), offers promising avenue for addressing this issue. Future work will aim to overcome these limitations and expand the method's applicability to more complex multi-agent systems, such those involving multiple embodied agents (Liu et al. 2024b).
(OOD) segments
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work is supported by the National Science Foundation of China(62495093, U24A20324), the Natural Science Foundation of Jiangsu (BK2024119, BK20243039), We thank Tencent AI Arena for their support, and the anonymous reviewers for their support and helpful discussions on improving the paper.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
Johannes Ackermann, Volker Gabler, Takayuki Osa, and Masashi Sugiyama. Reducing over- estimation bias in multi-agent domains using double centralized critics. arXiv preprint arXiv:1910001665 2019.
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B Tenenbaum, Tommi S Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? In International Conference on Learning Representations, 2022.
Stefano Albrecht, Filippos Christianos, and Lukas Schäfer. Multi-agent reinforcement learning: Foundations and modern approaches. MIT Press, 2024.
Jose A Medina, Michael Gillhofer, Michael Widrich, Thomas Unterthiner, Johannes Brand- stetter, and Sepp Hochreiter. Rudder: Return decomposition for delayed rewards. In Advances in Neurul Information Processing Systems, pp. 13566-13577 2019.
Alex Beeson and Giovanni Montana. Balancing policy constraint and ensemble size in uncertainty- based offline reinforcement learning. Machine Learning, 1331):443-488, 2024.
Avinandan Bose, Simon Shaolei Du and Maryam Fazel. Offline multi-task transfer rl with representational penalization. arXiv preprint (tt.:2402.1157000 2024.
Ian Char, Viraj Mehta, Adam Villaflor, John M Dolan, and Jeff Schneider. Bats: Best action trajectory stitching. arXiv preprint arXiv:2204 12026, 2022.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch, Decision transformer: Reinforcement learning via sequence modeling. In Advances in Neural Information Processing Systems, pp. 15084-155 15097, 2021.
Xiaocong Chen, Siyu Wang, Julian McAuley, Dietmar Jannach, and Lina Yao. On the opportunities and challenges of offline reinforcement learning for recommender systems. ACM Transactions on Information Systems, 42(6):1-26, 2024.
Jifci Deng, Seppo Sierla, Jie Sun, and Valeriy Vyatkin. Offline reinforcement learning for industrial process control: case study from steel industry. Information Sciences, 632:221-231, 2023.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In
Advances in Neural Information Processing Systems, pp. 8780-8794, 2021. 2021.
Wei Du and Shifei Ding. A survey on multi-agent deep reinforcement learning: from the perspective of challenges and applications. Artificial Intelligence Review, 54(5):3215-3238 2021.
Benjamin Ellis, Jonathan Cook, Skander Moalla, Mikayel Samvelyan, Mingfei Sun, Anuj Mahajan, Jakob Foerster, and Shimon Whiteson. Smacv2: An improved benchmark for cooperative multi- agent reinforcement learning. Advances in Neural Information Processing Systems, volume 36, 2023.
Claude Formanek, Asad Jeewa, Jonathan Shock, and Amu Pretorius. Off-the-gris grid marl: Datasets with baselines for offline multi-agent reinforcement learning. arXiv preprint arXiv:2302.00522 00521, 2023.
11

---

Published as a conference paper at ICLR 2025
Scott Fujimoto, David Meger, and Doina Precup. Off Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning, pp. 2052-2062) 2019.
Sven Gronauer and Klaus Diepold. Multi-agent deep reinforcement learning: A survey, Artificial Intelligence Review, 55(2))895-9944 2022.
Cong Guan, Feng Chen, Lei Yuan, Zongzhang Zhang. and Yang Yu. Efficient communication via self-supervised information aggregation for online and offline multiagent reinforcement learning. IEEE Transactions Neural Networks and Learning Systems, 2024.
Jun Guo, Yonghong Chen, Yihang Hao, Zixin Yin, Yin Yu, and Simin Li. Towards comprehensive testing on the robustness of cooperative multi-agent reinforcement learning. arXiv preprint arXiv:2204.07932, 2022.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2019.
Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, and Xue- long Li. Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning. Advances in neural information processing systems, 2024.
Haoyang He. A survey on offline model-based reinforcement learning. arXiv preprint arXiv:2305.033: 033660 2023.
Charles A Hepburn and Giovanni Montana. Model-based trajectory stitching for improved offline reinforcement learning. arXiv preprint arXiv:2211.11603, 2022.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NewrlPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, pp. 6840-6851, 2020.
Zhang-Wer Wei Hong, Aviral Kumar, Sathwik Karnik, Abhishek Bhandwaldar, Akash Srivastava, Joni Pajarinen, Romain Laroche, Abhishek Gupta, and Pulkit Agrawal. Beyond uniform sampling: Offline reinforcement learning with imbalanced datasets. In Advances in Neural Information Processing Systems, pp. 4985-5009, 2023.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, pp. 9902-9915, 2022.
Jiechuan Jiang and Zongqing Lu. Offline decentralized multi-agent reinforcement learning. In European Conference on Artificial Intelligence, pp. 1148-1155, 2021.
Amirhossein Kazerouni, Ehsan Khodapanah Aghdam, Moein Heidari, Reza Azad, Mohsen Fayyaz, llker Hacihaliloglu, and Dorit Merhof. Diffusion models in medical imaging: A comprehensive survey. Medical Image Analysis, 88: 102846, 2023.
Maurice G Kendall. The treatment of ties in ranking problems. Biometrika, 33(3):239-255 1945.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-learning for offline reinforcement learning. In Advances in Neural Information Processing Systems, pp. 1 1199-1191, 2020.
Nathan Lambert, Markus Wulfmeeirr William Whitney, Arunkumar Byravan, Michael Bloesch, Vibhavari Dasagi, Tim Hertweck, and Martin Riedmiller. The challenges of exploration for offline reinforcement learning. arXiv preprint 2rXiv:2201. 2022.
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. In Advances in neural information processing systems, pp. 19884-19895, 2020.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives open problems. arXiv preprint arXiv:2005.01643, 2020.
12

---

Published as a conference paper at ICLR 2025
Guanghe Li, Yixiang Shan, Zhengbang Zhu, Ting Long, and Weinan Zhang. Diffstitch: Boosting of- fline reinforcement learning with diffusion based trajectory stitching. In International Conference Machine Learning, pp. 28597-28660 2024a.
Xinran Li, Zifan LIU, Shibo Chen, and Jun Zhang. Individual contributions as intrinsic exploration scaffolds for multi-agent reinforcement learning. In International Conference on Machine Learning. 2024b.
Xinran Li, Ling Pan, and Jun Zhang. Kaleidoscope:
Learnable masks for heterogeneous multi-agent reinforcement learning. In Advances in Neural Information Processing Systems, 2024c.
Zhuoran Li. Ling Pan, and Longbo Huang. Beyond conservatism: Diffussio policies in offline
multi-agent reinforcement learning. arXiv preprint arXiv:2307.01442, 2023.
Tenglong Liu, Yang Li, Yixing Lan, Hao Gao, Wei Pan, and Xin Xu. Adaptive advantage-guided policy regularization for offine reinforcement learning. arXiv preprint arXiv:2405. 19909, 2024a.
Yang Liu, Weixing Chen, Yongjie Bai, Jingzhou Luo, Xinshuai Song, Kaixuan Jiang, Zhida Li, Ganlong Zhao, Junyi Lin, Guanbin Li, et al. Aligming cyber space with physical world: A comprehensive survey on embodied ai. arXiv preprint arXiv:2407.06886 2024b.
Ryan Lowe, Yi Wu, Aviv Tamar, Jean Harb Pieter Abbeel, and Igor Mordatch. Multi-agent actor-critic for mixed cooperative-conpettiie environments. In Advances in Neural Information Processing Systems. pp. 6379-639999 2017.
Fan-Ming Luo, Tian Xu, Hang Lai, Xiong-Hur Chen, Weinan Zhang, and Yang Yu. survey on model-based reinforcement learning. Science China Information Sciences, 67(2):12110, 2024.
Xueguang Lyu, Yuchen Xiao, Brett Daley, and Christopher Amato. Contrasting centralized and decentralized critics in multi-agent reinforcement learning. In Proceedings of the International Joint Conference on Autonomous Agents and Multiagent Systems, pp. 844-852, 2021.
Yixiu Mao, Hongchang Zhang, Chen Chen, Yi Xu, and Xiangyang Ji. Supported value regularization for offline reinforcement learning. In Advances in Neural Information Processing Systems, 2024.
Linghui Meng, Jingqing Ruan, Xuantang Xiong, Xiyun Li, Xi Zhang, Dengpeng Xing, and Bo Xu. M3: Modularization for multi-task and multi-agent offline pre-training. In Proceedings of the International Conference Autonomous Agents and Multiagent Systems, pp. 1624-1633, 2023a.
Linghui Meng, Muning Wen, Chenyang Le, Xiyun Li, Dengpeng Xing, Weinan Zhang, Ying Wen, Haifeng Zhang, Jun Wang, Yaodong Yang, et al. Offline pre-trained multi-agent decision transformer. Machine Intelligence Research, 20(2):233-248. 2023b.
Frans A Oliehoek, Christopher Amato, al. A concise introduction to decentralized POMDPs, volume 1. Springer, 2016.
OpenA1(2024). Video generation models world simulators. https:// /ttp:///ppa...... com/ ressarchh/video-eeneeteee generation models as world simulators] 2024.
Afshin Oroojlooy and Davood Hajinezhad. A review of cooperative multi-agent deep reinforcement learning. Applied Intelligence, pp. 1-46, 2022.
Ling Pan, Longbo Huang, Tengyu Ma, and Huazhe Xu. Plan better amid conservatism: Offline multi-agent reinforcement learning with actor rectification. In International Conference on Machine Learning. pp. 17221-17237, 2022.
Bei Peng, Tabish Rashid, Christian Schroeder de Witt, Alexae-Ae Kamienny, Philip Torr, Wendelin Böhmer, and Shimon Whiteson. Facmac: Factored multi-agent centralised policy gradients. In Advances in Neural Information Processing Systems, volume 34, pp. 12208-12221, 2021.
Rafael Figueiredo Prudencio, Marcos ROA Maximo, and Esther Luna Colombini. A survey offline reinforcement learning: Taxonomy, review, and open problems. IEEE Transactions on Neural Networks and Learning Systems, pp. 1-0, 2023.
13

---

Published as a conference paper at ICLR 2025
Yuhang Ran, Yi-Chen Li, Fuxiang Zhang, Zongzhang Zhang, and Yang Yu. Policy regularization with dataset constraint for offline reinforcement learning. In International Conference on Machine Learning. pp. 28701-28711 2023.
Tabish Rashid, Mikayel Samvelyan, Christian Schroeder, Gregory Farquhar, Jakob Foerster, and Shimon Whiteson, Qmix: Monotonic value function factorisation for deep multi-agent reinforce- ment learning. In International Conference on Machine Learning, pp. 4295-4304, 2018.
Erick Rosete- Beas, Oier Mees, Gabriel Kalweit, Joschka Boedecker, and Wolfram Burgard. Latent plans for task-agnostic offline reinforcement learning. In Conference on Robot Learning, pp- 1838-1849, 2023.
Mikayel Samvelyan, Tabish Rashid, Christian Schroeder de Witt, Gregory Farquhar, Nantas Nardelli, Tim GJ Rudner, Chia-Man Hung. Philip HS Torr, Jakob Foerster, and Shimon Whiteson. The starcraft multi-agent challenge. In Proceedings of the 18th International Conference on Autonomous Agents and MultiAgent Systems, pp. 2186-2188 2019.
Guillaume Sartoretti, Justin Kerr, Yunfei Shi, Glenn Wagner, TK Satish Kumar, Sven Koenig, and Howie Choset. Primal: Pathfinding via reinforcement and imitation multi-agent learning. IEEE Robotics and Automation Letters, 4(3):2378-2) 2019.
Kajetan Schweighofer, Andreas Radler, Corius-Constantin Dinu, Markus Hofmarcher, Vihang Patil, Angela Bitto-Nemling, Hamid Eghbal-zadeh, and Sepp Hochreiter. A dataset perspective on offline reinforcement learning. arXiv preprint arXiv:2111.0477 2021.
Jianzhun Shao, Yun Qu, Chen Chen, Hongchang Zhang, and Xiangyang Ji. Counterfactual conservative q learning for offline ageltt reinforcement learning. In Advances in Neural Information Processing Systems, 2023.
Ruizhe Shi, Yuyao Liu, Yanjie Ze, Simon Shaolei Du, and Huazhe Xu. Unleashing the power of pre-trained language models for offline reinforcement learning. In The Twelfth International Conference on Learning Representations, 2023.
Siddharth Singi, Zhanpeng He, Alvin Pan, Sandip Patel, Gunnar A Sigurdsson, Robinson Piramuthu, Shuran Song, and Matei Ciocarlie. Decision making for numan-in-the-lone robotic agents via uncertainty-awart reinforcement learning. In 2024 IEEE International Conference on Robotics and Automation (ICRA), pp. 7939-7945. IEEE, 2024.
Jiaming Song, Hongyu Ren, Dorsa Sadigh, and Stefano Ermon. Multi-agent generative adversarial imitation learning. Advances in neural information processing systems, 2018.
Yang Song, Jascha Sohl-Dickstein. Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference Learning Representations, 2021.
Chuanneng Sun, Songjun Huang, and Dario Pompili. Llm-based multi-agent reinforcement learn- ing: Current and future directions. arXiv preprint arXiv:2405.1 106, 2024.
Mukund Sundararajan, Ankur Taly, and Qiqi Yan. Axiomatic attribution for deep networks. In International Conference Machine Learning, pp. 3319-3328, 2017.
Peter Sunehag, Guy Lever, Audrunas Gruslys. Wojciech Marian Czarnecki, Vinícius Flores Zam- baldi, Max Jaderberg, Marc Lanctot, Nicolas Sonnerat, Joel Leibo, Karl Tuyls, and Thore Graepel. Value-decompositioon networks for cooperative multi-agent learning based on team reward. In International Conference on Autonomous Agents and Multiagent Systems, pp. 2085- 2087, 2018.
2018.
Qi Tian, Kun Kuang, Furui Liu, and Baoxiang Wang. Learning from good trajectories in offline multi-agent reinforcement learning. In Association for the Advancement of Artificial Intelligence, pp. 11672-11680. 2023.
Wei-Cheng Tseng, agent reinforcement learning with knowledge distillation. In Advances Neural Information Processing Systems, pp. 226-237, 2022.
Johnson Wang, Yen-Chen Lin, and Phillip Isola. Offline multi-
14

---

Published as a conference paper at ICLR 2025
Elise van der Pol, Herke Hoof, Frans A Oliehoek, and Max Welling. Multi-agent mdp homomorphic networks. In International Conference on Learning Representations, 2021.
Jianhong Wang, Wangkun Xu, Yunjie Gu, Wenbin Song, and Tim C Green. Multi-agent reinforce- ment learning for active voltage control on power distribution networks. In Advances in Neural 2021.
Information Processing Systems. pp.
Jianhong Wang, Yang Li, Yuan Zhang, Wei Pan, and Samuel Kaski. Open ad hoc teamwork with cooperative game theory. International Conference on Machine Learning, 2024a.
Xiangsen Wang and Xianyuan Zhan. Offline multi-agent reinforcement learning with coupled value factorization. In Proceedings of the International Conference on Autonomous Agents and Multiagent Systems, pp. 2781-2783, 2023.
Xiangsen Wang, Haoran Xu, Yinan Zheng, and Xianyuan Zhan. Offline multi-agent reinforcement learning with implicit global-to-local value regularization. Advances Neural Information Processing Systems, pp. 52413-52429, 2023.
Xiaofeng Wang, Zheng Zhu, Guan Huang, Xinze Chen, Jiagang Zhu, and Jiwen Lu. Drivedreamer: Towards real-world-driven world models for autonomous driving. In European Conference on Computer Vision, 2024b.
Xihuai Wang, Zhicheng Zhang, and Weinan Zhang. Model-based multi-agent reinforcement learning: Recent progress and prospects. arXiv preprint arXiv:2203. 10603, 2022.
Zhuohan Wang and Carmine Ventre. A financial time series denoiser based on diffusion models. In Proceedings of the 5th ACM International Conference on Al in Finance, pp- 72-80, 2024.
Kasun Weerakoon, Adarsh Jagan Sathyamoorthy, Mohamed Elnoor, and Dinesh Manocha. Vapor: Legged robot navigation in unstructured outdoor environments using offline reinforcement learn- ing. In 2024 IEEE International Conference on Robotics and Automation (ICRA), pp. 10344- 10350. IEEE, 2024.
Muning Wen, Jakub Grudzien Kuba, Runji Lin, Weinan Zhang, Ying Wen, Jun Wang, and Yaodong Yang. Multi-agent reinforcement learning is a sequence modeling problem. In Advances in Neural Information Processing Systems 35, pp. 16509-16521, 2022.
Chengjie Wu, Pingzhong Tang, Jun Yang, Yujing Hu, Tangjie Lv, Changjie Fan, and Chongjie Zhang. Conservative offline policy adaptation in multi-agent games. Advances in Neural Information Processing Systems, 2024.
Ke Xue, Jiacheng Xu, Lei Yuan, Miqing Li, Chao Qian, Zongzhang Zhang, and Yang Yu. Multi- agent dynamic algorithm configuration. In Advances in Neural Information Processing Systems, pp. 20147-20161, 2022.
Qisong Yang. Thiago D Simão, Simon H Tindemans, and Matthijs TJ Spaan. Wesac: Worst-case soft actor critic for contt-ccontt reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 10639-10646, 2021a
Sherry Yang, Ofir Nachum, Yilun Du Jason Wei, Pieter Abbeel, and Dale Schuurmans. Foun- dation models for decision making: Problems, methods, and opportunities. arXiv preprint 1.Xi:22303.4419.. 2023.
Tianpei Yang, Weixun Wang, Hongyao Tang, Jianye Hao, Zhaopeng Meng, Hangyu Mao, Dong Li, Wulong Liu, Yingfeng Chen, Yujing Hu, et al. An efficient transfer learning framework for multiagent reinforeement learning. In Advances in neural information processing systems, pp- 17037-17048, 2021b.
Yaodong Yang, Jianye Hao, Guangyong Chen, Hongyao Tang, Yingfeng Chen. Yujing Hu, Changjie Fan, and Zhongyu Wei. Q-value path decomposition for deep multiagent reinforcement learning. In International Conference on Machine Learning, pp. 10706-10715, 2020.
15

---

Published as a conference paper at ICLR 2025
Yigin Yang, Xiaoteng Ma, Chenghao Li, Zewu Zheng, Qiyuan Zhang, Gao Huang, Jun Yang, and Qianchuan Zhao. Believe what you see: Implicit constraint approach for offline multi-agent reinforcement learning. In Advances in Neural Information Processing Systems, pp. 10299- 10312, 2021c.
Denis Yarats, llya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International conference on learning representations, 2021.
Chao Yu, Akash Velu, Eugene Vinitsky, Jiaxuan Gao, Yu Wang, Alexandre Bayen, and Yi Wu. The surprising effectiveness of PPO cooperative multi-agent games. In Advances in Neural Information Processing Systems, pp. 24611-24624 2022.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. In Advances in Neural Information Processing Systems, pp. 14129-14143 2020.
Tianhe Yu, Aviral Kumar, Yevgen Chebotar, Karol Hausman, Sergey Levine, and Chelsea Finn. Conservative data sharing for multi-task offline reinforcement learning. In Advances in Neural Information Processing Systems, pp. 11501-11516, 2021.
Yang Yu. Towards sample efficient reinforcement learning. In IJCAL, pp. 5739-5743, 2018.
Lei Yuan, Ziqian Zhang, Lihe Li, Cong Guan, and Yang Yu. A survey of progress on cooperative multi-agent reinforcement learning in open environment. arXiv preprint
2023.
Fuxiang Zhang, Chengxing Jia, Yi-Chen Li, Lei Yuan, Yang Yu, and Zongzhang Zhang. Discovering generalizable multi-agent coordination skills from multi-task offline data. International Conference on Learning Representations, 2023.
Kaiqing Zhang, Zhuoran Yang, Han Liu, Tong Zhang, and Tamer Basar. Finite-sample analysis for decentralized batch multiagent reinforcement learning with networked agents. IEEE Transactions on Automatic Control, 66( 56(12):5925-594) 2021a.
Shenao Zhang, Li Shen, and Lei Han. Learning meta representations for agents in multi-agent reinforcement learning. arXiv preprint arXiv:2108. 12988, 2021b.
Yang Zhang, Shixin Yang, Chenjia Bai, Fei Wu, Xiu Li, Xuelong Li, and Zhen Wang. Towards efficient llm grounding for embodied multi-agent collaboration. arXiv preprint arXiv:2405. 14314, 2024.
Ziqian Zhang. Offpymarl: Benchmarked implementations of offline reinforcement learning algo- rithms. hhttps:// com/ bo2g-boooofflinne-eeet -ramorr--fffpyaaaaar 2023.
Changxi Zhu, Mehdi Dastani, and Shihan Wang. survey of multi-agent reinforcement learning with communication. arXiv preprint arXiv:2203.0899 2022.
Zhengbang Zhu, Minghuan Liu, Liyuan Mao, Bingyi Kang. Minkai Xu, Yong Yu, Stefano Ermon, and Weinan Zhang. Madiff: Offline multi-agent learning with diffusion models. arXiv preprint arXiv:2305. 173300 2023a.
Zhengbang Zhu, Hanye Zhao, Haoran He, Yichao Zhong, Shenyu Zhang, Yong Yu, and Weinan Zhang. Diffusion models for reinforcement learning: A survey, arXiv preprint arXiv:2311.01223, 2023b.
A MORE DETAILS ABOUT RELATED WORK
Multi-agent Reinforcement Learning (MARL) has made significant progress, making it well- suited for addressing large-scale, complex, real-time, and uncertain real-world problems. Model- ing such problems as single-agent systems often inefficient and fails to align with real-world conditions. In contrast, formulating them Multi-Agent Systems (Abrrcet et al. 2024 is
16

---

Published as a conference paper at ICLR 2025
typically more appropriate. Moreover, when agents share common objective, the problem becomes a cooperative MARL task(Oroojlooy & Hajinezhad| 2022). which has seen substantial advancements across diverse domains such as pathfinding (Sartoretti et all, 2019). active voltage control (Wang et at at (021), dynamic algorithm configuration (Xue et al. 2022), and Large Language Model (L.LM) applications (L.iu et al. 2024b). Numerous approaches have been developed enhance agent coor- dination, including volicy-based methods like MADDPG (Lowe et al.. 2017) and MAPPO (Yu et al. 2022). value-based techniques such as VDN (Sunchag et 2018)) and QMIX (Rashid al. 2018). as well as approaches leveraging transformer architectures (Wen et al. (022). These methods have demonstrated impressive coordination capabiliitiss across various tasks, including SMAC, Hanabi, and GRF (Yu et al.) 2022). Beyond these approaches and their variants, other methods have been proposed to explore cooperative MARL, such as efficient communication polices (Zhu et al. 2022).
offline policy deployment (Zhang et al. (023), model learning in MARI. (Wang et al. 2022).
robustness to perturbations (Guo ct al., 2022). sparse reward environments (Li et al. 2024b). ad
hoc teamwork (Wang et al., [0224)) and training paradigms like CTDE (Centralized Training with Decentralized Execution) (Lyu et al. 2021). Additionally, several methods have been introduced to design testbeds for MARL algorithms, such as SMAC(Samvelyan et al.] 2019). offline data
environments (Formanek et al. 2023), and communication testbeds (Guan et al. 20224) etc.
Offiin Multi-agent Reinforcement Learning Offiin reinforcement learning (Revne
eet al. 2020 Prudencio et al., 2023) has garnered significant research interest recently, focusing
on data-driven training paradigm that eliminates the need for direct interaction with the aivironment(rrueecioo et al. 2023). Early works (Fujimoto et al., 2019) primarily addressed the challenges of distributional shift in offline learning, focusing on eehavior-consrriine policies to reduce extrapolation errors when estimating unseen data (Kumar et al.) 2020]. Offline Multi-Age Agent Reinforcement Learning (MARL) is an emerging and promising area of research (Zhang et al.) 2021a). One class of offline MARI. approaches focuses on learning policies from offiie data with policy constraints. For example, ICQ (Yang et al.) 2021c) effectively mitigated extrapolation errors by leveraging only offline data in MARL scenarios. MABCQ (Jiang & Lu 2021) introduced fully decentralized setting for offline MARL, employing techniques like value bias and transfer normalization for more efficient learning. OMAR (Pan al., 2022) combined first-order policy gradients with zero-order optimization to circumvent local optima. MADT (Meng et al., (023b) utilized the sequential modeling power of transformers, applying it both offline and online MARL tasks. Research (Tian et al. (023)) further explored offline MARL by explicitly accounting for the diversity of agent trajectories and proposed new framework called Shared Individual Trajectories (SIT). In another work (Teeng et al., 2022), a teacher policy is first trained using full observation, action, and reward data, after which student policies are distilled from the teacher policy, capturing structural relationships between the teacher and agent behaviors. ODIS (Zhang eet al. 2023) introduced a novel algorithm for offline MARL, focusing on discovering cooperative skills from multi-task data. Recently, the Off-the-Grid MARL (OG-MARL) framework was released (Formanek et al., (023), offering a new benchmark for offline MARL dataset generation and algorithm evaluation. M3 (Meng et al.) 2023a) innovatively proposed multi-task and multi- agent offline pre-trane pre-training modules to learn higher-Ievel transferable policy representations. Lastly, OMAC (Wang & Zhan 2023) introduced
offline MARL algorithm based on coupled value decomposition, decomposing global value functions into local and shared components while ensuring consistent credit assignment between global state values and Q-value functions.
Sample Efficient Reinforcement Learning tackles the challenge of minimizing the number of interactions an agent requires with its environment learn an optimal policy. In traditional rein- forcement learning (RL), agents often need extensive exploration, which can be costly or impractical in real-worid applications like robotics, autonomous driving, and healthcare, where data collection is expensive or risky. The goal of sample-efficient RL is to maximize learning performance while minimizing environmental interactions (Yu) 2018). In single-agent RI, common techniques in- clude image-based transformations such as cropping, rotation, and flipping, which are particularly effective in visual tasks, as demonstrated by RAD (Laskin et al., 2020) and DrQ (Yarats et al.) 2021 Other approaches, such as perturbing or resampling from the experience replay buffer, encourage exploration. Model-based methods like Dreamer (Hafner et al., 2019) and MOPO (Yu et al., 2020) leverage environmental dynamics models to generate synthetic data. In multi-agent reinforcement learning (MARL), the complexity of agent interactions presents unique challenges for data augmentation, as improper methods can lead system non-stationarity. Some approaches
17

---

Published as a conference paper at ICLR 2025
address this by exploiting the structure of multi-agent systems, utilizing symmetry (van der Pol ee al. (021). policy similarity (Yang et al. 2021b). parameter sharing (Li et al. (024c), or learning meta-policy for multi-task environments Zhang et al. 2021b). Inspired by the rapid development of foundation models Yang et al. 2023). J recent work has begun to leverage diffusion models (Zhu et al. 2023a Li al. 2023] Zhu al. (023b), Transformers (Wen al. 2022). and large language models (LLMs) for MARL (Zhang et al.) 2024), showing improvements in some settings, While
these approaches contribute to sample efficiency, no current methods focus on generating trajectories
and performing trajectory stitching for enhanced policy learning.
B MORE COMPARISONS WITH OTHER OFFLINE MARL METHODS
Our method focuses on addressing sample efficiency issues by leveraging diffusion models to per- form trajectory stitching for data augmentation in temporally and spatially imbalanced datasets. This approach allows offline MARL algorithms to achieve better performance by learning from the enhanced dataset. In this section, we outine other prominent methods in the field of offline MARL and compare their differences with MADiTS.
SIT TTian et al. 2023) is offline MARL method that learns an effective joint policy from agent-wise imbalanced multi-agent datasets, addressing spatial imbalance in particular. It uses an attention-based reward decomposition network to perform offline credit assignment, identifying high-quality individual trajectories for sharing among agents. Additionally, it employs a graph attention network for conservative policy training. Unlike SIT, which directly learns a joint policy from imbalanced datasets, our approach focuses on addressing sample efficiency issues and introduces a data-oriented augmentation pipeline. MADiTS employs a diffusion model-based trajectory stitching mechanism to enhance dataset quality. The augmented datasets can be used by any agnostic offline MARL algorithm, offering flexibility and modularity.
DOM2 (Li et al. 2023 is diffusion model-based offline MARL algorithm designed to improve policy expressiveness and diversity. It achieves significant gains in performance, generalization, and data efficiency by employing
accelerated solver for diffusion-based policy construction and policy regularizer, while also scaling up the dataset size to enhance policy learning. DOM2 demon- strates outstanding results on balanced datasets, outperforming traditional conservatism-baase of- fline MARL methods. In contrast, MADITS focuses on addressing the complex data structures inherent in imbalanced datasets. By tackling temporal and spatial imbalances, MADITS enhances dataset quality, providing better support for other offline MARL algorithms and improving overall performance in imbalanced scenarios.
MADiff (Zuu et al. 2023a) is another diffusion model-based offline MARL algorithm, designed to predict future joint actions for -eakingn-making by modeling teammates. It employs an attention- based diffusion model to capture joint observation sequences and infers actions for planning. MAD- iff achieves superior performance on balanced datasets, excelling in standard offline MARL settings. By contrast, MADITS specifically targets the data imbalance problem in offline MARL, addressing sample efficiency issues. While MADiff improves learning efficiency in balanced settings, MADITS innovatively focuses sample efficiency issues for imbalanced datasets. To emphasize the differ- ences, we compare the performance of MADiff extended to data augmentation with MADiTS across various environments. The results demonstrate the superior effectiveness of MADITS in addressing data imbalance challenges.
C THEORETICAL PROOF
Before presenting the proof of Equation 7 we first introduce Integrated Gradient (IG) (Sundararajan et al. 2017), which is an interpretability technique for neural networks, initially proposed in the field of computer vision. IG leverages path integrals of gradients along a path between a baseline input and the actual input to summarize how each feature affects the output of the deep neural network as the model's prediction moves from F(b) to F F(x) along the line connecting the b and x.
Formally, given an input x, IG computes the attribution of each feature by accumulating the gradients of function F with respect to feature i, using a straight line as the integration path, defined by 7(a)- = x' + a(x x²), where x' is the baseline and is the input. Then Integrated Gradient for
18

---

Published as a conference paper at ICLR 2025
feature i is then given by:
(G) = (I; - 24)) f
0F(z' + +a(x - -x))==)))) da,
(10)
o==0
dz;
where @F(x) 0z, represents the gradient of function F at point with respect to the i-th feature.
Lemma 1. [andararajan et al. 2017 If : : + R is differentiable almost everywhere then
-((x) F()) F(b),
(11)
where (x)(); represents the integrated gradient of function F. along the path T w.r. the j-th dimension of the input x over a straight-line path and b is the baseline for the straight-line path.
Lemma T shows how to compute the straight-line path integral when the function is almost every- where differentiable. Due to excellent interpretability, IG is also used in reinforcement learning. et al. (019) applies integrated gradient to reinforcement
For instance, RUDDER learning, addressing the problem sparse delayed rewards in single-agent reinforcement learning and demonstrating excellent performance. By applying IG, researchers show:
Lemma 2. (Yngg et al. 2020) For any joint observation-acctio trajectory segment (o,, a,, -. ... 0r+h-1, a(th-1) from timestep to timestep t 1, we have:
T-1 (1Gg(enn) R=tjEx,
PathIG
(12)
JEX;
As is shown: the trajectory from timestep timestep & h 1 is split according to the timestep. 1+2
that is, T (T!*1 T -(th-2 where T!+1 represents the trajectory from (o,, a.) to (Ot+1, a.+1). The integral of the entire path segment can be obtained by summing the straight- line paths of each adjacent timestep. Lemma 2 shows how to calculate the path integral on the trajectory path in the field of reinforcement learning.
Based on Lemma T and Lemma 2 the proof of Equation can be given below:
Proof. According to the definition of return-to-go Å(x.. TT) = E+ r(xr):
id Tr) - r(x++h1))
RR(w,, Ti) -fi(-i)))-( R(12+1, Tr) + R(0 Tr) - M(124227T))
(13)
+...+ + A(xx41-2, T.) R(Z++h-1:71)
For given o and a,, the denoising process generates т, conditioned on a fixed expected return that adheres to a predefined distribution, there exists randomized function R such that the reward function can be expressed as r(x:) = R(xx)) R(x+11). By Lemma we have
R(x, Tr) r(21+h-1) R(x,)) (((+++1) R(2441) M(x+11 + ((()4412) - R(1z4n-1)
Σ IG, (z) 7+++1) Σ IG, (x: +... Σ
IG, (I;
jEv
jE
JE
Σ IG (z) (1) +G (G) (112)))))))))))))) + +IG, (Z:
jE
(14)
By Lemma have
Å(x T() r(0(4h-1)
Σ (G,((: (x; (Gy(;; X,42) + IG, (G;
jEa
PathIG (I; (1)=) + PathIG2 (zn + PathIGa(2; ( (t
(15)
19

---

Published as a conference paper at ICLR 2025
Classify 1.2, ... ,d according to the agents they belong to, we have:
R(z T() r()
PathIGi (I; TH 1) (z 1) +... PathIGa(m:77) 7111 1)
Σ
PathIG (ng (,g---1)) +
PathIG (1; Topta-1)
JEX1
JEX2
+ +...
PathIGay (ri (444t-1))
jEăn
PathIGy y
JEX,
(16)
This equation shows that the return-to-oo can be decomposed into the sum of integrated gradients of each agent's respective features, providing possibility of offline credit assignment between agents.
D CIRCULAR SHIFT
To facilitate cross-agent integration of excellent trajectories by allowing excellent individual behav- iors of one agent be learned by other homogeneous agents, we propose circular shift method to allows the excellent individual behaviors of one agent to be uniformly distributed across all identical agent positions. Therefore, we define the cyclic shift operator C C:T" + T such that:
a H H-1
o
of +
a ++
--111 H(+H-1
a
0
ai+
H
C
(17)
14H-1
044
0g
14
Specifically, we first compare the cumulative discounted rewards of different trajectory segments at the same timestep and select the top Teir % the joint observation-attio trajectory segments to form the set D good seg- For all trajectories in this set, assuming all agents are homogeneous, apply the circular shift operator n - times to each trajectory, adding each newly generated trajectory segment to the original dataset:
Dug + Degg (( T E D Dgood_seg 1 1sk n E,
(18)
Here, Deey denotes the dataset of trajectory segments obtained by partitioning the original dataset D into segments of length H. For hetergeneous settings, circular shift is conducted within groups of agents of the same state and action space. By applying circular shift, we can increase the proportion of excellent trajectory segments in the original imbalanced dataset, which facilitates the cross-agent stitching of excellent trajectories and thus improves the optimization effect for suboptimal agents.
E IMPLEMENTATION DETAILS AND HYPERPARAMETERS
We introduce the model implementation details of generative model Ga. inverse dynamics model faw. the forward dynamics model find and the team reward prediction model godd, The implemen- tation of offline MARL algorithms and values of hyperparameters are provided subsequently.
Specifically, we use MADiff (Zhu et al. 2023a for the implementation of our diffusion model, where U-Net is used to model the individual trajectories of agents and an attention layer is applied before every decoder block in the U-Net of all agents. Additionally, three-layer multi-layer perceptrons are used to implement the inverse dynamics model, forward dynamics model, and team reward prediction model, respectively, with ReLU used as the activation function between layers.
For the training stages of these models, the loss function is defined as:
(( Φ. th, w) = (k)
[||c ca(x (1 R(x)) + 36, A)]
(19)
+ Oldynamics Ldynamics (b, +b + Oreeward Creward (w),
20

---

Published as a conference paper at ICLR 2025
where (01, og.... ,0 is the observation sequence sampled from trajectory segments dataset D, k is sampled from uniform distribution (k)., p is the dropout rate for conditional diffusion
model, Ldynamics and Lrwardd losses in Equation and 8 Oldynamies and Areward hyperparameters.
For the implementation of offline MARL algorithms for evaluation, we use behavior cloning, OMIGA (Wang et al. 2023 and CFCQL (Shao et al. 2023 for evaluating the quality of datasets before and after data augmentation. Specifically, we use the 'Offline MARL framework
OfiPyMARL' codebase (Zhang, 2023) and default hyperparameters from GitHub to implement these algorithms, with training steps set to 100k.
The experiments
conducted servers outfitted with GeForce RTX 2080 Ti. We compare the computational costs of MADiTS and other data augmentation method on the Cooperative Navigation task. For MADiTS, the model training phase takes about 36 GPU hours of single GeForce RTX 2080 Ti and the trajectory stitching phase takes about 4 GPU hours. As comparison, MADiff takes about 36 GPU hours on the same device and MA-MBTS takes about 48 hours for looking for valid transitions for stitching. We can see that MADITS achieves better performance under comparable computational costs. including MBTS, MADiff, and MADiTS.
The normalized scores in Table and are calculated Smom 100 (S Smoor) (S.eil Snoor). ) where Seeil and Stoor represent the ceiling and floor scores for each task, determined by the deciles of episode returns in the expert datasets. Specifically, for the tasks Spread, Tag, World, and 4ant, the respective ceiling and floor scores are {1592, 971}. {801, 320}, {400, 182]. and (5976, 1033].
The hyperparameters used in our work can be categorized into two groups: environment-specific shown in Table 2 and vviroommnt-inddeeedeen shown in Table
Table 2: Envirommenntinndeeenddntt hyperparameters
Hyperparameter
Value
Hyperparameter
Value
10%
fond hidden dim
256
Taug
1:1
Ir
2e-4
w
1.2
training steps
le6
K
200
optimizer
Adam
condition dropout rate
0.25
batch size for stitching
256
batch size for training
32
Aimit
3
Oddynamics
10
clay hidden dim
256
Okewwar
0.01
hidden dim
256
IG step size
10
Table 3: Enviromment-inddpeenddn hyperparameters
Hyperparameters
CN
PP
World
2m_vs_L
3m
&recon
0.01
0.03
0.02
0.2
0.2
limic
5
5
5
10
10
Srank
2.33
2.33
2.33
1.5
2.33
R,
1800
1000
500
20
20
Hyperparameters
2s3z
12m
terran 5
zerg 5
5
4ant
Seeon
5
5
5
5
2
h limit
10
20
20
20
30
8rank
2.33
8
2.33
2.333
2.33
R.
20
20
20
20
4000
21

---

Published as a conference paper at ICLR 2025
F PSEUDOCODE
Algorithm Data Augmentation Process
Input: Original dataset D, generative model Gg. inverse and forward dynamics models f ffwd team reward prediction model y/dd
1: Train generative model Go, inverse dynamics model fln forward dynamics model ffrd team
reward prediction model gir according to Algorithm
2: Dang =
3: while (Daugl I Faug |D] do
4: Sample an initial joint observation 0₁ from D as the current observation.
5: Initialize Tgen = }or}. ). total regeneration count regen = O, single regeneration count
sigel
-.
6: Using {01} as the initial joint observation, a joint trajectory Tgen generated according to
Algorithm 3
7: Daug Daug (Tgen)
8: end while
9: D* + DDDD
10: return D
Algorithm 2 Model Training Process
Input: Original dataset DD generative model Go, inverse and forward dynamics models fin ffwd team reward prediction model g
Split D into trajectory segments Daeg lengh H, and find the top Teir % joint observation-action trajectory segments D sood based on the return ranking.
2: Apply circular shift described in Appendix D Doag + Dsag (T) T E D good seg. 1 k
n 1,k N].
3: Train diffusion model Gg based on Equation 2
4: Train inverse dynamics model f forward dynamics model findd and team reward prediction
model g.nd based on Equation and 8
The data augmentation process (Algorithm 1 contains conducting model training first (Algorithm 2 and generate trajectories by stitching (Algorithm B) During the model training process, first identify high-quality trajectory segments from the original dataset and apply circular shift (Lines 1-2). Then we train the diffusion model and other models required for the augmentation process (Lines 3-4). During data augmentation (Algorithm [T] line uses the diffusion model and other models obtained from the training process, and then lines generate a certain number of joint trajectories to expand the original dataset.
The process generating a single trajectory is detailed Algorithm 3 We generate a trajectory segment first (Lines 2-3), and apply dynamics constraint, truncating the trajectory to retain only parts that pass the checks (Lines 4-13). Then we discard trajectory segments that consistently fail to meet the constraints after multiple generation attempts (Lines 14-16). If identify any underperforming agents (Lines 17-20), we will optimize their behavior by partial noising (Line 21). Finally, apply dynamies constraint again the optimized joint trajectories, truncating them to retain only the valid parts (Lines 22-32).
G DETAILED DESCRIPTION OF THE EVALUATION ALGORIHTMM AND BASELINES
We provide more detailed introduction to the evaluation algorithms and baselines the experiment in this section.
OMIGA (Wang et al. 2023 OMIGA provides a structured approach to transform global value regularization into local value regularization, which supports efficient in-sample learning. This
22

---

Published as a conference paper at ICLR 2025
Algorithm 3 Single Trajectory Generation Process
Input: Original dataset D, generative model Gg. inverse dynamics model fer, forward dynamics gowd team reward prediction model godd
1: while A trajectory generation is not finished do
2: Initialize the valid trajectory segment Tvalid
3: Condition on the current observation O, to generate joint observation trajectory segment Tobs
length H according to Equation H do
4: for k=
= t, 1,
5:
Detect whether (Ōx, Ok+1) in Tobs satisfies the dynamics consistency, and obtain the joint action åд and the team reward rk- Valid then
6:
7:
Tvalid Tvalid U {āx; rk: ōx+1}.
8:
regen sings =
9:
else
10:
O, = =
regen f wt + 1, regen regen 12: end if
13:
end for
14:
if regen limiiimi regen Jimit then
15:
BREAK
16:
end if
17:
Inspect underperforming agents B B (Tvalid)
18:
if B 0 then
Tgen Tgen
Tvalid
20:
else
Use the diffusion model resampling obtain Tregenooss described in Equation for k = t,t 1...+........ h do
22:
Detect whether (Ōg, Õ ) in Tregen satisfies the dynamics consistency, and obtain the joint action àR and the team reward rk.
24:
Valid then
25:
Tgen Tgen U {āz, TR, 1x+1).
26:
else
27:
0, 0
28:
end
29:
end for
30:
end if
end while
32:
return Tgen-
method bridges the gap between multi-agent value decomposition and policy learning by integrating offline regularization techniques. facilitating principled multi-agent reinforcement learning.
CFCQL (Shao et al., 2023). CFCQL computes conservative value estimates for each agent in counterfactual manner and then aggregates them into an overall conservative value function. This approach avoids treating the multi-agent system as single high-dimensional entity and improves upon directly applying single-agent methods to multi-agent settings.
MBTS (Hepburn & Montana, 2022). MBTS is a model-based trajectory stitching technique that learns a state transition model and value function. By connecting high-quality trajectory segments, MBTS generates new trajectories that replace suboptimal data, effectively enhancing multi-agent performance through state-action exploration and improved dataset composition.
MADifT (Zhu et al. 2023a). MADiff is another diffusion-based offline MARL algorithm, designed to predict future joint observations for decision masiing by modeling teammates. It employs an attention-baal based diffusion model to capture joint observation sequences and infers actions for plan- ning. In our experiments, we compare the performance of MADiff extended to data augmentation with MADITS across various environments.
23

---

Published as a conference paper at ICLR 2025
Table 4: Evaluation results on additional multi-agent imbalanced datasets. The mean and standard eлrоr are computed based on the normalized average return or average battle won rate of the evaluation algorithms trained on the datasets, with 5 different random seeds. We bold the highest scores on exp-m and exp-s datasets, respectively.
Balanenl
Original
MA-METS
MADIIT
MADITS (Curs)
Riss
Algs
exp
exp-m
exp-m
exp-5
exp-s
BC
0.95 + 0.00
0.33 = 0.11
0.63 = 0.28
0.81 + 0.11
0.56 0.22
0.5 # 0.12
0.66 :: 0.21
1L58 IJ17
0.65 :: 0.12
OMIGA
0.97 d: 0.03
0.9 = 0.02
0.55 : 0.02
0.85 : 0.03
0.55 0.03
0.96 2: 0.03
0.86 =: 0.04
0.95 0.02
0.87 2: 0.02
CTCQL
0.90 1: 0.04
0.78 = 0.09
0.50 a 0.11
0.70 0.11
0.46 0.00
0.79 1: 0.07
0.53 11 0.09
0.50 0.05
0.63 2: 0,08
Average
BC
0.55 + 0.08
0.46 = 0.02
0.38 + 0.27
0.46 + 0.05
0.41 0.13
0.51 + 0.03
0.42 + 0.08
11:4 0JR
0.46 +: 0.09
5
OMIGA
0.59 # 0.06
0.55 = 0.15
0.54 0.11
0.58 + 0.08
0.56 0.04
0.6 # 0.06
0.57 = 0.04
0.70 0J06
0.61 :: 0.08
CECQL
0.63 1: 0.09
0.34 = 0.12
0.49 = 0.09
0.54 0,30
0.49 119::00.1 0.19
0.53 4: 0.06
0.50 =1 0.07
0.55 0.11
0.52 2: 0.06
0.30=
0.27
0.26 =i
0.33
OMIGA
=
0.31
0.36
0.33
CTCQL
0.51
0.44
0.37
0.45
0.47
0.51
0,46
6.00
12.58
3.68
10.58
20.29 6.57
Aask
52,30
33.12 =
10.99
8.72
30.61
34.35
4.47
48,02
CTCQL
7.30
32.41
24.00
8.07
48.43
20.20
32.40
16.44
30.72
15.85
36.92
26.22
43,00
33,42
H ADDITIONAL EXPERIMENT RESULTS
H.1 EVALUATION ON ADDITIONAL ENVIRONMENTS
In this section, we analyze MADITS's effectiveness on augmenting offline MARI datasets of more tasks, including SMAC 12m, SMACv2 terran_5_vs_5, zerg-5_vs_5, and MAMuloCo 4ant. The performance of various learned policies on these datasets are presented Table From the results, we can observe that MA-MBTS demonstrates only slight improvement over Original. MADiff improves upon the Original, highlighting the need for diffusion in data modeling. but overlooks dy- namics consistency and data imbalance, leading to performance gap with MADiTS on imbalanced datasets. On the contrary, MADITS outperforms the baseline method MA-MBTS and achieves better performance compared to MADiff most of datasets, highlighting its broad applicability. We can find that MADITS demonstrates enhanced sample efficiency across various methods, even as the number of agents increases.
H.2 EVALUATION ON BALANCED OFFLINE DATASETS
To thoroughly evaluate the general effectiveness of MADITS, we also utilize the multi-agent offline dataset in MPE with continuous action space built OMAR (Pan et al. 2022) for data augmentation, testing the effectiveness of our method on balanced datasets. This dataset comprises four different quality datasets collected from policies of varying quality trained by MATD3 (Ackermann et al.) 2019), including Expert, Medium, Medium-Replay, and Random. The Expert, Medium and Random datasets consist of one million samples each, generated by fully trained policies, perfoni-erforniing policies, or random policy in online environment. The Aediuum-Replay dataset records all data in the replay buffer before the policy reaching medium performance level. We assess the enhancement in dataset quality using behavior cloning policy, measuring the average return of the BC policy in the online environment as the quality metric.
From the results shown in Table 5 improvement over Original, while MADITS achieves the best performance despite the varying quality of datasets, which proves the robustness and generality of MADITS. Notably, in the mixed-
first observe that MA-MBTS demonstrates only slight
24

---

Published as a conference paper at ICLR 2025
Table 5: Evaluation results multi-agent balanced datasets. The mean and standard error are computed based on the average return of BC policies trained on the datasets, with 5 random seeds.
Envs
Dataset
Original
MA-MBTS
MADITS
expert
100.42 5.67
101.70 + 3.13
127.91 = 5.63
medium
80.46 = 5.05
81.20 8.21
91.31 2.24
CN
md-replay
32.17
84.47 + 8.04
random
2.69
2.79
22.38 4.59
expert
77.54
79.47 = 9.96
97.04 + 11.81
medium
54.14 10.19
53.19 7.99
63.25 10.15
Pp
md-replay
-6.72 2.39
25.04 + 11.45
random
-8.94 1.55
-9.12 1.36
-7.38 + 2.49
expert
76.89 + 23.08
73.26 + 17.45
138.10 34.01
medium
54.99 # 14.27
62.62 14.99
96.27 11.78
World
md-replay
19.28 + 4.06
20.49 2.09
29.46 4.65
random
6.19 = 3.08
5.92 = 2.95
14.01 + 2.95
quality md-replay datasets, the enhanced datasets after stitching can rival the quality of the medium dataset. This highlights that MADITS effectively leverages small amount of high-quality data to improve the quality of other suboptimal data.
H.3 FULL PARAMETER STUDIES
Apart from the sensitivity analysis of the reconstruction threshold 8 in Section 5.4 we also test the sensitivity of two key hyperparameters, augmentation ratio Taug and regeneration limit limit-
50
50
Episade Return
45
Epicade
35
20
02
0.4
0.5
1.0
20
0.089
0.005
0.01
0.02
0.55
20
Fagg
Sone
Jwet
(a)
(b)
(c)
Figure 5: Average returns on exp-m dataset in Cooperative Navigation of BC policies the hyperparameters change. The errors bars represent the standard error different random seeds.
From the results shown in Figure 5 we observe that as Taug gradually decreases, the performance of the new dataset slowly declines; however, still shows considerable improvement compared to the original dataset. This indicates that the trajectories generated through stitching can overcome the imbalances in the original dataset, with only small portion effectively enhancing the dataset's quality. Furthermore, the impact of hiii remains the same &reson As limi increase, the value first rise and then fall, which also reflects the negative impact of either overly lenient or overly strict constraint like what observe d
H.4 EVALUATION ON MORE ADDITIONAL PERTURBATION SETTINGS
Besides the experiments conducted on datasets with different timesteps of perturbations timbs we also explore the impact of qualities of behavior policy to the augmentation effects by applying
25

---

Published as a conference paper at ICLR 2025
Table Evaluation results on more additional perturbation settings on Cooperative Navigation task. The mean and standard error are computed based on the average return of BC policies trained on the datasets, with different random seeds.
Dataset
fisb
before
after
before
after
before
after
hefore
after
3.76
-36.12
-42.39
perturbations to medium-performee policies and collect more imbalanced datasets, denoted as md- and md-s. The average returns of the BC policy on these datasets are shown in Table from which we can observe that the trend of enhancement is consistent with the results of that in imbalanced datasets collect by sepert-performing behavior policies.
H.S ANALYSIS OFFLINE CREDIT ASSIGNMENT ACCURACY INTEGRATED GRADIENT
0.8
0.6
Density
0.4
0.2
0.0
-1.0
-0.5
0.0
0.5
1.0
Figure 6: Distribution visualization by kernel density estimation of the average Kendall correlation coefficients for trajectory segments compared with true data in the exp-m dataset the CN task. A coefficient less than 0 indicates negative correlation, greater than 0 indicates positive correlation, and 0 indicates independence.
In this section, analyze the accuracy of the credit assignment method in our method. In Coop- erative Navigation, the closer a landmark to an agent, the higher the return the team can achieve. Based understanding of the environment's true reward function, actual contribution at every time step as ground truth (this oracle information is not used in our method and only for analysis). Therefore, calculate the true individual contributions of agents from the offline dataset D and rank the contributions for each agent every time step rank) onaele-
derive each agent's
The similarity between the true ranking and the estimated ranking measured using the Kendall cor- relation coefficient (Kendall 1945). This coefficient evaluates the monotonic relationship between two ordinal variables, ranging from [-1, 1). A coefficient less than indicates a negative correlation, greater than 0 indicates positive correlation, and value of 0 implies independence. The closer
26

---

Published as a conference paper at ICLR 2025
the absolute value of the coefficient is to 1, the stronger the relationship, making it well-suited for assessing the consistency between the estimated and true rankings. The mean Kendall correlation coefficient for each trajectory segment is computed as:
T
mean Kendall
(T)
Tkendall
(rank. -. rank, ) (rank, oraalee poracle
ranky (oracte)
T
P=t
where Kendall is directly implemented using the 'scipy' library in Python3. Considering that offline credit assignment is performed on high-return trajectory segments generated by the diffusion model, we apply integrated gradient (IG) to compute the contribution of each agent for the trajectory segment dataset in the Cooperative Navigation exp-m dataset. We then calculate the average Kendall correlation coefficient for each trajectory segment and plotted its distribution
histogram.
The distribution of the computed Kendall correlation coefficients is presented in Figure 6 It can be observed that the majority of trajectory segments have average Kendall correlation coefficient greater than 0, demonstrating that the offline credit assignment method based on Integrated Gradients achieves high level of accuracy in practice.
H.6 DIVERSITY VISUALIZATION OF SYNTHESIZED TRAJECTORIES
Some previous works (Tian et al. 2023 et al. 2023) have emphasized the importance of diverse trajectories in offline multi-agent reinforcement learning. Here, we examine how our data augmen- tation method impacts the diversity of trajectories in the augmented dataset. Figure 7 visualizes 8 generated trajectories with identical agents starting positions of and landmarks positions in the augmented CN exp-m dataset, comparing results from w/o DC+PN and MADiTS, respectively.
As observed in Figure ((a) the inherent diversity of the diffusion models (Ho et al. 2020) allows w/o DC + PN to exhibit considerable diversity, despite issues of dynamics consisteney and suboptimal performance. In contrast, after incorporating our dynamics constraints and behavior correction, each agent's s paths to the landmarks still show significant diversity, indicating the capabilities of MADITS maintain impressive diversity while simultaneously enhancing trajectory quality.
I FURTHER DISCUSSIONS ON MADITS
I.1 APPLICATION TO DOOPERATIVE-COMPERR
SETTINGS
Our paper mainly addresses the cooperative multi-agent reinforcement learning (MARL) problem, widely studied setup where all agents share global reward. Through extensive experiments across various offline MARL environments, we demonstrate that MADITS significantly improves MARL performance. While our current focus on cooperative settings, the method can be naturally extended to cooperative-compeiity scenarios by equipping each team with its own MADiTS model. In this setup, teams can independently learn model parameters using tailored buffers for training. We leave the further exploration of this topic to future work.
I.2 CONCERNS ON THE USE OF DIFFUSION MODELS
Biases in the original offline dataset can propagate into the diffusion model, potentially affecting the quality of the generated trajectories. To mitigate this, our method employs bidirectional dynamics constraint, ensuring that the generated trajectories remain consistent with the environmental dynam- ics. Additionally, we integrate an offline credit assignment technique to identify and optimize the performance of underperforming agents within the generated trajectory segments, further enhancing the overall quality and utility of the augmented data.
On one hand, generative models like ChatGPT (Achiam et al. 2023) and SORA (OpenAl 2024) rely on large-scale datasets to train architectures such as Transformers or iiffussio-bbaed models. These models exhibit exceptional generative capabilities across domains like language and video, aligning with scaling laws that link performance data size. Recognizing the importance of data, these methods often use autoregressive training or advanced techniques to optimize data fitting. For real- world applications such autonomous driving (Wang al. 2024b), healthcare (Kazerouní al.
27

---

Published as a conference paper at ICLR 2025
(a)
(b)
Figure 7: Diversity visualization of generated trajectories with identical starting positions of agents and positions of landmarks in the augmented CN exp-m dataset. Lighter-colored trajectories are the extracted from observations of other agents. (a) Trajectories generated by w/o DC+PN. (b) Trajectories generated by MADITS.
2023), and finance (Wang & Ventre 2024), diffusion models require extensive and diverse datasets to ensure robust performance. Recent advancements have highlighted the potential of diffusion models to transform these domains.
On the other hand, despite their capabilities, these methods face challenges such as unreliable or unrepresentative data. ensure reliability real-world applications, techniques like human-in- the-loop testing (Singi et al. 2024) and risk control mechanisms (Yang et al. 202 are crucial.
To address the issue synthetic data deviating from real-world distributions, our method MADITS introduces bidirectional dynamics constraint to align generated trajectories with environmental dynamics. Moreover, the offline credit assignment technique enhances robustness by identifying and improving underperforming agents in generated segments. Experimental results validate the effectiveness of MADITS overcoming these challenges.
28


Reviews:
[{"review_id": "W6PteyPNhS", "reviewer": "Reviewer_Af9P", "summary": "The paper investigates the use of diffusion models for offline multi-agent reinforcement learning (MARL), extending certain diffusion-based approaches originally developed for single-agent offline RL to a multi-agent setting. Specifically, the authors combine Denoising Diffusion Probabilistic Models (DDPM), stitching techniques (to merge high-reward and low-reward data points and address data imbalance), and Integrated Gradients (IG) for identifying underperforming agents. This results in a data augmentation process that can enhance offline datasets and other MARL algorithms. Experiments are conducted using MPE and SMACv1, with comparisons to BC and OMIGA.", "questions": "- How does your approach compare to MADiff?\n- Can these methods be applied to cooperative-competitive settings?\n- Would the approach work effectively with balanced but bad datasets? What would happen if the data were balanced but consisted primarily of transitions with low rewards?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The diffusion approach makes sense. Using a diffusion model to enrich offline datasets could support several other MARL algorithms.\n- The employed techniques align with recent advancements in diffusion-based RL.\n- The paper is well-written and easy to follow.", "weaknesses": "1. The techniques seem to be a straightforward extension from diffusion-based single-agent RL.\n2. Most techniques, except those in Section 4.2, appear to be directly adapted from single-agent settings, with single-agent states and actions replaced by global observations and actions.\n3. There is already existing work on diffusion models for offline MARL. For instance, how does this approach compare to MADiff [1], which also develops diffusion models to enhance offline datasets in offline MARL? Notably, MADiff compares its methods against several MARL baselines, which seems lacking in this work.\n4. In addition to MPE and SMACv1, more challenging MARL tasks like MAMuJoCo and SMACv2 should be considered for comparisons.\n5. It’s unclear how this approach compares to other data-augmentation techniques in offline MARL. Additional citations, discussions, and experimental comparisons would be needed.\n\n\n**Reference:**\n[1] Zhu, Z., Liu, M., Mao, L., Kang, B., Xu, M., Yu, Y., Ermon, S., & Zhang, W. (2023). MADiff: Offline Multi-agent Learning with Diffusion Models. arXiv preprint arXiv:2305.17330.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "22iMU1mFJp", "reviewer": "Reviewer_onPK", "summary": "The paper proposes the MADiTS method to address the imbalance of offline datasets in both temporal and spatial dimensions. Specifically, it uses diffusion models to generate trajectory segments, then employs bidirectional environmental dynamics constraint to ensure the consistency of trajectories with environmental dynamics, and finally utilizes offline credit assignment technique to identify and optimize the behavior of underperforming agents in the generated segments.", "questions": "1.\tHow does the proposed method address the imbalance of the dataset in terms of time and space? For example, which technique solves the time imbalance problem, and which solves the spatial imbalance problem?\n2.\tHow are environment-independent hyperparameters like $\\delta_{recon}$ determined? What effects do different values have?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThe overall structure of the paper is complete, the content is substantial, and the experimental evidence is convincing.\n2.\tGenerating trajectories using diffusion models and ensuring the generated segments conform to environmental dynamics with bidirectional dynamic constraints is similar to learning a world model and using it to verify the rationality of trajectories.\n3.\tAs a data augmentation method, MADiTS is compatible with any offline MARL algorithm and has good versatility.", "weaknesses": "1.\tThe paper uses several certain threshold hyperparameters (environment-independent), which may make the method too sensitive to the setting of hyperparameters. Moreover, the paper does not conduct experimental analysis on methods with different hyperparameters.\n2.\tMADiTS involves multiple models and steps, which may lead to computational complexity. Additionally, the paper does not include experimental analysis on this aspect.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LmNHEXYJsI", "reviewer": "Reviewer_qWY9", "summary": "The paper presents a novel data augmentation method, MADiTS, to generate high-quality multi-agent trajectories for offline policy learning. A multi-agent trajectory diffusion model is trained to generate augmented future trajectories conditioned on the first-step joint observations. Multiple techniques are proposed to improve the quality of augmented trajectories. First, a forward dynamics model is used together with the inverse dynamics model to filter out dyanmic inconsist segments. Then, the interagted gradient is adopted to determine underperforming agents, and regenerating those agents' trajectories while keeping others' fixed. Experimental results on MPE and SMAC datasets demonstrate superior performence of MADiTS over baseline augmentation methods.", "questions": "See weaknesses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper is well-written and easy to follow.\n\n2. The idea of using the interagted gradient to find underperforming agents is novel to me, and I believe it might have border potential in multi-agent learning.\n\n3. The experiment results are impressive, achieving SOTA performance on all datasets.", "weaknesses": "1. The experiments are limited to environments with a small number of agents (up to five). The effectiveness of MADiTS on datasets with a larger number of agents has not been demonstrated.\n\n2. Given the known issues [1] with SMAC (v1), why were results not included for SMACv2?\n\n3. Unlike DiffStitch, which conditions on both the first and last trajectory segments, the diffusion model in MADiTS generates the subsequent trajectory based solely on the first state, which makes the naming of \"Stitch\" a little confusing.\n\n4. A key highlight of this paper is the use of integrated gradients to identify underperforming agents. However, this part lacks sufficient analysis. It would be beneficial to quantitatively demonstrate a general correlation between the value of integrated gradients and underperformance degrees across different environments.\n\n5. I have doubts about the extent to which regeneration can improve the overall reward. The multi-agent diffusion model learns from the joint trajectory distribution, meaning it is constrained by the combination of cooperative capabilities existed in the training data. Suppose, in a two-agent environment (agent 0 and agent 1), the training set includes two joint trajectories, $\\tau_1$ and $\\tau_2$. In $\\tau_1$, agent 0 significantly contributes to the team reward, while agent 1’s contribution is minimal. In $\\tau_2$, the situation is reversed. Now, suppose the model’s first generation produces a trajectory where agent 0 has high contribution and agent 1 has low contribution. If we successfully use integrated gradients to identify agent 1 and attempt to condition on agent 0’s trajectory to regenerate agent 1’s trajectory, the model should be unable to produce a high-contribution trajectory for agent 1, as the training dataset lacks joint trajectories where both agents have high contributions. Therefore, in this case, MADiTS may not effectively combine individual trajectories to yield better joint trajectories. The authors should discuss similar cases, clarifying the method’s limitations and maybe requirements for the training dataset.\n\n[1] Ellis, Benjamin, et al. \"Smacv2: An improved benchmark for cooperative multi-agent reinforcement learning.\" *Advances in Neural Information Processing Systems* 36 (2024).\n\nUpdate: Since the authors' responses have adequately addressed my concerns, I recommend accepting the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gyo5plZBmL", "reviewer": "Reviewer_bitr", "summary": "The authors propose a method called MADiTS to enhance the quality of imbalanced offline datasets in Multi-Agent Reinforcement Learning (MARL). By applying diffusion-based data augmentation, MADiTS generates high-return trajectories while enforcing dynamics constraints to ensure they align with the environment dynamics. Additionally, a behavior correction mechanism correct the actions of suboptimal agents, improving their trajectories within shared reward contexts. Experiments on MPE and SMAC datasets show that MADiTS can improve the performance of offline algorithms like Behavior Cloning (BC) and other offline algorithms.", "questions": "1. Why does this approach generate an imbalanced dataset by perturbing expert trajectories? Why is this better than using a dataset collected by mixed policies and combine their trajectories in a dataset?\n2. How does the proposed method impact the diversity of the trajectories in the augmented dataset? Do the dynamics constraints or behavior correction mechanism hurt the diversity of the trajectories?\n3. What is the computational cost of the proposed method compared to the baselines? How does the proposed method scale with the number of agents, particularly with the integrated gradient computation?\n4. How does MADiTS compare to standalone MaDiff, DOM2 and SIT? What are the key differences, and why were these methods not used as baselines in the primary experiments?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The authors address an important issue of dataset quality in MARL, particularly relevant for offline applications where data imbalance often impacts policy performance in real-world scenarios.\n- The paper is well-written, clear and easy to follow. \n- The authors apply their method to datasets collected on SMAC and MPE, two popular MARL environments, and apply it to behavior cloning (BC) and two modern offline MARL algorithms - OMIGA and CFCQL.\n- The authors provide an ablation study to show the impact of different components of their MADiTS method.", "weaknesses": "Methodology:\n- The authors provide details on how they generate the imbalanced dataset, using random perturbations on expert trajectories. However, it is not clear why this is better than using a dataset collected by mixed policies, e.g. taking medium and expert policies and combining their trajectories in a dataset. This appears to be more realistic and relevant to the problem of imbalanced datasets in MARL. A discussion on this choice would be relevant.\n- Work such as [2,3] highlights the importance of diverse trajectories in offline MARL. This impact of MADiTS on the diversity of the trajectories is not discussed in the paper.  \n- Clarity is needed regarding the computational cost of MADiTS. Details such as runtime and comparison with other trajectory augmentation methods (e.g., MA-MBTS or baseline models) would be relevant, especially since the text includes vague descriptors like “until a satisfactory trajectory is obtained.” Further, insights into how MADiTS scales with agent numbers or in larger MARL environments would enhance applicability.\n\nComparisons with Related Diffusion Models:\n- In the appendix, the authors mention they use MaDiff [1] architecture for the diffusion model. However, the original MaDiff model is not included as a baseline in the main experiments, nor is Diffusion Offline Multi-agent Model (DOM2) [3], another relevant diffusion-based MARL trajectory augmentation approach. Furthermore, Shared Individual Trajectories (SIT) [2], another work that aims at helping with imbalanced datasets in offline MARL is only briefly mentioned in the appendix, and not mentioned in the main text.  A comparison with MaDiff (or clarification if “w/o dc + pn” in Figure 4a is equivalent to MaDiff) would strengthen the evaluation. Including DOM2 and SIT as baselines or discussing the differences would clarify MADiTS's unique contributions to diffusion-based MARL augmentation. \n\n1] Zhu Z, Liu M, Mao L, Kang B, Xu M, Yu Y, Ermon S, Zhang W. Madiff: Offline multi-agent learning with diffusion models. arXiv preprint arXiv:2305.17330. 2023 May 27.\n\n2] Tian Q, Kuang K, Liu F, Wang B. Learning from good trajectories in offline multi-agent reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence 2023 Jun 26 (Vol. 37, No. 10, pp. 11672-11680).\n\n3] Li Z, Pan L, Huang L. Beyond conservatism: Diffusion policies in offline multi-agent reinforcement learning. arXiv preprint arXiv:2307.01472. 2023 Jul 4.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)