

\section*{Learning to Communicate in Multi-Agent Systems via Transformer-Guided Program Synthesis} 

\section{Extending to Multi-Round Communications}
\label{sec:multihop}

The formulation in Section \ref{sec:approach} can be extended to multiple rounds of communications per time step. For the transformer architecture with two rounds of communications, first there is an \emph{internal network} $\pi_\theta^H: \mathbb{R}^{d_{\text{in}}} \to \mathbb{R}^{d_{\text{out}}}$ that combines the state and the cumulative message into an internal vector $h^i = \pi_{\theta}^H\left(s^i,~\sum_{j=1}^N\alpha^{j\to i}m^{j\to i}\right)$. Next, we compute the next round of messages as ${m'}^{i\to j}=\pi_{\theta}^{M'}(h^i,o^{i,j})$ which replaces the state $s^i$ in the original equation with the internal state $h^i$. New keys and queries are also generated as $k'^{i,j}=\pi_{\theta}^{K'}(s^i,o^{i,j})$ and $q'^{i}=\pi_{\theta}^{Q'}(s^i)$, but these still use the original state $s^i$. Finally, the Equations~\ref{eq:tf_softmax} and~\ref{eq:tf_out} are repeated to compute the action. This architecture can be extended similarly to an arbitrary number of rounds of communications. A programmatic policy for $R$ rounds of communications will have $R$ different programs (one for each round). We synthesize these programs independently. To synthesize the communication program $P_r$ for the $r$-th round of communication, we use the hard attention weights $\alpha^{P_r}$ for the $r$-th round and use the original soft attention weights for the other rounds $r' \neq r$ to compute the synthesis objective $\tilde{J}(P_r)$. 



\section{Experimental Details}
The code and a short video illustrating the different tasks used in the paper can be found in \url{https://github.com/jinala/multi-agent-neurosym-transformers}. Figure \ref{fig:covering} shows the initial and goal positions for the unlabeled goals task, along with attention maps produced by our program policies for the two rounds of communications at a particular timestep. 

There are four main hyper-parameters in our synthesis algorithm. 
\begin{itemize}
    \item $\tilde{\lambda}$ in Section \ref{sec:synthesis}: This parameter strikes a balance between minimizing the difference in the actions (with and without hard attention) and  minimizing the maximum communication degree. We use $\tilde{\lambda} = \{0.3, 0.5, 0.7, 1.0\}$.
    \item The number of rules in the program $= \{2, 3, 4, 5\}$.
    \item The depth of the Boolean conditions in the filter expressions $= 2$.
    \item The feature map $\phi$ used in the filter predicates and the map functions. We have 2 versions: 1) for every vector $(x, y)$ in the state $s$ and the observations $o$, we also encode the norm $\sqrt{x^2 + y^2}$ and the angle $\text{tan}^{-1}(y/x)$ as part of the features; 2) on top of 1, we add quadratic features $(x_s x_o, x_s y_o, y_s x_o, y_s y_o )$ where $(x_s, y_s)$ is the state and $(x_o, y_o)$ is the observation.
\end{itemize}
We used cross validation to choose these parameters. In particular, we chose the ones that produced the lowest cumulative reward on a validation set of rollouts; if the cumulative rewards are similar, we chose the ones that reduced the communication degree.

\begin{figure*}[h]
    \centering
    \begin{subfigure}[b]{0.24\textwidth}
         \centering
         \includegraphics[width=\textwidth]{prog-retrained-covering_0.pdf}
         \caption{}
         \label{fig:covering-init}
     \end{subfigure}
     \begin{subfigure}[b]{0.24\textwidth}
         \centering
         \includegraphics[width=\textwidth]{prog-retrained-covering_200.pdf}
         \caption{}
         \label{fig:covering-final}
     \end{subfigure}
     \begin{subfigure}[b]{0.48\textwidth}
         \centering
         \includegraphics[width=\textwidth]{prog-retrained-attn-covering_10.pdf}
         \caption{}
         \label{fig:covering-attn}
     \end{subfigure}
    \caption{Unlabeled goals task: (a) Initial positions of the agents and the locations of the goals to cover (b) Final configuration of the agents where 8 out of  the 10 goals are covered (c) Attention maps of \texttt{prog-retrained} for the two rounds of communication. }
    \label{fig:covering}
\end{figure*}

\section{Additional Baselines}

We compare to two additional baselines: (i) an ablation of our approach that learns only deterministic rules---i.e., rules with \texttt{random} are excluded from the search space (\texttt{det-prog} and \texttt{det-prog-retrained}), and (ii) a learned communication policy in the form of a decision tree (\texttt{dt} and \texttt{dt-retrained}). For (ii), to train the decision tree, we constructed a supervised dataset by (i) collecting the soft-attentions from the transformer model, and (ii) solving the global hard-attention problem at each timestep to ensure that the maximum degree (both in-degree and out-degree) is at most $k$, where $k$ is chosen as described in Section~\ref{sec:experiments} (i.e., to match the number of rules in our programmatic communication structure). Then, we train the decision tree using supervised data on this dataset.

Figure~\ref{fig:new_baselines} shows the performance using both the loss and the maximum communication degree for these two baselines. The decision tree baselines (\texttt{dt} and \texttt{dt-retrained}) perform poorly in-terms of the communication degree for all the tasks, demonstrating that  domain-specific programs that operate over lists are necessary for the communication policy to reduce communication.

The deterministic baseline (\texttt{det-prog-retrained}) achieves a similar loss as \texttt{prog-retrained} for the \texttt{random-cross} and \texttt{random-grid} tasks; however, it has worse out-degrees of communication. For these tasks, it is most likely difficult for a deterministic program to distinguish the different agents in a group; thus, all agents are requesting messages from a small set of agents. For the \texttt{unlabeled goals} task, the deterministic baseline has a lower degree of communication but has higher loss than \texttt{prog-retrained}. Again, we hypothesize that the deterministic rules are insufficient for an agent to distinguish the other agents, which led to a low in-degree (and consequently low out-degree), which is not sufficient to solve the task. 

\begin{figure*}
    \centering
    \begin{subfigure}[b]{0.32\textwidth}
         \centering
         \includegraphics[width=\textwidth]{cross-stats-new.pdf}
         \caption{}
         \label{fig:covering-init}
     \end{subfigure}
     \begin{subfigure}[b]{0.32\textwidth}
         \centering
         \includegraphics[width=\textwidth]{grid-stats-new.pdf}
         \caption{}
         \label{fig:covering-final}
    \end{subfigure}
    \begin{subfigure}[b]{0.32\textwidth}
         \centering
         \includegraphics[width=\textwidth]{unlabeled-stats-new.pdf}
         \caption{}
         \label{fig:covering-final}
    \end{subfigure}
    \caption{Statistics of cumulative loss and communication graph degrees for the  additional baselines, for (a) \texttt{random-cross}, (b) \texttt{random-grid}, and (c) \texttt{unlabeled-goals}. }
    \label{fig:new_baselines}
\end{figure*}


\iffalse
\begin{figure*}
\centering
\begin{tabular}{cc}
\includegraphics[width=0.4\textwidth]{together_traj_prog_fix.pdf} &
\includegraphics[width=0.4\textwidth]{together_traj_prog.pdf} \\
(a) prog-retrained, fix left group & (b) prog-retrained \\
\includegraphics[width=0.4\textwidth]{together_traj_tf_fix.pdf} &
\includegraphics[width=0.4\textwidth]{together_traj_tf.pdf} \\
(c) tf-full, fix left group & (d) tf-full 
\end{tabular}
\caption{}
\label{fig:grid}
\end{figure*}
\fi

\section{Comparison to Communication Decisions as Actions}

The multi-agent communication problem can be formulated as an MDP where decisions about which agents to communicate with are part of the action.
\begin{figure}
    \begin{subfigure}[b]{0.32\textwidth}
         \centering
         \includegraphics[width=\textwidth]{random-cross-pg-stats.pdf}
        
         \label{fig:pg-stats-cross}
     \end{subfigure}
     \begin{subfigure}[b]{0.32\textwidth}
         \centering
         \includegraphics[width=\textwidth]{random-grid-pg-stats.pdf}
        
         \label{fig:pg-stats-grid}
     \end{subfigure}
     \begin{subfigure}[b]{0.32\textwidth}
         \centering
         \includegraphics[width=\textwidth]{unlabeled-pg-stats.pdf}
        
         \label{fig:pg-stats-grid}
     \end{subfigure}
    
    \caption{Comparing program policy with RL policy that treats communications as actions.  RL1 and RL2 correspond to two different hyper-parameters in the policy gradient algorithm.}
    \label{fig:pg-comp}
\end{figure}
We performed additional experiments to compare to this approach. Since the action space now includes discrete actions, we use the policy gradient algorithm to train the policy. We tuned several hyper-parameters including (i) weights for balancing the reward term with the communication cost, (ii) whether to use a shaped reward function, and (iii) whether to initialize the policy with the pre-trained transformer policy.

Results are shown in Figure~\ref{fig:pg-comp}. Here, rl1 is the baseline policy that achieves the lowest loss across all hyper-parameters we tried; however, this policy has a very high communication degree. In addition, rl2 is the policy with lowest communication degree; however, this policy has very high loss. 


As can be seen, our approach performs significantly better than the baseline. We believe this is due to the combinatorial blowup in the action space---i.e., there is a binary communication decision for each pair of agents, so the number of communication actions is $2^{N-1}$ per agent and $2^{N(N-1)}$ for all agents (where $N$ is the number of agents).
Our approach addresses this challenge by using the transformer as a teacher.

\section{Case Study with Noisy Communication}

We consider a new benchmark based on the random grid task, but where the communication link between any pair of agents has a 50\% probability of failing. 
\begin{figure}
\centering
\includegraphics[width=0.4\textwidth]{random-grid-noisy-stats.pdf}
\caption{Random grid task with noisy communications.}
\label{fig:noisy-stats}
\end{figure}
The results are shown in Figure~\ref{fig:noisy-stats}. As can be seen, the programmatic communication policy has similar loss as the transformer policy while simultaneously achieving lower communication degree. Here, the best performing policy has four rules (i.e., $K=4$), whereas for the previous random grid task, the programmatic policy only has 2 rules. Intuitively, each agent is attempting to communicate with more of the other agents to compensate for the missing communications.

\section{Neurosymbolic Transformer Policies}
\label{sec:approach}

Our algorithm has three steps. First, we use reinforcement learning to train an oracle policy based on transformers (Section~\ref{sec:oracle}), which uses soft attention to prioritize edges in the communication graph. However, attention is soft, so every agent communicates with every other agent. Second, we synthesize a programmatic communication policy (Section~\ref{sec:dsl}) by having it mimic the transformer (Section~\ref{sec:synthesis}). Third, we combine the programmatic communication policy with the transformer policy by overriding the soft attention of the transformer by the hard attention imposed by the programmatic communication policy (Section~\ref{sec:combined}), and re-train the transformer network to fine-tune its performance (Section~\ref{sec:retrain}). We discuss an extension to multi-round communications in Appendix~A.

\subsection{Oracle Policy}
\label{sec:oracle}

\newcommand{k}{k}
\newcommand{v}{v}
\newcommand{q}{q}

\newcommand{K}{K}
\newcommand{V}{V}
\newcommand{Q}{Q}
\newcommand{O}{O}

We begin by training an \emph{oracle policy} that guides the synthesis of our programmatic communication policy. Our oracle is a neural network policy $\pi_{\theta}=(\pi_{\theta}^C,\pi_{\theta}^M,\pi_{\theta}^A)$ based on the transformer architecture~\cite{vaswani2017attention}; its parameters $\theta$ are trained using reinforcement learning to optimize $J^R(\pi_{\theta})$. At a high level, $\pi_{\theta}^C$ communicates with all other agents, $\pi_{\theta}^M$ is the value computed by the transformer for pairs $(i,j)\in[N]^2$, and $\pi^A$ is the output layer of the transformer. While each agent $i$ receives information from every other agent $j$, the transformer computes a soft attention score $\alpha^{j\to i}$ that indicates how much weight agent $i$ places on the message $m^{j\to i}$ from $j$.

More precisely, the communication policy is $c^i=\pi_{\theta}^C(s^i,o^i)=[N]$ and the message policy is $m^{i\to j}=\pi_{\theta}^M(s^i,o^{i,j})$. The action policy is itself composed of a \emph{key network} $\pi_{\theta}^K:\mathcal{S}\times\mathcal{O}\to\mathbb{R}^{d}$, a \emph{query network} $\pi_{\theta}^Q:\mathcal{S}\to\mathbb{R}^d$, and an \emph{output network} $\pi_{\theta}^O:\mathcal{S}\times\mathcal{M}\to\mathcal{A}$; then, we have
\begin{align}
\pi_{\theta}^A(s^i,o^i,m^i)=\pi_{\theta}^O\left(s^i,~\sum_{j=1}^N\alpha^{j\to i}m^{j\to i}\right), \label{eq:tf_out}
\end{align}
where the attention $\alpha^{j\to i}\in[0,1]$ of agent $i$ to the message $m^{j\to i}$ received from agent $j$ is
\begin{align}
(\alpha^{1\to i}, ..., \alpha^{N\to i}) = \text{softmax} \left(\frac{\langleq^{i},k^{i,1}\rangle}{\sqrt{d}}, ..., \frac{\langleq^{i}, k^{i,N}\rangle}{\sqrt{d}} \right), \label{eq:tf_softmax}
\end{align}
where $k^{i,j}=\pi_{\theta}^K(s^i,o^{i,j})$ and $q^{i}=\pi_{\theta}^Q(s^i)$ for all $i,j\in[N]$. Since $\pi_{\theta}$ is fully differentiable, we can use any reinforcement learning algorithm to train $\theta$; assuming $f$ and $r$ are known and differentiable, we use a model-based reinforcement learning algorithm that backpropagates through them~\cite{deisenroth2011pilco,bastani2020sample}.

\subsection{Programmatic Communication Policies}
\label{sec:dsl}

\newcommand{R}{R}
\newcommand{R^f}{R^f}
\newcommand{R^m}{R^m}

\newcommand{\text{filter}}{\text{filter}}
\newcommand{\text{map}}{\text{map}}
\newcommand{\text{reduce}}{\text{reduce}}
\newcommand{\text{random}}{\text{random}}
\newcommand{\text{argmax}}{\text{argmax}}

Due to the combinatorial nature of the communication choice and the communication objective, we are interested in training communication policies represented as programs. At a high level, our communication policy $\pi_P^C$ is a parameterized  program $P:\mathcal{S}\times\mathcal{O}^N\to\mathcal{C}^K$, which is a set of $K$ rules $P=(R_1,...,R_K)$ where each rule $R:\mathcal{S}\times\mathcal{O}^N\to\mathcal{C}$ selects a single other agent from whom to request information---i.e., $\pi_P^C(s^i,o^i)=P(s^i,o^i)=(R_1(s^i,o^i),...,R_K(s^i,o^i))$.

When applied to agent $i$, each rule first constructs a list
\begin{align*}
\ell=(x_1,...,x_N)=((s^i,o^{i,1},1),...,(s^i,o^{i,N},N))\in\mathcal{X}^N,
\end{align*}
where $\mathcal{X}=\mathcal{S}\times\mathcal{O}\times[N]$ encodes an observation of another agent. Then, it applies a combination of standard list operations to $\ell$---in particular, we consider two combinations
\begin{align*}
R ~::=~ \text{argmax}(\text{map}(F, \text{filter}(B, \ell))) \mid \text{random}(\text{filter}(B, \ell)).
\end{align*}
Intuitively, the first kind of rule is a deterministic aggregation rule, which uses $F$ to score every agent after filtering and then chooses the one with the best score, and the second kind of rule is a nondeterministic choice rule which randomly chooses one of the other agents after filtering.

More precisely, filter outputs the list of elements $x\in\ell$ such that $B(x)=1$, where $B:\mathcal{X}\to\{0,1\}$. Similarly, map outputs the list of pairs $(F(x),j)$ for $x=(s^i,o^{i,j},j)\in\ell$, where $F:\mathcal{X}\to\mathbb{R}$. Next, argmax inputs a list $((v^{j_1},j_1),...,(v^{j_H},j_H))\in(\mathbb{R}\times[N])^H$, where $v^j$ is a score computed for agent $j$, and outputs the agent $j$ with the highest score $v^j$. Finally, $\text{random}(\ell)$ takes as input a list
$((s^i,o^{i,j_1},j_1),...,(s^i,o^{i,j_H},j_H))\in\mathcal{X}^H$, and outputs $j_h$ for a uniformly random $h\in[H]$.

Finally, the filter predicates $B$ and map functions $F$ have the following form:
\begin{align*}
B  ~::=~  \langle\beta,\phi(s^i,o^{i,j})\rangle \ge 0 \mid B \wedge B \; \vert \; B \vee B
\qquad\qquad
F ~::=~ \langle\beta,\phi(s^i,o^{i,j})\rangle,
\end{align*}
where $\beta\in\mathbb{R}^{d'}$ are weights and $\phi:\mathcal{S}\times\mathcal{O}\to\mathbb{R}^{d'}$ is a feature map.

\subsection{Combined Transformer \& Programmatic Communication Policies}
\label{sec:combined}

A programmatic communication policy $\pi_P^C$ only chooses which other agents to communicate with; thus, we must combine it with a message policy and an action policy. In particular, we combine $\pi_P^C$ with the transformer oracle $\pi_{\theta}$ to form a combined policy $\pi_{P,\theta}=(\pi_P^C,\pi_{\theta}^M,\pi_{P,\theta}^A)$, where (i) the oracle communication policy $\pi_{\theta}^C$ is replaced with our programmatic communication policy $\pi_P^C$, and (ii) we use $\pi_P^C$ to compute hard attention weights that replace the soft attention weights in $\pi_{P,\theta}^A$.

The first modification is straightforward; we describe the second in more detail. First, we use $\pi^C_P$ as a mask to get the hard attention weights: 
\begin{align*}
\alpha^{P,j\to i} =
\begin{cases}
\alpha^{j\to i}/Z &\text{if}~j \in \pi_P^{C}(s^i,o^i) \\
0 &\text{otherwise}
\end{cases} \qquad \text{where} \qquad Z = \sum_{j \in \pi_P^{C}(s^i,o^i) } \alpha^{j\to i}
\end{align*}
Now, we can use $\alpha^{P,j\to i}$ in place of $\alpha^{j\to i}$ when computing the action using $\pi_{\theta}$---i.e.,
\begin{align*}
\pi_{P,\theta}^A(s^i,m^i) = \pi_\theta^O\left(s^i,~\sum_{j=1}^N \alpha^{P,j\to i} m^{j\to i}\right),
\end{align*}
where the messages $m^{j\to i}=\pi_{\theta}^M(s^j,o^{j,i})$ are as before, and $\pi_{\theta}^O$ is the output network of $\pi_{\theta}^A$.

\subsection{Synthesis Algorithm}
\label{sec:synthesis}

To optimize $P$ over the search space of programs, we use a combinatorial search algorithm based on MCMC~\cite{metropolis1953equation,hastings1970monte,schkufza2013stochastic}. Given our oracle policy $\pi_{\theta}$, our algorithm maximizes the surrogate objective
\begin{align*}
\tilde{J}(P;\theta)=\tilde{J}^R(P;\theta)+\tilde{\lambda}\tilde{J}^C(P)=-\mathbb{E}_{s_0\sim\mathcal{P}_0}\left[\sum_{t=0}^T\|a_t-a_t^P\|_1 \right ] - \tilde{\lambda} \mathbb{E}_{s_0\sim\mathcal{P}_0}\left[\sum_{t=0}^T\max_{i\in[N]}\text{deg}(i;c_t^P)\right],
\end{align*}
where $\tilde{\lambda}\in\mathbb{R}_{>0}$ is a hyperparameter, the surrogate reward objective $\tilde{J}^R$ aims to have the actions $a_t^P$ output by $\pi_{P,\theta}$ match the actions $a_t$ output by the $\pi_{\theta}$, and the surrogate communication objective $\tilde{J}^C$ aims to minimize the degree of the communication graph $c_t^P$ computed using $\pi_{P}^C$.

Finally, a key to ensuring MCMC performs well is for sampling candidate programs $P$ and evaluating the objective $\tilde{J}(P)$ to be very efficient. The former is straightforward; our algorithm uses standard choice of neighbor programs that can be sampled very efficiently.
For the latter, we precompute a dataset of tuples $D=\{(s,o,\alpha,a)\}$ by sampling trajectories of horizon $T$ using the oracle policy $\pi_{\theta}$. Given a tuple in $D$ and a candidate program $P$, we can easily compute the corresponding values $(a^P,\alpha^P,c^P)$. Thus, we can evaluate $\tilde{J}(P)$ using these values. 

Note that, we sample trajectories  using the  transformer policy rather than using a program policy. The latter approach is less  efficient  because we have to sample trajectories at every iteration  of the MCMC algorithm and we cannot batch  process the objective metric across timesteps.  
  The only potential drawback of sampling using the transformer policy is that the objective $\tilde{J}(P)$  can be affected by the shift in the trajectories. However, as our experiments demonstrate, we achieve good results despite any possible distribution shift; hence, the efficiency gains far outweigh any cons. 



\subsection{Re-training the Transformer}
\label{sec:retrain}
Once our algorithm has synthesized a program $P$, we can form the combined policy $\pi_{P,\theta}$ to control the multi-agent system. One remaining issue is that the parameters $\theta$ are optimized for using the original soft attention weights $\alpha^{j\to i}$ rather than the hard attention weights $\alpha^{P,j\to i}$. Thus, we re-train the parameters of the transformer models in $\pi_{P,\theta}$. This training is identical to how $\pi_{\theta}$ was originally trained, except we use $\alpha^{P,j\to i}$ instead of $\alpha^{j\to i}$ to compute the action at each step.

\section*{Broader Impact}

Broadly speaking, reinforcement learning has the promise to significantly improve the usability of robotics in open-world settings. Our work focuses on leveraging reinforcement learning to help solve complex decentralized multi-agent planning problems, specifically by helping automate the design of communication policies that account for computational and bandwidth constraints. Solutions to these problems have a wide range of applications, both ones with positive societal impact---e.g., search and rescue, disaster response, transportation, agriculture, and constructions---and ones with controversial or negative impact---e.g., surveillance and military. These applications are broadly true of any work that improves the capabilities of multi-agent systems such as self-driving cars or drones. Restricting the capabilities of these systems based on ethical considerations is a key direction for future work.

Beyond communication constraints, security and robustness are important requirements for multi-agent systems. While we do not explicitly study these properties, a key advantage of reduced communication is to improve the resilience and robustness of the system and reduce the probability of failure, since there are fewer points of failure. Furthermore, communication policies that include stochastic rules are typically more robust since they can replace a broken communication link with another randomly selected link without sacrificing performance.

Furthermore, our research may have applications in other areas of machine learning. In general, there has been growing interest in learning programmatic representations to augment neural network models to improve interpretability, robustness, and generalizability. Along these dimensions, our work could potentially impact other applications such as NLP where transformers are state-of-the-art. In particular, our work takes a step in this direction by replacing soft attention weights in transformers with programmatic attention rules. The programmatic nature of these weights makes them much easier to interpret, as does the fact that the weights are hard rather than soft (since we now have a guarantee that parts of the input are irrelevant to the computation). 





\begin{ack}

We gratefully acknowledge support from ONR Grant N00014-17-1-2699, IBM Research,
DARPA Grant HR001120C0015, Boeing,
NSF  Grant 1917852,
ARL Grant DCIST CRA W911NF-17-2-0181, NSF Grants CNS-1521617 and CCF-1910769, ARO Grant W911NF-13-1-0350, ONR Grant N00014-20-1-2822, ONR Grant N00014-20-S-B001, DARPA Grant FA8750-19-2-0201, ARO Grant W911NF-20-1-0080,
IBM Research,
and Qualcomm Research.
The views expressed are those of the authors and do not reflect the official policy or position of the Department of Defense or the U.S. Government.

\end{ack}
\section{Introduction}

Many real-world robotics systems are distributed, with teams of agents needing to coordinate to share information and solve problems. Reinforcement learning has recently been demonstrated as a promising approach to automatically solve such multi-agent planning problems~\cite{tan1993multi,littman1994markov,foerster2016learning,lowe2017multi,foerster2018counterfactual,khan2019learning}.

A key challenge in (cooperative) multi-agent planning is how to coordinate with other agents, both deciding whom to communicate with and what information to share. One approach is to let agents communicate with all other agents; however, letting agents communicate arbitrarily can lead to poor generalization~\cite{khan2019graph,paulos2019decentralization}; furthermore, it cannot account for physical constraints such as limited bandwidth. A second approach is to manually impose a communication graph on the agents, typically based on distance~\cite{khan2019graph,tolstaya2019learning,paulos2019decentralization,shishika2020cooperative}. However, this manual structure may not reflect the optimal communication structure---for instance, one agent may prefer to communicate with another one that is farther away but in its desired path. A third approach is to use a transformer~\cite{vaswani2017attention} as the policy network~\cite{das2019tarmac}, which uses attention to choose which other agents to focus on. However, since the attention is soft, each agent still communicates with every other agent.

We study the problem of learning a communication policy that solves a multi-agent planning task while minimizing the amount of communication required. We measure the amount of communication on a given step as the maximum degree (in both directions) of the communication graph on that step; this metric captures the maximum amount of communication any single agent must perform at that step. While we focus on this metric, our approach easily extends to handling other metrics---e.g., the total number of edges in the communication graph, the maximum in-degree, and the maximum out-degree, as well as general combinations of these metrics.

A key question is how to represent the communication policy; in particular, it must be sufficiently expressive to capture communication structures that both achieve  high reward and has low communication degree, while simultaneously being easy to train. Neural network policies can likely capture good communication structures, but they are hard to train since the maximum degree of the communication graph is a discrete objective that cannot be optimized using gradient descent. An alternative is to use a structured model such as a decision tree~\cite{breiman1984classification} or rule list~\cite{wang2015falling} and train using combinatorial optimization. However, these models perform poorly since choosing whom to communicate with requires reasoning over sets of other agents---e.g., to avoid collisions, an agent must communicate with its nearest neighbor in its direction of travel.

We propose to use programs to represent communication policies. In contrast to rule lists, our programmatic polices include components such as filter and map that operate over sets of inputs. Furthermore, programmatic policies are discrete in nature, making them amenable to combinatorial optimization; in particular, we can compute a programmatic policy that minimizes the communication graph degree using a stochastic synthesis algorithm~\cite{schkufza2013stochastic} based on MCMC sampling~\cite{metropolis1953equation,hastings1970monte}. 

A key aspect of our programs is that they can include a random choice operator. Intuitively, random choice is a key ingredient needed to minimize the communication graph degree without global coordination. For example, suppose there are two groups of agents, and each agent in group $A$ needs to communicate with an agent in group $B$, but the specific one does not matter. Using a deterministic communication policy, since the same policy is shared among all agents, each agent in group $A$ might choose to communicate with the same agent $j$ in group $B$ (e.g., if agents in the same group have similar states). Then, agent $j$ will have a very high degree in the communication graph, which is undesirable. In contrast, having each agent in group $A$ communicate with a uniformly random agent in group $B$ provides a near-optimal solution to this problem, without requiring the agents to explicitly coordinate their decisions.


While we can minimize the communication graph degree using stochastic search, we still need to choose actions based on the communicated information. Thus, we propose a learning algorithm that integrates our programmatic communication policy with a transformer policy for selecting actions. We refer to the combination of the transformer and the programmatic communication policy as a \emph{neurosymbolic transformer}. This algorithm learns the two policies jointly. At a high level, our algorithm first trains a transformer policy for solving the multi-agent task; as described above, the soft attention weights capture the extent to which an edge in the communication graph is useful. Next, our algorithm trains a programmatic communication policy that optimizes both goals: (i) match the transformer as closely as possible, and (ii) minimize the maximum degree of the communication graph at each step. In contrast to the transformer policy, this communication policy makes hard decisions about which other agents to communicate with. Finally, our algorithm re-trains the weights of the transformer policy, except where the (hard) attention weights are chosen by the communication policy.


We evaluate our approach on several multi-agent planning tasks that require agents to coordinate to achieve their goals. We demonstrate that our algorithm learns communication policies that  achieve task performance similar to the original transformer policy (i.e., where each agent communicates with every other agent), while significantly reducing the amount of communication. Our results demonstrate that our algorithm is a promising approach for training policies for multi-agent systems that additionally optimize combinatorial properties of the communication graph~\footnote{The code and a video illustrating  the different tasks are available at \url{https://github.com/jinala/multi-agent-neurosym-transformers}.}

\begin{figure}
\centering
\includegraphics[width=0.32\textwidth]{01_traj.pdf}
\includegraphics[width=0.32\textwidth]{2g_tf-attn.pdf}
\includegraphics[width=0.32\textwidth]{2g_prog-retrained-attn.pdf}
\caption{Left: Two groups of agents (red vs. blue) at their initial positions (circles) trying to reach their goal positions (crosses). Agents must communicate both within group and across groups to choose a collision free path to take (solid line shows a path for a single agent in each group). Middle: The soft attention weights of the transformer policy computed by the agent along the $y$-axis for the message received from the agent along the $x$-axis for the initial step.
Right: The hard attentions learned by the programmatic communication policy to imitate the transformer.}
    \label{fig:run_ex}
\end{figure}
\textbf{Example.}
Consider the example in Figure~\ref{fig:run_ex}, where agents in group $1$ (blue) are navigating from the left to their goal on the right, while agents in group $2$ (red) are navigating from the right to their goal on the left. In this example, agents have noisy observations of the positions of other agents (e.g., based on cameras or LIDAR); however, they do not have access to internal information such as their planned trajectories or even their goals.
Thus, to solve this task, each agent must communicate both with its closest neighbor in the same group (to avoid colliding with them), as well as with any agent in the opposite group (to coordinate so their trajectories do not collide). The communication graph of the transformer policy (in terms of soft attention weights) is shown in Figure~\ref{fig:run_ex} (middle); every agent needs to communicate with all other agents. The programmatic communication policy synthesized by our algorithm is
\begin{align*}
\text{argmax}(\text{map}(-d^{i,j}, \text{filter}(\theta^{i,j} \ge -1.85, l))),\qquad
\text{random}(\text{filter}(d^{i,j} \ge 3.41, l)).
\end{align*}
Agent $i$ uses this program to choose two other agents $j$ from the list of agents $\ell$ from whom to request information; $d^{i,j}$ is the distance between them and $\theta^{i,j}$ is the angle between them. The first rule chooses the nearest agent $j$ (besides itself) such that $\theta^{i,j}\in[-1.85, \pi]$, and the second chooses a random agent in the other group. The communication graph is visualized in Figure~\ref{fig:run_ex} (right).




\textbf{Related work.}
There has been a great deal of recent interest in using reinforcement learning to automatically infer good communication structures for solving multi-agent planning problems~\cite{khan2019graph,tolstaya2019learning,paulos2019decentralization,das2019tarmac,shishika2020cooperative}. Much of this work focuses on inferring what to communicate rather than whom to communicate with; they handcraft the communication structure to be a graph (typically based on distance)~\cite{khan2019graph,tolstaya2019learning,paulos2019decentralization}, and then use a graph neural network~\cite{scarselli2008graph,kipf2016semi} as the policy network. There has been some prior work using transformer network to infer the communication graph~\cite{das2019tarmac}; however, they rely on soft attention, so the communication graph remains fully connected. Prior work~\cite{singh2018learning} frames the multi-agent communication problem as a MDP problem where the decisions of  when to communicate are  part of  the action space.  However, in  our case,  we want to learn who to communicate with in-addition to when to communicate. This  results in a large discrete action space and we found that RL algorithms perform poorly on this space. Our proposed approach addresses  this challenge by using the transformer as a teacher. 

There has also been a great deal of interest using an oracle (in our case, the transformer policy) to train a policy (in our case, the programmatic communication policy)~\cite{levine2013guided,montgomery2016guided}. In the context of multi-agent planning, this approach has been used to train a decentralized control policy using a centralized one~\cite{tolstaya2019learning}; however, their communication structure is manually designed.

Finally, in the direction of program synthesis, there has been much recent interest in leveraging program synthesis in the context of machine learning~\cite{ellis2015unsupervised,ellis2018learning,valkov2018houdini,pu2018selecting,young2019learning,liu2019learning}. Specifically in the context of reinforcement learning, it has been used to synthesize programmatic control policies that are more interpretable~\cite{verma2018programmatically,verma2019imitation}, that are easier to formally verify~\cite{bastani2018verifiable,sadraddini2019polytopic}, or that generalize better~\cite{inala2020synthesizing}; to the best of our knowledge, none of these approaches have considered multi-agent planning problems.






\section{Introduction}

Real-world robotics systems are often distributed in nature, with teams of agents needing to coordinate to share information and solve problems. Reinforcement learning has recently been demonstrated as a promising approach to automatically solve such multi-agent planning problems~\cite{tan1993multi,littman1994markov,foerster2016learning,lowe2017multi,foerster2018counterfactual,khan2019learning}.

A key challenge in (cooperative) multi-agent planning is how to coordinate with other agents, both deciding whom to communicate with and what information to share. One approach is to let agents communicate with all other agents; however, letting agents communicate arbitrarily can lead to poor generalization~\cite{khan2019graph,paulos2019decentralization}; furthermore, it cannot account for physical constraints such as limited bandwidth. A second approach is to manually impose a communication graph on the agents, typically based on distance~\cite{khan2019graph,tolstaya2019learning,paulos2019decentralization,shishika2020cooperative}. However, this manual structure may not reflect the optimal communication structure---for instance, one agent may prefer to communicate with another one that is farther away but in its desired path. A third approach is to use a transformer~\cite{vaswani2017attention} as the policy network~\cite{das2019tarmac}, which uses attention to choose which other agents to focus on. However, since the attention is soft, each agent still communicates with every other agent.

We study the problem of learning a communication policy that solves a multi-agent planning task while minimizing the amount of communication required. We measure the amount of communication on a given step as the maximum degree (in both directions) of the communication graph on that step; this metric captures the maximum amount of communication any single agent must perform at that step. While we focus on this metric, our approach easily extends to handling other metrics---e.g., the total number of edges in the communication graph, the maximum in-degree, and the maximum out-degree, as well as general combinations of these metrics.

The challenge with optimizing the maximum degree is that it is a global property of the communication graph. For instance, suppose each agent independently chooses whose information they want to receive; then, they may all choose to communicate with one agent that has highly desired information, but then this agent has very high degree in the communication graph. Thus, the agents must implicitly coordinate their decisions to minimize communication.

We propose an approach to addressing this problem. At a high level, our algorithm first trains a transformer policy for solving the multi-agent task; as described above, the soft attention weights capture the extent to which an edge in the communication graph is useful. Next, our algorithm trains a communication policy that optimizes both goals: (i) match the transformer as closely possible, and (ii) minimize the maximum degree of the communication graph at each step. In contrast to the transformer policy, this communication policy makes hard decisions about which other agents to communicate with. Finally, our algorithm re-trains the weights of the transformer policy, except where the (hard) attention weights are chosen by the communication policy.

A key question is how to train the communication policy. Training a neural network policy is hard since the maximum degree of the communication graph is a discrete objective that cannot be optimized using gradient descent. An alternative is to train a structured model such as a decision tree~\cite{breiman1984classification} or rule list~\cite{wang2015falling} using combinatorial optimization. However, these models perform poorly since choosing whom to communicate with requires reasoning over sets of other agents---e.g., to avoid collisions, an agent must communicate with its nearest neighbor in its direction of travel.

Instead, our algorithm synthesizes \emph{programmatic} communication policies; in contrast to rule lists, these polices include components such as filter and map that operate over sets of inputs. In particular, we define a domain-specific language (DSL) encoding programs that each select a fixed number of other agents to communicate with. Then, we use a stochastic synthesis algorithm~\cite{schkufza2013stochastic} based on MCMC sampling~\cite{metropolis1953equation,hastings1970monte} to synthesize programs that optimize both objectives (i) and (ii) above.

In addition to standard list-processing components, an important construct in our DSL is a random choice operator. Intuitively, random choice is a key ingredient needed to minimize the communication graph degree without global coordination. For example, if suppose there are two groups of agents, and each agent in group $A$ needs to communicate with an agent in group $B$, but the specific one does not matter. Using a deterministic communication policy, each agent in group $A$ might choose to communicate with the same agent in group $B$ (assuming it is hard for the policy distinguish agents within a group). In contrast, having each agent in group $A$ communicate with a uniformly random agent in group $B$ provides a near-optimal solution to this problem.

We evaluate our approach on several multi-agent planning tasks that require agents to coordinate to achieve their goals. We demonstrate that our algorithm learns communication policies that  achieve task performance similar to the original transformer policy (i.e., where each agent communicates with every other agent), while significantly reducing the amount of communication. Our results demonstrate that our algorithm is a promising approach for training policies for multi-agent systems that additionally optimize combinatorial properties of the communication graph.

\begin{figure}
\centering
\includegraphics[width=0.32\textwidth]{01_traj.pdf}
\includegraphics[width=0.32\textwidth]{2g_tf-attn.pdf}
\includegraphics[width=0.32\textwidth]{2g_prog-retrained-attn.pdf}
\caption{Left: Two groups of agents (red vs. blue) at their initial positions (circles) trying to reach their goal positions (crosses). Agents must communicate both within group and across groups to choose a collision free path to take (solid line shows a path for a single agent in each group). Middle: The soft attention weights of the transformer policy computed by the agent along the $y$-axis for the message received from the agent along the $x$-axis for the initial step. Group 1 attends to messages from all agents, whereas Group 2 only attends to messages from Group 1. Right: The hard attentions learned by the programmatic communication policy to imitate the transformer.}
    \label{fig:run_ex}
\end{figure}
\textbf{Example.}
Consider the example in Figure~\ref{fig:run_ex}, where agents in group $1$ (blue) are navigating from the left to their goal on the right, while agents in group $2$ (red) are navigating from the right to their goal on the left. In this example, agents have noisy observations of the positions other agents (e.g., based on cameras or LIDAR); however, they do not have access to internal information such as their planned trajectories or even their goals.
Thus, to solve this task, each agent must communicate both with its closest neighbor in the same group (to avoid colliding with them), as well as with any agent in the opposite group (to coordinate so their trajectories so they do not collide). The communication graph of the transformer policy is shown in Figure~\ref{fig:run_ex} (middle); as can be seen, the graph is not too dense, but it nevertheless has high degree since every agent in group $1$ communicates with every agent in group $2$ and vice versa. The programmatic communication policy synthesized by our algorithm is
\begin{align*}
\text{argmax}(\text{map}(-d^{i,j}, \text{filter}(\theta^{i,j} \ge -1.85, l))),\qquad
\text{random}(\text{filter}(d^{i,j} \ge 3.41, l)).
\end{align*}
Agent $i$ uses this program to choose two other agents $j$ from the list of agents $\ell$ from whom to request information; $d^{i,j}$ is the distance between them and $\theta^{i,j}$ is the angle between them. The first rule chooses the nearest agent $j$ (besides itself) such that $\theta^{i,j}\in[-1.85, 3.14]$, and the second chooses a random agent in the other group. The communication graph is visualized in Figure~\ref{fig:run_ex} (right).




\textbf{Related work.}
There has been a great deal of recent interest in using reinforcement learning to automatically infer good communication structures for solving multi-agent planning problems~\cite{khan2019graph,tolstaya2019learning,paulos2019decentralization,das2019tarmac,shishika2020cooperative}. Much of this work focuses on inferring what to communicate rather than whom to communicate with; they handcraft the communication structure to be a graph (typically based on distance)~\cite{khan2019graph,tolstaya2019learning,paulos2019decentralization}, and then use a graph neural network~\cite{scarselli2008graph,kipf2016semi} as the policy network. There has been some prior work using transformer network to infer the communication graph~\cite{das2019tarmac}; however, they rely on soft attention, so the communication graph remains fully connected.

There has also been a great deal of interest using an oracle (in our case, the transformer policy) to train a policy (in our case, the programmatic communication policy)~\cite{levine2013guided,montgomery2016guided}. In the context of multi-agent planning, this approach has been used to train a decentralized control policy using a centralized one~\cite{tolstaya2019learning}; however, their communication structure is manually designed.

Finally, in the direction of program synthesis, there has been much recent interest in leveraging program synthesis in the context of machine learning~\cite{ellis2015unsupervised,ellis2018learning,valkov2018houdini,pu2018selecting,young2019learning,liu2019learning}. Specifically in the context of reinforcement learning, it has been used to synthesize programmatic control policies that are more interpretable~\cite{verma2018programmatically,verma2019imitation}, that easier to formally verify~\cite{bastani2018verifiable,sadraddini2019polytopic}, or that generalize better~\cite{inala2020synthesizing}; to the best of our knowledge, none of these approaches have considered multi-agent planning problems.





\section{Problem Formulation}
\label{sec:problem}

\newcommand{N}{N}
\newcommand{K}{K}
\newcommand{\mathcal{S}}{\mathcal{S}}
\newcommand{\mathcal{A}}{\mathcal{A}}
\newcommand{\mathcal{O}}{\mathcal{O}}

\newcommand{s}{s}
\newcommand{o}{o}

\newcommand{\pi}{\pi}
\newcommand{a}{a}
\newcommand{f}{f}
\newcommand{h}{h}
\newcommand{\mathcal{P}_0}{\mathcal{P}_0}
\newcommand{\mathbb{R}}{\mathbb{R}}
\newcommand{\mathcal{C}}{\mathcal{C}}
\newcommand{c}{c}
\newcommand{\mathcal{M}}{\mathcal{M}}
\newcommand{m}{m}

\newcommand{S}{S}
\newcommand{O}{O}
\newcommand{C}{C}
\newcommand{M}{M}
\newcommand{A}{A}
\newcommand{R}{R}

We formulate the multi-agent planning problem as a decentralized partially observable Markov decision process (POMDP). We consider $N$ agents $i\in[N]=\{1,...,N\}$ with states $s^i\in\mathcal{S}\subseteq\mathbb{R}^{d_S}$, actions $a^i\in\mathcal{A}\subseteq\mathbb{R}^{d_A}$, and observations $o^{i,j}\in\mathcal{O}\subseteq\mathbb{R}^{d_O}$ for every pair of  agents ($j\in[N]$). Following prior work~\cite{lowe2017multi, foerster2018counterfactual, das2019tarmac}, we operate under the premise of centralized training and decentralized execution. Hence, during training the POMDP has global states $\mathcal{S}^N$, global actions $\mathcal{A}^N$, global observations $\mathcal{O}^{N\timesN}$, transition function $f:\mathcal{S}^N\times\mathcal{A}^N\to\mathcal{S}^N$, observation function $h:\mathcal{S}^N\to\mathcal{O}^{N\timesN}$, initial state distribution $s_0\sim\mathcal{P}_0$, and reward function $r:\mathcal{S}^N\times\mathcal{A}^N \to \mathbb{R}$.



The agents all use the same policy $\pi=(\pi^C,\pi^M,\pi^A)$ divided into a \emph{communication policy} $\pi^C$ (choose other agents from whom to request information), a \emph{message policy} $\pi^M$ (choose what messages to send to other agents), and an \emph{action policy} $\pi^A$ (choose what action to take). Below, we describe how each agent $i\in[N]$ chooses its action $a^i$ at any time step. 

\textbf{Step 1 (Choose communication).}
The communication policy $\pi^C:\mathcal{S}\times\mathcal{O}^N\to\mathcal{C}^K$ inputs the state $s^i$ of current agent $i$ and its observations $o^i=(o^{i,1},...,o^{i,N})$, and outputs $K$ other agents $c^i=\pi^C(s^i,o^i)\in\mathcal{C}^K=[N]^K$ from whom to request information. The \emph{communication graph} $c=(c^1,...,c^N)\in\mathcal{C}^{N\timesK}$ is the directed graph $G=(V,E)$ with nodes $V=[N]$ and edges $E=\{j\to i \mid (i,j) \in[N]^2 \wedge j\in\pi^C(s^i,o^i)\}$. For example, in the communication graph in Figure~\ref{fig:run_ex} (right), $c^0 = (1, 19)$---i.e., agent $0$ in group $1$ receives messages from agent $1$ in group $1$ and agent $19$ in group $2$.

\textbf{Step 2 (Choose and send/receive messages).}
For every other agent $j\in[N]$, the message policy $\pi^M:\mathcal{S}\times\mathcal{O}\to\mathcal{M}$ inputs $s^i$ and $o^{i,j}$ and outputs a message $m^{i\to j}=\pi^M(s^i,o^{i,j})$
to be sent to $j$ if requested. Then, agent $i$ receives messages $m^i=\{m^{j\to i}\mid j\inc^i\}\in\mathcal{M}^K$.

\textbf{Step 3 (Choose action).}
The action policy $\pi^A:\mathcal{S}\times\mathcal{O}^N\times\mathcal{M}^K\to\mathcal{A}$ inputs $s^i$, $o^i$, and $m^i$, and outputs action $a^i=\pi^A(s^i,o^i,m^i)$ to take.

\textbf{Sampling a trajectory/rollout.}
Given initial state $s_0\sim\mathcal{P}_0$ and time horizon $T$, $\pi$ generates the trajectory $(s_0,s_1,...,s_T)$, where $o_t=h(s_t)$ and $s_{t+1}=f(s_t,a_t)$, and where for all $i\in[N]$, we have $c^i_t=\pi^C(s^i_t,o^i_t)$, $m^i_t=\{\pi^M(s^j_t,o^{j,i}_t)\mid j\inc^i_t\}$, and $a^i_t=\pi^A(s^i_t,o^i_t,m^i_t)$.

\textbf{Objective.}
Then, our goal is to train a policy $\pi$ that maximizes the objective
\begin{align*}
J(\pi)=J^R(\pi)+\lambda J^C(\pi)=\mathbb{E}_{s_0\sim\mathcal{P}_0}\left[\sum_{t=0}^T\gamma^tr(s_t,a_t)\right]-\lambda\mathbb{E}_{s_0\sim\mathcal{P}_0}\left[\sum_{t=0}^T\max_{i\in[N]}\text{deg}(i;c_t)\right],
\end{align*}
where $\lambda\in\mathbb{R}_{>0}$ is a hyperparameter, the reward objective $J^R$ is the time-discounted expected cumulative reward over time horizon $T$ with discount factor $\gamma\in(0,1)$, and the communication objective $J^C$ is to minimize the degree of the communication graph,
where $c_t$ is the communication graph on step $t$, and $\text{deg}(i;c_t)$ is the sum of the incoming and outgoing edges for node $i$ in $c_t$. Each agent computes its action based on just its state, its observations of other agents, and communications received from the other agents; thus, the policy can be executed in a decentralized way.

\textbf{Assumptions on the observations of other agents.}
We assume that $o^{i,j}$ is available through visual observation (e.g., camera or LIDAR), and therefore does not require extra communication. In all experiments, we use $o^{i,j} = x^j - x^i + \epsilon^{i,j}$---i.e., the position $x^j$ of agent $j$ relative to the position $x^i$ of agent $i$, plus i.i.d. Gaussian noise $\epsilon^{i,j}$. This information can often be obtained from visual observations (e.g., using an object detector); $\epsilon^{i,j}$ represents noise in the visual localization process.

The observation $o^{i,j}$ is necessary since it forms the basis for agent $i$ to decide whether to communicate with agent $j$; if it is unavailable, then $i$ has no way to distinguish the other agents. If $o^{i,j}$ is unavailable for a subset of agents $j$ (e.g., they are outside of sensor range), we could use a mask to indicate that the data is missing. We could also replace it with alternative information such as the most recent message from $j$ or the most recent observation of $j$.

We emphasize that $o^{i,j}$ does not contain important internal information available to the other agents---e.g., their chosen goals and their planned actions/trajectories. This additional information is critical for the agents to coordinate their actions and the agents must learn to communicate such information.

\section{Experiments}
\label{sec:experiments}

\begin{figure*}
    \centering
    \begin{subfigure}[b]{0.32\textwidth}
         \centering
         \includegraphics[width=\textwidth]{all-cross.pdf}
         \caption{}
         \label{fig:stats-cross}
     \end{subfigure}
     \begin{subfigure}[b]{0.32\textwidth}
         \centering
         \includegraphics[width=\textwidth]{all-grid.pdf}
         \caption{}
         \label{fig:stats-grid}
     \end{subfigure}
     \begin{subfigure}[b]{0.32\textwidth}
         \centering
         \includegraphics[width=\textwidth]{all-covering.pdf}
         \caption{}
         \label{fig:stats-grid}
     \end{subfigure}
    \caption{Statistics of cumulative loss and communication graph degrees across baselines, for (a) \texttt{random-cross}, (b) \texttt{random-grid}, and (c) \texttt{unlabeled-goals}. We omit communication degrees for \texttt{tf-full}, since it requires communication between all pairs of agents.}
    \label{fig:formation-stats}
\end{figure*}

\textbf{Formation task.}
We consider multi-agent formation flying tasks in 2D space~\cite{khan2019graph}. Each agent has a starting position and an assigned goal position. The task is to learn a decentralized policy for the agents to reach the goals while avoiding collisions. The agents are arranged into a small number of groups (between 1 and 4): starting positions for agents within a group are close together, as are goal positions.
Each agent's state $s^i$ contains its current position $x^i$ and goal position $g^i$. The observations $o^{i,j}=x^j - x^i + \epsilon^{i,j}$ are the relative positions of the other agents, corrupted by i.i.d. Gaussian noise $\epsilon^{i,j}\sim \mathcal{N}(0,\sigma^2)$. The actions $a^i$ are agent velocities, subject to $\|a^i\|_2\le v_{\text{max}}$.
The reward at each step is $r(s,a) = r^g(s,a) - r^c(s,a)$, where the goal reward $r^g(s,a) = - \sum_{i\in[N]} \| x^i - g^i \|_2$ is the negative sum of distances of all agents to their goals, and the collision penalty $r^c(s,a) = \sum_{i,j\in [N], i\neq j} \max\{ p_c(2 - \| x^i - x^j \|_2/d_c), 0\}$ is the hinge loss between each pair of agents, where $p_c$ is the collision penalty weight and $d_c$ is the collision distance.


We consider two instances. First, \texttt{random-cross} contains up to 4 groups; each possible group occurs independently with probability $0.33$. The starting positions in each group (if present) are sampled uniformly randomly inside 4 boxes with center $b$ equal to $(-\ell, 0)$, $(0, -\ell)$, $(\ell, 0)$, and $(0, \ell)$, respectively, and the goal positions of each group are sampled randomly from boxes with centers at $-b$.
The challenge is that agents in one group must communicate with agents in other groups to adaptively choose the most efficient path to their goals.  Second, \texttt{random-grid} contains 3 groups with starting positions sampled in boxes centered at $(-\ell, 0)$, $(0,0)$, and $(\ell,0)$, respectively, and the goal positions are sampled in boxes centered at randomly chosen positions $(b_x,b_y)\in\{-\ell,0,\ell\}^2$ (i.e., on a $3\times3$ grid), with the constraint that the starting box and goal box of a group are adjacent and the boxes are all distinct. The challenge is that each agent must learn whom to communicate with depending on its goal.

\textbf{Unlabeled goals task.} 
This task is a cooperative navigation task with unlabeled goals~\cite{lowe2017multi} that has $N$ agents along with $N$ goals at positions $g_1,...,g_N$ (see Figure~5 in the appendix). The task is to drive the agents to cover as many goals as possible. We note that this task is not just a navigation task. Since the agents are not pre-assigned  to goals, there is a combinatorial aspect where they must communicate to assign themselves to different goals. The agent state $s^i$ is its own position $x^i$ and the positions of the goals (ordered by distance at the initial time step). The observations $o^{i,j}$ are the relative positions to the other agents, corrupted by Gaussian noise. The actions $a^i = (p_1^i, \cdots, p_l^i, \cdots, p_N^i)$ are the weights (normalized to $1$) over the goals; the agent moves in the direction of the weighted sum of goals---i.e., its velocity is $a^i=\sum_{k\in[N]}p_k^i(g_k - x^i)$. The reward is $r(s, a) = \sum_{k\in[N]} \text{max}_{i\in[N]} p^i_k - N$---i.e., the sum over goals of the maximum weight that any agent assigns to that goal minus $N$.

\begin{figure*}
\centering
    \begin{subfigure}[b]{0.49\textwidth}
         \centering
         \includegraphics[width=0.49\textwidth]{0123_traj.pdf}
         \includegraphics[width=0.49\textwidth]{2_traj.pdf}
         \caption{}
         \label{fig:traj}
     \end{subfigure}
     \begin{subfigure}[b]{0.5\textwidth}
         \centering
         \includegraphics[width=0.49\textwidth]{hard-attn-attn.pdf}
         \includegraphics[width=0.49\textwidth]{prog-retrained-attn.pdf}
         \caption{}
         \label{fig:attn}
     \end{subfigure}
\caption{(a) For \texttt{random-cross}, trajectories taken by each group (i.e., averaged over all agents in that group) when all four groups are present (left) and only one group is present (right), by \texttt{prog-retrained} (solid) and \texttt{dist} (dashed). Initial positions are circles and goal positions are crosses. (b) Attention weights for \texttt{hard-attn} and \texttt{prog-retrained} at a single step near the start of a rollout, computed by the agent along the $y$-axis for the message from the agent along the $x$-axis.}
\label{fig:fourgroup}
\end{figure*}

\textbf{Baselines.}
We consider four baselines. (i) Fixed communication (\texttt{dist}): A transformer, but where other agents are masked based on distance so each agent can only attend to its $k$ nearest neighbors. We find this model outperforms GCNs with the same communication structure~\cite{khan2019graph}, since its attention parameters enable each agent to re-weight the messages it receives.
(ii) Transformer (\texttt{tf-full}): The oracle transformer policy from Section~\ref{sec:oracle}; here, each agent communicates with all other agents. (iii) Transformer + hard attention (\texttt{hard-attn}): The transformer policy, but where the communication degree is reduced by constraining each agent to only receive messages from the $k$ other agents with the largest attention scores. Note that this approach only minimizes the maximum in-degree, not necessarily the maximum out-degree; minimizing both would require a centralized algorithm. %
(iv) Transformer + program (\texttt{prog}): An ablation of our approach that does not retrain the transformer after synthesizing the programmatic communication policy (i.e., it skips the step in Section~\ref{sec:retrain}). (v) Transformer + retrained program (\texttt{prog-retrain}): Our full approach.

The tasks \texttt{random-cross} and \texttt{random-grid} perform $1$ round of communications per time step for all the baselines, while the $\texttt{unlabeled-goals}$ task uses $2$ rounds of communications.
For all approaches, we train the model with 10k rollouts.
For synthesizing the programmatic policy, we build a dataset using 300 rollouts and run MCMC for 10000 steps. We retrain the transformer with 1000 rollouts.
We constrain the maximum in-degree to be a constant $d_0$ across all approaches (except \texttt{tf-full}, where each agent communicates with every other agent); for \texttt{dist} and \texttt{hard-attn}, we do so by setting the communication neighbors to be $k=d_0$, and for \texttt{prog} and \texttt{prog-retrain}, we choose the number of rules to be $K=d_0$.
This choice ensures fair comparison across approaches.


\textbf{Results.}
We measure performance using both the loss (i.e., negative reward) and maximum communication degree (i.e., maximum degree of the communication graph), averaged over the time horizon. Because the in-degree of every agent is constant, the maximum degree equals the in-degree plus the maximum out-degree. Thus, we report the maximum in-degree and the maximum out-degree separately. Results are in
Figure \ref{fig:formation-stats}; we report mean and standard deviation over 20 random seeds. 

For \texttt{random-cross} and \texttt{random-grid} tasks, our approach achieves loss similar to the best loss (i.e., that achieved by the full transformer), while simultaneously achieving the best communication graph degree. In general, approaches that learn communication structure (i.e., \texttt{tf-full}, \texttt{hard-attn}, and \texttt{prog-retrained}) perform better than having a fixed communication structure (i.e., \texttt{dist}). In addition, using the programmatic communication policy is more effective at reducing the maximum degree (in particular, the maximum out-degree) compared with thresholding the transformer attention (i.e., \texttt{hard-attn}). Finally, retraining the transformer is necessary for the programmatic communication policy to perform well in terms of loss. 
For \texttt{unlabeled-goals} task, our approach performs almost similar to \texttt{dist}  baseline and slightly worse than \texttt{tf-full} baseline, but achieves a smaller maximum degree. Moreover, the loss is significantly lower than the loss of 4.13 achieved when no communications are allowed. 
We give additional results and experimental details in the appendix.




\textbf{Learning whom to communicate with.}
Figure \ref{fig:fourgroup}a shows two examples from \texttt{random-cross}: all four groups are present (left), and only a single group is present (right). In the former case, the groups must traverse complex trajectories to avoid collisions, whereas in the latter case, the single group can move directly to the goal. However, with a fixed communication structure, the policy \texttt{dist} cannot decide whether to use the complex trajectory or the direct trajectory, since it cannot communicate with agents in other groups to determine if it should avoid them. Thus, it always takes the complex trajectory. In contrast, our approach successfully decides between the complex and direct trajectories. 

\textbf{Reducing the communication degree.}
Figure \ref{fig:fourgroup}b shows the attention maps of \texttt{hard-attn} and \texttt{prog-retrained} for the \texttt{random-cross} task at a single step. Agents using \texttt{hard-attn} often attend to messages from a small subset of agents; thus, even if the maximum in-degree is low, the maximum out-degree is high---i.e., there are a few agents that must send messages to many other agents. In contrast, \texttt{prog-retrained} uses randomness to distribute communication across agents.

\textbf{Understanding the learned program policy.}
Figure \ref{fig:rules} visualizes a programmatic communication policy for \texttt{random-grid}. Here, (a) and (b) visualize the nondeterministic rule for two different configurations. As can be seen, the region from which the rule chooses an agent (depicted in orange) is in the direction of the goal of the agent, presumably to perform longer-term path planning. The deterministic rule (Figure \ref{fig:rule-2-rlu}) prioritizes choosing a nearby agent, presumably to avoid collisions. Thus, the rules focus on communication with other agents relevant to planning.

\begin{figure*}
\centering
    \begin{subfigure}[b]{0.33\textwidth}
         \centering
         \includegraphics[width=0.98\textwidth]{quadratic_500_rule_1_rlu_24.pdf}
         \caption{}
         \label{fig:rule-1-rlu}
     \end{subfigure}
     \begin{subfigure}[b]{0.33\textwidth}
         \centering
         \includegraphics[width=0.98\textwidth]{quadratic_500_rule_1_url_24.pdf}
         \caption{}
         \label{fig:rule-1-url}
     \end{subfigure}
     \begin{subfigure}[b]{0.32\textwidth}
         \centering
         \includegraphics[width=\textwidth]{quadratic_500_rule_2_rlu_24.pdf}
         \caption{}
         \label{fig:rule-2-rlu}
     \end{subfigure}
\caption{Visualization of a programmatic communication policy for \texttt{random-grid}, which has two rules---one determinitsic and one nondeterministic. The red circle denotes the agent making the decision, the red cross denotes its goal, and the green circle denotes the agent selected by the rule. (a,b) Visualization of the nondeterministic rule for two configurations with different goals; orange denotes the region where the filter condition is satisfied (i.e., the rule chooses a random agent in this region). (c) Visualization of the deterministic rule, showing the score output by the map operator; darker values are higher (i.e., the rule chooses the agent with the darkest value).}
\label{fig:rules}
\end{figure*}

\section{Conclusion}

We have proposed an approach for synthesizing programmatic communication policies for decentralized control of multi-agent systems. Our approach performs as well as state-of-the-art transformer policies while significantly reducing the amount of communication required to achieve complex multi-agent planning goals. We leave much room for future work---e.g., exploring other measures of the amount of communication, better understanding what information is being communicated, and handling environments with more complex observations such as camera images or LIDAR scans.
