MA2CL:Masked Attentive Contrastive Learning for Multi-Agent Reinforcement Learning

1 Introduction

Recent advances in reinforcement learning (RL) and multiagent reinforcement learning (MARL) have led to remarkable progress in developing artificial agents that can cooperate to solve complex tasks. While the performance is encouraging, the agents require extensive training time and millions of interactions with the environment, especially in a vision-based setting (agents learn from visual observation). Therefore, enhancing the sample efficiency has become a challenge in both RL and MARL communities.

Various techniques have been proposed to improve the sample efficiency of RL in single-agent environments through joint learning, which combines the RL loss with auxiliary tasks in the form of self-supervised learning (SSL).

Some methods utilize data augmentation to generate multiple views representation for constructing SSL learning objectives. Some other methods employ a dynamic model to predict future states given the current state and future action sequences, then use the prediction results and the ground truth to construct SSL loss. These dynamic models aim to learn temporally aware representations and thus improve the agent's ability to predict the potential outcomes of its actions over time. The SSL auxiliary tasks provide additional representation supervision and facilitate the acquisition of informative representations, which better serve policy learning.

The dynamic model has been demonstrated effective to learn temporally aware representations in single-agent RL, where agents have a global observation of the environment. However, in the multi-agent setting, the situation becomes vastly different, as each agent only receives a partial observation from the environment that is influenced by others agents. This makes building a dynamic model for each agent with incomplete information challenging. The observation of different agents may be interrelated and contain agent-level information. Furthermore, in cooperative tasks, all agents need to collaborate to achieve a common goal, and the agent-level information in their observation should be considered when making decisions. So it is necessary for agents to learn representations with team awareness in MARL. MAJOR [Feng et al., 2023] extends the dynamic model in SPR [Schwarzer et al., 2020] for multi-agent settings, but only focuses on temporal awareness, without specifically taking advantage of the correlation among the agents. At present, it appears that recent works in MARL have not explicitly incorporated representation with agent-level information as a learning objective.

In this paper, we introduce a novel representation learning framework called Multi-Agent Masked Attentive Contrastive Learning (MA2CL). MA2CL aims to encourage representations to be both temporal and agent-level predictive, achieved by reconstructing the "masked agent's" observation in latent space. We sample observations of all agents at the same time step from the replay buffer and treat them as a sequence with contextual relationships in the agent dimension. Next, we randomly mask several agents with information from the previous time step, then map the masked observations into latent space. We utilize an attention reconstruction model truct the masked agent in latent space, generating another view of the masked agent's representation. We construct a contrastive learning objective for training, based on the intuition that the reconstructed representations should be similar to the ground truth while dissimilar to others. In this way, we build a representation learning objective optimized together with the policy learning objectives, as shown in Fig. 1.

It is worth noting that our algorithm can be easily integrated as an auxiliary task module in many multi-agent algorithms. In order to assess the effectiveness of MA2CL, we implement it based on the state-of-the-art MARL baselines and compared their performance with our approach across various multi-agent environments involving both vision-based and state-based scenarios. Our method outperformed the baselines in these evaluations.

The contributions of our work are summarized as follows:

1. We propose MA2CL, an attentive contrastive representation learning framework for MARL algorithms, which encourages agents to learn effective representations.
2. We implement MA2CL on both MAT and MAPPO, demonstrating its flexibility and ability to be incorporated into various off-policy MARL algorithms.
3. Through extensive experiments, we demonstrate that MA2CL outperforms previous methods and achieves state-of-the-art performance in both vision-based and state-based multi-agent environments.

2 RelatedWorks

Sample-Efficient Reinforcement Learning Sample efficiency assesses how well interaction data are utilized for model training [Huang et al., 2022]. Sample-efficient RL tries to maximize the expected return of the policy during training by interacting with the environment as little as feasible [Ma et al., 2022]. To improve the sample efficiency of RL that learns a policy network from high-dimensional inputs in single-agent settings, recent works design auxiliary tasks to explicitly improve the learned representations [Yarats et al., 2021; Srinivas et al., 2020; Zhu et al., 2022; Lee et al., 2020; Schwarzer et al., 2020; Zhao et al., 2022; Ye et al., 2021; Yu et al., 2021b; Yu et al., 2022], or adopt data augmentation techniques, to improve the diversity of data used for training [Laskin et al., 2020; Yarats et al., 2020]. In multi-agent environments, the observations of various agents already provide a diverse range of observations, and the need for additional data augmentation may be less pressing.

Representation Learning in MARL To the best of our knowledge, there are few works that have investigated the promotion of representation in the context of multi-agent reinforcement learning (MARL). In [Shang et al., 2021], it utilizes an agent-centric predictive objective, where each agent is tasked to predict its future location and incorporates this objective into an agent-centric attention module. However, although an auxiliary task was first introduced to MARL by [Shang et al., 2021], it is only designed to predict the position of agent in a 2D environment and is not adaptable. Another work focus representation learning in MARL is MAJOR [Feng et al., 2023], which employs a joint transition model to predict the future latent representations of all agents. However, constructing the auxiliary task using predictions from all agents results in focusing on the entire team's information in the temporal sequence rather than correlation in the agent level. In this work, our auxiliary task encourages agents to take full advantage of agent-level information which is important in various cooperative tasks.

3 Backgroud

Dec-POMDP Cooperative MARL problems are often modeled using decentralized partially observable Markov decision processes (Dec-POMDPs) [Oliehoek and Amato, 2016], ⟨N,O,A,R, P, γ⟩. N = {1, . . . , n} is the set of agents, O = ∏ni=1 Oi is the product of local observation spaces of the agents, namely the joint observation space, A = ∏ni=1 Ai is the joint action space, R : O ×A → R is the joint reward function, P : O ×A×O → R is the transition probability function, and γ ∈ [0, 1) is the discount factor. At each time step t ∈ N, each agent observes a local observation oit ∈ Oi and takes an action ait according to its policy πi . The next set of observations ot+1 is updated based on a transition probability function, and the entire team receives a joint reward R(ot). The goal is to maximize the expected cumulative joint reward over a finite or infinite number of steps. In our framework, we use the base MARL algorithm as a policy learning part.

Multi-Agent Proximal Policy Optimization (MAPPO) MAPPO [Yu et al., 2021a] is a method for applying the Proximal PolicyOptimization (PPO) [Schulman et al., 2017] algorithm to multi-agent reinforcement learning (MARL). Each agent in MAPPO has a representation encoder and a policy network. The representation encoder process the observation of one agent, and the policy network generates an action based on the representation. The parameters of the representation encoder and policy network are shared by all agents for training, MAPPO updates the parameters using the aggregated trajectories of all agents.

Multi-Agent Transformer (MAT) MAT [Wen et al., 2022] is a transformer encoder-decoder architecture that changes the joint policy search problem into a sequential decisionmaking process. The encoder maps an input sequence of observations to latent representations and the decoder generates a sequence of actions in an auto-regressive manner. MAT simplifies the sequence modeling paradigm for multi-agent reinforcement learning by treating the team of agents as a single sequence. The learning objectives for MAT are represented by LEncoder and LDecoder.

The learning objectives of MAPPO and MAT serve as the MARL loss and are detailed in the supplementary materials.

Contrastive Learning Contrastive learning is an optimization objective for self-supervised algorithms and is used in many auxiliary tasks for representation learning in RL. Contrastive learning can be regarded as an instance classification problem, where one instance should be distinguished from others. Given a query q and keys K = {k0, k1, . . .} and an explicitly known partition of K (different view of q ), P(K) = ({k+} ,K\{k+}). the goal of contrastive learning is to ensure that q matches with k+ relatively more than any of the keys in K\{k+}. We utilize the InfoNCE loss in [Oord et al., 2018] as the contrastive loss function. It can be mathematically defined as follows:

where ω(q, k) = qT Wk is a bi-linear inner-product similarity function in [Srinivas et al., 2020], andW is a learned matrix.

4 Method

Multi-Agent Masked Attentive Contrastive Learning (MA2CL) is an auxiliary task that aims to enhance representation learning in MARL by reconstructing masked agent with information from other agents and historical trajectory. This allows for better utilization of temporal and agent-level information in observation, further improving the sample efficiency in MARL tasks. Besides, MA2CL can be integrated easily into various off-policy MARL algorithms. The framework of MA2CL is shown in Fig. 2, and each component will be introduced in the following subsections.

4.1 Observation and ActionMasking

We sample the observations and actions of all agents at the same time step t from the replay buffer B. Denote the observation and action as a sequence:

where N represents the total number of agents. Then, we will randomly select Nm agents for masking, as indicated by:

If Mi = 1 then the agent will be masked. In our masking strategy, if agent im is selected, its observation and action (oim t ,aim t ) will be modified as follows: oim t is replaced by oim t-1 and aim t is replaced by aim t-1 , where (oim t-1 , aim t-1 ) is the observation and action of agent im at time step t-1, then we get the masked observation and action sequence as :

We will utilize the masked observation sequence during the encoding and reconstruction stages. The masked actions are only employed in the reconstruction stage. Our aim is truct the masked observation by leveraging information from both agent and time domains. The number of masked agents, denoted as Nm(1 ≤ Nm ≤ N) is a hyperparameter that would be set prior to training.

4.2 Encoding Observation Sequence

We use two encoders to obtain representations from masked and original observation sequences respectively. The online encoder, denoted as ψ, is the same as the representation(encoder) network in the MARL agent. It can be a centralized encoder that processes all agent's observations simultaneously or a decentralized encoder set: ψ =

{ψ1 , ψ2 , . . . , ψN }, where each agent's observation is processed by its own encoder. For simplicity, we use ψ to denote the online encoder. This encoder maps the masked observation ˜ot into latent space, represented as ˜zt. This process is formulated as ˜zt = ψ(˜ot). Similar to [Srinivas et al., 2020], we employ a target encoder, denoted as ¯ ψ, to encode the original observation sequence, formulated as zt = ¯ ψ(ot). The target encoder has the same architecture as the online encoder, and its parameters are updated by an exponential moving average (EMA) of the online encoder parameters. Given online encoder ψ parameterized by θ , target encoder ¯ ψ parameterized by ¯ θ and the momentum coefficient τ ∈ [0, 1), the target encoder will be updated as follows:

Previous works [Chen et al., 2020b; Chen et al., 2020a] have experimentally demonstrated that introducing a nonlinear layer before the contrastive loss can significantly improve performance. Therefore, we also include a non-linear projection layer ϕ before reconstruction in our model. This layer is implemented as an MLP network with GELU activation.

We use the online projector ϕ to process the latent state feature sequence ˜zt, and we get obtain the following result: ˆzt = {ˆ z1 t , ˆ z2 t , . . . , ˆ zN t }, where ˆ zi t = ϕ(˜ zi t ). For the encoded results obtained from original observations, denoted as z, we use a target projector ¯ ϕ for further processing. target projector has the same structure as the online projector and its parameters are updated with online projector following the same EMA update strategy in the encoder. We get the target latent feature sequence as follows:

As shown in Fig. 2, we apply a stop-gradient operation on target projector to avoid model collapse, following [Grill et al., 2020; Srinivas et al., 2020].

4.3 Attentive Reconstruction

We propose an attentive reconstruction model Ψ to recover the representation of the masked observations. The model Ψ consists of L identical blocks, each of which is an attention layer. To add the identity information of each agent in the team, we adopt relative positional embedding, which are commonly used in the standard Transformer [Vaswani et al., 2017]. truct the masked information in latent space, we leverage three types of information: (a) the representation sequence ˆzt, which is obtained from the masked observation sequence ˜ot through processed by online encoder and online projector sequentially. This sequence contains information from other agents and historical trajectory; (b) the action information ˜at, which corresponds to ˆzt provides decision information from each agent, including the masked ones; (c) the identity information p = (p1 , p2 , . . . , pN ), which is obtained from the positional embedding and reflects the identity of each agent in the sequence. To integrate these three types of information, we concatenate ˆ zi t with the masked action ˜ait , and add identity information pi to the concatenating vector, as the identity information is related to both ˆ zi t and ˜ait . Thus, the inputs tokens of the predictive model can be expressed as:

Algorithm 1 Training Process for MA2CL

Input: number of agents N, number of masking agents NM Parameter: parameters in online encoder ψ, online projector ϕ, similarity function W,policy network, target encoder ¯ ψ , target projector ¯ ϕ.

Determine EMA coefficient τ

• 1: while Training do
• 2: Interact with the environment and collect the transition: B ← B ∪ (ot, at, rt, ot+1)
• 3: Sample a minibatch (ot, at, rt, ot+1) fromB.
• 4: Calculate RL loss Lrl based on a given basic MARL algorithm (e.g. MAT, MAPPO)
• 5: Sample another minibatch (ot, at) fromB
• 6: Randomly masks NM agent: ˜ot, ˜at ←Mask(ot, at)
• 7: Process ˜ot and ot based on Eq. (4),Eq. (5) and Eq. (7), obtain ˆ yt and yt.
• 8: Calculate contrastive loss Lcl based on Eq. (8)
• 9: Calculate total loss: Ltoatal = Lrl +Lcl
• 10: Update online encoder ψ, online projector, similarity function W, policy network with Ltotal
• 11: Update target encoder and projector based on Eq. (3)
• 12: end while

The input token sequence is processed by L attention layers in the attentive reconstruction model. The l-th layer processes the input token sequence according to the following steps:

The Layer Normalization (LN), Multi-Headed Self-Attention (MSHA) layer, and Feed-Forward Network(FFN) are the same as those used in the Transformer [Vaswani et al., 2017].

Given the input x, we try to recover a sequence of all agent observations in the latent space using the attentive reconstruction model, then we get:

Using the reconstructed feature sequence ˆ yt and the target feature sequence y from Eq. (4), we construct a contrastive loss to improve the accuracy of the reconstruction model and the ability of the encoder, which will be described in detail in the following section.

4.4 Contrastive Loss

Inspired by the success of [Oord et al., 2018; Srinivas et al., 2020], we use contrastive learning to optimize the parameters of online encoder, projector, and reconstruction model. The query set Q = {qi|qi = ˆ yi } is the reconstructed feature of the masked agent according to Eq. (7). The key set K = {ki|ki = yi } is encoded from the non-masked observation. As defined in Eq. (4), (qi, ki) is a query-key pair from the same agent. We use the function ω(qT , k) from Eq. (1) to measure the similarity between query and key. We can formulate the contrastive loss as follows:

The contrastive loss Lcl in Eq. (8) is designed based on the following intuition: the reconstructed feature ˆ yi (i.e., qi) should be similar to its corresponding original feature yi (i.e., ki) while being distinct from the others. Mi (as specified in Eq. (2) ) emphasizes the focus on the reconstructed features of the masked agents rather than those of the unmasked ones. By minimizing the contrastive loss Lcl, we aim to train an accurate attentive reconstruction model and a powerful encoder. The powerful encoder is able to extract an informative representation from observation input, allowing the reconstruction model to recover the sequence in the latent space even when some agents are masked.

4.5 Overall Training Objective

The agent in MARL algorithms consists of two parts: the representation encoder and the policy network. The online encoder of MA2CL is the same as the representation encoder in the MARL algorithm for policy learning. MA2CL provides a powerful encoder for basic MARL algorithm, which helps agents learn cooperative policy more quickly. In MA2CL, representation learning is optimized together with policy learning. Thus, the overall training objective of our method is as below:

where Lrl is the loss functions of the base MARL algorithm (e.g. MAT [Wen et al., 2022], MAPPO [Yu et al., 2021a]), and Lcl is the loss functions of the MA2CL we introduced. λ is a hyperparameter for balancing the two terms. Following [Srinivas et al., 2020], we fix λ as 1. A detailed example of MA2CL is provided in Algorithm 1.

5 Experiments

To evaluate the performance of MA2CL, we use different multi-agent reinforcement learning benchmarks, including both vision-based and state-based scenarios. We set mask agent number Nm = 1, attention layer L = 1. Other hyperparameters settings can be found in supplementary materials.

5.1 Setup

Vision-based MARL Environments In single-agent settings, representation learning techniques are commonly used in vision-based reinforcement learning (RL) to improve sample efficiency. We demonstrate the effectiveness of our proposed MA2CL model in vision-based MARL settings on the Multi-Agent Quadcopter Control benchmark (MAQC) [Panerati et al., 2021].

In the MAQC environment, each agent receives an RGB video frame ∈ R64×48×4 as an observation, which is captured from a camera fixed on the drone. MAQC allows for continuous action settings of varying difficulty, including RPM and PID. When operating under the RPM setting, the agent's actions directly control the motor speeds. In the PID setting, the agent's actions directly control the PID controller, which calculates the corresponding motor speeds. In this work, we evaluate MA2CL on two cooperative tasks Flock and LeaderFollower in both RPM and PID settings. Some detail about the task can be found in supplementary materials.

State-based MARL Environments In order to demonstrate the applicability of our method in both vision-based and state-based scenarios within MARL environments, and due to the lack of vision-based benchmarks, we conduct an addtional evaluation on commonly utilized state-based MARL benchmarks. We use a variety of state-based MARL scenarios, such as the StarCraft Multi-Agent Challenge (SMAC) [Samvelyan et al., 2019] and Multi-Agent MuJoCo [de Witt et al., 2020], where the performance of the baselines represents the current state-of-the-art performance in MARL.

Baseline Different current state-of-the-art baselines are selected for comparison, including MAT [Wen et al., 2022] and MAJOR [Feng et al., 2023], as they have been shown to perform well in both continuous and discrete MARL environments. The hyperparameters are kept consistent with those used in the original papers for a fair comparison. A total of five random seeds were used to obtain the mean and standard deviation of various evaluation metrics such as episode rewards in Multi-Agent MuJoCo, and the winning rate measured in the SMAC.

Additionally, we play four drone agents in the MAQC environment, incorporating different action settings such as RPM and PID. In the Multi-Agent MuJoCo environment, agents were able to observe the state of joints and bodies from themselves and their immediate neighbors.

(b) StarCraftII Multi-Agent Challenge

Figure 4: Results on state-based MARL environment: (a) Multi-Agent MuJoCo, (b) StarCraftII Multi-Agent Challenge, The results indicate that MA2CL outperforms MAJOR and MAT in both performance and sample efficiency across different state-based scenarios. Due to the page limit, additional results on other scenarios can be found in the supplementary materials.

5.2 Result and Discussion

To ensure a fair comparison, we implement MA2CL on MAT, as the MAJOR algorithm is also proposed based on MAT.

Result in vision-based Environment As shown in Fig. 3, the performance of the MA2CL is evaluated in the MAQC environment by comparing it to the MAJOR and MAT algorithms. We can observe that MA2CL demonstrates superior performance and sample efficiency in both RPM and PID action settings. Specifically, MA2CL achieves the highest episode reward in three out of four tasks. Additionally, MA2CL requires fewer steps to obtain high episode rewards compared to other baseline algorithms. The superior performance of MA2CL can be attributed to its ability to effectively capture task-relevant information from high-dimensional inputs (RGB image in MAQC), which enables it to generate more informative representations that incorporate both temporal and team awareness. This enhances the ability of the drones to learn good policy and effectively collaborate and complete cooperative tasks.

Result in state-based Environment In the state-based MARL environment, MA2CL also exhibits strong performance when compared to MAJOR and MAT. Fig. 4(a) and Fig. 4(b) show the performance of MA2CL and baselines in different tasks of Multi-Agent MuJoCo and SMAC respectively. The results indicate that MA2CL demonstrates robust performance, improving the performance of MAT in most tasks and surpassing that of MAJOR. These findings suggest that the MA2CL algorithm not only possesses the ability to extract task-relevant information from high-dimensional (image) input but also to effectively integrate information from agent-level in state-based observations, which provide more direct information. For instance, in the Multi-Agent MuJoCo environment, the decision-making process of the agents, in terms of the joints of the body, must take into account the velocities and positions of other joints to produce humanlike joint actions. Our MA2CL algorithm enables agents to fully leverage the information of neighboring agents present in their observations, allowing for decisions to be made based on a representation of team awareness and resulting in superior performance.

Effect on other MARL Algorithms To exhibit the plugand-play nature of our proposed MA2CL framework across various MARL algorithms, we implement MA2CL on the MAPPO algorithm, denated as MAPPO+MA2CL. We select MAPPO as a representative example of the actor-critic algorithm in MARL due to its utilization of a CNN encoder (for vision-based tasks) or an MLP encoder (for state-based tasks) to encoder observation from every single agent, in contrast to the transformer-based encoder processing observations from all agents in MAT. To ensure fairness in comparison, we also choose MAPPO+MAJOR, which is the implementation of MAJOR on MAPPO, and the standard MAPPO algorithm as baselines for evaluation. Our experiments are conducted in both visual and state-based MARL environments.

As shown in Fig. 5(a), Fig. 5(b), and Fig. 5(c), MA2CL is found to significantly enhance the performance of MAPPO in both visual and state-based tasks and outperform the MAPPO+MAJOR in these tasks. The representation encoder in MAPPO only utilizes observations from a single agent for decision-making, as opposed to the sequential decision model used in MAT, where agents have explicit agent-level information from previous agents' observations in decision-making. We can observe that, with the introduction of MA2CL, we obtain a more notable improvement effect on the MAPPO than on the MAT. This suggests that the MA2CL framework enhances the representation encoder in MAPPO to extract more task-relevant information and exploit the agent-level information from observation. This enhanced informative representation enables agents to effectively learn cooperative strategies, resulting in higher performance and increased sample efficiency in MAPPO+MA2CL. This highlights the potential of our MA2CL to promote various MARL algorithms.

Figure 5: Result for MAPPO+MA2CL, MAPPO+MAJOR, and MAPPO in both vision and state-based MARL environments. Due to the page limit, additional results on other scenarios can be found in supplementary materials.

5.3 Ablation Studies

To study the impact of masked agent numbers and masking strategy on MA2CL, we conduct ablation studies in which we varied the number of masked agents or masking strategy while holding other parameters constant. Additional ablation results can be found in the supplementary materials.

Variation in masked agent numbers. In MA2CL, the number of masked agents, Nm, determines the ratio of accurate agent-level information from the same time step. As Nm increases, more information comes from the historical trajectory. Therefore, the agent-level information utilized by the reconstruction model decreases and becomes outdated. As shown in Fig. 6(a), when we perform masking on three or four agents in a task involving four drones, the performance decreases slightly but still surpasses the baseline. This is because the agent-level information from the history is not accurate but also contains information in the time dimension, which is also beneficial for agent decision-making.

Different masking strategy. On the other hand, the masking strategy in MA2CL is also a crucial factor in the reconstruction process. One consideration is that in MA2CL, we only use oim t-1 to mask the agents, which may be viewed as "replacement" rather than "masking". Therefore, we explored two other masking strategies:

Figure 6: Result for ablation study on different mask agent number Nm and masking strategy

• (1) Use oim t-1 plus gaussian noise to mask oim t , namely Mask(oim t ) = oim t-1 +N (0, 1), denoted as MA2CL-PG;
• (2) Mask oim t with full gaussian noise, namely Mask(oim t ) =N(0, 1), denoted as MA2CL-FG.
• (3) Mask oim t with zero, namelyMask(oim t ) = 0, denoted as MA2CL-Z.

As shown in Fig. 6(b), even when using zero to completely mask the agent's observation, in which case only agentlevel information can be used for reconstruction, MA2CL-Z still demonstrates superior performance. This indicates that MA2CL indeed learns informative representation with agentlevel information and helps the agent learn better cooperation strategies. Additionally, the performance of MA2CL-PG is slightly worse than MA2CL, but better than MA2CL-FG and MA2CL-Z, indicating that using ot-1 for masking actually introduces temporal information. Therefore, MA2CL simultaneously considers valuable information in the time dimension and agent level, thus learning more informative representations to promote agent policy learning.

5.4 Conclusion

In this work, we present the Multi-Agent Masked Attentive Contrastive Learning (MA2CL) framework which utilizes an attentive reconstruction model that encourages the learning of temporal and team awareness representations, thereby enhancing the sample efficiency and performance of MARL algorithms. Extensive experiments on both vision-based and state-based cooperative MARL benchmarks show MA2CL set the state-of-the-art performance.

In terms of future work, there are several avenues that warrant further exploration. One potential direction is to investigate the use of more elaborate masking strategies to allow for attention to different types of information in MARL settings. Additionally, given the importance of sample efficiency in reinforcement learning, future efforts could be directed toward methods that are specifically designed to address this issue in vision-based multi-agent environments.What's more, we only consider the encoder of actor, and it is also worth thinking about the approach of treat the encoder of both critic and actor as the target of representation learning.

Acknowledgments

This work was supported in part by NSFC under Contract 61836011, and in part by the Fundamental Research Funds for the Central Universities under contract WK3490000007.