SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xue Yan  $^{1,2}$ , Jiaxian Guo  $^{3}$ , Xingzhou Lou  $^{1,2}$ , Jun Wang  $^{4}$ , Haifeng Zhang  $^{1,2}$ , Yali Du*5
$^{1}$ Institute of Automation, Chinese Academy of Science, Beijing, China
$^{2}$  School of Artificial Intelligence, University of Chinese Academy of Sciences, China
<sup>3</sup> University of Tokyo, Japan <sup>4</sup>University College London, UK <sup>5</sup>King's College London, UK {yanxue2021, louxingzhou2020, haifeng.zhang}@ia.ac.cn
jiaxian.guo@weblab.t.u-tokyo.ac.jp, jun.wang@cs.ucl.ac.uk, yali.du@kcl.ac.uk
The goal of zero-shot human-AI coordination is to develop an agent capable of collaborating with humans without relying on human data. Prevailing two-stage population-based methods require a diverse population of mutually distinct policies to simulate diverse human behaviors. The necessity of such populations severely limits their computational efficiency. To address this issue, we propose E3T, an Efficient End-to-End Training approach for zero-shot human-AI coordination. E3T employs a mixture of ego policy and random policy to construct the partner policy, making it both skilled in coordination and diverse. This way, the ego agent is trained end-to-end with this mixture policy, eliminating the need for a pre-trained population, and thus significantly improving training efficiency. In addition, we introduce a partner modeling module designed to predict the partner's actions based on historical contexts. With the predicted partner's action, the ego policy can adapt its strategy and take actions accordingly when collaborating with humans exhibiting different behavior patterns. Empirical results on the Overcooked environment demonstrate that our method substantially improves the training efficiency while preserving comparable or superior performance than the population-based baselines. Demo videos are available at https://sites.google.com/view/e3t-overcooked.
Cooperation between AI agents and humans has gained significant attention, spanning various fields such as robotics [5, 26], automatic driving [41], Human-AI dialogue [14] and Human-AI coordination games [30]. Due to the high cost of collecting human data and involving humans during training, recent studies [32, 15] have focused on zero-shot human-AI coordination problems, with the aim to train an ego agent capable of collaborating with humans without relying on human data.
To train such an ego agent, the prevailing approach [15, 21] involves training the ego agent with diverse partner agents that can simulate human behavior without the need of human data. One method for achieving this is through self-play [33, 29], where the ego agent plays against itself as the partner agent. However, empirical evidence has shown that self-play can lead to agents getting stuck in a single cooperative pattern during training [15], hindering their ability to adapt to diverse human behaviors. Alternative approaches like other-play [15] aim to introduce diversity by breaking the symmetry of self-play policies. Nevertheless, this method relies on the assumption of strict symmetry in the action or state space, limiting its practical applicability.
In contrast to relying solely on a single partner policy, advanced population-based methods [21, 42, 32, 18, 19] construct a population with multiple policies to simulate diverse human behaviors, and then train the ego agent to collaborate with this population. Among these methods, FCP [32] obtains a diverse population by training several self-play policies with different random seeds and saving multiple past checkpoints reflecting various coordination levels. TrajeDi [21] promotes diversity by maximizing the divergence of trajectory distributions within the partner population. MEP [42] achieves a cooperative and diverse population by maximizing self-play rewards and population entropy as auxiliary rewards.
However, these population-based methods suffer from two major limitations. Firstly, they are highly computationally inefficient due to their two-stage framework, which involves training a population of policies before training the ego agent. For example, training a MEP agent for one environment in overcooked [8] can take over 15 hours on a single 3090Ti GPU. By contrast, a self-play agent can be trained in less than 1 hour. Secondly, these methods only learn a single response for all partners in the population, neglecting the diverse behavior patterns among them. As a result, the learned ego agent lacks the capability to adjust its policy based on each partner's unique behavior pattern, leading to unsatisfactory performance in zero-shot human-AI coordination.
Figure 1: Illustration of training time and zero-shot coordination performance of baselines. The right  $y$ -axis is the required hours to train one model. The left  $y$ -axis represents the average rewards of collaborating with AI baselines across 5 layouts.
In this paper, we introduce an Efficient End-to-End Training approach, called E3T, which aims to improve both efficiency and the ability to recognize behavior patterns in zero-shot human-AI coordination. To achieve efficient training, E3T utilizes a blend of ego and random policies as partners during training, and integrates a partner modeling module into the ego policy. This approach is motivated by the idea that the partner policies should be both skilled in coordination and diverse [42]. By breaking down these two objectives into two policies i.e. ego policy and random policy, and merging them as the partner policy, E3T is able to be end-to-end trained without the need for pre-training a population, thus enjoying higher efficiency and a simpler implementation compared to population-based methods.
In order to improve the collaboration capabilities of of ego agents when working with partners exhibiting
different behaviors, we introduce a partner modeling module that predicts the partner's actions based on historical contexts. By incorporating this module into the ego policy, the agent is able to know not only the current observation, but also the predicted actions of the partner. This enables the ego agent to adapt its behavior more effectively, aligning it better with the partner's preferences and resulting in improved collaboration performance. To the best of our knowledge, our paper is the first to utilize a partner action prediction module in the context of a two-player human-AI coordination task.
Our contributions are three-fold. Firstly, we introduce an efficient end-to-end ego training approach, named E3T, for zero-shot Human-AI coordination, which offers higher training efficiency compared to population-based methods. Secondly, we propose a partner modeling module that predicts partner actions from the context, allowing the ego policy to adapt when collaborating with humans exhibiting different behavior patterns. Lastly, empirical results on Overcooked show that our method, E3T, significantly improves training efficiency while achieving comparable or superior performance to existing methods as Figure 1 shows. Specifically, E3T demonstrates a significant improvement in training efficiency when compared to state-of-the-art population-based methods.
Our work investigates how to train an ego agent for zero-shot Human-AI coordination, a field with broad applications in real-world scenarios such as robotics [5, 26], service allocation [36, 37, 35] and Human-AI coordination games [30]. In zero-shot Human-AI coordination tasks, reinforcement learning (RL) algorithms [15] are tasked with training a robust ego agent without relying on human data, due to the high cost of collecting real human data.
Self-Play Training There are a number of methods to tackle zero-shot coordination, broadly categorized into two classes: self-play and population-based training. Self-play has demonstrated impressive capability in competitive games such as Go [29], Dota [38], StarCraft [34] and imperfect information games [13]. Self-play can also master high-level skills in cooperative games [8]. However, self-play agents often struggle to coordinate effectively with previously unseen partners due to their tendency to adopt rigid conventions formed during training against themselves [32]. Other-play [15] attempts to mitigate these conventions by breaking the symmetries of self-play policies. However, it heavily relies on the prior knowledge of strict symmetries in the environment, limiting its applicability in general environments lacking such symmetries. In this work, E3T follows the self-play training framework, which avoids falling into conventions by encouraging the partner policies to be more diverse. In addition, our method offers greater generality compared to other-play, making it applicable in a wider range of environments.
Population-Based Training Methods Population-based methods train a robust AI by collaborating with a diverse population. Experimental results show that population-based methods achieve superior zero-shot coordination performance compared to self-play and other-play [15]. For example, PBT [8] employs an online evolutionary algorithm, which iteratively updates policy parameters and performs policy substitution within the population. FCP [32] achieves diversity within a population by training several self-play policies with different random seeds and saving multiple past checkpoints to obtain agents with various coordination levels. In a similar vein, TrajeDi [21] regulates the partner population to be diverse by maximizing the divergence of trajectory distributions between policies within the population. Furthermore, MEP [42] introduces an auxiliary reward by designing the entropy of the average of policies in the population, thus promoting population diversity. More recently, COLE [18, 19] reformulate two-player coordination as a game-theoretic problem, where the trainer's objective is to estimate the best responses to the most recent population based on the cooperative incompatibility distribution. Our work improves the partner behavior diversity by maximizing the entropy, but E3T directly increases the randomness of a single-agent partner policy, making it more efficient than MEP and TrajeDi.
Agent Modeling Agent modeling is an important ability for autonomous agents, enabling them to reason about various properties of other agents such as actions and goals [3]. Typical agent modeling methods focus on reasoning about types of other agents [1, 2, 31], plan recognition [9, 7], group modeling [23, 11, 20] and especially policy reconstruction [6, 25, 4, 27], which is also adopted in E3T to improve zero-shot human-AI coordination performance. However, a key difference is that E3T's ego agent's partner modeling module is trained by pairing with agents and aims to predict the actions of partners. In addition, we directly employ the predicted partner actions for decision-making, instead of using the embedding of partner behaviors [10].
3 Preliminary
In this section, we introduce the two-player cooperative Markov Game and a representative pipeline of population-based training for two-player zero-shot coordination.
Two-player Cooperative Markov Game A two-player cooperative Markov Game is described by a tuple  $\langle S, \mathcal{A}, \mathcal{A}, \mathcal{T}, R \rangle$ , where  $\mathcal{S}$  is a finite set of states,  $\mathcal{A}$  is the finite set of actions, and suppose two agents have the same action space.  $p_0: \mathcal{S} \to \mathbb{R}$  is the initial state distribution.  $\mathcal{T}: \mathcal{S} \times \mathcal{A} \times \mathcal{A} \times \mathcal{S} \to [0,1]$  is the transition function and assumes the transition dynamic is deterministic.  $R: \mathcal{S} \times \mathcal{A} \times \mathcal{A} \to \mathbb{R}$  is the reward function and assume rewards are bounded by  $\alpha$ , i.e.  $R(s, a^1, a^2) \leq \alpha, \forall s, a^1, a^2$ . We assume the transition and reward functions are unknown. In this work, we solve the two-policy cooperative Markov Game via single-agent reinforcement learning and the self-play framework as previous studies [8, 15]. Specifically, the ego policy  $\pi_{\mathrm{e}}$  is trained by pairing with a fixed partner policy  $\pi_{\mathrm{p}}$  with the object defined as  $\pi_{\mathrm{e}} = \arg \max_{\pi} V(\pi, \pi_{\mathrm{p}})$ . The value function of joint policy  $(\pi, \pi_{\mathrm{p}})$  is defined as  $V(\pi, \pi_{\mathrm{p}}) = E_{p_0, \mathcal{T}, a^1 \sim \pi, a^2 \sim \pi_{\mathrm{p}}} \sum_{t=0}^{\infty} \gamma^t R(s_t, a_t^1, a_t^2)$ , where  $\gamma \in (0,1)$  is the discounted factor.
Population-Based Training Population-based methods [21, 42, 32] aim to construct a population of policies to simulate diverse human behaviors. Intuitively, we expect that the policy in the population should be both skilled in coordination and diverse. Taking Maximum Entropy Population (MEP) [42] as an example, it constructs the partner population by maximizing the self-play rewards and simultaneously increasing the entropy of policies in the population. The objective of Maximum
Figure 2: (a). The illustration of the self-play training framework, which trains the ego policy by pairing it with the copied partner policy. (b). The illustration of E3T. The green box shows the decision process of the ego agent. The blue box shows that of the partner agent.
Entropy Population training is defined as:
$$
J (\bar {\pi}) = \mathbb {E} _ {\pi_ {i} \sim \mathcal {P}} \left[ \sum_ {t} \mathbb {E} _ {\left(s _ {t}, a _ {t} ^ {1}, a _ {t} ^ {2}\right) \sim \left(\pi_ {i}, \pi_ {i}\right)} \left[ R \left(s _ {t}, a _ {t} ^ {1}, a _ {t} ^ {2}\right) + \lambda \mathcal {H} (\bar {\pi} (.) | s _ {t})) \right] \right]. \tag {1}
$$
Here  $\mathcal{P} = \{\pi_1, \pi_2, \dots, \pi_n\}$  denotes a population of policies,  $\bar{\pi}(.|s_t) = \frac{1}{n} \sum_{i=1}^{n} \pi_i(|s_t)$  is the mean policy of the population  $\mathcal{P}$ , and  $\lambda$  determines the relative importance of the population entropy term against the task reward. The task reward under the joint action  $\mathbf{a}_t$  is obtained by self-play with a uniformly sampled policy from the population, and the population entropy term acts as an auxiliary reward. With this competent and diverse population as the partner population, an ego policy is trained by collaborating with the population via prioritized sampling that assigns a higher priority to the partner policy that is difficult to cooperate with. The objective of ego policy training is defined as:
$$
J \left(\pi_ {\mathrm {e}}\right) = \sum_ {t} \mathbb {E} _ {\left(s _ {t}, a _ {t} ^ {1}, a _ {t} ^ {2}\right) \sim \left(\pi_ {\mathrm {e}}, \rho (\mathcal {P})\right)} \left[ R \left(s _ {t}, a _ {t} ^ {1}, a _ {t} ^ {2}\right) \right], \tag {2}
$$
where  $\rho(.)$  as the prioritized sampling mechanism. However, the training process of MEP is inefficient due to the complex two-stage training and the necessity of maintaining a large population for behavioral diversity. Our work draws inspiration from MEP, which achieves behaviorally diverse and cooperative partners by maximizing the self-play reward and the entropy of partner policy and then trains a robust ego policy by cooperating with those partners. In the next section, we will discuss how to solve zero-shot coordination using an efficient end-to-end training process.
In this section, we introduce the proposed method, E3T, an efficient end-to-end training approach, which trains an ego policy  $\pi_{\mathrm{e}}$  through collaborating with a partner policy  $\pi_{\mathrm{p}}$ . To take into account both training efficiency and partner diversity, we employ a mixture of the copied ego policy and a random policy as the partner policy  $\pi_{\mathrm{p}}$ . The objective of E3T is defined as follows:
$$
J \left(\pi_ {\mathrm {e}}\right) = \sum_ {t} \mathbb {E} _ {\left(s _ {t}, a _ {t} ^ {1}, a _ {t} ^ {2}\right) \sim \left(\pi_ {\mathrm {e}}, \pi_ {\mathrm {p}}\right)} \left[ R \left(s _ {t}, a _ {t} ^ {1}, a _ {t} ^ {2}\right) \right] \tag {3}
$$
The proposed method contains two key components, a diverse partner policy and a partner modeling module for adapting to unseen partners, which will be introduced in Section 4.1 and Section 4.2 respectively. An overview of E3T is presented in Figure 2.
4.1 A Mixture Policy as Partner
We begin with the motivation of leveraging the proposed mixture partner policy. Recall that MEP trains a partner population by maximizing self-play rewards of policies within it and maximizing the population entropy as shown in Equation (1). The first term on maximizing self-play rewards tends to learn a population with cooperative ability. The second entropy term can force the partner population
to be diverse. However, optimizing the objective of Equation (1) suffers very high computational complexity due to the necessity of training a number of policies to collaborate effectively with themselves. Moreover, it exhibits unstable convergence properties, requiring meticulous hyperparameter tuning for additional entropy term [12].
To circumvent computational cost and non-stable training, we construct the partner policy via dividing Equation (1) into two parts with different partner policies learned, i.e.  $\pi_{\mathrm{p1}}$  and  $\pi_{\mathrm{p2}}$ , which aim to maximize the rewards  $\sum_t \mathbb{E}_{(s_t, \mathbf{a}_t) \sim \pi_{\mathrm{p1}}, \pi_{\mathrm{p1}}} [R(s_t, \mathbf{a}_t)]$  and the entropy of the policy output  $\sum_t \mathbb{E}_{(s_t, \mathbf{a}_t) \sim \pi_{\mathrm{p2}}} [\mathcal{H}(\pi_{\mathrm{p2}}(.|s_t))]$ , respectively. We can see that maximizing the former objective function is aligned with the goal of learning coordination skills, akin to the self-play training objective. In this way, we can directly use the ego policy  $\pi_{\mathrm{e}}$  as  $\pi_{\mathrm{p1}}$ . On the other hand, if we directly learn the latter one objective, i.e., the entropy of the mean policy, policy  $\pi_{\mathrm{p2}}$  will lead to the random policy  $\pi_{\mathrm{r}}$ , i.e., a uniform distribution over the action space.
To construct a partner who has both skills in coordination and diverse behaviors, we directly mix the above-discussed partner policies  $\pi_{\mathrm{p1}}$  and  $\pi_{\mathrm{p2}}$  as the partner policy  $\pi_{\mathrm{p}}$ . As previously discussed,  $\pi_{\mathrm{p1}}$  can directly apply ego policy  $\pi_{\mathrm{e}}$  to achieve high reward for coordination.  $\pi_{\mathrm{p2}}$  is a random policy  $\pi_{\mathrm{r}}$ , which has the maximum policy entropy. In this way, the partner policy  $\pi_{\mathrm{p}}$  is the mixture between the self-play ego policy and the random policy with the balance parameter  $\epsilon$ , formally defined as:
$$
\pi_ {\mathrm {p}} = \epsilon \pi_ {\mathrm {r}} + (1 - \epsilon) \pi_ {\mathrm {e}} \tag {4}
$$
We can see that the computation cost for obtaining this mixture partner policy is significantly lower than population-based methods since the conversion eliminates the need to train a population of policies. Additionally, the partner policy achieves behavioral diversity by simply leveraging a random policy with maximum entropy, whereas population-based approaches rely on more complex mechanisms, such as gradient descent [21] or reward shaping [42], to enforce the diversity. Both diversity and coordination ability are two keys to constructing a partner policy. Intuitively, the diversity reflected by entropy increases, and the coordination performance decreases with the random coefficient epsilon increasing. We attempt to quantitatively verify these intuitions through Propositions 1 and 2. Firstly, we analyze the behavior diversity of the partner policy of E3T. The entropy of partner policy reflects the diversity of partner behaviors. Proposition 1 shows that the entropy of the mixture partner policy is related to that of the random policy with maximum entropy.
Proposition 1 (Entropy of Mixture Partner Policy). The entropy of the mixture partner policy, defined as  $\pi_p = (1 - \epsilon)\pi_e + \epsilon \pi_r$ , satisfies that:
$$
\mathcal {H} \left(\pi_ {p} (.. | s _ {t})\right) \geq \epsilon \mathcal {H} \left(\pi_ {r} (.)\right) - C _ {1}, \forall s _ {t} \tag {5}
$$
where  $C_1 = \epsilon \log \epsilon + |\mathcal{A}|(1 - \epsilon)$ , and  $|\mathcal{A}|$  is the cardinality of the action space  $\mathcal{A}$ . In addition, the lower bound of  $\mathcal{H}(\pi_p(.|s_t))$  is increased as the increasing of  $\epsilon$ . Proof. See Appendix.
Secondly, we demonstrate the benefits of improving coordination skills when pairing with the partner policy of E3T. According to Proposition 2, although  $\pi_{\mathrm{e}}^{\prime}$  is the best response to the mixture partner policy rather than the ego policy itself, the coordination performance of  $\pi_{\mathrm{e}}^{\prime}$  with the ego policy is lower bounded by that of ego policy self-play with a punishment term  $\zeta$ . The  $\zeta$  is related to the balance parameter  $\epsilon$  and smaller  $\epsilon$ , which means less weight on the random policy, leads to lower  $\zeta$ .
Proposition 2 (Coordination Skill Improvement). Let the  $\pi_e'$  as the best response to  $\pi_p = (1 - \epsilon)\pi_e + \epsilon \pi_r$ , i.e  $\pi_e' = \arg \max_{\pi} V(\pi, \pi_p)$ , then we have  $V(\pi_e', \pi_e) \geq V(\pi_e, \pi_e) - \zeta$ , where  $\zeta = \frac{4\alpha\epsilon}{(1 - \gamma)^2}$ . Proof. See Appendix.
4.2 Partner Modeling
In this section, we describe the details of the proposed partner modeling module, designed to extract and analyze the behavior pattern of the partner from historical information. This module equips the ego policy with the ability to consider not only the current state but also the partner's behavior pattern, empowering the ego agent to adapt to previously unseen partners
Partner Action Prediction As in previous works [6, 25, 4], we model the partner agent by policy reconstruction. If the actual partner action distribution  $\mathbf{a}_t^{\mathrm{p}}$  is given, the two-player cooperative Markov Game will reduce to a single-agent MDP about the ego agent, in which the ego agent can easily find the best response to the current state. Formally, under the partner policy  $\pi_{\mathrm{p}}$  with  $\mathbf{a}_t^{\mathrm{p}} = \pi_{\mathrm{p}}(.|s_t)$ , the two-player MDP with transition function  $\mathcal{T}$  can be converted in to a single-agent MDP with transition
Algorithm 1 E3T: An efficient end-to-end training approach with mixture policies and partner modeling
Input: Parameters of context encoder  $E_{\mathrm{c}}$ , action prediction network  $F$ , ego policy  $\pi_{\mathrm{e}}$ , hyperparameter  $\epsilon$ ,  $k$ .
Output: Well-trained context encoder  $E_{\mathrm{c}}$ , action prediction network  $F$  and the ego policy  $\pi_{\mathrm{e}}$ .
1: while not done do
2: Set data buffer  $D = \emptyset$ .
3: for  $t = 1,2,\dots$  , step do
4: Encode context of the partner:  $z_{t}^{\mathrm{c}} = E_{\mathrm{c}}(\tau_{t})$
5: Predict partner action distribution:  $\hat{\mathbf{a}}_t^{\mathrm{p}} = F(z_t^{\mathrm{c}})$
6: Randomly select an action distribution:  $\hat{\mathbf{a}}_t^e$
7: Mixture partner policy as:  $\pi_{\mathrm{p}}(\cdot |s_t,\hat{\mathbf{a}}_t^{\mathrm{e}}) = (1 - \epsilon)\pi_{\mathrm{e}}(\cdot |s_t,\hat{\mathbf{a}}_t^{\mathrm{e}}) + \epsilon \pi_{\mathrm{r}}(\cdot)$
8: Gather data from  $\pi_{\mathrm{e}}(\cdot |s_t,\hat{\mathbf{a}}_t^{\mathrm{p}})$  and  $\pi_{\mathrm{p}}(\cdot |s_t,\hat{\mathbf{a}}_t^{\mathrm{e}})$ , then  $D = D\cup \left\{(\tau_{t + 1},s_t,a_t^{\mathrm{e}},a_t^{\mathrm{p}},r_t)\right\}$
9: end for
10: Update parameters of partner modeling module by minimizing Equation (7).
11: Update parameters of the ego policy  $\pi_{\mathrm{e}}$  via PPO.
12: end while
function denoted as  $\mathcal{T}_{\pi_{\mathrm{p}}}(s'|s_t,a) = \sum_{a^{\mathrm{p}}}\pi_{\mathrm{p}}(a^{\mathrm{p}}|s_t)\mathcal{T}(s'|s_t,a^{\mathrm{p}},a)$ . However, because the agents act simultaneously, the ego agent can only infer the partner's action from historical information instead of directly observing the actual partner's action.