Encouraging metric-aware diversity in contrastive representation space
Authors: Tianxu Li
Kun Zhu
Venue: NeurIPS 2025 poster ()
Keywords: Multi-agent reinforcement learning
Exploration
Multi-agent diversity
Wasserstein distance
Primary area: reinforcement_learning

Abstract: In cooperative Multi-Agent Reinforcement Learning (MARL), agents that share policy network parameters often learn similar behaviors, which hinders effective exploration and can lead to suboptimal cooperative policies. Recent advances have attempted to promote multi-agent diversity by leveraging the Wasserstein distance to increase policy differences. However, these methods cannot effectively encourage diverse policies due to ineffective Wasserstein distance caused by the policy similarity. To address this limitation, we propose Wasserstein Contrastive Diversity (WCD) exploration, a novel approach that promotes multi-agent diversity by maximizing the Wasserstein distance between the trajectory distributions of different agents in a latent representation space. To make the Wasserstein distance meaningful, we propose a novel next-step prediction method based on Contrastive Predictive Coding (CPC) to learn distinguishable trajectory representations. Additionally, we introduce an optimized kernel-based method to compute the Wasserstein distance more efficiently. Since the Wasserstein distance is inherently defined for two distributions, we extend it to support multiple agents, enabling diverse policy learning. Empirical evaluations across a variety of challenging multi-agent tasks demonstrate that WCD outperforms existing state-of-the-art methods, delivering superior performance and enhanced exploration.

Encouraging metric-aware diversity in contrastive representation space
Tianuu Li Kun Zhu"
College of Computer Science and Technology, Nanjing University of Aeronautics and Astronautics, China (tianxuli, (aaruu))Ouuaaeee edu. en
Abstract
In cooperative Multi-Agent Reinforcement Learning (MARL), agents that share policy network parameters often learn similar behaviors, which hinders effective exploration and can lead to suboptimal cooperative policies. Recent advances have attempted to promote multi-agent diversity by leveraging the Wasserstein distance increase policy differences. However, these methods cannot effectively encour- age diverse policies due to ineffective Wasserstein distance caused by the policy similarity. To address this limitation, we propose Wasserstein Contrastive Diversity (WCD) exploration, a novel approach that promotes multi-agent diversity by maxi- mizing the Wasserstein distance between the trajectory distributions of different agents in a latent representation space. To make the Wasserstein distance meaning- ful, we propose novel next-step prediction method based on Contrastive Predictive Coding (CPC) to learn distinguishable trajectory representations. Additionally, we introduce an optimized kernel-based method to compute the Wasserstein dis- tance more efficiently. Since the Wasserstein distance is inherently defined for two distributions, we extend it to support multiple agents, enabling diverse policy learning. Empirical evaluations across variety of challenging multi-agent tasks demonstrate that WCD outperforms existing state-of-the- methods, delivering superior performance and enhanced exploration.
Introduction
Multi-Agent Reinforcement Learning (MARL) has shown promise in addressing various multi-agent challenges, such as multiplayer video games [Vinyals al. 2019 and autonomous cars [Cao et al.) 2012], attracting growing interest in recent years. MARL facilitates efficient collaboration by training multiple agents together towards maximizing team rewards. Yet. there are still many challenges such as partial observation constraints and high scalability requirements, when learning effective cooperative policies for agents in complex multi-agent tasks. To resolve these issues, recent works commonly employ the Centralized Training with Decentralized Execution (CTDE) framework (Lowe) af. (017) where agents make decisions based on local observations using a decentralized policy jointly trained with global information, ensuring robust and stable performance.
The CTDE framework develops distinct decentralized policies for each agent, but training numerous policy network parameters can be inefficient. Thus, parameter sharing has become universal, allowing agents to share the same policy network parameters for action srciisionmmakine significantly reduces the number of parameters, leading to lower computational cost and speeding up training. Additionally, parameter sharing promotes experience sharing during centralized training. fostering robust policy learning and improving overall efficiency (Wang et al. 2020b). Despite these benefits, sharing policy parameters can lead homogeneous behaviors among agents, hindering
This practice
multi-agent diversity and efficient exploration (Hu et al. 2022
Terry et al. 2020]. In challenging
Corresponding author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

multi-agent tasks, extensive exploration and diverse policies are crucial (Jiang and Lu 2021 Li et al. 202T]. For example, in a football game, agents must adopt varied roles and strategies for effective collaboration and goal scoring.
Prior works primarily adopt the mutual information objective to promote diversity among agents (Jiang and Lu 2021 Li et al. 202 Jo et al. [022]. These methods encourage agents to explore diverse individual agent trajectories [Li al. [021]. observations [Jiang and Lu 2021] or formations (e.g., observation differences) [Jo et al. 20244. While these methods do lead to trajectories that are mutually different among agents, the mutual information objective cannot measure how different the induced trajectories are. Slight differences between trajectories are enough to maximize the mutual information objective, which does not necessarlly encourage exploration (see Section 3 for more details). To encourage exploration, some recent advances (Hu et al., 2024 Bettini et al.) [024]] are based on the Wasserstein distance, a metric-aware measure that quantifies the distance between two distributions. These methods achieve diversity among agents through dynamic parameter sharing (Hu et al. 2024] or policy distribution fusion (Bettini et al.) 2024) according to the Wasserstein distance metric. Unfortunately, these methods overlook the similarity of agents' initial policies resulting from shared policy network parameters, which causes the Wasserstein distance, intended to measure policy differences, to converge to zero, leading to ineffective diversity incentives.
To overcome this limitation and enjoy the metric-aware benefit of the Wasserstein distance, we propose a novel Wasserstein Contrastive Diversity (WCD) exploration method, which enlarges the Wasserstein distance in a latent contrastive trajectory representation space. To generate effective diversity incentives based on the Wasserstein distance, we propose novel next-step prediction method based on Contrastive Predictive Coding (CPC) to learn distinguishable trajectory representations, and calculate the Wasserstein distance between distributions of trajectory representations of different agents.
Our contributions can be summarized as follows: First, because of the similarities among agents' initial policies, the Wasserstein distance cannot provide effective feedback to learn diverse policies. To solve this issue, we consider a trajectory representation space in order to make the Wasserstein distance meaningful. To construct the representation space, we propose next-step prediction method based on Contrastive Predictive Coding (CPC) [Oord et al., 2018 to Iearn distinguishable trajectory representations. Second, due to the high computation cost of calculating the Wasserstein distance, we propose a novel Gaussian kernel method to optimize dual functions of the Wasserstein distance, significantly reducing the computational cost. Third, we extend the Wasserstein distance to multiple policy learning by introducing nearest neighbor intrinsic reward. We further integrate our method with QMIX. Fourth, we show the outperformance of our method against existing state-of-the-art methods by testing in various challenging multi-agent tasks.
2 Backgrounds
2.1 Multi-Agent System
We consider modeling the fully cooperative multi-agent Decentralized Partially Observable Markov Decision Process (Dec-POMDP) [Olichock and Amato 2015), defined as tuple (A, S,U, U, P, R. O, 7). Here, A denotes set of |A| agents, E S represents the global state of the environment, and U stands for the set of agents' actions. At each time step, each agent receives an observation o" E Ω drawn from the function O(s,a) and subsequently selects an action u² E U. All agents' actions collectively form a joint action u, leading the environment to transition to the next state s' based on the probability drawn from the transition function P ((s' 8, u. Simultaneously, the environment provides the agents with a shared team reward r = R(ss, u). E [0, 1) is the reward discount factor. The observation-action pairs (0", u²) of agent a during an episode constitute its trajectory T" E TT Each agent a learns its individual policy x" (u² ra), contributing to the formation of a joint policy π, aimed at maximizing the joint action-value function Q (s, u) = E Bg Mu:co IE 7 So = s, Ug = u, TT..
2.2 Wasserstein Distance
The Wasserstein distance formulates an optimal transport problem that measures the distance or discrepancy between two probability distributions (Villani et al., 2009]. Given two probability

---

(( Gomma
Predict
Netwock
a
Autoregresäre Madel
Auteregroaiie Mudel
e. (o)
= miinban Pa.
where 4 -P
Encoder
Encuder
(
QMIX
Trajectury
Figure I: Architecture WCD. On the left shows the proposed next-step prediction method which learns trajectory representations by predicting the next-step observation-actoe embedding with Contrastive Predictive Coding (CPC) to make the Wasserstein distance effective. On the right shows the combination of our method with QMIX. We update the policies of agents towards maximizing the intrinsic rewards Tw based on the Wasserstein Distance between the current agent and its nearest neighbor by introducing an additional intrinsic utility network Q
distributions p and q over domains X S 1h and y R²" respectively, the Wasserstein distance with a cost function e(x, y): X' x y + IR is defined as:
inf / c(x, y)dr(x,))
Wc(p.))
7EL((p.)) Xxx
where ['(p, q) is set of all possible couplings of distributions and q over the product space X V. The probability distributions p and q are the marginals of the coupling 7(x, y) over space X and V respectively, i.e., SM (2, yydy p(x) and JM (a,, y)dr q(v).
In practice, we adopt tractable smoothed Wasserstein distance Wo(p. q) based on Fenchel- Rockafellar duality [Villani et al. 2009].
(p(p q) sup (aormptaaaadll plx) - v(v) - exp (u(x) v(y) 1.0
-V)) -(0) v) )
(2)
B
where μ X - E and v y R are dual functions on continuous domains. B is a smoothing parameter. The dual form of the Wasserstein distance allows for the parametrization of dual functions, thereby mitigating the computational complexity of optimizing the optimal transport problem.
3 Quantitative comparison between the Wasserstein distance and the KL divergence
The mutual information objective is based the KL divergence, which is metric-agnostic. We next compare the KL divergence and the Wasserstein distance to illustrate that our Wasserstein distence-based method can more efficiently encourage multi-agent diversity compared to the KL divergence. To illustrate the difference between the Wasserstein distance and the KL divergence, take Gaussian distribution example. Let p A ("y:0²) and q N' (My, (2), As the probability mass of and converges to their means, thus we can achieve the KI. divergence between two distributions p and as lim,-+o DKL (pll))
oo, which independent of the specific means Mр and μt. The Wasserstein distance between and is lim W (p, q) [M, - Mql. We note that the Wasserstein distance provides an explicit measurement of distance, whereas the KL divergence focuses solely on distinguishability and has no relevance to the metric of the underlying data distribution. Thus, the KL divergence can easily maximized with slight differences between different agents' trajectories, which does not necessarily encourage visitations of diverse trajectories with large distances. However, due to the metric-aware property of the Wasserstein distance, our method can not only encourage the visitations of different trajectories, as in the KI. divergence, but also maximize the distance between diverse trajectories that leads to better trajectory space coverage and more sufficient exploration.
3

---

Wasserstein Contrastive Diversity
In this section, we detail our proposed Wasserstein Contrastive Diversity (WCD). First, we present how learn distinguishable trajectory representations to generate effective feedback for the Wasserstein distance. Then, we show how to maximize the Wasserstein distance between different trajectory distributions the latent representation space.
4.1 Contrastive Predictive Trajectory Representations
In the framework of CTDE, agents that share the same policy network parameters have similar initial policies. Therefore, the Wasserstein distance between any two agents' policy distributions tends to approach zero, i.e., W(X, Y) - 0, where X and Y respectively represent the policy distributions of two agents. Our motivation is propose a next-step prediction method based on Contrastive Predictive Coding (CPC) [Oord et al. 2018] to learn distinguishable trajectory representations, enabling trajectories induced by different agents' policies collapse to diverse distributions latent representation space. In this space, the Wasserstein distance can effectively encourage the learning of diverse policies.
Initially, as shown in the left of Figure I we encode the cntevaationaaction pairs x; = (of uf') with non-linear encoder ge, into latent embedding z = ge. (aî). Then, we use an autoregressive model 9e, to summarize all the latent embeddings and output the trajectory representation of = ge, (2
at timestep t. We simply denote ge {ga.- 9e, to represent the overall trajectory encoder. For simplicity, we adopt standard architectures such as MLPs for ge. and GRUs for ge,-
To train to learn distinguishable trajectory representations, we model a density ratio that preserves the underlying information between the trajectory representation of and the next-step observation- action 7;41:
P p(2f+1 (
(1(+1
(2
where of (2(+1, = exp(ge. (17+1) We)) = exp(22++ T " ct) calculates the similarity between the next-step observation-action embedding 2 and a linear transformation W of with the parame- ter W used for the next-step prediction. Compared to modeling p (xf+1 cq) directly by generative method that requires reconstruct every detail in 2 modeling the density ratio has lower compu- tation cost and is more effective in extracting shared information between 2 and cf. Moreover, we infer the latent embedding #[+1 instead of the raw x++1, which avoids modeling high-dimensional observation-action space. To let f (If+1, <t) be proportional to the density ratio, inspired by CPC, given a set of next-step observation-action pairs of all agents C = faft2 (of (ett))()))) (4 we minimize a InfoNCE loss [Oord et al. 2018
f ( G) Exerec (Iq c7)
CN
E (ef C)
log
By using the next-step observation-action pairs of other agents as noisy samples in Equation 4 and contrasting the trajectory representation of with these noises, the trajectory representation of stays close to its associated next-step abservation-action embedding while being far away from other noisy embeddings, As result, the trajectory encoder go trained by minimizing the InfoNCE loss to learn distinguishable trajectory representations.
4.2 Wasserstein Distance between Trajectory Representations
We then encourage the exploration of diverse trajectories by maximizing the Wasserstein distance between the trajectory distributions of different agents in latent representation space. Let Pr, and Pr, be the trajectory representation distributions of agent and agent 2, respectively. The Wasserstein distance between Pa; and Pxa is defined as follows:

---

(m(ct) - v(c v(22) c(c,, dc) )
)
WelPr: Pro))= = suu Ee Mop e8 p22
plc) - v(c²) exp
(5)
p.u
where the cost function f ot) is represented by the Fuclidean distance between the points of and c². i.e. e(c,, ef) ||c| og ||- notable that to compute the Wasserstein distance, we may
simply parameterize dual functions with neural networks like previous works [Pacchiano et al. [2020] Dadashi et al.) [222]] I He et al.) 2022 Park et al. 2024]. However, this may lead to high computational costs in our multi-agent settings, as we need to compute the Wasserstein distance for each pair of agents. To learn optimal dual functions μ and to compute the Wasserstein distance with low computational costs, resort to the kernel method (Hearst et al. [998| that has been widely used in machine learning. Specifically, we consider representing dual functions with linear combinations of Gaussian kernel functions approximated by the random feature map [Rahimi and Recht 2007).
For example, let the dual function j has the following form: μ(() (X#) <(x). For E R <(x) Vm cos (Gx + b) represents m-dimenssonna random feature map, where G E mR is Gaussian with entries sampled from a normal distribution W 1) and b E IR" with entries sampled from uniform distribution (0, 2x). This means that when we optimize the dual function μs, we only need to learn the dual vector X E IRm, which significantly reduces the computational cost compared with parameterizing dual functions with computationally intensive neural networks.
To learn optimal dual functions, we perform stochastic gradient descent (SGD) over the Wasserstein distance objective in Equation Given dual functions µ and that are modeled by kernels and E, respectively, and trajectory representation samples {cl, c ~ (Pa, Pra). we apply the chain rule to Equation and the gradients with respect to X and X" are
(A) We Eetree, (ppperr--ppp (1 (1-v)
v)
( )}
dn(c))
-cx((ct)
(6)
(AN) @x((c) (20) - C(ct,
where y = exp
3
We approximate the expectation by averaging the function values over a batch of trajectory represen- tation samples from the replay buffer that is used store agent experiences during training.
As we have computed the value of the Wasserstein distance, we can view the Wasserstein distance as an intrinsic reward r.. W (Pm PPg)) which enables us to deploy our method in MARL algorithms to maximize the Wasserstein distance. When the number of agents |A] is more than two, the trajectory of arbitrary agent should keep distance with any other agent. In practice, we empirically find that employing an intrinsic reward Far !in =', W (Pr.: PT.,) for each agent to keep the trajectory of the current agent a to away from its nearest neighbor trajectory in a latent representation space can lead to better performance. The pseudocode for our method can be found in Appendix
4.3 Practical Learning Algorithm
We next show how to integrate our method with QMIX [Rashid et al. 2018). state-of-the-arr MARL algorithm. QMIX learns optimal individual policies, that maximizes shared team rewards, for agents through optimizing the joint action-value function QT approximated by Qiot, output of mixing network that monotonically mixes the agent utilitics (where the policies are derived) of all agents. In QMIX, in order maximize the Wasserstein distance- based intrinsic rewards, cannot simply add each agent's 's intrinsic rewards to the shared team reward. More detailed explanations can be found in Appendix A To integrate method with QMIX, additionally introduce intrinsic utility network Q which takes as input the agent utility Qa uz) and the trajectory representation et We update Q" towards maximizing the intrinsic rewards by minimizing the TD loss follows
DTD Biog of of [ (ef, Qalof, uf)) w)²
where y == rw + 2 (c+t Qa (oi+1, 4(+1)
where Q and Q. are target networks employed to stabilize training and D the replay buffer for storing trajectory samples. LTD can be seen as regularizer that introduces auxiliary gradient
5

---

Emakes
(a) Illustration of(b) Average rewards(c) Visitation heatmap(d) Visitation heatmap(e) Pac-Men
Wasserstein lance
dis-
of MAPD
of WCD
Figure 2: Performance comparison between our proposed WCD and baselines in Pac-Men. We report both the mean and standard deviation of the performance tested across five random seeds.
the agent utility network Qu in order to learn diverse trajectories. We can thus get the total loss function
Liotal = CTD + aCTD
(8)
where LTD is the TD loss of QMIX train Qtat and is coefficient that changes the weight of CTD- As a -> O, our method converges to QMIX. Through minimizing Ltotal, we train the overall framework of our method end-to-end in a centralized manner. As a result, agents learn their policies towards maximizing both team rewards and the Wasserstein distance between different agent's trajectory representation distributions. Someone may question whether the regularizer Ly conflicts with the global optimum, where agents sometimes need to behave similarly. For example, in SMAC scenarios, agents may be required to fire at the same enemy to quickly defeat it. Empirical results demonstrate that our method does not hinder agents from learning such similar behaviors when they yield greater environmental rewards. This is because our method is essentially an exploration technique, similar to Soft Actor-Crite (Haarnoja et al., 2018 that employs entropy regularization to maximize policy entropy, thereby encouraging sufficient exploration. Numerous studies [Ladosz et al. 2022 have shown that efficient exploration is beneficial for learning optimal policies.
For policy gradient methods, we refer the reader to Appendix D where we integrate method with the policy gradient-based method MAPPO.
Experiments
In this section, we use challenging multi-agent tasks from Pac-Men, SMAC, and SMACv2 to demonstrate the outperformance of our method. We show comparison of our method against the state-of-the-art methods such as decooposccoioiitioonn methods (QMIX [Rashid et al.) 2018 and
QTRAN [Son et al., 2019)), role-based diversity methods (RODE [Wang et al. [020cc ). mutual information- diversity methods (MAVEN [Mahajan et al. 2019 EOI Jiang and Lu 2021], J,
SCDS Li et al. 2021], PMIC (Li et al., 2022), LIPO [Charakorn et al. 2023 and FoX eJo et al. 2024 ) and Wasserstein distance-baut basedd diversity methods (MAPD Hu et al. 2024 and DiCo [Bettini al. 2024]). ). Without loss of generality, the comparison results are shown with both the mean and standard deviation of the performance tested across five random seeds. For a fair comparison, we adopt the same hyperparameters and policy network architecture across all methods. More training details and hyperparameters are provided Appendix K
5.1 Pac-Men
We first test our method in Pac-Men, as illustrated in Figure 2a to investigate the effectiveness of our method in encouraging multi-agent diversity. Pac-Men foraging game, where four agents initialized at the center of the maze try to eat the dots randomly distributed in four edge rooms. Agents can move to these rooms along paths of different lengths. Each agent only has a partial observation of 4 x4 grid around them. The goal of the agent is to collect as many dots as possible to achieve more rewards. Notably, agents arriving at the same edge room may result in inefficient competition. They expected to behave differently and move to different rooms.

---

OME ORTX
STRMW
ECDE
LPO
K0D1
MAPO
100
60
Test 300-00
Test Ra0-7
60-
10
10
6.0
Timestopa
1.4
20
0.0
Timestoce
20
0.0
Timestaps
(a) 3s5% (easy)
(b) 2c_ vs _s__v__e_e_ 647g (hard)
(c)7sz. (hard)
00-
60-
Test Me RasaS
ToP RoooS
20-
20
Timesteps
3es
(d) _1_____________ vS 8z (super hard)
(e) corridor (super hard)
(f)3s5z_vs Sas__vs_3scc _36z (super hard)
Figure 3: Performance comparison between our proposed WCD and baselines in the SMAC scenarios.
The results shown in Figure (6) demonsste the outperformance of our method compared to baselines. Through maximizing the Wasserstein distance between different trajectory distributions in latent space, agents respectively move to the four edge rooms, as depicted by Figure 2d leading to diverse policies and efficient cooperation. MAPD fails learn diverse policies. As shown in Figure 2c some agents adopt the same policy and move to the same edge room, resulting in poor performance. From Figure [e] we note that MAPD does not provide effective diversity measures in terms of the average Wasserstein distance, which does not necessarily encourage exploration of diverse policies. However, by learning a contrastive trajectory representation space through the proposed next-step prediction method, our method produces effective Wasserstein distance as intrinsic rewards to learn diverse policies. Some mutual information- based baselines such as EOI and SCDS, employing the metric- agnostic variational intrinsic reward, achieve similar performance. They may not find the edge room with the longest path, leading to sub-optimal performance. This is because the variational intrinsic reward converges quickly due to its metric-agnostic property, leading to insufficient incentives for exploration. Conversely, our Wasserstein distance-based metric-aware intrinsic reward can continuously provide effective reward signals for agents to encourage sufficient exploration.
5.2 SMAC
We then test our method on the StarCraft Multi-Agent Challenge (SMAC) (Samvelyan et al., 2019). a commonly used benchmark for evaluating cooperative MARL algorithms, consisting of various combat scenarios with different difficulties. We evaluate our method in 6 scenarios of SMAC including 3s5z (easy), 2c_ vs 64zg (hard), 7sz (hard). 6h_vs, vs 8z (super hard), corridor (super hard), and 3s5z_vs vs_ 3s6z (super hard). The version of SMAC adopted in our experiments is SC2.4.10. The performance comparison are not applicable across different SMAC versions.
As shown in Figure 3 our method maintains its outperformance in both easy and hard scenarios and significantly outperforms all baselines in the super hard scenarios. QMIX struggles to learn optimal cooperative policies in the super hard scenarios. However, our method can efficiently improve the performance of QMIX by encouraging multi-agent diversity. Moreover, our method shows a substantial performance boost compared to MAPD, highlighting the effectivenesss of the Wasserstein distance in the contrastive trajectory representation space. MAPD even achieves worse performance than the mutual information-baut methods in the super hard scenarios. This phenomenon due to the fact that the Wasserstein distance metric employed in MAPD may not work properly under the policy network parameter-sharinn setting. Compared to mutual information-based methods, our method achieves better performance due to the maximization of the metric-aware Wasserstein distance in the contrastive trajectory representation space, leading to more sufficient exploration. We further present visualization examples of diverse policies learned by our method in the super hard
7

---

WAD WCD--QMI
MAVEN
201
MAC
PM
MAP
DICK
80
Ies: Min tarehs
Tos (M Rale/M
600
20
Timestens
Loc
Timestops
Timesteps
100
(a) terran vs
(b) protoss_5_vs_0
5_vs_
Figure 4: Performance comparison between WCD and baselines in the SMACv2 scenarios.
(a) MAVEN (b) EOI
(c)SCDS
(d) LIPO
(e) MAPD
WCD
Figure 5: Visitation heatmaps of different algorithms the terran_5_vs__ 5_vs_5 scenario.
scenarios in Appendix | The mutual information based methods may not enable agents learn trajectories with large variations. HOI does not result in satisfactory performance as the trajectory classifier employed
EOI overfits the agent identity information, impeding further exploration. Moreover, it is notable that our method also achieves satisfactory performance in the easy 3s5z. scenario where agents sometimes need to behave in the same way to master the trick of "focus fire', demonstrating that our method would not prevent the homogeneous behaviors that can lead to more environmental rewards. More experimental results related to such homogeneous behaviors can be found in Appendix E.2] These results reveal that our method efficiently balances exploration and exploitation, resulting in the learning of optimal cooperative policies.
Stochasticity and Exploration Although SMAC consists of many challenging scenarios, the agents may overfit the timesteps regardless of real environmental states Ellis et al. [2022] since the team compositions and the initial positions of units are the same in each episode. We further adopt the SMACv2 benchmark EElli al.] [2022]- SMACv2 introduces stochasticity by deploying random team compositions and random initial positions, which challenges agents to continuously explore optimal policies. The performance comparison of our method against baselines are shown in Figure Our method achieves superior performance in all scenarios compared to the baselines. Our method significantly improves the performance QMIX by introducing effective Wasserstein distance objective regularizer to encourage multi-agent diversity. The mutual information based methods do not yield satisfactory performance. We believe this because the variational intrinsic reward adopted in these methods converge quickly when the trajectories of different agents identified. As a result, cannot provide effective feedback for agents to continuously explore. Moreover, MAPD fails to yield sufficient exploration to adapt to the environment stochasticity due the ineffective Wasserstein distance incentives. Instead, our method can continuously provide efficient Wasserstein distance-based intrinsic rewards to encourage exploration. This be verified by the visitation heatmaps of agents trained by various methods shown in Figure We observe that agents trained by our method achieve more extensive environmental exploration compared those trained using baselines distributed only in partial areas.
5.3 Ablation Study
We conduct ablation studies to evaluate the contributions of the main components our method. To test the contribution of the autoregressive model employed to learn trajectory representations, we ablate the autoregressive model and only use the non-linear encoder go, regardless of the trajectory context. To measure the contribution of the representation learning method, we design five variants: (i) employing a randomly initialized encoder with fixed parameters for encoding trajectories, (ii) learning trajectory representations by directly predicting the agent identities of various trajectories

---

Meurca metgiaez Wisserssen dasinto
Fandom apent Wosperstenn distance
8
Average Wateerstein distnne
1.50
1.00
102
10
candaa
MCD
VIDD oi crod sting inslectory reprenettatira
MCD food
randem Miaserstin dasance sarrage Mor sendew sidtane
0:5
Timestaps
1et
(a)
(b)
Figure 6: (a) Performance comparison of different variants in the scenarios of SMAC. (b) Different kinds of Wasserstein distances.
instead of employing the InfoNCE loss, (ii) learning trajectory representations by adopting a generative method to model [2]+1 | et) instead of modeling the density ratio, (iv) predicting the trajectory representation G4+1 instead of the latent emhedding 1 and (v) directly predicting the raw observation-action 2:+11- To test the Wasserstein distance objective, we ablate the nearest neighbor intrinsic reward Tw and use the Wasserstein distance between trajectory representation distributions of the current agent and another randomly selected agent, and the average Wasserstein distance as intrinsic rewards, respectively.
We test these variants in the scenarios from SMAC, and the results shown in Figure 6a We note that the absence of any of the components employed in our method results in significant performance degradation. Encoding trajectory representations with a fixed encoder leads poor performance, demonstrating the importance of learning distinguishable trajectory representations. Moreover, learning trajectory representations by minimizing the identity prediction loss or learning generative model is less efficient than our method. These methods do not necessarily learn distinguishable trajectory representations with large variations, thus the representations may not work properly in the Wasserstein distance objective to produce efficient feedback. Also, using the generative method leads lower learning efficiency due high computational cost. Predicting the trajectory representations or the raw observation-aciinn adopted in our method. The average Wasserstein distance does not yield satisfactory performance and even achieves worse performance than the random agent Wasserstein distance. As shown in Figure 6b the average Wasserstein distance intrinsic rewards do not provide effective incentives to encourage multi-agent diversity. Instead, our nearest neighbor Wasserstein distance is more sensitive to the trajectory representation variations. Despite the performance degradation caused by different kinds of Wasserstein distances, these Wasserstein distance methods also lead to significant performance improvement over QMIX, demonstrating the robustness of our representation learning method. As the difficulty of the task increases, we note obvious performance degradation caused by the ablation of the autoregressive model, indicating that learning trajectory representations results in more robust performance, especially in hard tasks.
does not lead to better performance than predicting the latent embeddings
6 Related Works
Diversity within MARL aims to learn diversified policies among agents to encourage efficient explo- ration. To achieve this goal, numerous diversity-driven methods have proposed different intrinsic motivations or regularizers. RODE [Wang et al. 2020c promotes diversity by assigning distinct actions to predefined roles; however, its effectiveness may decrease in scenarios with continuous actions and extensive action spaces. MAVEN [Mahajan et al.) [0199] introduces a value-based ap- proach that conditions agents' joint behaviors on a shared latent variable controlled by a hierarchical policy. EOI (Jiang and Lu 2021 utilizes supervised learning approach to promote agent individu- ality, employing a probabilistic classifier to predict agents' probability distributions based on their observations. SCDS (Li et al. 2021] concentrates on enhancing multi- diversity by optimizing mutual information between agent identities and trajectories. PMIC [Li et al. 2022] adopts a unique approach by maximizing the mutual information concerning superior cooperative behaviors while minimizing it regarding inferior behaviors. LIPO [Charakorn et al.] 2023) uses policy compatibility proxy to learn diverse policies and diversifies agents' behaviors through the mutual information

---

proposes formation-bus based exploration. encouraging visitations of diverse formations by guiding agents to fully understand their current formations. Although these approaches show promise in enhancing multi-agent diversity, the KI. divergence derived from the mutual information objective may lead insufficient exploration.
objective. FoX OJo et al. 2024
Wasserstein distance, emerging as an advanced measure of distribution dissimilarity, has garnered attentions of researchers from the community of machine learning. Many generative models (Arjovsky et al. 2017 Ambrogioni et al. 2018 Patrini et al. 2020 Tolstikhin et al. 2018) have incorporated the Wasserstein distance objective and demonstrate the effectiveness Wasserstein distance scenarios where distributions become degenerate on sub- monimalnf within pixel space. Some recent advances (Hu al. 2024 Bettini et al. [024] rely on the maximization of the Wasserstein distance to enlarge the policy differences among agents. MAPD [Hu et al., 2024 introduces the use of the Wasserstein distance as a metric for measuring policy differences. It normalizes the action distributions of different agents and computes the Wasserstein distance between them. Different from MAPD, our method learns a contrastive trajectory representation space by contrasting the trajectories of different agents learn distinguishable trajectory representations to make the Wasserstein distance meaningful. Empirical results shown in Section 5 demonstrates the superiority of our proposed representation learning method compared to MAPD. DiCo [Bettini et al., 2024] proposes using metric based on the average Wasserstein distance to control the diversity among agents. It performs well a simple multi-agent navigation task. However, it may suffer from insufficient exploration in challenging multi-agent tasks due to the controlled diversity. Furthermore, these methods fail to capture the similarities between agent policies, resulting in an ineffective of the Wasserstein distance, which ultimately impacts the performance of the proposed methods.
7 Limitations and Future Directions
The cost function of the Wasserstein distance determines how the probability mass is transferred. For simplicity, we resort to the Euclidean distance the cost function of the Wasserstein distance in the experimental settings. However, selecting an optimal cost function for the Wasserstein distance to address specific multi-agent tasks still remains a challenge.
8 Conclusion
In this paper, we propose a novel WCD exploration method. Unlike previous Wasserstein distance- based methods, our method maximizes the Wasserstein distance between the trajectory distributions of different agents in a contrastive trajectory representation space learned by a next-step prediction method, leading to sufficient exploration. We deploy our method in MARL by introducing nearest neighbor intrinsic reward based on the Wasserstein distance. The experimental results demonstrate that our method learns more diverse policies and leads to more sufficient exploration compared to the bascline methods. This simple yet effective method provides novel idea of learning useful representations to promote exploration, which shows promising results in learning cooperative policies for challenging multi-agent tasks.
Acknowledgments and Disclosure of Funding
This work was supported in part by the Fundamental Research Funds for the Central Universities (Grant No. NS2024055), in part by National Natural Science Foundation of China 62061 146002).
References
L. Ambrogioni, U. Güçlüü Y. Güçlütürk, M. Hinne, M. A. van Gerven, and E. Maris. Wasserstein variational inference. Advances in Neural Information Processing Systems, 31, 2018.
M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pages 214-223. PMLR, 2017.
M. Bettini, R. Kortvelesy, and Prorok. Controlling behavioral diversity in multi-agent reinforce- ment learning. arXiv preprint 4rXhv:2405.15054, 2024.
10

---

Y. Cao, W. Yu, W. Ren, and G. Chen. An overview of recent progress in the study of distributed multi-agent coordination. IEEE Transactions on Industrial informaties, (1):427-438, 2012.
R. Charakorn, P. Manoonpong, and N. Dilokthanakul. Generating diverse cooperative agents by learn- ing incompatible policies. In The Eleventh International Conference on Learning Representations, 2023.
R. Dadashi, L. Hussenot, M. Geist, and O. Pietquin. Primal wasserstein imitation learning. In ICLR 2021-Ninth International Conference on Learning Representations, 2021.
B. Ellis, S. Moalla, M. Samvelyan, M. Sun, Mahajan, J. Foerster, and S. Whiteson. Smaev2: An improved benchmark for cooperative multi-agent reinforcement learning. arXiv preprint arXiv:2212.07489, 2022.
A. Genevay, M. Cuturi, G. Peyré, and F. Bach. Stochastic optimization for large-scale optimal transport. Advances in neural information processing systems, 29, 2016.
Haamoja, A. Zhou, P. Abbeel, and Levine. Soft actor-critic: Off Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International conference on machine learning. pages 1861-1870. Pmir, 2018.
S. He, Y. Jiang, H. Zhang, J. Shao, and X. Ji. Wasserstein unsupervised reinforcement learning. In Proceedings of the AAAI Conference Artificial Intelligence, volume 36, pages 6884-6892, 2022.
M. A. Hearst, S. T. Dumais, E. Osuna, Platt, and B. Scholkopf. Support vector machines. IEEE Intelligent Systems and their applications, 13(4):18-28 1998.
S. Hu, C. Xie, X. Liang, and X. Chang. Policy diagnosis via measuring role diversity in cooperative multi-agent rl. In International Conference on Machine Learning, pages 9041-9071. PMLR, 2022.
T. Hu, Z. Pu, X. Ai, T. Qiu, and J. Yi. Measuring policy distance for multi-agent reinforcement learning. arXiv preprint 12X1 2024.
Jiang and Z. Lu. The emergence of individuality. International Conference on Machine Learning. pages 4992-5001. PMLR, 2021.
Y. Jo, S. Lee, J. Yeom, and S. Han. Fox: formation-awar exploration in multi-ageen reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 12985-12994, 2024.
Ladosz, L. Weng, M. Kim, and H. Oh. Exploration in deep reinforcement learning: survey. Information Fusion, 85:1-22 2022.
C. Li, T. Wang, C. Wu, Q Zhao, J. Yang. and C. Zhang. Celebrating diversity in shared multi-agent reinforcement learning. Advances Neural Information Processing Systems, 34:3991-4002, 2021.
Li, H. Tang, T. Yang. X. Hao, T. Sang, Y. Zheng, J. Hao, M. E. Taylor, W. Tao, Z. Wang, et al. Pmic: improving multi-agent reinforcement learning with progressive mutual information collaboration. arXiv preprint ar arXiv:2203.0855
2022.
R. Lowe, Y. I. Wu, A. Tamar, J. Harb, O. Pieter Abbeel, and I. Mordatch. Multi-agent actor-critic for mixed cooperativecooppettie environments. Advances in neural information processing systems, 30, 2017.
X. Ma. Y. Yang, Li, Y. Lu, Q. Zhao, and J. Yang. Modeling the interaction between agents in cooperative multi-agent reinforcement learning. In Proceedings of the 20th International Conference Autonomous Agents and MultiAgent Systems, pages 861, 2021.
A. Mahajan, T. Rashid, M. Samvelyan, and S. Whiteson. Maven: Multi-agent variational exploration. Advances in Neural Information Processing Systems, 32, 2019.
K. K. Ndousse, D. Eck, S. Levine, and N. Jaques. Emergent social learning via multi-agent rein- forcement learning. International conference on machine learning, pages 7991-8004. PMLR, 2021.
11

---

F. A. Olichoek and C. Amato. A concise introduction to decentralized pomdps, 2015.
A. v. d. Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXh:: 2018.
A. Pacchiano, J. Parker-Holder, Y. Tang, K. Choromanski, A. Choromanska, and M. Jordan. Learning to score behaviors for guided policy optimization. In International Conference on Machine Learning, pages 7445-7454. PMLR, 2020.
Park, O. Rybkin, and S. Levine. METRA: Scalable unsupervised RL with metric-aware abstraction. In The Twelfth International Conference on Learning Representations, 2024. URL https:/7
G. Patrini, R. Van den Berg, P. Forre, M. Carioni, S. Bhargav, M. Welling, T. Genewein, and F. Nielsen. Sinkhorn autoencoders. In Uncertainty in Artificial Intelligence, pages 733-743. PMLR, 2020.
A. Rahimi and B. Recht. Random features for large-scale kernel machines, Advances in neural information processing systems, 20, 2007.
Rashid, C. De Witt, G. Farquhar, Foerster, S. Whiteson, and M. Samvelyan. Qmix: Monotonic value function factorisation for deep multi-agent reinforcement learning. In 35th International Conference on Machine Learning, ICML 2018, pages 6846-6859, 2018.
M. Samvelyan, T. Rashid, C. S. De Witt, G. Farquhar, N. Nardelli, T. G. Rudner, C.-M. Hung, P. H. Torr, J. Foerster, and S. Whiteson. The starcraft multi-agent challenge. arXiv preprint arXiv: 1ee.119020044.4 2019.
K. Son, D. Kim, W. J. Kang, D. E. Hostallero, and Y. Yi. Qtran: Learning to factorize with transformation for cooperative multi-agent reinforcement learning. In International conference on machine learning, pages 5887-5896. PMLR, 2019.
Sunehag, G. Lever, A. Gruslys, W. M. Czamecki, V. Zambaldi, M. Jaderberg, M. Lanctot, N. Son- nerat, J. Z. Leibo, K. Tuyls, et al. Value-decompositioo networks for cooperative multi-agent learning based on team reward. In Proceedings of the I 17th International Conference on Autonomous Agents and MultiAgen Systems, pages 2085-2087, 2018
J. K. Terry, N. Grammel, S. Son, B. Black, and A. Agrawal. Revisiting parameter sharing in multi- ageet- deep reinforcement learning. arXiv preprint 201i:20005.366220 2020.
I. Tolstikhin, O. Bousquet, S. Gelly, and B. Schölkopf. Wasserstein auto-encoders. In 6th International Conference on Learning Representations (ICLR 2018). OpenReview. net, 2018.
C. Villani et al. Optimal transport: old and new, volume 338. Springer, 2009.
O. Vinyals, I. Babuschkin, W. M. Czarnecki, M. Mathieu, A. Dudzik, J. Chung, D. H. Choi, R. Powell, T. Ewalds, Georgiev, al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782):350-3354 2019.
J. Wang, Z. Ren, T. Liu, Y. Yu, and C. Zhang. Qplex: Duplex dueling multi-agent q-learning, arXiv preprint arXiv:2008.000 010062 2020a.
T. Wang, H. Dong, V. Lesser, and C. Zhang. Roma: Multi-agent reinforcement learning with emergent roles. arXiv preprini arXiv:2003.0m 08039, 2020b.
T. Wang, T. Gupta, A. Mahajan, B. Peng, S. Whiteson, and C. Zhang. Rode: Learning roles to decompose multi-agen tasks. arXiv preprint arXiv:2010.0/ 523, 2020c.
Y. Wang, B. Han, T. Wang, H. Dong, and C. Zhang. Dop: Off -ofl-o multi-agent decomposed policy gradients. In International conference on learning representations, 2020d.
Y. Yang, X. Ma, C. Li, Zheng, Q. Zhang, G. Huang, Yang, and Q. Zhao. Believe what you see: Implicit constraint approach for offline multi-agent reinforcement learning, Advances in Neural Information Processing Systems, 34: 10299-103312 2021.
T. Zhang, Y. Li, C. Wang, G. Xie, and Z. Lu. Fop: Factorizing optimal joint policy of maximum- entropy multi-agent reinforcement learning. In International Conference on Machine Learning, pages 12491- 12500.. PMLR, 2021.
12

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our main claims made in the abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of work in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
13

---

Justification:
We provide the full set of assumptions and a complete (and correct) proof.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: Yes]
Justification: We provide the training details and hyperparameter settings needed to repro- duce the experimental results in Appendix K
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
14

---

Answer: [Yes]
Justification:
We provide the source code in the supplemental material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips.ccm httpsscccpppbccccceeceeeeee wwwpc//wwgeeeeeeeeeeeeeee Seweeeeeeoooeemiieeeeeeeeeeeeeee ssiniiiicyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
We provide the training details in Appendix K
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We report both the mean and standard deviation of performance results, averaged over five random seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (c.g., Normally distributed errors).
15

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide the computation platform used our work Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
Our research follows the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
Our work aims at advancing the field of Machine Learning and has no societal
impact.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
16

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We have properly cited the benchmarks used in work.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name the license (c.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
17

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes],
Justification: provide the source code in the supplemental material with the documenta- tion.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
18

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standarr components.
Please refer to our LLM policy (https: /ttpp/////euips..... www.errpssscccooofeeeee ///Cooofeeenceee22255//////////e for what should or should be described.
19

---

A The TD loss of QMIX
The TD loss of QMIX to learn the optimal Qtot defined as:
(
LTD Σ
(++ max Qeot (St+1, Me+1) Qroe (Se; uz) uc+1
where Qiat is the target network and b is the size of transition samples from the replay buffer D. the global reward shared among agents. Note that since all agent's policies are jointly trained by minimizing the TD loss, we cannot simply apply each agent's 's intrinsic reward Tw to the global reward r to formulate a swwad--haping independently train each agent's individual policy, That why we need to learn an additional intrinsic utility network to maximize the intrinsic reward Far-
B Theoretical analysis of convergence guarantee of our method
We next provide a theoretical analysis to guarantee the convergence of the MARL when our Wasser- stein distance-based intrinsic reward introduced. To achieve this, we theoretically analysis the boundeness of the intrinsic rewards, which is a standard and critical condition for proving conver- gence. The bounded rewards state that the rewards received by the agents at each time step are limited within finite range. For our intrinsic reward, boundedness ensures that the diversity-seeking objective does not overwhelm the primary goal of maximizing the team reward. the intrinsic reward were unbounded, it could lead to policies that are arbitrarily diverse but suboptimal in terms of task performance. Specifically, for our proposed WCD method, we prove that the Wasserstein distance-based intrinsic reward is bounded. The boundedness of this intrinsic reward is guaranteed by properties of the learned representation space and the Wasserstein distance itself.
1. Boundedness of the Compact Representation Space
The trajectory representations guaranteed be in compact (i.e., closed and bounded) set due to the architecture of the trajectory encoder. The encoder is composed multi-layer perceptron (MLP) and Gated Recurrent Unit (GRU), both of which use activation functions that constrain their outputs.
Let the trajectory encoder be represented by the function go X' C, where X' is the space of trajectory histories and is the latent representation space. The encoder is a composition of the non-linear encoder ge, and the autoregressive model ga,-
ge, (go, (=i)), ... 90, (x(-1))
(10)
The boundedness of the output of can be shown by analyzing the components:
Bounded Activation Functions: The MLP and GRU components of the encoder activation functions such as the sigmoid (o(x) = r=-->) or the hyperbolic tangent (lanh(z)), which are
inherently bounded. For any input æ, their outputs are confined to a specific range (e.g., [0, 1] for sigmoid, [-1, 1] for tanh).
Bounded Weights and Biases: During training, the weights and biases of the neural networks are typically regularized (e.g., through weight decay), which keeps them within finite range.
Given that the inputs to each layer are transformations of the outputs from previous layers with bounded activation functions, the final output of the network, which is the trajectory representation et, will also bounded. we assume the maximum possible value for any dimension of the representation vector is C Cmaz. then the representation space subset of hypercube in RRd;
C c [-Cmaax Cmax]" d
(11)
where d is the dimension of the representation space. Since this set is closed and bounded, the representation space compact.
2. Boundedness of the Wasserstein Distance
20

---

The Wasserstein distance is used to measure the cost of optimal transport between two probability distributions. When these distributions are defined over compact set, the Wasserstein distance is also bounded.
The I-Wasserstein distance has dual representation known the Kantorovich-Rubinste duality, which is particularly useful for understanding its properties:
(p. = sup 11
Beup (f()
(12)
where the supremum is taken over all I-Lipschitz functions f.
Let and be two probability distributions over the compact representation space C. Since C is bounded, there exists a constant M such that for any two points c1 02 E C, the distance between them is bounded: d(c1, 02) & M.
For any 1-Lipschitz function f and any two points c1, E C, have:
If(c)) - f(c2) - d(c 02) S M
(13)
Now, consider the expectation of under the distributions p and q-
1m==p[((1) Ey [f(]]] 1/6 [(x)dp(()) Jo s())d((v))
(14)
By choosing a reference point co E C, we can write:
((oo)dd(()))) - (f(0) ((o))dllo)))
(15)
Since [f(c) f(co)] S d(c, < M for any E C, we have:
[xzplf(x)) - Eywglf(() (f())] & Mdp(t) + Mdq(y) = 2M
(16)
This shows that the Wasserstein distance is bounded by the diameter of the compact set on which the distributions are supported. The intrinsic reward Tws being the minimum of these bounded distances. is therefore also bounded. This formalizes the boundedness assumption, which is essential for the stability and convergence of the WCD algorithm.
C Pseudocode for WCD
The pseudocode for WCD is given in Algorithm
!
D Integrating WCD with the policy-based method
We have implemented our method with the value-based method QMIX. Here, we illustrate the integra- tion of proposed WCD with policy-based methods. Specifically, integrate WCD with MAPPO, a state-of-the-at policy-based MARL algorithm measured by SMAC. In MAPPO, all agents share an actor network and critic network. As each agent learns its own critic, we can straightforwardly incor- porate a shaped reward Tene + arw (where Tenv represents the environmental reward and Tw denotes the Wasserstein distance-based intrinsic reward) when computing the reward-te each agent's critic network. The remaining components of MAPPO do not require modification. conduct experiments on Pac-Men, SMAC, and SMACv2 to test the performance of WCD+MAPPO. The results, presented in Table demonstrate the superior performance of WCD+MAPPO compared to the baselines.
R for updating
21

---

Algorithm 1 Wasserstein Contrastive Diversity (WCD)
Initialize dual functions je and v. Initialize the joint policy = (ma )
Randomly initialize Qtot for QMIX. repeat
for each episode do
Collect the trajectories of all agents induced by the joint policy π. Store them into replay buffer D.
end for
Sample a batch of trajectories T from the replay buffer D.
Train the trajectory encoder go to learn trajectory representations by minimizing the InfoNCE
loss given by Equation 4
Train dual functions is and by SGD with the gradient given by Equation 6
Compute the intrinsic reward rw = min
at al W (Pv., Px.,) for each agent.
|A]
Jointly train the policy Ta for each agent by minimizing Ltotal LTD + aLTD- until Quot is converged
E Environmental details and Additional experimental results
E.1
Environmental details and experimental results
In Pac-Men, four agents are initialized in the central room of a maze. Each agent is restricted to observing a 4x4 grid surrounding them. Randomly distributed dots are present each edge room. The objective for the agent is to gather as many dots as possible in each edge room. We vary the lengths of paths to evaluate the exploration of environments. Specifically, path lengths for downward, leftward, rightward, and upward directions are set to 3, 6, 6, and 10, respectively. Only one path falls within the agent's observation scope. Dots in each room will respawn once all have been consumed by agents. Agents receive environmental reward equal to the total number of dots consumed in each time step.
The SMAC benchmark includes many cooperative tasks based on Blizzard's 's real-time strategy game StarCraft II, designed evaluate the efficacy of different Multi-Agent Reinforcement Learning (MARL) algorithms. Agent-level control in SMAC utilizes the Machine Learning APIs provided by StarCraft and DeepMind's PySC2. Each task presents combat scenario with two armies: one led by allied RL agents and the other by a non-learning game Al. The game ends when all units from any army perish or a predefined time limit reached. The objective for allied agents is to maximize the game's 's win rate. To achieve this, agents must learn sequence of actions to effectively collaborate with allies in vanquishing enemy forces. An illustrative example of such collaboration involves mastering kiting skills, where agents organize formations based armor types lure enemy units into pursuit while maintaining a safe distance to minimize damage. The SC2.4.10 10 version of StarCraft II is utilized, and performance comparison across different versions not applicable. Experiments are conducted across six scenarios, including 3s5z, 2c_ vs "s__4gz, 7sz, 6h vs_ 8z, corridor, and 3s52_ vs ss_3s6 spanning various difficulty levels.
SMAC greatly limited by its lack of stochasticity. To remedy this, the newly released SMACv2 proposes modifications such as incorporating random team compositions and random start positions. These adjustments aim to inject more stochastic clements into the environment to effectively evaluate the exploration capabilities of MARL algorithms. We conduct experiments in three SMACv2 scenarios: terran_5 5______________ protoss 5_ 5_vs_5, and 5erg___vs__ 5. In SMACv2, each race in the game of StarCraft II employs three unit types, with units algorithmically assembled into teams. The probability of each unit type appearing in each episode remains fixed throughout training and testing phases. Allied agents have the same unit types their adversaries. In each episode, allied agents randomly deployed on the map using either a reflect or surround style.
We present the average returns of all algorithms in Pac-Men, SMAC, and SMACv2, along with their standard deviation over five random seeds, in Table [] The results indicate the significant performance superiority of our method over bascline methods.
22

---

Table Average returns of all algorithms in Pac-Men, SMAC, and SMACv2. denotes the standard deviation over five random seeds.
SMAC
SMACV2
Mechod
Pac-Men
60_
corridor
SerranL5 5_1
profoss, 5_
serg. serg_55_v._5
QMIX
021110.04
0.7210.13
0.85_.008 0.08
0.17.00.2
0.23.0.0
0.5.10.07
0.2610012
0.68.00.33
0,53.0006
0,41_0004
MAPPO
0.4910.03
0.8110005
0.83_0004
0.5240.06
0.53
0.6210.055
0.57.0008
0.5210004
0.47.0000
0,37
MAVEN
0.32
0.512:0:21
0.72-0006
0000000000
0.0240.04
0.36:10.88
0.18:10.15
0.58:10.04
0.31
0.29
B0I
0.1
0.87200077
0.33-0002
0.37
0.08
0.25
0.2210113
0.651:0.55
0.42:0103
0.47:0004
QTRAN
0.28:00088
0.2110.19 19
0,75
0.00:00000
0.02:00.0
0000
0.00:10001
0.42:10.02
0.40:0004
0.25
SCDS
0.37.00055
0,7610.00
0.57
0.11100033
0.03.0.01
0.56.10.06
0,0010.000
0.52.10.03
0,47.0006
0.38
PMIC
0.34100033
0,8210.03
0.79
0.58.10.02
0.5110.08
0.37100033
0.18.10.06
0.47
0.36.0..0
0.42.000
LIPO
0.4310.02
0.711.003
0.76
0.3911004
0.361.0.66
0.7:00033
0.2110003
0.4310002
0,46.0.00
0,37-000
FOX
0.39
0.74
0.64:0.55
0.5665003
0.45
0.5240004
0.310004
0.5410005
0.56
0.49-0002
RODE
0.37.Q012
0.72600.3
0.69:0002
0.5486003
0.13
0.37
0.3:50022
0,40:100
0,340000
MAPD
0.23:0013
0.634:0044
0.550003
D.00000222
0.26
0.12#008
0.10:0006
0.442000
0,38=1002
0.36-000
I6Co
0.2700022
0.82+0003
0.71-003
0.56++000
0,34+0.04
0.05+0004
0.14+008
0A5+0004
0,36-0.0
0.31-0.0
WCD+OMIX
0.87.6.03
0,9510.03
0,06
0.8710.04
0,830003
0.85110.0
0:5210003
0.8510005
0.90
0,84
WCD+MAPPO
0.3220.22
0,5210.02
0,89
0.941:.03
0,7910.04
0.57.1005
0,5910.04
0,8910.03
0.82-0.02
0.91
E.2
Additional results
Google Research Football We also evaluate our method in three scenarios from Google Research Football (GRF), challenging, physics- based environment that simulates football game. In this setting, agents must learn strategic planning, coordination, and precise timing to succeed. The players on the left side (excluding the goalkeeper) serve agents, trained to develop cooperative policies, while the right-side players are controlled by the game engine. The agents operate within a discrete action space of 19 options, which includes moving in eight directions, sliding, shooting, and passing. Observations for each agent consist of the positions and movement directions of the agent itself, other agents, and the ball. The results of the comparison are presented in Table 2 where our method consistently outperforms the baseline methods across all scenarios.
Table 2: Performance comparisons our method against the baseline methods in Google Research
Football.
Method
academy
3_ 3___yy_________ _eith_epppperrrr
academy 4__rsi____________ $_vs_2__i
with with__eepperr
academy__oouue cocpller__ooeeerrreeee atlack ahad
QMIX
0,2310.05
0.13110.09
0.1710.03
MAPPO
0.31+00.99
0.18++0.09
0.23+00.07
MAVEN
0.18:+0.06
0.08=10.06
0.13=0.09
EOI
0.17.10.05
0.0510.03
0.07
QTRAN
0,25++0.03
03
0.11=00.05
SCDS
0,42=0.13
0.25
0,47=0.06
PMIC
0.23
0.11100.7
0.16:10.07
LIPO
0.1910.05
0007.0.03
0.12
FoX
0,57+00.55
0.41+0.13
0.33-0008
RODE
0.37+
0.16+00.10
0.28
MAPD
0.23:10.11
0.11
0.19:10007
DiCo
0.42 #0.06
0.29
0.21=0.12
WCD+QMIX
0.82+0.11
0.68+
0.76=0009
WCD+MAPPO
0.78110.15
0.75
0.13
Homogeneous behaviors Agents may sometimes desire to behave in the same way. For instance, allied agents in the scenarios of SMAC might take the same action to fire at the same enemy in order to rapidly defeat it. In this section, to demonstrate the effectiveness of our method in learning such behaviors, we evaluate focus fire. The results are shown in Table 3 Our method outperforms QMIX across all scenarios, demonstrating that our method would not prevent the homogeneous behaviors if they can lead to more environmental rewards. In contrast, our method encourages sufficient exploration to search for such optimal cooperative behaviors.
method in four homogeneous scenarios of SMAC that require the trick of
Scalability The scalability the MARL algorithms refers to their ability to effectively handle the growing number of agents in the environment. The action space grows exponentially with the number of agents, highlighting the urgent need for exploration. In this section, we evaluate the scalability of our method in four scenarios of SMACv2 with an increasing number of agents: terran 5_ vs 5,
23

---

Table 3: Performance of our method and QMIX in homogeneous scenarios.
Method
8m
5m vs 6m
8m vs 9m
10m vs. 11m
WCD+QMIX
0.94
0.95:10.03
0.93:10.04 0.58=00005
0.91 0.03 0.43++0.04
QMIX 0.87+0.03
0.65++0.04
terran_10_vs__ 110_v__10, terran_15 1_____________________________ vs_ and terran_22 perran_20__v vs__0. We present the results in Table f Our method maintains its outperformance over QMIX across all scenarios. QMIX suffers from poor scalability due to limited exploration, while our method scales well with an increasing number of agents, demonstrating that our method can lead to sufficient exploration of action space by enlarging the Wasserstein distance between trajectory distributions of different agents in the latent representation space.
Table Performance of our method and QMIX in scenarios of SMACv2 with different number of agents
Method WCD+QMIX
terrn ___ vs_
terran 10_ vs_10
terran 15 vs_ vs_15
terran 20 vs vs_20
0.83+0.03
0.86 +0.02
0.83 +0.04
0.81 +0.05
QMIX
0.68+0003
0.39-0.04
0.24 +0.06
0.1 1-1.50
F
Comparison with e-Greedy
The e-greedy method is a commonly used exploration strategy in many RL algorithms. Typically, increasing the value of enhances exploration. In this section, we compare our Wasserstein distance- based method with e-greedy to highlight its effectiveness in promoting exploration within MARL. For this comparison, we set the values to 0.05, 0.075, and 0.1 for QMIX, and evaluate these settings in the challenging scenarios including corridor, 3s5z vs 3s6z, terran vs__,, and protoss 5 5_vs vs_5. 5, The results, presented in Table 5 show that our Wasserstein distance-based method is more effective in fostering exploration compared to simply increasing c. Notably, increasing values does not lead to significant performance gains. In multi-agent settings, higher values primarily increase randomness in an individual agent's s action selection without enhancing diversity or coordination among agents, as they fail to consider the trajectories of other agents, resulting in inefficient exploration.
Table 5: Comparison of performance between our method and QMIX using various e values
Method
corridor
3s5z_vs_3sd0
terran_ _eaa____v__
protoss_5_vs__ _s_
= 0.05 (QMIX)
0.57 +0.07
0.36 +0.12
0.68 +0.03
0.53 +0.05
0.075 (QMIX)
0.61 +0.04
0.39 +0.11
0.72 +0.04
0.62 +0.07
0.1 (QMIX)
0.63 #0.06
0.44 0.15
0.74 10.03
0.69 10.06
Wasserstein distance (our method)
0.85 0.04
0.82 0.03
0.85 +0.03
0.90 :0.02
G Evaluations of different kernel functions
We use the Gaussian kernel by default in our paper. We may also use a linear kernel to parameterize dual functions. To evaluate the effectiivnees of using the linear kernel for dual functions, design linear kernel variant and test it in the super hard scenarios of SMAC. The results are shown in Table We note that using the linear kernel to parameterize dual functions leads to significant performance decline. We suspect this is because the dual function may not be linear functions. Using the linear kernel constraints the representation ability of the dual function.
H Evaluations of different values for the weight of the intrinsic reward a
The values for the weight of the intrinsic reward in different scenarios are listed in Table 5 in our paper. To investigate the effect of different weights of the intrinsic reward, we evaluate different weight values in the easy scenario 355z and the super hard scenario corridor. The results are shown in Table 7] The results demonstrate that our method not very sensitive the values of the weight. Sub-optimal weights do not result in significant performance drop even in the super hard scenario.
24

---

Table 6: Performance comparisons of WCD with different kernel functions in the scenarios of SMAC
Method
6h_vs vs 8z
corridor
3s57_ vs_3s6z
WCD (Linear Kernel)
0.37 0.07
0.39 + 0.05
0.32 = 0.03
WCD (Ours)
0.85 + 0.03
0.90 + 0.03
0.87 0.04
Table 7: Performance comparisons of WCD with different values for the weight of the intrinsic
reward (t.
3s5z
corridor
Method a 0.02
a 0.05
Q = 0.1
= 0.02
a = 0.05
a = 0.1
WCD 0.89 + 0.03
0.91 + 0.02
0.93 0.03
0.82 + 0.07
0.85 + 0.04
0.81 + 0.05
I Evaluations of different cost functions
In our paper, we mainly use the Wasserstein distance to encourage sufficient exploration and simply adopt the Euclidean distance the cost function in many prior works. We may also use cosine similarity as the cost function, which measures the direciton differences between data points. We test the cosine similarity in Pac-Men, where agents need to move to different directions. The resutls are shown Table We note that the Wasserstein distance based on the cosine similarity achieves higher rewards in Pac-Men. In our work, we do not specifically discuss different cost functions and use the default Euclidean distance because we want to be consistent with prior works using the Wasserstein distance to ensure a fair comparison.
Table 8: Performance comparisons WCD using different cost functions.
Method
Pac-Men
WCD (Cosine Similarity)
94 + 0.05
WCD (Euclidean Distance)
87 0.03
Additional computational overhead analysis
We note that compared QMIX. our method needs to train the trajectory encoder by minimizing the InfoNCE loss and the dual functions to calculate the Wasserstein distance objective, which brings additional computational overhead. We next present the comparisons of training time between QMIX and our method in the three super hard scenarios of SMAC under the same computation platform in Table 9 The results show that compared to QMIX, the addition of our proposed method does not consume lot of extra training time.
K Training Details and Hyperparameters
In this section, we provide the training details and hyperparameters adopted in our experiments. To implement the one-step prediction method, we use two-layer MLP with hidden size of 64 for the encoder ga, followed by the batch normalization and a GRU unit for the autoregressive model ga,- We adopt a dual vector with dimension m of 64 to parameterize the dual function. To integrate our method with QMIX, the intrinsic agent utility network is implemented with two-layer MI.P with hidden size of 64. We keep other components the same as in QMIX.
The policy networks of all agents are implemented with Deep Recurrent Q-Networks. At each time step, an agent's policy network processes local observation input, which is then forwarded through a fully-connected hidden layer, followed by a GRU unit, and ultimately fully-connected layer generating U outputs, where U the number of actions. Furthermore, all agents' policies share the same policy network parameters to accelerate training. We set the evaluation interval to 10K steps followed by test episodes. We run all methods for 5 million steps We employ hard updates to update target networks every 200 episodes in SMAC and SMACv2. In Pac-Men, we utilize soft updates for updating target networks with a momentum of 0.01. The common hyperparameters consistent across various methods for each multi-agent task. Detailed
tested tasks.
25

---

Table 9: Comparisons of training time between QMIX and our method in the three super hard scenarios of SMAC
Training time
Methods
6h vs_ 87
corridor
3s57_vs_3sse 35___v__s66
QMIX
8h 35m 29s
7h 17m 30s
10h 42m 18s
WCD+QMIX
8h 43m 21s
7h 36m 53s
10h 49m 37s
hyperparameters are provided in Table 10 The replay buffer size set to 5K. We implement our
method using NumPy and PyTorch. All experiments are performed on NVIDIA GeForce RTX 4090 GPU.
Table 10: Hyperparameters
Pac-Men
SMAC
SMACY2
hidden dimension learning rute
64
128
0.0003
0.005
optimizer
target update
0.01(soft)
200thard)
batch size
32
64
0.03
0.05
0,005 for 3s5z, 2c_vs_ 2c_xs_64xg, Sm, 5m_ 6m, $m_vs_9m and 10m_vs_llmm
WCD+QMIX
0.03
0.05 for 7ez, 6h vs 82, corridor, and 3652 vs. sr__v_______ 3562
0.005 for 3s5z. 2c 642g, 8m, 5m, 6m, 8m. vs. 9m, and 10m_ Ilm.
for WCD+MAPPO
0.01
0.03 for 7sz, vs 8z, corridor, and 3852_ vs_662
0.03
200,000 for 3s5z, 2c 642g. 8m, 5m vs. 6m, 8m vs. 9m, and 10m_vs Hm,
epsilon anneul time
200,000
500,000
500,000 for 7sz, 6h_ _s_ _z, corridor and $s5z_vs_3s66
L
Visualizations
Challenging tasks typically necessitate complex cooperative behaviors requring agents to learn diverse policies. We next present
visualization examples of diverse policies learned by method the super hard scenarios (6h_vs_He _o__v,__ corridor, and 3s5z_vs (ir__vs_3se vs_3s62) in Figure 7 the dh_vs_8z scenario, one agent first leaves the team, causing most enemies to follow the lone agent's movements. The agent continues moving away to draw the enemies' fire and cover other agents. Other agents then quickly surround the few remaining enemies. Through learning such cooperative tacties, agents successfully scatter the enemies' powerful attacks. If all agents behave similarly and directly move towards enemies, they would be killed by enemies immediately. Similar tactics can also be observed in the other two scenarios, demonstrating the effectiveness of our method in encouraging multi-agent diversity.
We also present the visitation heatmaps of baseline methods and our method in the protoss_5. __v_ss___________ and the zerg_5 5 vs_5 scenarios in Figures 8 and respectively. The visitation heatmaps reveal that our proposed WCD leads to more sufficient exploration compared to the baselines. We believe this is because the baseline methods does not provide effective incentives for exploring the environment. As a result, the agents trained by baseline methods are slow to search for randomly appearing enemies on the map. In contrast, method enables sufficient exploration by enlarging the Wasserstein distance.
26

---

corridor
3s5x_ys vs 3336
Figure 7: Visualization examples of diverse policies emerging in 6h_ vs 8z (top), corridor (medium), and 3s5z_vs, vs_3s6z (bottom) from initial (left) to final (right). Green and red shadows represent agents and enemies, respectively. Green and red arrows represent the moving directions of agents and enemies, respectively.
(a) MAVEN
(b) BOI
(c)SCDS
(d) LIPO
(e) MAPD
(l WCD
Figure 8: Visitation heatmaps of different algorithms in the protoss vs 5 scenario.
(a) MAVEN
(b) BOI
(c)SCDS
(d) LIPO
(e) MAPD
(f WCD
Figure 9: Visitation heatmaps of different algorithms in the zerg_ 5 vs scenario.
27


Reviews:
[{"review_id": "kidvfeGtXb", "reviewer": "Reviewer_WaHo", "summary": "This paper explores a novel approach to promoting diversity in Multi-Agent Reinforcement Learning (MARL). Specifically, the work proposes a method for enhancing trajectory diversity based on the Wasserstein distance. By integrating Contrastive Predictive Trajectory Representations, kernel-based methods, and nearest neighbor techniques, the authors develop a practical algorithmic framework. Experimental results on Pac-Men and SMAC environments demonstrate that the proposed method outperforms several existing diversity-oriented approaches.", "questions": "Ethical Concerns:\n\nIn the domain of MARL diversity, several prior works have addressed similar problems using similar methodologies. However, the absence of citations to these works is unacceptable. The authors must clarify the distinctiveness of their approach in relation to these existing studies.\n\nBased on the publicly available information on OpenReview, relevant publications include:\n\n1.Li T, Zhu K, Li J, et al. Learning distinguishable trajectory representation with contrastive loss[J]. Advances in Neural Information Processing Systems, 2024, 37: 64454-64478.\n\n2.Li T, Zhu K. Toward Efficient Multi-Agent Exploration With Trajectory Entropy Maximization[C]//The Thirteenth International Conference on Learning Representations.\n\n3.Li T, Zhu K. Unsupervised Multi-Agent Diversity With Wasserstein Distance[J].", "limitations": "Yes.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**: \n\n(*Major*) From a writing perspective, the paper is well-written and easy for readers to understand. \nThis paper has a clear structure with corresponding evidence supporting relevant arguments. \nMoreover, the figures in the article (Fig.1 and Fig.2) are clearly drawn, facilitating expression and readability. \n\n(*Major*) The experiments in the paper are comprehensive, with relatively complete baseline comparisons (comparing classical methods, diversity-based methods, and WD-based methods), and sufficient ablation studies have been conducted. \n\n(Minor) The experiment corresponding to Figure 6 in the paper provides a thorough comparison of various variants of WCD, particularly the comparisons among generative models, autoregressive models, contrastive models, and fixed encoders, which can provide certain insights. \n\n\n---\n\n\n**Weaknesses**: \n\n(*Major*) The authors lack detailed comparisons with two works. The methods and experimental results of these two papers are highly similar to this work, and I hope the authors can conduct in-depth comparisons with these two methods in both methodological analysis and experimental analysis. \n\n1. Li T, Zhu K, Li J, et al. Learning distinguishable trajectory representation with contrastive loss[J]. Advances in Neural Information Processing Systems, 2024, 37: 64454-64478. \n\n2. Li T, Zhu K. Toward Efficient Multi-Agent Exploration With Trajectory Entropy Maximization[C]//The Thirteenth International Conference on Learning Representations.", "was_revised": true, "final_justification": "The authors have addressed my ethical concerns. This was the sole reason for my initial rejection of the paper. I have rewritten my review and rescored accordingly.", "rebuttal": "Thank you for the in-depth review. Here are our responses to your questions and concerns:\n\n**1. Comparison with Previous Works**\n\nWe did not initially compare our work with [1,2] because they are very recent publications. Moreover, a previous version of our paper was also submitted during the ICLR 2025 review period. Therefore, our work and [1,2] were developed contemporaneously. Following the reviewer's suggestion, we have now added a comparison with these two works from both methodological and experimental perspectives.\n\nThe main contribution of our work is making the Wasserstein distance an effective measure for encouraging exploration in multi-agent settings. To achieve this, we propose a next-step prediction method based on Contrastive Predictive Coding (CPC) to learn a trajectory representation space, and we then employ the Wasserstein distance within this latent space. This approach is fundamentally different from the contributions of [1] and [2]. [1] uses a contrastive lower bound to maximize the mutual information between trajectories and learnable identity representations, thereby encouraging multi-agent diversity. This addresses the limitation of previous mutual information methods that tend to overfit to fixed agent identities. In contrast, [2] encourages exploration by maximizing trajectory entropy within a representation space learned by an improved version of SimCLR in a multi-agent setting.\n\nWhile the representation learning methods of all three works involve contrastive learning, the specific methods employed are distinct. Our work proposes a next-step prediction method based on CPC; [1] uses an improved InfoNCE loss to predict learnable identity representations; and [2] is based on an improved version of SimCLR.\n\nWe would also like to clarify that the reviewer's claims (b) and (c) are incorrect. Neither [1] nor [2] involves the Wasserstein Distance (WD), contrary to the reviewer's assertion.\n\nFrom an experimental perspective, all three works achieve satisfactory performance in certain StarCraft Multi-Agent Challenge (SMAC) scenarios. To further compare their performance, we have evaluated all three methods on two additional challenging, sparse-reward SMAC scenarios: so_many_banelings and offence_hard. In so_many_banelings, a reward is given only when the agents win the game. The offence_hard scenario includes advanced environmental obstacles that agents must navigate to strategically engage enemies; rewards are provided only for damaging or eliminating an enemy. The comparison results are shown below:\n\n| Method | so_many_banelings | offence_hard |\n| :--- | :--- | :--- |\n| CTR+QMIX [1] | $0.13 \\pm 0.06$ | $0.26 \\pm 0.11$ |\n| TEE+QMIX [2] | $0.21 \\pm 0.08$ | $0.37 \\pm 0.06$ |\n| WCD+QMIX (ours) | $0.63 \\pm 0.09$ | $0.75 \\pm 0.13$ |\n\nThese results show that our method outperforms the other two, demonstrating the effectiveness of maximizing the Wasserstein distance in the learned trajectory representation space.\n\n[1] Li T, Zhu K, Li J, et al. Learning distinguishable trajectory representation with contrastive loss. *Advances in Neural Information Processing Systems*, 2024, 37: 64454-64478.\n\n[2] Li T, Zhu K. Toward Efficient Multi-Agent Exploration With Trajectory Entropy Maximization. *The Thirteenth International Conference on Learning Representations*, 2025.\n\n**2. Ethical Concerns**\n\nWe note that the reviewer raises significant concerns about the similarity between our work and [3]. Due to the double-blind review policy, we cannot disclose our authorship. The only information we can publicly provide is that, like our submission, [3] is also a paper that was not accepted at ICLR 2025. To fully address this concern, we respectfully request that the Area Chair (AC) confidentially confirm the authorship of our submission and of [3] with the Program Chairs (PC).\n\n[3] Li T., Zhu K. Unsupervised Multi-Agent Diversity With Wasserstein Distance. *Submitted to ICLR 2025.*"}, {"review_id": "kdFLXGQIFr", "reviewer": "Reviewer_scNh", "summary": "To effectively encourage diverse policies in cooperative MARL, the paper presented Wasserstein Contrastive Diversity (WCD) exploration to promote multi-agent diversity by maximizing the Wasserstein distance between the trajectory distributions of different agents in a latent representation space, where CPC is adopted to learn different trajectory representations. Experimental results show the effectiveness of the proposed approach.", "questions": "1. Although experimental results have been reported to verify the effectiveness of the proposed method, there are no theoretical guarantees for the performance improvements. \n2. In the paper, the authors claimed that by employing the proposed intrinsic reward function, better performance can be achieved experimentally by keeping the current trajectory away from its nearest neighbor trajectory. However, for tasks that benefit from homogeneity, it may not be true. \n3. In the literature, DiCo aims to address the same diversity issue in MARL while theoretically proving the desired diversity can be achieved. Can the proposed method have the same theoretical guarantee? Moreover, more discussions and comparisons about the performances between WCT+QMIX and DiCo can be provided. It is also difficult for me to distinguish the experimental curves representing WCT+QMIX and the experimental curves representing DiCo, they almost have the same color in the experimental figures. \n4. How does DiCo perform in terms of exploration, which appears to be missing from the analysis in Figures 8 and 9?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths: \nThe paper is well-motivated and clearly written,  Wasserstein Contrastive Diversity (WCD) exploration is proposed to promote multi-agent diversity by maximizing the Wasserstein distance between the trajectory distributions of different agents in a latent representation space, and extensive experiments conducted on several multi-agent tasks show the effectiveness of the proposed framework. \n\nWeakness:\nThe presented method is mainly the direct combination of existing techniques to address the diversity issue, e.g., CPC for next-step prediction/Trajectory Representation, Wasserstein distance for diversity measurement, the contribution is not significant. More importantly, while experimental results are competitive, no theoretical guarantees of the performance are provided. Additionally, more experimental results, especially the comparison with DiCo, can be reported.", "was_revised": true, "final_justification": "The authors have addressed the concerns on the contributions and theoretical guarantees of the proposed framework.", "rebuttal": "Thank you for your careful review and for providing us with detailed and helpful feedback. We respond to your concerns and questions below.\n\n**Weakness 1:** The presented method is mainly the direct combination of existing techniques to address the diversity issue, e.g., CPC for next-step prediction/Trajectory Representation, Wasserstein distance for diversity measurement; the contribution is not significant.\n\nPrevious methods based on the Wasserstein distance overlook the similarity of agents’ initial policies resulting from shared policy network parameters, which causes the Wasserstein distance to be an ineffective diversity incentive. To solve this, the main contribution of our work is to enlarge the Wasserstein distance in the latent representation space, making the Wasserstein distance measure effective in multi-agent settings. To learn meaningful representations, we propose employing CPC to contrast different agents' trajectories. Next, we propose a novel kernel method to calculate the Wasserstein distance more efficiently, since directly computing the Wasserstein distance in multi-agent settings is computationally expensive. Moreover, we design a novel intrinsic reward based on the nearest-neighbor Wasserstein distance to efficiently encourage multi-agent diversity. These contributions have not been proposed in previous works.\n\n**Weakness 2:** More importantly, while experimental results are competitive, no theoretical guarantees of the performance are provided.\n\nDiCo provides a theoretical analysis to guarantee that the policies defined by Equation 4 can achieve the desired diversity. Unlike DiCo, we do not directly control diversity. Instead, we enlarge the Wasserstein distance to learn diverse behaviors, solving the issue of inefficient exploration caused by shared policy network parameters. However, such an exploration technique may affect the convergence of the MARL method. Next, we provide a theoretical analysis to guarantee the convergence of the MARL method when our Wasserstein distance-based intrinsic reward is introduced. To achieve this, we theoretically analyze the boundedness of the intrinsic rewards, which is a standard and critical condition for proving convergence. The bounded rewards condition states that the rewards received by the agents at each time step are limited to a finite range. For our intrinsic reward, boundedness ensures that the diversity-seeking objective does not overwhelm the primary goal of maximizing the team reward. If the intrinsic reward were unbounded, it could lead to policies that are arbitrarily diverse but suboptimal in terms of task performance. Specifically, for our proposed WCD method, we prove that the Wasserstein distance-based intrinsic reward is bounded. The boundedness of this intrinsic reward is guaranteed by the properties of the learned representation space and the Wasserstein distance itself.\n\n**1. Boundedness of the Compact Representation Space**\n\nThe trajectory representations are guaranteed to be in a compact (i.e., closed and bounded) set due to the architecture of the trajectory encoder. The encoder is composed of a multi-layer perceptron (MLP) and a Gated Recurrent Unit (GRU), both of which use activation functions that constrain their outputs.\n\nLet the trajectory encoder be represented by the function $g_\\theta: \\mathcal{X} \\rightarrow \\mathcal{C}$, where $\\mathcal{X}$ is the space of trajectory histories and $\\mathcal{C}$ is the latent representation space. The encoder is a composition of the non-linear encoder $g_{\\theta_e}$ and the autoregressive model $g_{\\theta_g}$:\n\n$$c_t^a = g_{\\theta_g}(g_{\\theta_e}(x_0^a), \\dots, g_{\\theta_e}(x_{t-1}^a))$$\n\nThe boundedness of the output $c_t^a$ can be shown by analyzing the components:\n\na. **Bounded Activation Functions**: The MLP and GRU components of the encoder use activation functions such as the sigmoid ($\\sigma(x) = \\frac{1}{1 + e^{-x}}$) or the hyperbolic tangent ($\\tanh(x)$), which are inherently bounded. For any input $x$, their outputs are confined to a specific range (e.g., $[0, 1]$ for sigmoid, $[-1, 1]$ for tanh).\n\nb. **Bounded Weights and Biases**: During training, the weights and biases of the neural networks are typically regularized (e.g., through weight decay), which keeps them within a finite range.\n\nGiven that the inputs to each layer are transformations of the outputs from previous layers with bounded activation functions, the final output of the network, which is the trajectory representation $c_t^a$, will also be bounded. If we assume the maximum possible value for any dimension of the representation vector is $C_{max}$, then the representation space is a subset of a hypercube in $\\mathbb{R}^d$:\n\n$$\\mathcal{C} \\subset [-C_{max}, C_{max}]^d$$\n\nwhere $d$ is the dimension of the representation space. Since this set is closed and bounded, the representation space is compact.\n\n**2. Boundedness of the Wasserstein Distance**\n\nThe Wasserstein distance is used to measure the cost of optimal transport between two probability distributions. When these distributions are defined over a compact set, the Wasserstein distance is also bounded.\n\nThe 1-Wasserstein distance has a dual representation known as the Kantorovich-Rubinstein duality, which is particularly useful for understanding its properties:\n\n$$W\\_1(p, q) = \\sup\\_{\\|f\\|\\_L \\le 1} \\left| \\mathbb{E}\\_{x \\sim p}[f(x)] - \\mathbb{E}\\_{y \\sim q}[f(y)] \\right|$$\n\nwhere the supremum is taken over all 1-Lipschitz functions $f$.\n\nLet $p$ and $q$ be two probability distributions over the compact representation space $\\mathcal{C}$. Since $\\mathcal{C}$ is bounded, there exists a constant $M$ such that the distance between any two points $c_1, c_2 \\in \\mathcal{C}$ is bounded: $d(c_1, c_2) \\le M$.\n\nFor any 1-Lipschitz function $f$ and any two points $c_1, c_2 \\in \\mathcal{C}$, we have:\n$$|f(c_1) - f(c_2)| \\le d(c_1, c_2) \\le M$$\n\nNow, consider the expectation of $f$ under the distributions $p$ and $q$:\n$$\\left| \\mathbb{E}\\_{x \\sim p}[f(x)] - \\mathbb{E}\\_{y \\sim q}[f(y)] \\right| = \\left| \\int\\_{\\mathcal{C}} f(x) dp(x) - \\int_{\\mathcal{C}} f(y) dq(y) \\right|$$\n\nBy choosing a reference point $c_0 \\in \\mathcal{C}$, we can write:\n$$\\left| \\int (f(x) - f(c_0)) dp(x) - \\int (f(y) - f(c_0)) dq(y) \\right|$$\n\nSince $|f(c) - f(c_0)| \\le d(c, c_0) \\le M$ for any $c \\in \\mathcal{C}$, we have:\n$$\\left| \\mathbb{E}\\_{x \\sim p}[f(x)] - \\mathbb{E}\\_{y \\sim q}[f(y)] \\right| \\le \\int M dp(x) + \\int M dq(y) = 2M$$\n\nThis shows that the Wasserstein distance is bounded, as the distributions are supported on a compact set. The intrinsic reward $r_w$, being the minimum of these bounded distances, is therefore also bounded. This formalizes the boundedness assumption, which is essential for the stability and convergence of the WCD algorithm.\n\n**Weakness 3:** Additionally, more experimental results, especially the comparison with DiCo, can be reported.\n\nThe numerical results of the comparisons between WCD+QMIX and DiCo are provided in Tables 1 and 2 in the Appendix, which are clearer than the learning curves. In our work, we compare WCD+QMIX with DiCo on a total of 13 challenging tasks from Pac-Men, SMAC, SMACv2, and GRF.\n\n\n**Q1:** Although experimental results have been reported to verify the effectiveness of the proposed method, there are no theoretical guarantees for the performance improvements.\n\nPlease see our response to Weakness 2.\n\n**Q2:** In the paper, the authors claimed that by employing the proposed intrinsic reward function, better performance can be achieved experimentally by keeping the current trajectory away from its nearest neighbor trajectory. However, for tasks that benefit from homogeneity, it may not be true.\n\nIn our work, we additionally conduct experiments to investigate whether our method prevents the learning of homogeneous behaviors in the second paragraph of Appendix F.2. This issue is also discussed in lines 253-258 of the main paper. The results demonstrate that our method does not prevent homogeneous behaviors that can lead to higher environmental rewards. Essentially, as discussed in our response to Weakness 2, our method is an exploration technique. In our response to Weakness 2, we theoretically prove that our intrinsic reward is bounded, which ensures the diversity-seeking objective does not overwhelm the primary goal of maximizing the team reward. If the intrinsic reward were unbounded, it could lead to policies that are arbitrarily diverse but suboptimal in terms of task performance. Moreover, as shown in Equation 8, our method only introduces a regularizer to the policy network training, which is weighted by a coefficient $\\alpha$.\n\n**Q3:** In the literature, DiCo aims to address the same diversity issue in MARL while theoretically proving the desired diversity can be achieved. Can the proposed method have the same theoretical guarantee? Moreover, more discussions and comparisons about the performances between WCT+QMIX and DiCo can be provided. It is also difficult for me to distinguish the experimental curves representing WCT+QMIX and the experimental curves representing DiCo; they almost have the same color in the experimental figures.\n\nPlease see our responses to Weakness 2 and Weakness 3.\n\n**Q4:** How does DiCo perform in terms of exploration, which appears to be missing from the analysis in Figures 8 and 9?\n\nIn our work, we do not examine DiCo's performance in terms of exploration. Instead, we compare another similar Wasserstein distance-based method, MAPD. The results demonstrate that MAPD does not yield sufficient exploration of the environment space. Thank you for your valuable suggestion; we have additionally added the visitation heatmap of DiCo to Figures 8 and 9. We note that DiCo also does not lead to sufficient exploration because of its constrained diversity."}, {"review_id": "bmHsCkUbxh", "reviewer": "Reviewer_kYgf", "summary": "This paper proposes Wasserstein Contrastive Diversity (WCD), a method to promote diverse behaviors in cooperative multi-agent reinforcement learning (MARL), especially under shared policy networks. WCD learns contrastive trajectory representations via Contrastive Predictive Coding (CPC) and maximizes the Wasserstein distance between agent trajectories in this latent space. This encourages meaningful behavioral diversity. The method integrates efficiently with algorithms like QMIX and MAPPO using a nearest-neighbor intrinsic reward. Experiments on challenging MARL benchmarks show that WCD significantly improves exploration and overall performance.", "questions": "- The authors adopt the Euclidean distance in the latent trajectory representation space as the cost function for the Wasserstein computation. Have you explored or considered task-specific or learned cost functions that might better reflect semantic differences in behavior? How sensitive is WCD to this choice?\n\n- Since the effectiveness of WCD heavily depends on CPC for producing distinguishable trajectory embeddings, how robust is the method when CPC representations are noisy or collapse in more complex observation spaces (e.g., continuous control with high-dimensional inputs)? \n\n- While your kernel-based approach improves efficiency over dual-network parametrization, computing Wasserstein distances across many agent pairs may still be costly. Can you provide runtime and memory analysis for scalability when the number of agents grows significantly (e.g., >50 agents)?\n\n- In scenarios where homogeneous behavior is optimal (e.g., focus fire in SMAC), how does WCD avoid hindering convergence to such strategies? Is the nearest-neighbor intrinsic reward turned off or decayed during training, or does it naturally diminish when agents converge?\n\n- WCD is evaluated only in fully cooperative settings. Can the proposed approach be extended to mixed or competitive MARL environments? If so, how would the diversity metric or reward structure need to be adapted?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strength\n\n- Clear motivation. The paper addresses a well-known limitation in cooperative MARL,l ack of agent diversity under shared policy network and motivates the use of metric-aware approaches like Wasserstein distance to overcome it.\n\n- The method is shown to integrate seamlessly with both value-based (QMIX) and policy-based (MAPPO) MARL algorithms, increasing its practical applicability.\n\n## Weakness\n\n- While the method is empirically sound, the paper lacks deeper theoretical analysis (e.g., convergence guarantees, diversity-exploration trade-offs), which could strengthen its foundation.\n\n- The representation learning relies heavily on CPC and InfoNCE loss. If CPC fails to capture meaningful distinctions in more complex domains, WCD’s effectiveness could degrade\n\n- Although kernel methods reduce cost, computing pairwise Wasserstein distances (even with nearest-neighbor selection) could still be expensive in very large-scale multi-agent systems.", "was_revised": true, "final_justification": "Most of my concerns are addressed in the rebuttal, however, the performance under more complex environments is unevaluated, thus I will maintain my score, a borderline accept", "rebuttal": "We sincerely appreciate your thorough review and valuable feedback on our paper. We answer your questions below:\n\n**Weakness 1:** While the method is empirically sound, the paper lacks deeper theoretical analysis (e.g., convergence guarantees, diversity-exploration trade-offs), which could strengthen its foundation.\n\nWe thank the reviewer for their valuable suggestions. We now provide a theoretical analysis to guarantee the convergence of MARL when our Wasserstein distance-based intrinsic reward is introduced. To achieve this, we theoretically analyze the **boundedness** of the intrinsic rewards, which is a standard and critical condition for proving convergence. The boundedness of rewards means that the rewards received by the agents at each time step are limited to a finite range. For our intrinsic reward, boundedness ensures that the diversity-seeking objective does not overwhelm the primary goal of maximizing the team reward. If the intrinsic reward were unbounded, it could lead to policies that are arbitrarily diverse but suboptimal in terms of task performance. Specifically, for our proposed WCD method, we prove that the Wasserstein distance-based intrinsic reward is bounded. The boundedness of this intrinsic reward is guaranteed by the properties of the learned representation space and the Wasserstein distance itself.\n\n**1. Boundedness of the Compact Representation Space**\n\nThe trajectory representations are guaranteed to lie in a compact (i.e., closed and bounded) set due to the architecture of the trajectory encoder. The encoder is composed of a multi-layer perceptron (MLP) and a Gated Recurrent Unit (GRU), both of which use activation functions that constrain their outputs.\n\nLet the trajectory encoder be represented by the function $g_\\theta: \\mathcal{X} \\rightarrow \\mathcal{C}$, where $\\mathcal{X}$ is the space of trajectory histories and $\\mathcal{C}$ is the latent representation space. The encoder is a composition of the non-linear encoder $g_{\\theta_e}$ and the autoregressive model $g_{\\theta_g}$.\n\n$c_t^a = g_{\\theta_g}(g_{\\theta_e}(x_0^a), ..., g_{\\theta_e}(x_{t-1}^a))$\n\nThe boundedness of the output $c_t^a$ can be shown by analyzing the components:\n\na.  **Bounded Activation Functions**: The MLP and GRU components of the encoder use activation functions such as the sigmoid ($\\sigma(x) = \\frac{1}{1 + e^{-x}}$) or the hyperbolic tangent ($\\tanh(x)$), which are inherently bounded. For any input $x$, their outputs are confined to a specific range (e.g., $[0, 1]$ for sigmoid, $[-1, 1]$ for tanh).\n\nb.  **Bounded Weights and Biases**: During training, the weights and biases of the neural networks are typically regularized (e.g., through weight decay), which keeps them within a finite range.\n\nGiven that the inputs to each layer are transformations of the outputs from previous layers with bounded activation functions, the final output of the network, which is the trajectory representation $c_t^a$, will also be bounded. If we assume the maximum possible value for any dimension of the representation vector is $C_{max}$, then the representation space $\\mathcal{C}$ is a subset of a hypercube in $\\mathbb{R}^d$:\n\n$\\mathcal{C} \\subset [-C_{max}, C_{max}]^d$\n\nwhere $d$ is the dimension of the representation space. Since this set is closed and bounded, the representation space is compact.\n\n**2. Boundedness of the Wasserstein Distance**\n\nThe Wasserstein distance is used to measure the cost of optimal transport between two probability distributions. When these distributions are defined over a compact set, the Wasserstein distance is also bounded.\n\nThe 1-Wasserstein distance has a dual representation known as the Kantorovich-Rubinstein duality, which is particularly useful for understanding its properties:\n\n$W\\_1(p, q) = \\sup_{\\|f\\|\\_L \\le 1} \\left| \\mathbb{E}\\_{x \\sim p}[f(x)] - \\mathbb{E}\\_{y \\sim q}[f(y)] \\right|$\n\nwhere the supremum is taken over all 1-Lipschitz functions $f$.\n\nLet $p$ and $q$ be two probability distributions over the compact representation space $\\mathcal{C}$. Since $\\mathcal{C}$ is bounded, there exists a constant $M$ (the diameter of the set) such that for any two points $c_1, c_2 \\in \\mathcal{C}$, the distance between them is bounded: $d(c_1, c_2) \\le M$.\n\nThe 1-Wasserstein distance has a dual representation known as the Kantorovich-Rubinstein duality, which is particularly useful for understanding its properties:\n\n$$W\\_1(p, q) = \\sup\\_{\\|f\\|\\_L \\le 1} \\left| \\mathbb{E}\\_{x \\sim p}[f(x)] - \\mathbb{E}\\_{y \\sim q}[f(y)] \\right|$$\n\nwhere the supremum is taken over all 1-Lipschitz functions $f$.\n\nLet $p$ and $q$ be two probability distributions over the compact representation space $\\mathcal{C}$. Since $\\mathcal{C}$ is bounded, there exists a constant $M$ such that the distance between any two points $c_1, c_2 \\in \\mathcal{C}$ is bounded: $d(c_1, c_2) \\le M$.\n\nFor any 1-Lipschitz function $f$ and any two points $c_1, c_2 \\in \\mathcal{C}$, we have:\n$$|f(c_1) - f(c_2)| \\le d(c_1, c_2) \\le M$$\n\nNow, consider the expectation of $f$ under the distributions $p$ and $q$:\n$$\\left| \\mathbb{E}\\_{x \\sim p}[f(x)] - \\mathbb{E}\\_{y \\sim q}[f(y)] \\right| = \\left| \\int\\_{\\mathcal{C}} f(x) dp(x) - \\int_{\\mathcal{C}} f(y) dq(y) \\right|$$\n\nBy choosing a reference point $c_0 \\in \\mathcal{C}$, we can write:\n$$\\left| \\int (f(x) - f(c_0)) dp(x) - \\int (f(y) - f(c_0)) dq(y) \\right|$$\n\nSince $|f(c) - f(c_0)| \\le d(c, c_0) \\le M$ for any $c \\in \\mathcal{C}$, we have:\n$$\\left| \\mathbb{E}\\_{x \\sim p}[f(x)] - \\mathbb{E}\\_{y \\sim q}[f(y)] \\right| \\le \\int M dp(x) + \\int M dq(y) = 2M$$\n\nThis shows that the Wasserstein distance is bounded, as the distributions are supported on a compact set. The intrinsic reward $r_w$, being the minimum of these bounded distances, is therefore also bounded. This formalizes the boundedness assumption, which is essential for the stability and convergence of the WCD algorithm.\n\n**Weakness 2:** If CPC fails to capture meaningful distinctions in more complex domains, WCD’s effectiveness could degrade.\n\nOur CPC-based next-step prediction method is robust in noisy, high-dimensional, and complex domains. This robustness stems from its core principles of focusing on shared, high-level information while discarding local, low-level details and noise. Furthermore, the process of contrasting positive and negative examples makes the model robust against noisy views and ensures that the distance between different representations is meaningful, which is key to the success of using the Wasserstein distance.\n\n\n**Weakness 3:** Computing pairwise Wasserstein distances (even with nearest-neighbor selection) could still be expensive in very large-scale multi-agent systems.\n\nThe additional computational overhead introduced by our method is analyzed in Appendix K. A comparison of the training times between QMIX and our method demonstrates that our method does not add significant extra training time. During the rebuttal phase, we evaluated the training time in the large-scale 27m_vs_30m scenario (a standard SMAC scenario with the largest number of agents). We ran each method for 5 million steps on the same computational platform. The comparison of training times is shown below:\n\n| Methods | Training time |\n|---|---|\n| QMIX | 11h 27m 18s |\n| WCD+QMIX (ours) | 11h 43m 39s |\n\nThe results demonstrate that our method does not significantly increase the training time and scales well with an increasing number of agents. We will further examine our method in large-scale multi-agent tasks with more agents in our future work.\n\n\n**Q1:** Have you explored or considered task-specific or learned cost functions that might better reflect semantic differences in behavior? How sensitive is WCD to this choice?\n\nWe explore the effect of different cost functions on the performance of agents in Appendix J, where we compare the Euclidean distance with cosine similarity, which measures the directional differences between data points. The Wasserstein distance based on cosine similarity achieves slightly better performance in Pac-Men due to the directional characteristics of cosine similarity. The results demonstrate that the final performance of agents is not highly sensitive to the choice of cost function. In our work, we do not specifically discuss different cost functions and use the default Euclidean distance because we want to be consistent with prior works that use the Wasserstein distance to ensure a fair comparison.\n\n**Q2:** See **Weakness 2**.\n\n**Q3:** See **Weakness 3**.\n\n**Q4:** In scenarios where homogeneous behavior is optimal (e.g., focus fire in SMAC), how does WCD avoid hindering convergence to such strategies? Is the nearest-neighbor intrinsic reward turned off or decayed during training, or does it naturally diminish when agents converge?\n\nAs discussed in the theoretical analysis in our response to **Weakness 1**, our intrinsic reward is bounded to prevent the diversity-seeking objective from overshadowing the main goal of maximizing the team reward. Thus, our method does not hinder the learning of homogeneous behaviors that are optimal in some scenarios. An unbounded intrinsic reward could result in policies that are highly diverse but ultimately ineffective for the task. Moreover, as shown in Equation 8, we control the influence of our diversity objective relative to the main objective by introducing a weighting coefficient, $\\alpha$.\n\n\n**Q5:** WCD is evaluated only in fully cooperative settings. Can the proposed approach be extended to mixed or competitive MARL environments? If so, how would the diversity metric or reward structure need to be adapted?\n\nThe Pac-Men and Google Research Football (GRF) environments used in our work are mixed MARL environments. In Pac-Men, agents may compete for dots if they move to the same room. In GRF, agent-controlled players may compete for the ball to earn more rewards. Our method successfully enables agents to learn optimal policies in these tasks."}, {"review_id": "5hAgXRKZ8D", "reviewer": "Reviewer_6Eyc", "summary": "This paper addresses the challenge of insufficient exploration in cooperative Multi-Agent Reinforcement Learning (MARL), where agents with shared policy parameters tend to learn similar behaviors, resulting in suboptimal cooperation. \nThe authors propose Wasserstein Contrastive Diversity (WCD) exploration, which maximizes the Wasserstein distance between agents’ trajectory distributions in a learned latent space. The approach leverages Contrastive Predictive Coding to obtain distinguishable trajectory representations and introduces an efficient, kernel-based computation of the Wasserstein distance. \nThe method also extends the Wasserstein distance to handle multiple agents. Experimental results on challenging MARL benchmarks show that WCD achieves superior performance and promotes more effective exploration compared to state-of-the-art methods.", "questions": "1. How should the value of $\\alpha$ be selected in practice? Does its optimal value depend heavily on specific cases or environments?\n2. What is the motivation for choosing QMIX as the base framework for this work?\n3. How is the standard error in the experiments calculated?", "limitations": "The discussion of limitations could be more candid and comprehensive.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths \n\n1. The proposed method is both well-motivated and novel. The authors clearly articulate the motivation for employing the Wasserstein Distance between trajectory representations.\n2. The paper is well-written and easy to follow.\n3. The experimental results demonstrate the effectiveness of the proposed approach.\n\nWeaknesses\n\n1. The paper adopts InfoNCE and utilizes next states from other agents to compute the loss. Could the authors provide more insight into the motivation and rationale behind this design choice?\n2. The method requires retraining the trajectory encoder at each step. Does this impact the efficiency of the algorithm? Are there alternative approaches that could allow for encoder reuse to improve efficiency?\n3. The discussion of limitations could be more candid and comprehensive.", "was_revised": true, "final_justification": "The author addressed my concern in the rebuttal and I maintain my score.", "rebuttal": "We sincerely appreciate your thorough review and valuable feedback on our paper. We answer your questions below:\n\n**Weakness 1: The paper adopts InfoNCE and utilizes next states from other agents to compute the loss. Could the authors provide more insight into the motivation and rationale behind this design choice?**\n\nAs we discussed in the first paragraph of Section 3.1, our goal is to distinguish the trajectories of different agents in a latent representation space, enabling the Wasserstein distance to be effective. To achieve this, the InfoNCE loss uses a **contrastive** operation to enable the trajectory representation to stay close to its associated next-step observation-action embedding while being far away from other agents' noisy embeddings. Moreover, predicting the next states helps learn meaningful trajectory representations that are robust to the **noise** in the trajectory sequence, making the learned representations focus on shared, high-level information while discarding local, low-level details and noise.\n\n**Weakness 2: The method requires retraining the trajectory encoder at each step. Does this impact the efficiency of the algorithm? Are there alternative approaches that could allow for encoder reuse to improve efficiency?**\n\nOur method does not retrain the trajectory encoder at each step. During **the** training interval, we train the trajectory encoder to adapt to the sampled trajectory batch before training the **agents' policies**. This helps stabilize training since the trajectories induced by the learning policies can **vary** in different training **phases**. Moreover, the additional computational overhead introduced by our method (including trajectory encoder training and Wasserstein distance computation) is analyzed in **Appendix** K. The **comparison** of training **times** between QMIX and our method demonstrates that our method does not consume **significant** extra training time.\n\n**Weakness 3: The discussion of limitations could be more candid and comprehensive.**\n\nThank you for your suggestions. We have added the following limitation analysis to the manuscript: We note that our method encourages multi-agent diversity by effectively enlarging the Wasserstein distance in a latent representation space. However, learning trajectory representations using contrastive learning can be **computationally expensive**. Further adaptations of contrastive learning in multi-agent settings need to be considered in future work to make trajectory representation learning more effective and less costly. Moreover, as the number of agents **increases**, the computational cost of the Wasserstein distance objective **grows** significantly, posing a challenge to the scalability of our method.\n\n**Q1: How should the value of $\\alpha$ be selected in practice? Does its optimal value depend heavily on specific cases or environments?**\n\nWe investigate the effect of the intrinsic reward weight, $\\alpha$, on final performance in Appendix I. The results show that our method is not highly sensitive to this value. Sub-optimal weights do not cause a significant performance drop, even in the **most difficult** scenarios. In our experiments, we set the same weight value for the **most difficult** tasks, as the optimal value does not appear to depend heavily on the specific task.\n\n**Q2: What is the motivation for choosing QMIX as the base framework for this work?**\n\nFirst, QMIX is widely used as the base method in baselines such as EOI, SCDS, and MAVEN. We use the same base method to help **guarantee** a fair **comparison**, which **reduces** the influence of the base method on agent performance. Second, QMIX is fundamental, simple yet efficient, and easy to **integrate**.\n\n**Q3: How is the standard error in the experiments calculated?**\n\nAll experimental results are presented as the mean and standard deviation of performance across five random seeds. To calculate the standard deviation, we first compute the mean performance (e.g., win rates in SMAC) over the five random seeds. Then, we find the squared differences from the mean. The standard deviation is the square root of the variance (the average of these squared differences)."}]

Author rebuttal:

Meta-review: This work sets off to address the challenge of exploration in cooperative multi-agent RL. This paper proposes a method for enhancing trajectory diversity based on a Wasserstein Contrastive Diversity (WCD) exploration, which maximizes the Wasserstein Distance between the agent's trajectory distributions in a latent space. These representations are obtained via contrastive predictive coding and introduces a method to compute these Wasserstein distances efficiently, extending it to the case of multiple agents. The reviewers agreed the method is well motivated and novel, the paper is well written, and the experimental section demonstrates the method's effectiveness. The reviewers raised the point that the work is thin in theoretical justification. The merits of this work outweigh the negative points raised in the reviewer discussion and it is thus clear this work can enrich the conference's program.
Decision: Accept (poster)