Scalable Communication for Multi-Agent Reinforcement Learning
via Transformer-Based Email Mechanism

Policy Gradient (PG) reinforcement learning has the advantage of learning a policy network explicitly, in contrast to value-based reinforcement learning methods. PG methods optimize the policy parameter $\theta$ to maximize its objective function $J(\theta)=E_{S}(V_{\pi_{\theta}}(s))$. However, due to the variance of environments, it is hard to choose a subtle learning rate in reinforcement learning. To resolve this problem and ensure the safe optimization of policy learning, the Trust Region Policy Optimization (TRPO) Schulman et al. (2015) increases the constraint of the parameter difference between policy updates.

Based on TRPO, a simplified version Proximal Policy Optimization Schulman et al. (2017) is carried out, maintaining the motivation to constrain the learning step while more efficient and easy to be implemented. The object function of PPO can be written as:

which forces the ratio of $\frac{\pi_{\theta}(a|s)}{\pi_{\theta_{k}}(a|s)}$ to locate in the interval $(1-\epsilon,1+\epsilon)$, so that the new $\theta$ is not too far away from old $\theta_{k}$.

Multi-agent PPO (MAPPO) introduces PPO into the multi-agent scenario Yu et al. (2021). MAPPO mainly considers decentralized partially observable Markov decision processes (DEC-POMDP). In an environment with $n$ agents, $s\in S$ denotes the state of the environment. The agent i only has a local observation of environment $o_{i}=O(s)$ and chooses its action based on its observation and policy $a_{i}=\pi_{i}(a_{i}|o_{i})$. The joint action $A=(a_{1},...,a_{n})$ denotes the set of actions of all agents. Then, the environment transits its state based on the transition probability $P(s^{\prime}|s,A)$. In MARL, all the agents will get rewards based on the transition of state and their actions (or more likely joint action) $r_{i}=\mathcal{R}(s,A)$. Each agent is supposed to get a higher accumulated reward $\sum_{t}{r_{i}^{t}}$. Therefore, the agents optimize their policy to maximize the discount accumulated reward $J(\theta)=E_{a^{t},s^{t}}[\sum_{t}{\gamma^{t}\mathcal{R}(s^{t},a^{t})}]$, where $\gamma\in(0,1]$ is the discount factor.

MAPPO utilizes parameter sharing within homogeneous agents, i.e., homogeneous agents share the same set of network structure and parameters during training and testing. MAPPO is also a CTDE framework, namely, each PPO agent maintains an actor network $\pi_{\theta}$ to learn the policy and a critic network $V_{\phi}(s)$ to learn the value function, where $\theta$ and $\phi$ are the parameters of policy network and value network, respectively. The value function requires the global state and only works during training procedures to reduce variance. In our work, we take MAPPO as our baseline and backbone, and add TEM as the communication mechanism into MAPPO.

We design a communication protocol following the way how humans communicate by email. The information flow is like a forwarding chain: the chain starts with an agent with key information to share, and the following agents merge their own information into the chain and then forward the new message to the next agent. The chain ends when the final agent finds the message useless for others, or there are no more potential communication objects.

When designing the communication protocol, we mainly consider the following questions: (1) Whether to communicate? (2) Whom to communicate with? (3) What to communicate? (4) How to utilize the messages?

(1) Whether to communicate? As shown in Fig 2, in every step of execution and training, the first stage is communication. When the communication stage is done, the actor networks for each agent will make the action decisions by the observations and messages. In the communication stage, each agent has the chance to decide whether to start a new chain and send a message. And the agents who receive messages can decide whether to continue forwarding the messages. Multiple message chains are allowed and the information from different chains is merged if a shared node exists.

(2) Whom to communicate with? For partial-observation (PO) problems, communication with all the agents is not reasonable and effective. The direct communication with the agent outside the observation range may not bring helpful information as the sender does not even know the receiver’s state. Therefore, we do not model all the other agents to decide whether to communicate with them like in some previous works Ding et al. (2020); Yuan et al. (2022). Instead, when the agent i chooses communication objects, we only consider the agents in the observation range $\mathcal{O}_{i}$ , and in our experiments, $\mathcal{O}_{i}$ includes the nearest several agents of the agent i. By training the message module, the agent can predict the impact of the message on other agents, and is more likely to choose the one with the highest impact to communicate.

We combine the two decisions (1) and (2) into one communication action. The communication actions of agent i $m\_a_{i}$ include not sending at all $m\_a_{i}=0$, and sending to one agent j in the observation range $m\_a_{i}=j,~{}(j\in\mathcal{O}_{i})$. Namely:

This way, the agent can decide when and whom to

communicate by one action, simplifying the modeling and learning.

(3) What to communicate? To keep the information from the head nodes in the chain, and merge the information from different chains, every agent maintains a message buffer to store the messages. In practice, the message buffer is implemented as a queue, with a fixed storage length, but can flexibly push in and pop out elements as the communication goes (First Input First Output, FIFO). We use $m\_b_{i}$ to denote all the messages inside the agent i ’s message buffer. When sending the new message, the agent i merges its own observation into the chain, then the message chain expands to $\langle m\_b_{i},o_{i}\rangle$. Here, the operation $\langle\cdot\rangle$ denotes pushing into the queue to concatenate the messages. The buffer is clear when every step starts.

(4) How to utilize the messages? Instead of some previous works Zhang et al. (2019); Yuan et al. (2022), we do not think the messages directly influence the value estimation of other agents is the natural way of communication. The information exchange should be separated from the information utilization. And the final effects of the messages should be determined by the receiver instead of the sender. Thus, in our model, messages are taken as a counterpart of the observation, serving as part of the inputs of the actor network.

The schematics of the network design in our model are shown as Fig 3. Each agent has an actor network to observe the environment and communicate with other agents. The actor network of the agent i will output the action to interact with the environment $a_{i}$, the action to communicate $m\_a_{i}$ (whether to communicate and whom to communicate with), and the corresponding message to be sent $m_{i}$. To better utilize the history information and get a smoother action sequence, an RNN is employed in the actor network. Thus the agent i also keeps a hidden state $h_{i}^{t}$, and updates it every time step.

The actor network consists of two sub-networks, the action network and the message network. The action network mainly concentrates on the task itself and tries to get better rewards by outputting reasonable actions. The message network concentrates on the communication to share information with other agents instead. The two sub-networks exchange the representation feature of the observations $o\_f_{i}$ and that of the messages $m\_f_{i}$ to merge the information.

In the action network, $o\_f_{i}$ is learned by a multi-layer perceptron (MLP), and then the action network concatenates $o\_f_{i}$ and $m\_f_{i}$ to input into the RNN together with the hidden state from the last time step $h_{i}^{t-1}$. Another MLP, in the end, processes the output of the RNN to get the action $a_{i}$.

On the other hand, the messages from other agents like $m_{j}\cdots m_{k}$ are stored in the message buffer, like the email inbox. The embedding layer (we implement it as a full connected (FC) layer by practice) converts the messages to fit the input dimensions of the encoders. $N_{e}$ sequential encoder modules and $N_{d}$ sequential decoder modules are followed by the embedding layer. The output of encoder modules $m\_f_{i}$ serves as the representation of all the messages in the buffer, with the key information emphasized by the attention mechanism. The decoder modules further combine the information from both of $m\_f_{i}$ and $o\_f_{i}$ to get the output $m\_dec_{i}$. Finally, one MLP produces the communication decision $m\_a_{i}$.

For each encoder module, it takes in $m\_enc_{i}^{n_{e}-1}$ from the embedding layer or the last encoder, then generates $m\_enc_{i}^{n_{e}}$ as the input for the next layer. The transformer in the module can model the sequential information and is flexible to fit message chains with different lengths. Also, the attention mechanism will help the agent to pick out the key information from the chain. $n_{e}$ implies the position of the layer in the encoder sequence. To prevent gradient vanishing, the encoder module employs the residual connections to link the self-attention mechanism and the MLP Wen et al. (2022). The structure of the self-attention mechanism is the same as the attention network in the decoder while $k,q$ and $v$ are generated from the same input $m\_enc_{i}^{n_{e}-1}$.

In the decoder module, the first $m\_dec_{i}^{0}$ is the representation of the observations $o\_f_{i}$. In the attention network, full connected layers generates key $k$ and query $q$ by $m\_dec_{i}^{n_{d}-1}$ and $m\_f_{i}$, respectively. Also, the third FC layer generates value $v$ from $m\_dec_{i}^{n_{d}-1}$. $k$ and $q$ are used for calculating the weights $\alpha$ of the value $v$ as Equation 3.

In fact, the weight $\alpha$ learns the correlations between the $m\_dec_{i}^{n_{d}-1}$ and $m\_f_{i}$. By multiplying $v$ and $\alpha$, then we get the weighted representation of $m\_dec_{i}^{n_{d}-1}$ from the ending FC layer. With a similar structure of the residual connections and MLP, we get $m\_dec_{i}^{n_{d}}$ as the input for the next layer.

The communication among the agents in a collaborative task aims to share the key information that one believes is useful for some specific agents. So the learning of the message network is driven by the impact of the message to be sent.

As the communication will not change either the action of other agents or the loss of the action network, an independent loss to model the influence of the messages on other agents’ actions is needed. We denote the communication loss as $\mathcal{L}^{(i)}_{m}(\theta)$, where $\theta$ is the parameters of the actor network. The action of an agent j is sampled from the categorical distribution $P(a_{j}|o_{j},m\_b_{j})$ learned by the action network. Then, when considering the new message from the agent i $m_{i}$, we can estimate the distribution $P(a_{j}|o_{j},\langle m\_b_{j},m_{i}\rangle)$ as the consequence of the communication. Kullback-Leibler (KL) divergence is widely used to measure the discrepancy between these two conditional probability distributions. Thus, the causal effect $\Gamma_{j}^{(i)}$ of the message from agent i on agent j can be defined as:

By considering all the possible agents to send the message to in the observation range, we can get the expectation of the causal effect of the message $\mathbb{E}\Gamma^{(i)}(\theta)$ by Equation 5:

where $\mathcal{O}_{i}$ denotes the observation range of the agent i. The communication decision of agent i is sampled form the categorical distribution $P_{\theta}(m\_a_{i}|o_{i},m\_b_{i})$ learned by the message network. $P_{\theta}$ denotes that the gradient of this item should be propagated when training.

The expectation $\mathbb{E}\Gamma^{(i)}(\theta)$ represents the overall effect the message $m_{i}$ can bring to the whole system, which we should maximize in the loss function. However, communication should also be sparse and efficient. If we do not control the communication times by the external guidance, the agents will tend to send as many messages as possible to get higher $\mathbb{E}\Gamma^{(i)}(\theta)$. Therefore, we also designed another item for communication loss to reduce the communication overhead. When the agent i chooses not to send the message to any agents in the observation range for most of the times, the probability $P_{\theta}(m\_a_{i}=0|o_{i},m\_b_{i})$ should be relatively high. So we need to maximize this probability at the same time.

So far, we can get the final communication loss $\mathcal{L}^{(i)}_{m}(\theta)$ by the following equation and maximize it when training.

where $\delta$ is the weight of the communication reduction.

The loss of the action network $\mathcal{L}^{(i)}_{a}(\theta)$ is defined followed by Equation 1 in MAPPO as:

where $r_{\theta}^{(i)}=\frac{\pi_{\theta}(a^{(i)}|o^{(i)})}{\pi_{\theta_{old}}(a^{(i)}|o^{(i)})}$, $A^{(i)}_{\pi_{\theta_{old}}}$ is the advantage function.

What’s more, to encourage more exploration when training, we adopt an entropy loss $\mathcal{L}^{(i)}_{e}(\theta)$ as Yu et al. (2021):

We can get the overall loss function for the actor network when training:

where $n$ is the number of the agents, and $\lambda_{m},\lambda_{e}$ are the coefficients to weight the corresponding losses.

The loss of critic network keeps the same as MAPPO.

StarCraft II is a real-time strategy game serving as a benchmark in the MARL community. In the Starcraft Multi-Agent Challenge (SMAC) task, N units controlled by the learned algorithm try to kill all the M enemies, demanding proper cooperation strategies and micro-control of movement and attack. We choose the hard map 5m vs. 6m to evaluate TEM. TEM controls 5 Marines to fight with 6 enemy Marines.

The baselines include MAPPO, MADDPG, Full Communication (FC) and Randomly-stop Communication (RC). MAPPO is the CTDE backbone we are using in the following experiments, which is proven to have state-of-the-art performance on several MARL cooperative benchmarks Yu et al. (2021). MADDPG is another classic CTDE approach for multi-agent cooperation tasks Lowe et al. (2020). FC and RC are two special cases of TEM. We keep the communication protocol the same, but disable the decoder in the message module, instead, the agents choose the communication targets by pre-defined rules. In FC, the agent will keep randomly choosing someone to communicate with, to extend the message chain until no one is available. In RC, the agent will randomly stop the message chain by a probability $p$, or keep forwarding to a random one.

We run the experiments over five seeds. For each seed, we compute the win rate over 32 test games after each training iteration as shown in Fig. 4. TEM gets the highest win rate over the baselines. FC and RC perform worse than MAPPO benchmark. One possible reason is that targeted communication by TEM could improve cooperation while random communication by FC and RC may bring redundant information for decision-making. The win rate of MADDPG remains zero, showing that it is hard to defeat an army with more units.

In the Predator Prey (PP) task, N predators try to chase, surround and finally capture M preys, as shown in Fig 1. The predators are the agents to be trained and the preys flee in the opposite direction of the closest predator at a faster speed following pre-defined rules. So the predators have to be grouped automatically and cooperate to surround each prey, and it is impossible for one predator to capture a prey itself. In practice, we set N as 7 and M as 3, denoted as 7-3 scenario.

Different from the PP task in some previous works, here, the agents can only partially observe the teammates and targets. The rewards are the sum of the agents’ negative distances to their closest preys or landmarks. In addition, the agents are penalized for collisions with other agents.

The baselines include MAPPO, I2C, MADDPG, DDPG, FC and RC. I2C proposes an individual communication mechanism Ding et al. (2020). DDPG is a classic deep reinforcement learning algorithm for continuous control Lillicrap et al. (2019). We apply DDPG independently to each agent as a baseline without considering cooperation.

As shown in Fig 5, while other baselines gradually converge at the last episodes, TEM keeps raising the rewards and improves the final reward by 17.2% compared with MAPPO.

In the Cooperative Navigation (CN) task, N agents try to occupy N stationary landmarks separately, as shown in Fig 7. The positions of landmarks and agents are randomly initialized. The best strategy is that each agent has a different target from the beginning through communication instead of rescheduling when collisions happen because of choosing the same target. In practice, we set N as 7, denoted as 7-7 scenario. The baselines and reward settings are the same as PP.

We compare TEM with the baselines on the training performance Fig 6. We can see that TEM converges to the highest reward than all the baselines. FC and RC are only slightly better than MAPPO, suggesting that the communication actions $m\_a$ learned by TEM are targeted, and the message chain brings helpful information to the ones that really need it.

We compare the illustrations on CN between TEM and MAPPO in Fig 7. In (a), the TEM agents Agent 1 and Agent 4 notice Landmark 1 by communication (pink message chain). Thus each agent moves straight forward to the corresponding landmarks. While in (b), the MAPPO agents miss Landmark 1, so for Agent 4, there will be nowhere to go. Agent 4 first tries to scramble with Agent 1 but fails, then turns to Agent 2. Agent 2 is forced to leave to avoid collision and turns to Agent 3. We can see that communication brought by TEM can improve cooperation and reduce internal strife.

We further examine the scalability of TEM on PP in Table 1. We take average episode rewards (R), successful capture times (S), collision times (C) as the metrics. For R and S, the performance is better when the values are higher, while for C, the performance is better when the values are lower. Note that the existing selective MARL communication approaches are not scalable due to the modeling of each agent, so the baselines are the transferred MAPPO and the specifically trained MAPPO. We directly transfer the learned model from the 7-3 scenario to 9-3 and 3-1 scenarios without further training. For 9-3 scenario, two new agents are included so the task will be easier. But more agents also increase the risks of collision, so the cooperation mode could be different and the agents need to communicate to suit the new scenario. For TEM, the average episode rewards rise from -40.5 to -17.9, and the gain is 55.8%, while for MAPPO, the gain of rewards is 52.3%. TEM does not only perform better after transferring, but also gains more.

For 3-1 scenario, both the numbers of the agents and preys change. The results show that TEM still keeps a better performance on all the metrics. Moreover, it shows that after TEM learns how to communicate in a complex scenario, it can successfully transfer to simple ones.

We also train MAPPO from scratch specifically on 9-3 and 3-1 (denoted as MAPPO (learned)), and the performance of transferred TEM (trained on 7-3) is close to MAPPO (learned) on 9-3 without training. But the transferred TEM works worse on 3-1, and we suggest that cooperation by communication may not play an essential role in such a simple environment. We further finetune TEM (7-3) on the new scenarios and the finetuned models even outperform the specially learned MAPPO.

Similar experiments are conducted on CN as shown in Table 2. TEM keeps the scalability when transferred from 7-7 scenario to 6-6 and 9-9, and outperforms the transferred MAPPO. Surprisingly, the transferred TEM even outperforms the MAPPO trained from scratch (denoted as MAPPO (learned)) on most metrics. It suggests that CN requires more communication to coordinate the agents to explore all the landmarks. The results also show that the communication pattern learned from 7-7 still works well in other scenarios. Similarly, the finetuned TEM gets even better performance.