SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jiangxing Wang
jiangxiw@stu.pku.edu.cn
Deheng Ye
Tencent Inc.
dericye@tencent.com
Zongqing Lu†
BAAI
zongqing.lu@pku.edu.cn
Despite the success of cooperative multi-agent reinforcement learning algorithms, most of them focus on a single team composition, which prevents them from being used in more realistic scenarios where dynamic team composition is possible. While some studies attempt to solve this problem via multi-task learning in a fixed set of team compositions, there is still a risk of overfitting to the training set, which may lead to catastrophic performance when facing dramatically varying team compositions during execution. To address this problem, we propose to use mutual information (MI) as an augmented reward to prevent individual policies from relying too much on team-related information and encourage agents to learn policies that are robust in different team compositions. Optimizing this MI-augmented objective in an off-policy manner can be intractable due to the existence of dynamic marginal distribution. To alleviate this problem, we first propose a multi-agent policy iteration algorithm with a fixed marginal distribution and prove its convergence and optimality. Then, we propose to employ the Blahut-Arimoto algorithm and an imaginary team composition distribution for optimization with approximate marginal distribution as the practical implementation. Empirically, our method demonstrates strong zero-shot generalization to dynamic team compositions in complex cooperative tasks.
The cooperative multi-agent reinforcement learning (MARL) problem has attracted the attention of many researchers for being a well-abstracted model for many real-world problems, such as traffic signal control (Wang et al., 2021), autonomous warehouse (Zhou et al., 2021), and even AutoML (Wang et al., 2022) as the feedback to the machine learning community. In a cooperative MARL problem, we aim to train a group of agents that can cooperate to achieve a common goal. Such a common goal is often defined by a global reward function that is shared among all agents. Although this objective is naturally centralized, we want agents to be able to execute in a fully decentralized manner. Under such a requirement, Kraemer and Banerjee (2016) proposed the centralized training with decentralized execution (CTDE) framework, where a centralized critic is learned to evaluate the performance of the joint policy in terms of the global reward and a group of decentralized individual policies are learned via the centralized critic to realize decentralized execution.
With a centralized critic, multi-agent policy gradient methods directly use it to guide the update of each decentralized individual policy. Based on this idea, a series of studies (Lowe et al., 2017; Kuba et al., 2022; Yu et al., 2022; Ye et al., 2020) have been proposed with different optimization
techniques for policy improvement. On the other hand, as the centralized critic is used to guide the learning of decentralized individual policies, many CTDE algorithms choose to factorize the centralized critic into decentralized individual utilities via the mixer network. This line of research is called value decomposition (Sunehag et al., 2018). Based on different design choices of the mixer network, a variety of value decomposition methods (Sunehag et al., 2018; Rashid et al., 2018; Wang et al., 2020; Zhang et al., 2021) has been proposed and achieved great success in cooperative MARL problems.
Despite the success of CTDE methods, previous research mainly focuses on training agents under a fixed team composition, which has been shown to exhibit serious overfitting issues (Wen et al., 2022), and leads to catastrophic performance when facing unseen team compositions. One natural way to solve this problem is to introduce multiple team compositions during training to alleviate overfitting. However, due to the existence of the mixer network and the network structure of individual utilities and policies, such a strategy cannot be simply applied to many existing CTDE methods.
REFIL (Iqbal et al., 2021) attempts to address this problem and proposes to use multi-head attention (Vaswani et al., 2017) in the mixer network and individual utilities to handle dynamic team compositions. It further proposes an imaginary objective based on random sub-group partitioning to accelerate the training process given a fixed set of team compositions. Although REFIL is able to handle dynamic team compositions, as it is still trained on a fixed set of team compositions, there is still a risk of overfitting to not a single team composition, but a set of team compositions, which has been demonstrated in several previous studies (Liu et al., 2021; Shao et al., 2022).
The overfitting issue can be attributed to the lack of robustness in different team compositions. When facing an arbitrary team composition during execution, one agent's observed information about team composition can arbitrarily vary from what it experienced during training. If the agent puts too much credit on this highly varied information to make decisions, it may fail to achieve robust behavior. As the variation of team composition is uncontrollable, one way to achieve robust behavior is to reduce the reliance on team-related information. Based on this intuition, we propose MIPI (Mutual-Information Regularized Multi-Agent Policy Iteration), minimizing the mutual information between the policy of the agent and the team-related information to encourage robust behavior of each agent. Inspired by SAC (Haarnoja et al., 2018), we combine the global reward of environment and mutual information of each agent as a new objective and learn individual policies to optimize them at the same time. As the incorporation of mutual information imposes a challenge on optimization due to the existence of dynamic marginal distribution, we first propose a multi-agent policy iteration algorithm with a fixed marginal distribution and prove its convergence and optimality. Then, we propose to utilize the Blahut-Arimoto algorithm (Cover, 1999) and an imaginary team composition distribution for optimization under an approximate dynamic marginal distribution as the practical implementation.
To empirically justify our algorithm, we first evaluate the performance of MIPI in a simple yet challenging matrix game. Compared with the other two baselines using pure environmental reward and entropy-augmented reward, using mutual information as an augmented reward can help the agent find the policy that can achieve consistent performance across different team compositions. Then, we move to a more complicated scenario, StarCraft Micromanagement Tasks. While having the same level of performance in the training set, MIPI achieves better zero-shot generalization results in unseen team compositions during evaluation.
2.1 Centralized Training with Decentralized Execution (CTDE)
CTDE methods can be categorized into value decomposition and multi-agent policy gradient, depending on whether a centralized critic is decomposed or not. For value decomposition methods, a centralized critic is decomposed into decentralized utilities through the mixer network. Different mixers have been proposed as different interpretations of the Individual-Global-Maximum (Rashid et al., 2018) (IGM) principle or its equivalence, which ensures the consistency between optimal local actions and optimal joint action. VDN (Sunehag et al., 2018) and QMIX (Rashid et al., 2018) give sufficient conditions for IGM by additivity and monotonicity, respectively. QPLEX (Wang et al., 2020) and FOP (Zhang et al., 2021) take advantage of duplex dueling architecture to guarantee IGM.
In multi-agent policy gradient, a centralized critic function is directly used to guide the update of each decentralized individual policy. Most multi-agent policy gradient methods can be considered as an extension of the policy gradient method from RL to MARL. For example, MAPPDG (Lowe et al., 2017) extends DDPG (Lillicrap et al., 2016), HATRPO (Kuba et al., 2022) extends TRPO (Schulman et al., 2015), MAPPO (Yu et al., 2022) and CoPPO (Wu et al., 2021) extend PPO (Schulman et al., 2017).
2.2 Dynamic Team Composition
While classical CTDE methods mainly focus on fixed team compositions, in real-world applications, agents that can adapt to dynamic team composition are preferable. To address this problem, REFIL (Iqbal et al., 2021) incorporates multi-head attention (Vaswani et al., 2017) into the networks and further introduces an imaginary objective based on random sub-group partitioning to accelerate the training process on a fixed set of team compositions. While REFIL learns policies that can handle dynamic team compositions, studies (Liu et al., 2021; Shao et al., 2022) suggest that it generalizes poorly on unseen team compositions. The necessity of training agents that can generalize to unseen team compositions is evident using the automated warehouse as an example, where it is very common to add more agents (more agents got purchased) or delete some agents (some agents got broken). Therefore, agents have to deal with different teams in the application scenario, which can not be fully covered during training. In order to learn policies that can adapt to unseen team compositions, many studies choose to sacrifice the requirement of decentralized execution. For example, a centralized agent is assumed in COPA (Liu et al., 2021), which has a global view of the environment and coordinates agents by distributing individual strategies. In SOG (Shao et al., 2022), a communication channel is assumed to elect conductors, so that the corresponding groups are constructed with conductor-follower consensus. Unlike these methods, we do not assume any kind of centralization during execution, such that the ability of decentralized execution is fully preserved. In CollaQ (Zhang et al., 2020), the decentralized utility function is decomposed into two terms: the self-term that only relies on the agent's own state, and the interactive term that is related to states of nearby agents. By using an additional MARA loss to constrain the contribution of the interactive term in the decentralized utility function, CollaQ solves the generalization on the dynamic team composition problem with CTDE being preserved. Unlike this method, MIPI uses mutual information to constrain the contribution of team-related information in the agent's policy.
2.3 Information-Theoretic Principles in RL
In the standard RL problem, the objective is to solely optimize the environmental reward. However, in many problems, we not only want to optimize the cumulative rewards but also want the learned policy to exhibit some other properties. Therefore, a line of research across single-agent and multi-agent domains has been proposed using different information-theoretic principles for policy optimization. For example, SQL (Haarnoja et al., 2017) and SAC (Haarnoja et al., 2018) incorporate the maximum entropy principle to encourage exploration and diverse behavior. FOP (Zhang et al., 2021) further extends this idea to MARL and proves its convergence and optimality. As the augmented entropy term distorts the original objective in the original MDP, which may lead to undesirable behavior in some scenarios (Eysenbach and Levine, 2019). To solve this problem, DMAC (Su and Lu, 2022) proposes to use the divergence between the current policy and previous policy to replace entropy, which yields a bound of the discrepancy between the converged policy and optimal policy in the original MDP. Using divergence to guide the policy optimization is also very popular in offline RL (Levine et al., 2020), for example, F-BRC (Kostrikov et al., 2021) and ICQ (Yang et al., 2021), where the divergence is used to control the similarity between learned policy and behavior policy. While the entropy and KL divergence can be seen as a measurement of the distance to a fixed policy, the mutual information is about the distance to a dynamic marginal policy. In MIRL (Grau-Moya et al., 2019) and MIRACLE (Leibfried and Grau-Moya, 2020), the environmental reward is combined with mutual information to encourage the learned policy to be close to an optimal prior policy, which is also dynamically learned during the RL process instead of being fixed throughout. Unlike (Grau-Moya et al., 2019; Leibfried and Grau-Moya, 2020) that aim at a generalized version of SAC in single-agent RL, the purpose of our work is to solve generalization on dynamic team composition in MARL.
Mutual information (MI) has been widely used in previous MARL research for various purpose, e.g., exploration (Mahajan et al., 2019; Wang et al., 2019; Zheng et al., 2021), coordination (Konan et al., 2021; Kim et al., 2023), individuality (Jiang and Lu, 2021), diversity (Li et al., 2021) and social
influence (Jaques et al., 2019). Unlike these works that mainly focus on the performance of agents in a fixed team, we focus on the generalization ability of agents over different or even unseen teams. Also, these works mainly try to increase the mutual information between two variables to enhance the dependency between variables. However, in our work, we try to decrease the mutual information between the agent's policy and team-related information to reduce the dependency between these two variables and avoid overfitting. Being an exception, PMIC (Li et al., 2022) maximizes the MI associated with the superior trajectories and minimizes the MI associated with the inferior trajectories at the same time. However, it still focuses on the training of a fixed team, while our work focuses on the training of a dynamic team and the generalization over unseen teams.
In this paper, we formulate cooperative MARL with dynamic team composition as a multi-agent Markov decision process (MMDP) with entities. MMDP with entities can be defined by a tuple  $\langle \mathcal{E}, S, A, U, P, r, \gamma \rangle$ .  $\mathcal{E}$  is the set of entities in the environment,  $S$  is the set of states, and each state  $s$  is composed by the state of each entity  $s = \{s_e\}$ . It is worth noting that except agents  $a \in A$ , there are also other entities in the environments (e.g., landmarks, obstacles, agents with fixed behavior).  $U = U_1 \times \dots \times U_{|A|}$  is the joint action space, where  $U_i$  is the individual action space for each agent  $i$ . For the rigorousness of proof, we assume full observability such that at each state  $s \in S$ , each agent  $i$  receives state  $s$ , chooses an action  $u_i \in U_i$ , and all actions form a joint action  $u \in U$ . The state transitions to the next state  $s'$  upon  $u$  according to the transition function  $P(s'|s, u): S \times U \times S \to [0,1]$ , and all agents receive a shared reward  $r(s, u): S \times U \to \mathbb{R}$ . The objective is to learn an individual policy  $\pi_i(u_i | s)$  for each agent such that they can cooperate to maximize the expected cumulative discounted return,  $\mathbb{E}[\sum_{t=0}^{\infty} \gamma^t r_t]$ , where  $\gamma \in [0,1)$  is the discount factor. In CTDE, from a centralized perspective, a group of local policies can be viewed as a joint policy  $\pi_{\mathrm{jt}}(u | s)$ . For this joint policy, we can define the joint state-action value function  $Q_{\mathrm{jt}}(s_t, u_t) = \mathbb{E}_{s_{t+1: \infty}, u_{t+1: \infty}}[\sum_{k=0}^{\infty} \gamma^t r_{t+k}|s_t, u_t]$ . Note that although we assume full observability for the rigorousness of proof, we use the trajectory of each agent  $\tau_i \in \mathcal{T}_i: (Y \times U_i)^*$  to replace state  $s$  for each agent to settle the partial observability in practice, where  $Y$  is the observation space.
Since we are discussing dynamic team composition in this paper, we further denote  $s_i^+$  as team-unrelated information for agent  $i$  (agent's own information), and use  $s_i^-$  to denote team-related information that varies along team composition (e.g., information of other agents, landmarks, obstacles). Although we assume full observability for each agent (i.e.,  $s$  is the same for all agents),  $s_i^+$  and  $s_i^-$  can be different as the circumstance of each agent per se is different. One can easily conclude that  $s = \{s_i^+, s_i^-\}$  for each agent  $i$ .
In this section, we present our method, MIPI, as follows. In Section 4.1, we introduce the mutual information (MI) augmented objective for regularizing the reliance on  $s_i^-$  for each agent. However, due to the existence of the dynamic marginal distribution, direct optimization on this objective can be intractable in practice. Therefore, in Section 4.2, we first discuss a multi-agent policy iteration with a fixed marginal distribution and prove its convergence and optimality. Then, in Section 4.3, we discuss how to use an imaginary team composition distribution to achieve an approximate dynamic marginal distribution and how to use the Blahut-Arimoto algorithm (Cover, 1999) to optimize the corresponding objective. Finally, in Section 4.4, we summarize the learning framework of MIPI.
4.1 MI-Augmented Objective
The learning objective of standard MARL can be formulated as follows,
$$
\underset {\pi_ {\mathrm {j} t}} {\arg \max } \mathbb {E} _ {\rho (\boldsymbol {s} _ {0}), \pi_ {\mathrm {j} t}, P} \left[ \sum_ {t = 0} ^ {T} \gamma^ {t} r (\boldsymbol {s} _ {t}, \boldsymbol {u} _ {t}) \right], \tag {1}
$$
where  $\rho (s_0)$  is the distribution of initial state. We can further rewrite it as:
$$
\underset {\pi_ {\mathrm {j t}}} {\arg \max } \sum_ {t = 0} ^ {T} \mathbb {E} _ {\rho_ {\pi_ {\mathrm {j t}}} (\boldsymbol {s} _ {t})} \left[ \mathbb {E} _ {\pi_ {\mathrm {j t}} (\boldsymbol {u} _ {t} | \boldsymbol {s} _ {t})} \left[ \gamma^ {t} r (\boldsymbol {s} _ {t}, \boldsymbol {u} _ {t}) \right] \right], \tag {2}
$$
where  $\rho_{\pi_{\mathrm{jt}}}(s_t)$  is the marginal distribution over states at timestep  $t$ . Recall that the conditional mutual information between  $x$  and  $y$  given  $z$  can be expressed as follows,
$$
\mathbf {M I} (x; y | z) = \mathbb {E} _ {p (y, z)} \left[ \mathbb {E} _ {p (x | y, z)} \left[ \log \frac {p (x | y , z)}{p (x | z)} \right] \right].
$$
The conditional mutual information can be used to measure the dependency between variable  $x$  and  $y$  with  $z$  given. Therefore, as our goal is to reduce the reliance of  $\pi_i(u_i|s) = \pi_i(u_i|s_i^+, s_i^-)$  on  $s_i^-$ , we can formulate conditional mutual information as follows,
$$
\begin{array}{l} \mathbf {M I} (u _ {i}; s _ {i} ^ {-} | s _ {i} ^ {+}) = \mathbb {E} _ {\rho (s _ {i} ^ {+}, s _ {i} ^ {-})} \left[ \mathbb {E} _ {\pi_ {i} (u _ {i} | s _ {i} ^ {+}, s _ {i} ^ {-})} \left[ \log \frac {\pi_ {i} (u _ {i} | s _ {i} ^ {+} , s _ {i} ^ {-})}{\hat {\pi} _ {i} (u _ {i} | s _ {i} ^ {+})} \right] \right] \\ = \mathbb {E} _ {\rho (s _ {i} ^ {+}), \rho (s _ {i} ^ {-} | s _ {i} ^ {+})} \left[ \mathbb {E} _ {\pi_ {i} (u _ {i} | \boldsymbol {s})} \left[ \log \frac {\pi_ {i} (u _ {i} | \boldsymbol {s})}{\pi_ {i} (u _ {i} | s _ {i} ^ {+})} \right] \right], \\ \end{array}
$$
where  $\pi_i(u_i|s_i^+) = \sum_{s_i^-}\rho (s_i^- |s_i^+)\pi_i(u_i|s_i^+,s_i^-)$ . Incorporating the conditional mutual information of all agents into the standard MARL objective, we now have the MI-augmented objective used in this paper:
$$
\underset {\pi_ {\mathrm {j t}}} {\arg \max } \sum_ {t = 0} ^ {T} \mathbb {E} _ {\rho_ {\pi_ {\mathrm {j t}}} (\boldsymbol {s} _ {t})} \left[ \mathbb {E} _ {\pi_ {\mathrm {j t}} (\boldsymbol {u} _ {t} | \boldsymbol {s} _ {t})} \left[ \gamma^ {t} \left(r (\boldsymbol {s} _ {t}, \boldsymbol {u} _ {t}) - \alpha \sum_ {i} \log \frac {\pi_ {i} (u _ {i , t} | \boldsymbol {s} _ {t})}{\pi_ {i} (u _ {i , t} | s _ {i , t} ^ {+})}\right) \right] \right] \tag {3}
$$
$$
\text {s . t .} \quad \pi_ {i} \left(u _ {i, t} \mid s _ {i, t} ^ {+}\right) = \sum_ {s _ {i, t} ^ {-}} \rho_ {\pi_ {\mathrm {j t}}} \left(s _ {i, t} ^ {-} \mid s _ {i, t} ^ {+}\right) \pi_ {i} \left(u _ {i, t} \mid s _ {i, t} ^ {+}, s _ {i, t} ^ {-}\right), \tag {4}
$$
where the coefficient  $\alpha$  is used to determine the trade-off between maximizing global reward and minimizing mutual information.
4.2 Multi-Agent Policy Iteration with a Fixed Marginal Distribution
As we can see in (4), the optimization of (3) is highly coupled with a dynamic marginal distribution  $\pi_i(u_{i,t}|s_{i,t}^+)$ . What's even worse is, this marginal distribution is determined by  $\rho_{\pi_{\mathrm{jt}}}(s_{i,t}^{-}|s_{i,t}^{+})$ , which is related to  $\pi_{\mathrm{jt}}$ , making the optimization problem even harder. However, one may notice that, if such a marginal distribution  $\pi_i(u_{i,t}|s_{i,t}^+)$  is given and fixed, this problem becomes much easier and can be solved in an off-policy manner. Therefore, in this section, we introduce multi-agent policy iteration with a fixed marginal distribution and prove its convergence and optimality, and in the next section, we discuss how to approximate the dynamic marginal distribution by integrating constraints similar to (4) into this multi-agent policy iteration. First, let us define the joint value function  $V_{\mathrm{jt}}$  and joint state-action value function  $Q_{\mathrm{jt}}$  as follows,
$$
V _ {\mathrm {j t}} ^ {\pi_ {\mathrm {j t}}} (\boldsymbol {s}) = \mathbb {E} _ {\pi_ {\mathrm {j t}}} \left[ \sum_ {t = 0} \gamma^ {t} \left(r (\boldsymbol {s} _ {t}, \boldsymbol {u} _ {t}) - \alpha \sum_ {i} \log \frac {\pi_ {i} (u _ {i , t} | \boldsymbol {s} _ {t})}{\pi_ {i} (u _ {i , t} | s _ {i , t} ^ {+})}\right) | \boldsymbol {s} _ {0} = \boldsymbol {s} \right]
$$
$$
Q _ {\mathrm {j t}} ^ {\pi_ {\mathrm {j t}}} (\boldsymbol {s}, \boldsymbol {u}) = r (\boldsymbol {s}, \boldsymbol {u}) + \gamma \mathbb {E} _ {\boldsymbol {s} ^ {\prime} \sim P} \left[ V _ {\mathrm {j t}} \left(\boldsymbol {s} ^ {\prime}\right) \right],
$$
where  $\pi_i(u_{i,t}|s_{i,t}^+)$  is the fixed marginal distribution for each agent  $i$ . With the above definition, we can further deduce that:
$$
V _ {\mathrm {j t}} ^ {\pi_ {\mathrm {j t}}} (\boldsymbol {s}) = \mathbb {E} _ {\pi_ {\mathrm {j t}}} \left[ Q _ {\mathrm {j t}} ^ {\pi_ {\mathrm {j t}}} (\boldsymbol {s}, \boldsymbol {u}) - \alpha \sum_ {i} \log \frac {\pi_ {i} (u _ {i} | \boldsymbol {s})}{\pi_ {i} (u _ {i} | s _ {i} ^ {+})} \right].
$$
We can then define the joint policy evaluation operator as
$$
\Gamma_ {\pi_ {j t}} Q _ {j t} (\boldsymbol {s}, \boldsymbol {u}) := r (\boldsymbol {s}, \boldsymbol {u}) + \gamma \mathbb {E} _ {\boldsymbol {s} ^ {\prime}} \left[ V _ {j t} \left(\boldsymbol {s} ^ {\prime}\right) \right] \tag {5}
$$
and have the following lemma.
Lemma 1 (Joint Policy Evaluation). Consider the modified Bellman backup operator  $\Gamma_{\pi_{\mathrm{jt}}}$  (5) and a mapping  $Q_{\mathrm{jt}}^{0}: S \times U \to \mathbb{R}$  with  $|U| < \infty$ , and define  $Q_{\mathrm{jt}}^{k+1} = \Gamma_{\pi_{\mathrm{jt}}} Q_{\mathrm{jt}}^{k}$ . Then, the sequence  $Q_{\mathrm{jt}}^{k}$  will converge to the joint  $Q$ -function of  $\pi_{\mathrm{jt}}$  as  $k \to \infty$ .
Using Lemma 1, we can get  $Q_{\mathrm{jt}}$  for any joint policy  $\pi_{\mathrm{jt}}$ . However, it is hard for us to use  $Q_{\mathrm{jt}}$  for individual policy improvement. To solve this problem, many value decomposition methods choose to factorize the joint state-action value function  $Q_{\mathrm{jt}}$  into the utility function  $Q_{i}$  of each agent and use  $Q_{i}$  to guide the individual policy improvement of  $\pi_{i}$ . In this paper, we factorize the joint state-action value function into the following form, which is shared by many value decomposition methods (Zhang et al., 2021; Su and Lu, 2022; Wang et al., 2023):
$$
Q _ {\mathrm {j t}} ^ {\pi_ {\mathrm {j t}}} (\boldsymbol {s}, \boldsymbol {u}) = \sum_ {i} w _ {i} (\boldsymbol {s}) * Q _ {i} ^ {\pi_ {i}} (\boldsymbol {s}, u _ {i}) + b (\boldsymbol {s}). \tag {6}
$$
After the evaluation of the joint policy and the decomposition of  $Q_{\mathrm{jt}}$ , we construct the following optimization problem for individual policy improvement.
$$
\pi_ {i} ^ {\text {n e w}} = \arg \max  _ {\pi_ {i} ^ {\prime}} \mathbb {E} _ {\pi_ {i} ^ {\prime}} \left[ Q _ {i} ^ {\pi_ {i} ^ {\text {o l d}}} (\boldsymbol {s}, u _ {i}) - \alpha \log \frac {\pi_ {i} ^ {\prime} (u _ {i} | \boldsymbol {s})}{\pi_ {i} (u _ {i} | s _ {i} ^ {+})} \right] \tag {7}
$$