Semi-Supervised Imitation Learning of Team Policies from Suboptimal Demonstrations

1 Introduction

Teamwork is essential for the success of human enterprise. As artificial agents increasingly become parts of human life, thus, they too are expected to reason about and contribute to human teams. At the same time, teamwork is highly challenging to perfect. Successful human teams employ a variety of training techniques to improve coordination and teamwork [Tannenbaum and Salas, 2020]. Analogously, spurred by the need to enable and enhance human-agent collaboration, there has been growing work on developing computational techniques for training artificial agents to support human teams [Thomaz et al., 2016]. These techniques build upon a variety of AI paradigms, such as planning under uncertainty, reinforcement learning, and imitation learning.

In this work, we consider the paradigm of imitation learning [Argall et al., 2009; Osa et al., 2018], wherein an agent learns from demonstrations and (in contrast to reinforcement learning) can thereby learn policies for teamwork without the need of unsafe exploration. By providing novel multi-agent imitation learning techniques that are inspired by real-world teaming considerations, this work aims to enable agents to model, assess, and improve both human-human and humanAI teamwork in sequential tasks. Mathematically, imitation learning techniques seek to learn a single-agent behavioral policy (π), a stochastic function that encodes probability of selecting an action (a) in a task-specific context (s), given a dataset of (s, a)-tuples provided by a demonstrator. Typically, the context features (s) and actions (a) are assumed to be fully observable and measured using sensors.

An extended version of this paper, which includes supplementary material (appendices and video supplements) mentioned in the text, is available at https://arxiv.org/abs/2205.02959

Imitation learning has been extended to model multi-agent systems by seeking to learn a set of behavioral policies {πi|i=1:n}, one corresponding to each member (i=1:n) of the multi-agent system from demonstrations of teamwork [Le et al., 2017; Bhattacharyya et al., 2019; Song et al., 2018; Lin et al., 2019]. We provide a brief survey of related multi-agent imitation learning (MAIL) techniques in Appendix A. MAIL is an emerging area of research, wherein the existing works either focus on learning policies corresponding to game-theoretic equilibria of multi-agent systems, assume homogeneity in agents' capabilities, or assume data of optimal teaming behavior as the training input. However, in contrast to most settings considered in prior art, teamwork observed in practice [Salas et al., 2018; Seo et al., 2021] often differs in three key ways: (1) it may not correspond to a game-theoretic equilibrium, (2) it can be suboptimal due to its dependence on latent performance-shaping factors (such as team members' mental models), and (3) it can involve team members with different capabilities.

Informed by these three considerations of teamwork observed in the real world, in Sec. 3.3, we provide an alternate problem formulation of MAIL for collaborative tasks. In particular, we consider that the teammembers' behavior (πi) depends not only on the context features (s) but also on their (time-varying) mental models (xi) pertaining to teamwork. As the mental models cannot be readily sensed and manually annotating them is resource-intensive, they are modeled as partially observable to the imitation learner. Further, in teaming scenarios where the mental models are misaligned (i.e., the team members do not maintain a shared understanding), the demonstrations available for learning can be suboptimal.

Thus, adding to the challenge, in our problem formulation the imitation learner also needs to infer which segments of the teaming demonstrations are (sub)-optimal and perform learning under partial observability of the dynamic state (s, x).

Summary of Contributions. Towards this problem setting, we present Bayesian Team Imitation Learner (BTIL, pronounced as "bee-tul"), an imitation learning algorithm that can learn decentralized team policies from both optimal and sub-optimal demonstrations. To effectively learn the team policies, BTIL explicitly models the time-varying mental states of each teammember (xi) and jointly learns their transition model (Tx). To enable sample- and label-efficient policy learning, BTIL utilizes a Bayesian perspective and is capable of learning from partial supervision of the mental states. We benchmark our solutions against two existing techniques [Pomerleau, 1991; Song et al., 2018]. In our evaluations, we emphasize the challenge of collecting teaming demonstrations by collecting a novel dataset of human-agent teamwork for settings where (a) the labels of mental models are only partially available, and (b) the number of demonstrations is small relative to the size of the task's state space. Our experiments show that BTIL can learn team policies from small semi-supervised datasets of suboptimal teamwork and outperform the baselines across relevant metrics.

Running Example. To help describe our approach, we utilize the collaboration scenario shown in Fig. 1 as a running example. In this scenario, a two-member team composed of Alice and Rob is tasked with moving all boxes together to the flag in the least amount of time. At each time step, each agent can choose to move in any one of four cardinal directions, attempt to pick up or drop a box, or perform no operation. As each box is heavy, it cannot be lifted by one agent alone. To effectively move all boxes to the goal location, the team needs to coordinate which box to pick next as well as the path along which to move the box. Thus, in this scenario, the pertinent latent state (x) corresponds to the team's next target location (i.e., one of the boxes or the flag), while the observable state (s) corresponds to the locations of the agents and boxes.

During task execution, the team members' mental models (x) can be misaligned. For example, Alice may target the box at the top left while Rob targets the box at the bottom left, thereby demonstrating suboptimal teamwork. Further, to improve the collaboration, teammembers may or may not choose to change their target at any point during the task based on behavior of the other teammate. Thus, as motivated earlier, the demonstrations of teamwork may not correspond to a game-theoretic equilibrium and can be suboptimal due to dependence on partially observable, dynamic mental models. The goal of the imitation learner is to learn the mental modeldependent policies of the team given these (potentially suboptimal) partially observable demonstrations of teamwork.

2 RelatedWork

Our work relates to the following three sub-areas of imitation learning: multi-agent imitation learning, learning from suboptimal demonstrations, and learning from partially observable demonstrations. Here, we summarize research from these sub-areas and relate it to our approach. Please see Appendix A for a more detailed discussion of MAIL techniques.

Multi-agent Imitation Learning (MAIL). Although multiple MAIL algorithms exist, the problem setting considered in prior art differs from the one considered in this paper. Reiterating from Sec. 1, prior MAIL techniques either learn behavior corresponding to game-theoretic equilibria of multi-agent systems [Song et al., 2018; Lin et al., 2019], assume homogeneity in agents' capabilities [Bhattacharyya et al., 2019], or do not consider latent performance-shaping factors (such as mental models or cognitive states). Approaches that do model latent states assume that the latent state is either shared between members [Wang et al., 2021; Ivanovic et al., 2018] or time-invariant [Le et al., 2017]. In contrast, inspired by real world teaming considerations, we seek to develop MAIL algorithms that both recognize the individual and dynamically changing latent states of each member and are capable of learning multi-agent policies from different levels of supervision over the latent states.

Learning from Partially Observable Demonstrations. Imitation learning with occlusions or missing features has also received increasing attention in the last decade. These techniques, while not directly valid for the multi-agent setting considered in this work, inform our work. [Torabi et al., 2018; Sun and Ma, 2019] consider learning an agent policy from demonstrations that may not include data of the demonstrator's actions. Similarly, [Choi and Kim, 2011; Gangwani et al., 2020] allow incomplete specification of states by utilizing a belief state. [Unhelkar and Shah, 2019] explicitly model an agent's latent decision factors (e.g., mental states), which can change dynamically within an episode. While related to our work, these techniques only consider single-agent tasks and do not model the interaction between multiple agents. [Bogert and Doshi, 2018] provide an approach for multi-robot inverse reinforcement learning from partially observable demonstrations. In contrast to our approach, their work does not model agents' mental states or seek to learn from demonstrations of suboptimal teamwork.

Learning from Suboptimal Demonstrations. While classical imitation learning assumes demonstrations are generated from experts who behave optimally, a few approaches admit that demonstrations can be suboptimal in practice. For example, with an assumption that the majority of demonstrations are optimal, [Choi et al., 2019; Zheng et al., 2014] focus on imitation learning that is robust to suboptimal outliers. Meanwhile, [Brown et al., 2019; Chen et al., 2021; Zhang et al., 2021] aim to incorporate demonstrations that come from demonstrators whose level of expertise is unknown in order to overcome the challenge of scarce expert demonstrations. [Yang et al., 2021] utilize a latent action representation while learning the optimal policy from potentially suboptimal demonstrations. While related to our approach, these solutions for imitation learning from suboptimal demonstrations neither consider demonstrators' mental states nor multi-agent tasks. In contrast, our goal is to develop an approach that can learn stochastic multi-agent policies from demonstrations that are both partially observable and suboptimal.

3 ProblemFormulation

To formalize the problem of learning team policies from suboptimal and partially observable demonstrations, we first provide models for the team task and team members' behavior.

3.1 TaskModel

Due to our focus on learning task-oriented team policies, we require a model to represent team tasks. Borrowing from prior research in multi-agent systems [Oliehoek and Amato, 2016], we build upon the framework of multi-agent Markov decision processes (MMDP) to describe the tasks of interest. An MMDP models sequential collaborative tasks and is specified by the tuple Mtask .= (n, S, A, T, R, γ), where

• n, is the number of agents i indexed 1 : n;
• s ∈ S, denotes the set of task states;
• ai ∈ Ai, is the set of actions ai available to the i-th agent;
• A=×iAi is the set of joint actions, where a=[a1, · · · , an] denotes the joint action;
• Ts(s′ |s, a) : S×A×S → [0, 1] denotes the state transition probabilities, i.e., the probability of the next task state being s′ after the team agents executed action a in state s
• R(s, a) : S × A → R is the joint reward that the team receives after execution action a in state s.
• γ is the discount factor.

The MMDP model assumes that all agents have a shared objective and each agent has full observability of the task state and reward. The shared objective of the set of n agents, whomwe jointly refer to as the team, is to maximize their expected cumulative discounted reward, E[Σtγ tR(st, at)]. In this work, we focus on team tasks that can be modeled as MMDP where the set of states S and the set of actions A are finite. Several real-world tasks can be modeled using MMDP. For instance, the scenario described in the running example can be described as an MMDP with n=2, S modeling the task-relevant features (namely, the agent and box locations), and A modeling the actions available to the agent.

The solution to the MMDP is a set of n decentralized agent policies π1:n, where πi is the policy of the i-th agent. In the running example, this corresponds to policies of Alice and Rob. Mathematically, πi(ai|s) is a probability distribution of the i-th agent's actions ai conditioned on the MMDP state. Since each agent has full state observability, in theory, an MMDP can be solved optimally in a centralized manner by the team using MDP solvers [Puterman, 1990] before the task begins. If each team member follows this optimal policy faithfully, coordination between team members in MMDP tasks is guaranteed during task execution.

3.2 Agent Model

In practice, however, one seldom observes perfect coordination among team members, including in tasks where the agents have complete or near-complete observability of the task state and complete knowledge of the team's objective (e.g., healthcare team in an operating room, or a team of basketball or soccer players). The potential causes of this imperfect coordination are varied. For instance, imperfect coordination can occur due to inability to compute a joint policy, lack of prior coordination, imperfect execution, and different individual preference. To design an imitation learning algorithm that can effectively recover team policies, it is essential to explicitly consider these imperfections and latent causes of suboptimal teamwork.

Hence, to model teamwork observed in practice, we provide a latent variable model for each team member's (potentially suboptimal) behavior. Our model extends the Agent Markov Model (AMM), which explicitly models latent states of a single agent's behavior [Unhelkar and Shah, 2019], to model teamwork. In particular, we model each team member's behavior as the tuple (Xi, bxi , Txi , πi), where

• xi ∈ Xi denotes the latent states that influence the i-th agent's behavior during the task. These may include mental models, methods to tie-break if multiple optimal policies exist, or preferences over different task components.
• bxi (xi) ∈ Xi → [0, 1] denotes the probability distribution of the latent state at the start of the task.
• Txi (x′i |s, xi, a, s′ ) ∈ S×Xi×A×S×Xi → [0, 1] denotes the transition model of the latent state.
• πi(ai|s, xi) ∈ S×Xi×Ai → [0, 1] denotes the teammember's policy, a probability distribution of each member's decision ai conditioned on their decision factors (s, xi).

Referring to the running example, behavior of Alice and Rob depends not only the task context (MMDP state) but also on their latent preferences over the next target location. For each team member, the agent model helps in modeling this latent preference (as xi ∈ Xi), their latent state-dependent policy (as πi), and the evolution of their latent preference (via bxi and Txi ). While the above model is expressive and can represent a variety of team behaviors (e.g., suboptimal policies, evolution of latent preferences based on past behavior), we assume the transition dynamics Tx to be Markovian for computational tractability. Notationally, we jointly refer to the behavioral models for the whole team as (X, bx, Tx, π), where X= ×i Xi, Tx={Tx1 , · · · , Txn }, and π=[π1, · · · , πn]. The latent state of whole team is denoted as x=[x1, · · · , xn].

3.3 Problem Statement

In the classical MAIL setting, the goal is to learn the team policy from a set of teamwork demonstrations 1 τ .= (s0:h , a0:h ), where h denotes the demonstration length. In our setting, however, team behavior is additionally influenced by the trajectories of team members' latent states χ .= (x0:h ), which are partially observable and resource-intensive to annotate. Hence, we focus on semi-supervised learning of team policies, where x-labels are available for only a subset of the demonstrations.

Formally, our problem corresponds to learning the team policy π, given the MMDP task model (n, S,A, T,R, γ), a set of d observable demonstrations, τ1:d .= {τm}dm=1 , and labels of x for a subset l(≤ d) of the demonstrations, χ1:l .= {χm}lm=1 . For our running example, the problem corresponds to recovering the behavioral policy of the twoagent team (Alice and Rob), given d observable trajectories of agent and box locations τ1:d and labels of each agents' preferred target locations for a subset of the trajectories χ1:l.

4 Solution: Static Latent States

For ease of exposition, we first derive the policy learning algorithm for the case wherein team members do not change their latent mental model (x) during task execution; mathematically, Tx .= 1(x=x′ ). In the next section, we build upon the solution derived for this special case to solve the overall problem of Sec. 3.3. We note that, despite the simplification of static latent states, the learner needs to reason under partial state observability to learn the team policy.

Generative Model. To enable policy learning from a small number of demonstrations, we utilize a Bayesian approach and provide a generative model of team behavior. The generative model, shown in Fig. 2 for a two-agent team, models the process of generating team demonstrations. Each agent (indicated by the subscript) selects its action, ati ∼ πi(·|st , xi), at time step t, based on the current task state st , her mental state xi, and policy πi. In this special case, the latent state xi does not change during task execution and, thus, is represented without a superscript. The dynamics of task state st+1 depend on the MMDP model Ts(·|st , at ).

1 Notation: We use superscript to denote the time step. The subscript is overloaded and, based on the context, is used to denote the i-th agent, m-th demonstration, task state s, agent's latent state x, or action a. Finally, for notational convenience, we use 1abc(a′ , b′ , c′ ) to denote 1(a′ =a, b′ =b, c′ =c).

To complete the generative process, the model additionally includes priors (not shown in Fig. 2) for the latent state and policy. In absence of any additional domain knowledge, we assume that the policy is given as a Categorical distribution. Hence, we define the prior of a policy as its conjugate prior, Dirichlet distribution: πi,sx ∼ Dir(uπ ), where uπ =(uπ1 , · · · , uπ|A| ) are hyperparameters. Similarly, we assume that the latent state is drawn from the uniform distribution unless additional information is given: xi ∼ Uni(X). Given the generative model and data of semi-supervised demonstrations (τ1:d, χ1:l), the policy can be learned by maximizing the posterior: p(π|τ1:d, χ1:l).

Policy Learning with Supervision (l=d). For the case when labels of x are available for the entire dataset, we can directly compute the posterior distribution of the policy as:

Policy Learning with Semi-supervision (l < d). When labels of the latent state are only partially available, the likelihood p(τ1:d, χ1:l|π) cannot be readily computed as it depends on unknown variables, namely, the subset of the data for which latent states x labels are unavailable: {xm}m>l. Hence, to calculate the posterior distribution in a computational tractable manner, we explore paradigms for approximate Bayesian computation. We note the computing the posterior through exact inference (i.e., brute-force) is intractable due to the high-dimensional nature of our problem. Inspired by prior work on single-agent modeling [Johnson and Willsky, 2014; Unhelkar and Shah, 2019], we utilize meanfield variational inference (MFVI) [Beal, 2003] and derive an MFVI algorithm for modeling team policies. In MFVI, the posterior of the team policy is approximated as the variational distribution q(π) that maximizes the evidence lower bound (ELBO). For our problem, the ELBO corresponds to:

The solution for the optimization problem, argmaxq L(q), corresponds to the iterative computation of the local q(x) and global q(π) variational distributions until convergence:

where, Z = ∑x′∈X exp [∑τ ln ˜πi,sx′ a] is the partition function, the operator ˜ A denotes exp (Eq′ (A) [lnA]), and

Given the posterior q(π), the policy is simply estimated as the maximum a posteriori (MAP) estimate: ˆπ = argmaxπ q(π).

5 Solution: Dynamic Latent States

In contrast to the special case considered above, in practice, each member's task-specific mental model may evolve during the task. Hence, next, we extend the solution presented for static latent states to solve the general problem of Sec. 3.3. Analogous to the previous section, we first provide a generative model of team behavior and provide a MFVI-based approach for computing the MAP estimate of policy.

Generative Model. For the general setting, the generative process of team behavior additionally needs to model the temporal evolution of the team members' mental models. Thus, as shown in Fig. 3, we augment the generative model of Fig. 2 to include each agent's time-varying latent states xti and their dynamics: Txi .= P(xt+1 i |st , xti , at , st+1 ). In general, how the mental model evolves may not be known a priori and, thus, requires specification of a prior distribution. Similarly to the policy prior, we utilize Dirichlet distribution as the prior for the latent state transition model, i.e, Txi,sxas′ ∼ Dir(uT1 , · · · , uT|X| ), where uTx are hyperparameters. Both Tx and π jointly influence the likelihood of the semi-supervised demonstrations of team behavior, p(τ1:d, χ1:l|Tx, π), where Tx specifies the distribution of the next latent state xt+1 and π that of the action at . Due to this dependence, to recover the team policy using MFVI, we need an approach to compute their joint posterior p(Tx, π|τ1:d, χ1:l).

Bayesian Team Imitation Learner (BTIL). Similar to Sec. 4, our solution to the overall problem - abbreviated as BTIL- includes iterative computation of variational distributions. However, in contrast to the previous section, the variational distributions additionally include the posterior of latent state dynamics Tx. The BTIL algorithm builds upon MFVI and is derived by maximizing the following ELBO:

Algorithm 1 Bayesian Team Imitation Learner (BTIL)

Input: τ1:d, χ1:l

Parameters: uπ , uTx , N, Ts

• 1: Initialize wπ i ,wTx i for i=1 : n
• 2: Initialize posterior of all unlabeled states q({x0:h m }m>l)
• 3: while L(q) converges do
• 4: Update the variational parameters wπ 1:n , wTx 1:n
• 5: for all τm do
• 6: Compute forward F and backward B messages
• 7: Update posterior of all unlabeled states q({x0:h m }m>l)
• 8: end for
• 9: end while
• 10: Compute the policy posterior q(π) ∼ Dir(wπ i )
• 11: return argmaxπ q(π)

Alg. 1 provides the pseudocode of BTIL, which approximates the posterior distribution p(Tx, π|τ1:d, χ1:l) as independent variational distributions q(π)q(Tx), where q(πi)=Dir(wπ i ) and q(Txi )=Dir(wTx i ). The estimates of posterior distributions are improved by iteratively updating the variational parameters wπ , wTx (line 4). Similar to Eq. 5, the variational parameters are updated as:

An estimate of the posterior distribution of unlabeled states q(x) is required to compute the expectations in Eq. 7. This local variational distribution is given as follows:

To compute Eq. 8 in a tractable manner, we define the following forward-backward messages. These messages are be computed in a recursive manner (line 6) as follows:

For notational convenience, the subscript m is omitted. Given the forward F and backward messages B, BTIL computes the required local probabilities (line 7) as follows:

The time complexity of forward and backward messaging passing subroutine is O(h|X|2n ). Given the converged posterior q(π) in line 10, the team policy is estimated as the MAP estimate (line 11).

6 Experiments

We evaluate BTIL using two domains, Movers and Cleanup, which include aforementioned features and challenges of real world teamwork. These domains build upon the cooperative box pushing task of [Oliehoek and Amato, 2016], include unambiguous latent preferences, and allow for opportunities of (mis)-alignment of team members' mental models. Video demonstrations of collaborative task execution in these domains are included in the supplementary material.

Due to the latent nature of mental models, labels of human teammembers' latent state cannot be ascertained without significant manual effort and annotation. Further, existing multiagent dataset, to our knowledge, have not recorded mental models at each time step of collaborative tasks. Hence, to conduct the proposed experiments and compute pertinent metrics, we create two novel datasets for each domain: one synthetically generated and the other collected via human subject experimentation.

6.1 Domains

Movers. This domain realizes the running example of Sec. 1 in a 7 × 7 grid world. The two member team of Alice and Rob is tasked with carrying boxes to the goal position (flag). Boxes cannot be picked up by one agent alone; hence, to efficiently complete the task, the agents should coordinate on their latent preference over which box to pick or drop next. Each box can be either on its original location, held by both agents, or on the goal location. As described in the running example, agents can take one of the six actions at each step. In this domain, there are 38988 observable states and five possible mental states (corresponding to the three box pickup locations and two drop off locations), resulting in around 200000 states affecting each team member's decisions.

Cleanup. This domain has a similar configuration as Movers, but the environment includes lighter trash bags are placed instead of heavy boxes. Trash bags can only be picked up by one agent; hence, to effectively complete this task, each agent should carry different trash bags as much as possible. Each trash bag can be either on its original location, held by one of the agents, or be dropped at the goal location. In this domain, there are in total over 450000 states affecting agent decisions. The set of primitive actions and latent preferences available to each agent are same as the Movers domain.

6.2 Baselines andMetrics

Due to the novel features of our problem setting, to the best of our knowledge, existing algorithms do not readily apply to the general version of our problem. Existing MAIL solutions either do not model mental states, assume them to be aligned across all team members, or model them as time invariant. Hence, we benchmark our approach against baselines on special cases of our problem. In our design of experiments, we divide our problem into four settings based on two criteria: (a) whether the transition model of latent states Tx is known a priori or not, and (b) whether the latent states are completely labeled or not. We apply the behavioral cloning (BC) and MAGAIL as baselines for the setting of complete labels and the known transition model Tx [Pomerleau, 1991; Song et al., 2018]. The implementation of BC and MAGAIL is similar to those used in [Ho and Ermon, 2016; Song et al., 2018] respectively but adapted to handle discrete states. For the settings where labels are only partially available, we cannot apply existing algorithms; instead, we compare the performance of the supervised (BTIL-Sup) and the semi-supervised (BTIL-Semi) version of our approach. Implementation details of BTIL and the baselines are provided in Appendix D.

We evaluate our approach's ability to effectively learn team policies using the weighted Jensen-Shannon divergence (JS Div.) between the true and learned policies. Like [Unhelkar and Shah, 2019], the policy divergence metric is weighted by the relative counts of states (s, x) observed in the training set. The policy learning performance can only be computed in experiments with synthetic data, where the ground truth policy is known. Hence, in addition, we compare the ability to decode the unlabeled latent states using learned policies. In particular, we utilize the normalized Hamming distance (Hamming) between the decoded ˆx0:h i,m and true x0:h i,m sequences of latent states as the decoding metric. To decode the teammember's latent states, as detailed in Appendix B, we extend the algorithm presented by [Seo et al., 2021]. Lastly, to have a better sense of the worst case values of these highly nonlinear metrics, we utilize the Random baseline, which models the team policy as a Uniform distribution.

6.3 Results on Synthetic Data

We first present results on the synthetic dataset, which are summarized in Table 1. These experiments evaluate BTIL in two settings: with and without prior knowledge of Tx. We present additional results in Appendix E.

Data of Multi-agent Teamwork. The first dataset is synthetically generated by simulating teamwork between two artificial agents. For each domain, we implement the Markovian task model Ts, specify ground truth policies πi, and transitions of the team members Txi . To arrive at the agent policies, we specify rewards associated with each latent state, utilize value iteration to compute Q-values, and derive stochastic policies πi using the softmax operation over Q-values. Execution sequences are created by first assigning initial latent states xi to each team member and, then iteratively, (a) sampling team members' action ai ∼ π(·|s, xi), (b) sampling the next state s′ ∼ T(·|s, a), and (c) sampling the next latent state x′i ∼ Txi (·|xi, s, a, s′ ) until the task termination criteria or 200 time steps are reached. For each domain, we generate 200 demonstrations for training and 100 for evaluation. The proportion of suboptimal training demonstrations (as defined in Appendix C) is 49% and 7% for Movers and Cleanup, respectively.

BTIL Outperforms Baselines in Fully Supervised Settings. The first four rows of Table 1 provide results with 20 labeled demonstrations (d=l=20) and Tx as inputs. For this setting in the Movers domain, we observe that the supervised version of our algorithm (BTIL-Sup) learns more accurate team policies than the baselines. Other baselines (BC, MAGAIL) performed no better than the Random baseline. In Cleanup, our algorithm also outperformed other baselines in the policy learning metric (JS Div.). Despite a small training set, BTIL can learn the team policy with low JS-Divergence by effectively leveraging the x-labels and knowledge of Tx.

In terms of the decoding metric (Hamming), BC showed better results than BTIL in Cleanup domain. We posit that this trend occurs due to a combination of two reasons. First, the Movers domain requires tighter coordination relative to the Cleanup domain. In Movers, teammates must agree on which object to pick next to achieve coordination while in Cleanup, they only need to ensure that they are not picking the same object next. Second, the decoding metric (Hamming) assesses learning performance only on a subset of the state space (i.e., the states encountered in the test set); while JS Div. assess learning performance for the entire state space and, thus, is a better indicator of generalizability. This explains why, even in Cleanup, BC outperforms BTIL only on the decoding metric. In the Movers domain, which requires tighter coordination, BC performs poorly in both metrics.

BTIL Is Capable of Learning Team Policies Without Prior Knowledge of Tx. As shown in the bottom section of Table 1, BTIL can maintain its policy learning performance (JS Div.) even when Tx is unknown. These results suggest that BTIL is capable of learning Tx along with π and that joint learning of (π, Tx) is essential to imitation learning of team policies. Somewhat unsurprisingly, when the decoding algorithm utilizes the learnt Tx, the latent state decoding performance degrades relative to the known Tx case.

BTIL Effectively Utilizes Unsupervised Demonstrations to Improve Team Policy Learning. Lastly, we compare the performance of BTIL under semi-supervision, i.e., the general setting of Sec. 3.3. For these trials (denoted as BTILSup), we provide the algorithm additional 180 demonstrations without latent state labels (d=200, l=20). Comparing performance of BTIL-Sup and BTIL-Semi, we observe improvement in policy learning performance, highlighting the ability of BTIL to effectively leverage available unsupervised data. This ability is particular critical in practice, where collecting data of (s, a)-tuples can be significantly less resource intensive than arriving at the labels of mental states. In additional experiments reported in Appendix E, which further investigate the effect of training set size and amount of semisupervision on the learning performance, we observe that the decoding performance is enhanced with more labeled data and semi-supervision provides the most benefit when the labeled training set is small.

6.4 Results on Data of Human-AI Teamwork

Synthetically generated data, while being useful in validating policy learning performance, cannot capture the variability in behavior demonstrated by humans and human-agent teams. Hence, to benchmark our approach in more realistic settings, we evaluate our algorithm on a novel dataset of human-agent teamwork through an ethics board-approved human subject experiment with 33 participants (16 female, 17 male, mean age: 26.7±5.3 years), who were recruited at Rice University.

Data Collection Procedure. To collect this novel dataset of human-AI teamwork, we designed a web-based interface shown in Fig. 4. Through this web-based interface participants completed tasks in the two domains (Movers and Cleanup) with an AI teammate. The human participant served the role of Alice, while the AI teammate served the role of Rob depicted as a robot avatar.

When the experiment starts, participants are asked to complete a short demographic survey. Further, before the participants conduct any tasks, they were provided with an interactive tutorial to make them familiar with the task and the interface. Upon completing the interactive tutorial, the experiment consisted of 9 sessions: 4 sessions on the Movers domains and 5 tasks for Cleanup. For each domain, the first two In this task,

• Your goal is to move all the boxes to the flag as soon as possible.
• You can only pick up the box circled in red.
• You cannot pick up the box alone.

Special instructions for the test session,

• You will be periodically prompted to select your destination (red circle).
• Since you will be selecting the destination, you can use this added flexibility to complete the task faster.

sessions were used as practice sessions. These practice sessions were designed for participants to gain further familiarity with the user interface. Through the tutorials and practice sessions, each participant was requested to complete tasks collaboratively with a robot avatar which behaves according to an AI policy. The AI policy is generated similar to the synthetic experiment, i.e, by specifying rewards associated with each mental model, running value iteration, and taking the softmax function over Q-values.

For each task, participants' actions (clicking action buttons) are logged along with the task state (location of Alice, Rob, and the boxes) and their chosen destination (mental state corresponding to which box to pick up or drop next). In general, collecting ground truth values of participants' mental states is challenging; in these experiments, we achieve this through a destination-selection interface detailed in Appendix C. By encouraging them to complete the task as soon as possible through the provided instructions and by showing "Your Best" score, we gamify the collaborative task and expect participants to take goal-oriented actions.

Through this experiment, we collect 66 trajectories for Movers, and 99 trajectories for Cleanup, of which every time step is labeled. The average lengths of trajectories (in terms of time steps) are 68.1 in Movers and 40.1 in Cleanup. The proportions of suboptimal demonstrations are 17% and 7% in Movers and Cleanup, respectively. We use two thirds of the collected trajectories for training and the rest for evaluation.

Performance of BTIL Observed with Data of Multi-agent Teamwork Translates to Learning Policies of Humanagent Teamwork. As we cannot ascertain the true policy of a human, we utilize only the state decoding metric in these experiments. Further, as would be the case in practice, the mental model dynamics are also unavailable; hence, Tx needs to be learned by our algorithm and the MAGAIL baseline cannot be applied. Table 2 summarizes the decoding performance computed using π and Tx learned by variants of BTIL. All variants of BTIL are supplied with same amount of (s, a)- demonstrations but varying amount of x-labels. BTIL, even with a small amount of supervision, significantly outperform the Random baseline. In conjunction with the results computed with synthetic data, these experiments provide proofof-concept in the ability of BTIL to learn team policies from small semi-supervised datasets of optimal and suboptimal teamwork.

7 Concluding Remarks

We provide BTIL, a Bayesian approach to learn team policies from demonstrations of suboptimal teamwork. In most collaboration scenarios, it is challenging to collect large labeled datasets of teamwork due to changes in team membership, adaptations in team policies, and the need for labeling latent states. Inspired by these and other aspects of teamwork observed in practice, BTIL includes multiple desirable features, including (a) the ability to learn from small sets of semi-supervised data, (b) explicit modeling of team members' mental models and model alignment and (c) the ability to jointly infer team policy, latent state dynamics, and latent states. We confirm the ability of our algorithm to learn team policies on two novel datasets of teamwork, including one of human-AI teamwork. Our work also offers several avenues of future work, including the ability to consider collaborative tasks where the task state s itself may be partially observable and consideration of communicative actions.

Acknowledgments

We thank the anonymous reviewers for their detailed and constructive feedback. Sangwon Seo was partially supported by the Army Research Office through Cooperative Agreement Number W911NF-20-2-0214. The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the Army Research Office or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein.