Multi-Agent Communication with Information Preserving Graph Contrastive Learning

1 Introduction

Multi-agent reinforcement learning (MARL) has attracted considerable interest and achieved impressive results in various complex real-world applications, including autonomous driving [Xu et al., 2024], traffic signal control [Zhang et al., 2024], and auction market [Qiu et al., 2021]. Recently, the centralized training with decentralized execution (CTDE) framework has gained widespread adoption as an effective solution to address the challenges of non-stationarity and scalability in MARL. This framework involves developing decentralized policies in a centralized manner, enabling the sharing of experiences and parameters throughout the training process. Building on the CTDE framework, several value decomposition methods have been proposed [Liu et al., 2023; Rashid et al., 2020; Du et al., 2024]. These methods apply various constraints or restrictions to factorize the global value function to a combination of individual value functions. While the CTDE framework presents multiple benefits, the partial observability and stochasticity encountered during the decentralized execution period can heighten agents' uncertainty about the states and actions of other agents, potentially leading to miscoordination in their actions.

∗Corresponding author

Recently, several MARL methods have utilized communication learning protocols to improve coordination among agents, enabling them to share information, such as feature embeddings, during the execution phase. By facilitating interagent communication, these methods significantly enhance coordination across a variety of tasks [Gilmer et al., 2017]. Graph Neural networks (GNNs) effectively integrate both the feature and topological information in graph-structured data, facilitating robust feature representation learning for subsequent tasks. Due to the effective representation learning capability of GNN, communication learning through GNN has garnered significant research interest in MARL. In this case, agents can typically be represented as nodes, while communication channels between them are denoted as edges in a graph. Various MARL methods utilize this GNN-based communication framework, such as DGN [Jiang et al., 2020], LSC [Sheng et al., 2022], and MAGIC [Niu et al., 2021].

While GNN-based MARL methods [Liu et al., 2020; Das et al., 2019; Sukhbaatar and Fergus, 2016]have achieved success, they overlook a key limitation of GNN it relies on: homophily assumption. This assumption suggests that nodes within the same class are more likely to be connected. However, in real-world multi-agent systems, this assumption often fails, as agents of the same type can be spatially distant from one another. Under this situation, GNN-based MARL methods struggle with poor performance because the propagation mechanism within graph neighborhoods of GCN becomes problematic, leading to the mixing of irrelevant information from agents of different classes. Moreover, GNNbased MARL methods tend to overlook the critical role that feature similarity between agents plays in action coordination, which limits their overall effectiveness.

In cooperative MARL, expressive message representations that incorporate comprehensive information are essential for efficient action coordination. Therefore, it is crucial to extract and preserve valuable information from neighboring agents to learn effective message representations. To address these issues, we present a Multi-Agent communication protocol with Information preserving graph contrastive Learning (MAIL), which comprehensively preserves agent feature information while exploiting topological information.

Consider a football game as illustrated in Figure 1, and assume that each agent chooses adjacent agents within its range to establish the original graph. In MAIL, we utilize the original graph as one of the contrastive views and consider the combination of the feature graph and the original graph as a second contrastive view. By maximizing the agreement between the two contrastive views, MAIL effectively preserves the agents' feature information. To maintain the topological information, we introduce an additional contrastive learning module specifically designed to preserve global topological details. We directly utilize the higher-order view of the graph as the global topological view and contrast it with the original graph view for preserving global topological information. MAIL acquires high-quality message representations by extracting crucial information at both the topological and feature levels. We select several baselines and conduct experiments across various environments, with the results confirming that MAIL outperforms the baselines. The key contributions of our work can be outlined as follows:

• To the best of our knowledge, our research is the first to utilize graph contrastive learning in the context of

MARL, enabling effective communication learning.

• We present two contrastive learning modules that can individually extract the feature and topological information of the input graph while training them in a coordinated manner.
• The proposed method encourages the agent to learn significant information from both feature and topological aspects, enabling it to generate high-quality message representations for downstream action coordination.

2 RelatedWork

2.1 Graph Contrastive Learning

Graph Neural Networks (GNNs) have become a highly effective method for learning robust node representations. The term GNN refers to a broad spectrum of methodologies. In this study, we specifically define it as the Message Passing Graph Neural Network [Gilmer et al., 2017], recognized as the most general architecture within the GNN framework. Prominent examples of this architecture include GCN [Duvenaud et al., 2015], GAT [Veliˇckovi´c et al., 2017], and GraphSAGE [Hamilton et al., 2017]. While GNNs have shown impressive performance in many scenarios, their effectiveness may be impaired when the homophily assumption fails [Chen and Kou., 2023].

Graph contrastive learning (GCL) has emerged as a promising research direction, typically involving the design of different views and aiming to maximize the agreement between the representations of these views. DGI [Velickovic et al., 2019] concentrates on maximizing the mutual information between global graph-level representations and local node-level representations. Building upon DGI, GMI [Peng et al., 2020] utilizes two discriminators to directly compute mutual information between the input and the embeddings of both edges and nodes. In contrast, MVGRL [Hassani and Khasahmadi, 2020] focuses on learning both node-level and graph-level embeddings by contrasting node representations and facilitating node diffusion using augmented graph summary representations.

Most previous GCL methods rely on GNNs as the backbone encoder, achieving impressive performance in homophilic scenarios. However, limited attention has been given to developing GCL methods for situations where the homophily assumption fails. Moreover, current contrastive learning methods primarily generate views by transforming existing graph data. However, these data transformation strategies are inadequate for learning comprehensive and robust node embeddings. In our research, we adopt the graph contrastive learning scheme introduced in ASP [Chen and Kou., 2023], which effectively retains both feature and topological information from the input built graph.

2.2 GNN-basedMARL

In our work, we utilize the following GNN-based communication learning methods as baseline methods: TarMAC [Das et al., 2019] employs GAT to learn communication through a fully constructed communication graph. MAGIC [Niu et al., 2021] also utilizes GAT to facilitate multiple rounds of communication between agents. CommNet [Sukhbaatar and Fergus, 2016] introduces a communication channel that allows agents to dynamically enter and exit each other's communication range, similar to GNN methods that utilize mean aggregation. IC3Net [Singh et al., 2019] incorporates a gating mechanism to regulate communication learning. DGN [Gilmer et al., 2017] establishes communication protocols on graphs derived from the environment. LSC [Sheng et al., 2022] presents a hierarchical GNN that enhances communication learning by enabling message exchanges within groups and between agents. G2ANet [Liu et al., 2020] integrates both soft and hard attention mechanisms to dynamically adapt communication. DICG [Li et al., 2021] features a module for inferring the topology of the coordination graph, using GNNs for implicit reasoning on joint actions. Although GNN-based MARL has achieved success, the network homogeneity assumption of GNN in MARL is often unrealistic, which limits the performance of GNN-based MARL methods. Furthermore, GNN-based MARL methods often neglect the crucial role of feature similarity between agents in action coordination, which hampers their overall effectiveness.

3 Method

3.1 Preliminary

Let G = (V, E) represents a graph, where V = {vi, ..., vn} denotes the node set and E ⊆ V ×V denotes the edge set. The feature matrix is represented as X = [x1, x2, . . . , xn]T ∈ Rn×f , where xi ∈ Rf denotes the feature of vi. A ∈ {0, 1}n×n represents the adjacency matrix. In our study, the multi-agent system is represented as a graph G = (V,E,X), where the set of agents is denoted by V = {vi, ..., vn}, the edge set is represented as E ⊆ V × V , and the agent/node features are indicated by X ∈ Rn×f . Adjacent agents are defined as those agents located within a specified range of each agent, which is used to construct the original graph. Given a graph G = (V, E, X), K distinct transformations F1, . . . , FK can be applied to obtain different views v1, . . . , vK of the graph:

Graph transformation techniques encompass various methods such as node feature masking, edge perturbation, etc. A set of encoders f1, . . . , fK receives the corresponding views as inputs and produces the representations h1, . . . , hK of the graph from each view:

A cooperative MARL problem could be formally modeled as Dec-POMDP [Oliehoek, 2012], which can be represented by a tuple < I, S, O, U, P, R >. I denotes the agents' set, which is indexed from 1 to n. S denotes the state set of the environment. O represents the observations set, where oi ∈ O represents observation of agent i. U denotes the finite space of joint actions. For agent i, at each timestep, agent i takes its action ai depending on its local observation oi, and forms a joint action a = (a1, . . . , an) ∈ U. Based on the Markovian transition function P : S × U → S, the state changes and the agent obtains reward rt i based on the reward function R : S ×U →R. The target of the agent i is to maximize its total discounted reward Ri = ∑Tt=0 γ t rt i , in which γ ∈ [0, 1] denotes a discount factor. In Dec-POMDPs, the target of all agents is to learn an optimal joint policy π(τ, a) to maximize the global value Qπtot (τ, a) = Es,a[∑∞t=0 γ tR(s, a)], where τ represents the joint observation history.

The framework of the MAIL is illustrated in Figure 2. For agent i, it receives local observation oi and employs multilayer perceptron (MLP) with the gated recurrent unit (GRU) to produce the agent feature xi. Next, xi is fed into the GCL module to generate the expressive message embedding hi.

Subsequently, the message representation hi and local history τi are concatenated to generate the input for the individual action-value function to select an action and compute individual value Qi (τi, ai, hi). Finally, all individual actionvalues are fed to the mixing network to calculate the estimation of global value Qtot. In this study, the mixing network of QMIX [Rashid et al., 2020] is adopted, and it could be substituted with any mixing network from current value function factorization approaches.

3.2 Graph Contrastive Learning module

View Generation

As depicted in Figure 1, the original built graph data, without any data augmentation techniques, is referred to as the original view vo, serving as an anchor for the other views. To capture the diverse feature similarity relationships among agents, we employ the feature matrix X to create a k-NearestNeighbor (kNN) graph Gf , which we designate as the feature view vf . The kNN graph Gf can be constructed using various distance metrics, including Euclidean distance, Jaccard distance, or cosine distance. In our framework, we utilize cosine distance. Given a node pair (vi, vj), their cosine distance is defined as follows:

where |·| represents the magnitude of a vector. The adjacency matrix is denoted as AF , and the corresponding degree matrix of the kNN graph is denoted as DF .

To extract global topological information, we create a straightforward yet effective graph view. In contrast to edge perturbation strategies that arbitrarily alter the graph topology, edge diffusion methods [Gilmer et al., 2017] effectively preserve the expressive global information contained in graphs. Current edge diffusion approaches primarily utilize the heat kernel [Kondor and Lafferty., 2002] to create a global view. However, this solution requires computing matrix exponentials or inverses, which can be computationally inefficient. While approximations can help reduce computational complexity, they often necessitate tedious and timeconsuming hyperparameter tuning.

In contrast to existing approaches, we directly employ a higher-order view of the graph as the global topological view vt. In this framework, each node gathers information from neighbors that are l hops away. As l increases, more global topological information is incorporated. Both the kNN graph Gf and the original graph G can be utilized to generate these higher-order views. Our findings indicate that simply contrasting the higher-order views with the original graph views can yield comparable performance.

Feature Preserving Contrastive Learning

Current graph contrastive learning methods typically use GNN as the foundational encoder, effectively extracting the topological information inherent in graphs. Despite their effectiveness, these methods often neglect the similarities between node features derived from the feature matrices. In this study, we utilize a more robust mechanism that retains feature knowledge by leveraging both the feature graph view and the original graph view. While our framework permits the use of various GNNs, we have selected SGC [Wu et al., 2019] as the base encoder because of its simplicity and competitive performance. SGC streamlines the topology of GCN by eliminating intermediate nonlinearities, resulting in all learnable parameters being consolidated into a single matrix:

where S represents the normalized adjacency matrix with added self-loops S = ˜D-1/2 ˜A˜D-1/2 and W represents a trainable weight matrix. SP represents the representations generated by propagating information from agents that are P -hop away.

Rather than directly comparing the original view with the feature view, we utilize the feature view to complement the original view. Specifically, we incorporate the node embeddings from the feature view into the original view, treating the sum of these embeddings as the final contrastive view. Formally, we express this as follows:

where SF = D-1/2 F AFD-1/2 F . The weight matrix Wf is shared between Hf and Ho .

After acquiring the message representations Ho and Hf , we employ InfoNCE [Gutmann and Hyv¨arinen, 2020] to estimate the lower bound of the MI between them. For node vi, the learned node embeddings hoi and hfi serve as positive samples, while the embeddings of all other nodes are considered negative samples. With these definitions of negative and positive samples, the loss function for the feature-preserving contrastive learning module can be formulated as follows:

where θ represents the temperature parameter, and D(·) denotes the discriminator that calculates the agreement score between two vectors. In this context, we utilize cosine similarity for D(·).

Topological Preserving Contrastive Learning

While feature information has been uncovered through feature preserving contrastive learning, global topological information remains unexplored. Therefore, to enhance the expressiveness and robustness of the proposed framework, we introduce a topological preserving contrastive learning module. We continue to use SGC as the base encoder in this module. To capture global topological information, we directly compare the original graph view with the global topological view:

where l is a positive integer significantly larger than p, and SG ∈ {S, SF}. It is important to note that we employ separate weight matrices for Ho and Hr to ensure that they do not interfere with one another.

Given the node embeddings hri and hti of node vi, we compute the contrastive loss for topological preserving learning:

Model Training

With the two main components of MAIL established, we introduce the cross-module loss, which we employ to align the representations of the original view across various modules. Following this, we will outline the overall objective loss for training of MAIL.

The purpose of the cross-module loss is to align the representations of Ho and Hr . We have observed that contrasting representations from the same view but across different modules can effectively enhance the quality of learned message representations.

Lc (vi) =

The overall GCL objective loss of MAIL is defined as the sum of the feature preserving loss, the topological preserving loss, and the cross-module loss:

where λ1 and λ2 denote the tuning parameters to weight the importance of Lt and Lc, respectively.

Through the above description, we can obtain high-quality message representations in MAIL. Except for the graph contrastive learning optimization constraints on learning message representations in the communication module, all the parameters across the remaining modules of the framework are updated by minimizing the TD loss LTD. In the end, TD loss is explicitly expressed as follows:

where θ denotes all the parameters in the remaining modules and θ- denotes the parameters of the target network. Therefore, the overall optimization target of MAIL is expressed as follows,

where β represents a hyper-parameter that can be fine-tuned to balance between the graph contrastive learning optimization loss LGCL and the TD loss LTD. A detailed description of our framework is provided in Algorithm 1.

Algorithm 1 MAIL

• 1: Initialize: the parameters of networks, the maximum size of the replay buffer, and the frequency of network updating.
• 2: for each timestep t ∈ T do
• 3: for each agent i ∈ N do
• 4: // During the decentralized execution period
• 5: Generate agent feature xi by GRU and MLP
• 6: Construct graph G = (V,E,X) based on xi
• 7: Receive node representations Ho , Hf , Hr , andHt
• 8: Calculate feature loss Lf , topological loss Lt and, cross-module loss Lc with Eq.6, Eq.8, and Eq.9, respectively
• 9: Update parameters according to the overall GCL objective loss LGCL in Eq.10
• 10: Obtain final message representation hoi
• 11: Calculate action-value Qi based on hi and τi
• 12: ati ← π (Qi) (ϵ- greed )
• 13: Store τi and ati to replay buffer
• 14: // During centralized training period
• 15: Fed Qi to mixing network and obtain Qtot
• 16: Minimize loss function according to Eq.12
• 17: Update weights of all networks
• 18: end for
• 19: end for

4 Experiments

To verify the effectiveness of MAIL, we perform a range of experiments across 4 benchmarks: Predator-Prey [Sukhbaatar and Fergus, 2016], Traffic Junction [Sukhbaatar and Fergus, 2016], Battle [Zheng et al., 2018], StarCraft Multi-Agent Challenge [Vinyals et al., 2019]. Experiments are conducted with a GPU NVIDIA RTX 4090. The hyperparameters that we adjust are as follows: (i) k ∈ {3, 5, 10}, for k nearest neighbors, (ii) aggregation hops l ∈ {3, 5, 7}, (iii) λ1 = 0.2, λ2 = 0.3, and β = 0.2 depending on the experimental results. For each environment, 4 GNN-based MARL baselines (introduced in Related Work) have been chosen for ease of comparison without losing generality. The detailed hyperparameters and some experiments are given in the Appendix.

4.1 Predator-Prey

As shown in Figure 3(a), the Predator-Prey environment [Singh et al., 2019] involves multiple predators with limited vision, aiming to capture prey. Predators can move in four directions: down, up, left, or right. An episode is considered successful if all predators find the prey within the given time limit. We define two difficulty levels: a 10 × 10 grid with 5 predators, and a 20 × 20 grid with 10 predators. A superior approach is characterized by its ability to minimize the mean number of steps necessary to accomplish an episode. As shown in Table 1, MAIL captures the prey faster than the baselines in both scenarios. Figure 4(a) illustrates the success rate of approaches in a 20 × 20 grid with 10 predators, MAIL performs significantly better than baselines. Compared to GNN-based MARL, MAIL avoids the mixing of irrelevant information from agents and, therefore learns high-quality message representations for action coordination.

4.2 Traffic Junction

Traffic Junction environment [Sukhbaatar and Fergus, 2016] consists of vehicles (agents with limited visibility) and intersecting routes. This environment is a valuable evaluation environment for assessing the efficiency of communication. In this environment, the main goal is to ensure effective communication between vehicles to prevent collisions. In this scenario, vehicles approach junctions from diverse entry points using a probability represented as p. There is a maximum limit, denoted as Nc, imposed on the number of vehicles permitted in the environment. At each time step, the vehicles can take one of two actions: "brake" or "gas". We evaluate the performance of MAIL and baseline methods in Traffic Junction with two difficulty levels.

As depicted in Figure 3(b), in the medium difficulty level, the traffic junction environment comprises 2 two-way roads sorted within a 14 × 14 grid. In this scenario, the supreme agents' number allowed is 10 (Nc= 10, p = 0.2). On the other hand, the scenario with a hard difficulty level entails 4 two-way roads within an 18 × 18 grid. In this scenario, the supreme agents' number is 20 (Nc= 20, p = 0.05). The objective in both scenarios with two difficulty levels is to maximize the mean success rate, which is specified as having no collisions occur in an episode.

Table 2 shows the mean success rate for each method upon reaching convergence. As shown in Table 2, in both scenarios, MAIL outperforms other baselines. Figure 4(b) illustrates the learning curves of methods in the scenario with a medium difficulty level. As shown in Figure 4(b), the mean number of steps needed to accomplish an episode of MAIL is significantly fewer than other baseline methods, which demonstrates the effectiveness of the graph contrastive learning module of the proposed MAIL.

4.3 Battle

We chose the Battle environment [Zheng et al., 2018] to further evaluate the effectiveness of MAIL. Battle scenario, which includes ally agents and enemy agents. Ally agent can select one action of two actions: attack or move. The target for the ally agents is to kill enemy agents. In Battle, MAIL and other baselines are trained with the identical configuration. The ally agent can obtain a positive reward of +5 if it successfully attacks an enemy. Conversely, a negative reward -2 is applied if an ally agent is killed by the enemy agent.

Figure 4(c) depicts the mean reward of different methods in the Battle environment. It can be observed that MAIL significantly outperforms baseline methods. MAIL trained agents can acquire various tactical skills, such as encircling and enveloping. Against a single enemy, MAIL trained agents can effectively learn to coordinate and surround it to kill it. When facing a group of enemies, the agents can adeptly learn to target and attack one flank. Other baselines trained agents initially adopt suboptimal strategies, such as clustering in the corner to evade attacks. Table 3 shows the performance of MAIL compared to baseline methods in the Battle environment. MAIL consistently surpasses baseline methods in terms of mean reward, kills, and kill-death (K/D) ratio.

4.4 StarCraft Multi-Agent Challenge

StarCraft Multi-Agent Challenge (SMAC) [Vinyals et al., 2019] is an established environment developed within the prevalent game StarCraft II, as shown in Figure 3(d). In SMAC environment, MARL methods are employed to train ally agents, while enemy agents are managed by the built-in AI. The features collected by the agent encompass the following attributes of both allied and enemy units within its field of vision: unit type, location, distance, health, and shield. To heighten the coordination challenge for ally agents, the default experimental settings have been fine-tuned, resulting in a reduction of the visual field for agents from 9 to 2.

Figure 4(d) shows the mean win rates of different methods in scenario 1o10b vs 1r. We can observe that MAIL significantly outperforms other baselines. Table 4 shows the performance of MAIL and baselines in the different scenarios: MMM2, 1c3s5z, and 27m vs 30m. As shown in Table 4, MAIL significantly outperforms baselines in these scenarios.

4.5 Ablation analysis

We further evaluate the contribution of each component in MAIL. Specifically, we design three ablations masking different components: (i) w/o FP is MAIL without the feature preserving contrastive learning component. (ii) w/o TP is MAIL without the topological preserving contrastive learning component. (iii) w/o CR is MAIL without cross-module loss. The results are summarized in Table 5. We can observe declines in performance when any of the components are removed, highlighting the effectiveness of each one. In particular, when the feature preserving contrastive learning module is removed, the performance in all three scenarios drops significantly, indicating that this module is essential for boosting performance. Meanwhile, the topological preserving contrastive learning module and the cross-module loss contribute to further improvements.

Figure 5 illustrates the average number of epochs needed for each method to reach convergence. Compared to GNNbased MARL, MAIL prevents mixing irrelevant information from agents of different classes, enabling more efficient communication learning. MAIL consistently outperforms baselines as the number of agents increases, highlighting its scalability for tackling large-scale communication learning tasks.

5 Conclusion

We have introduced the graph contrastive learning optimization concept for communication learning methods of MARL. MAIL efficiently preserves both feature and topological information and learns expressive message representations, therefore enhancing action coordination between agents. Experimental results on various environments validate that the proposed MAIL outperforms existing GNN-based communication approaches. The effectiveness of MAIL offers significant potential for advancing efficient multi-agent communication learning. Our research promotes further research into multi-agent communication methods and broader, robust forms of communication learning. In future research, we plan to explore more GCL mechanisms and apply them to realworld multi-agent communication learning scenarios.

Acknowledgements

This work was supported by the National Key Research and Development Program of China (No. 2023YFF0725500), the National Natural Science Foundation of China under Grant (No.92367202, 62276265, and 62376138), and the Innovative Development Joint Fund Key Projects of Shandong NSF (ZR2022LZH007).