SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Wentao Zhu $^{1*}$  Jason Qin $^{1*}$  Yuke Lou $^{1}$  Hang Ye $^{1}$
Xiaoxuan Ma<sup>1</sup> Hai Ci<sup>1</sup> Yizhou Wang<sup>1,2</sup>
Humans exhibit a remarkable capacity for anticipating the actions of others and planning their own actions accordingly. In this study, we strive to replicate this ability by addressing the social motion prediction problem. We introduce a new benchmark, a novel formulation, and a cognition-inspired framework. We present Wusi, a 3D multi-person motion dataset under the context of team sports, which features intense and strategic human interactions and diverse pose distributions. By reformulating the problem from a multi-agent reinforcement learning perspective, we incorporate behavioral cloning and generative adversarial imitation learning to boost learning efficiency and generalization. Furthermore, we take into account the cognitive aspects of the human social action planning process and develop a cognitive hierarchy framework to predict strategic human social interactions. We conduct comprehensive experiments to validate the effectiveness of our proposed dataset and approach.
Human beings are inherently social creatures. Concretely, individuals unconsciously anticipate the actions of others and make informed decisions about their own behaviors in social contexts [13, 19, 54], which enables individuals to cooperate and compete with others in a variety of scenarios, from pedestrian traffic [25, 44] to team sports [64, 50]. Notably, task experts demonstrate exceptional skill in predicting others' movements in advance [37, 70].
To better understand and replicate this ability, the research community has proposed the task of future prediction for multiple interacting agents given their historical observations. The majority of prior work in this area focuses on modeling and predicting agent interactions at the trajectory level [42, 39, 52, 5, 34, 58, 75, 59, 43, 72] and has demonstrated promising results in applications such as autonomous driving [38, 49, 17, 36, 16, 27]. However, trajectory-based approaches can only reflect coarse-grained interactions (e.g., collision avoidance, social distancing) and fail to capture the rich and fine-grained human actions. To this end, some studies have investigated the multi-person motion prediction problem, which aims to forecast both trajectories and poses for a group of people [2, 3, 67, 68, 22, 63, 48, 73]. Despite recent advancements, there remain several critical challenges in this field. Firstly, existing multi-person motion datasets are primarily designed for human pose estimation tasks [65, 1, 31], and consequently, do not place particular emphasis on human interactions. Individuals in these datasets tend to move casually and interact with others at random, making future predictions both difficult and less meaningful. Secondly, the majority of prior methods concentrate on developing neural network architectures for end-to-end supervised training while overlooking the cognitive aspects of human social action planning. These two challenges are closely related and necessitate a comprehensive solution.
Figure 1: Example sequences from Wusi dataset. Three players in purple form the offensive team, and two players in orange form the defensive team. Left: a successful pass for the offensive team. Middle: a successful steal for the defensive team. Right: a successful pass following a fake pass.
In this work, we propose novel perspectives on this problem to address the aforementioned limitations. We begin by constructing a large-scale multi-person 3D motion dataset featuring intense and strategic interactions among participants. To achieve this, we turn to team sports, which offer several inherent advantages: 1) Well-defined game rules and global rewards implicitly constrain and guide individual actions. 2) Participants develop intricate interaction strategies based on their roles, such as planning cooperative actions with teammates while acting adversarially towards opponents. Skilled players even employ sophisticated techniques like deception, etc. 3) Human motions exhibit greater dynamism in terms of pose diversity and motion intensity, making motion prediction more challenging than in previous datasets.
Additionally, we present a new formulation of the multi-person motion prediction task as a multi-agent reinforcement learning (MARL) problem. Specifically, we model the task using imitation learning, where the objective is to learn a policy from expert demonstrations. We employ behavioral cloning (BC) [7, 51] to imitate expert behaviors from the dataset. To improve learning efficiency and generalization, we utilize generative adversarial imitation learning (GAIL) [26], aiming to render agents' policy indistinguishable from experts' policy. Furthermore, we propose a framework for human social interactions based on the cognitive hierarchy theory [10]. In particular, we assume that people base their decisions on their predictions regarding the likely actions of others, while others engage in similar decision-making processes from their perspectives. By considering the reasoning steps recursively, we posit that a level  $k$  agent takes actions based on the level  $k - 1$  agent actions, where  $k$  represents the depth of strategic thought. Building on this insight, we develop a computation model that can be elegantly integrated with the MARL formulation.
We summarize our contributions as follows: 1) We present Wusi, the first large-scale multi-human 3D motion dataset featuring intense and strategic interactions. We demonstrate that our dataset has greater motion diversity than existing datasets and poses a more significant challenge for the social motion prediction problem. 2) We propose a novel MARL formulation for the problem and develop an imitation learning baseline that combines behavioral cloning and generative adversarial imitation learning. 3) We introduce a cognitive hierarchy framework to model the strategic and game-theoretic human social interactions. Our approach outperforms the state-of-the-art methods in challenging long-term social motion predictions.
2.1 Multi-agent trajectory prediction
The multi-agent trajectory prediction problem has been extensively studied, especially for interacting traffic participants, e.g., pedestrians and vehicles. There exist multiple industrial benchmarks including Argoverse [12], nuScenes [8], Waymo Open Motion Dataset [18]. Meanwhile, various generation methods have been explored to model the interaction among different agents and predict their future trajectories [42, 39, 52, 5, 34, 58, 75, 59, 43, 72]. In this work, we aim to predict intense and strategic interactions for a group of humans with fine-grained body motions.
2.2 Multi-person motion dataset
The existing multi-human motion datasets can be categorized into 2D and 3D. The PoseTrack dataset [6, 29] offers video sequences with manually annotated 2D keypoints. However, 2D keypoints fail to accurately represent the proximity of individuals in the real world, and manual annotation is difficult to scale. Meanwhile, 3D multi-person motion datasets can be effectively constructed with motion capture (mocap) systems, including marker-based solutions and markerless ones. 3DPW [65] uses the inertial measurement units (IMUs) to obtain high-quality motion reconstructions. UMPM [61] employs a set of reflective markers to identify the joint positions. However, wearable sensors may struggle to capture complex movements and become expensive as the number of individuals increases. Other datasets, e.g., CMU-Mocap [1], Panoptic [31], MuPoTs-3D [41], ExPI [22], employ markerless mocap systems with multi-view cameras to address the occlusion problem and obtain 3D human motion through triangulation. In this work, we construct a multi-person motion dataset featuring high motion diversity and intense social interactions using multi-view cameras. We provide a comparison with previous datasets in Section 3.3.
2.3 Multi-person motion prediction
Joo et al. [32] propose to predict human motion conditioned on other individuals' kinesic signals in a triadic haggling scenario. Adeli et al. [2] introduce the 2D social motion forecasting (SoMoF) benchmark, which aims to predict multi-person trajectory and pose on the PoseTrack dataset [6]. They also propose a baseline model using a shared GRU encoder and a pooling layer to incorporate social clues. TRiPOD [3] employs attention graphs to characterize the spatiotemporal social interactions. It also considers the joint occlusion and body invisibility issues arising from 2D observations. Futuremotion [67] provides a simple baseline using no social context, yet achieving competitive performance on SoMoF. As 2D multi-person motion prediction faces challenges such as data scarcity, depth ambiguity, and occlusions, more recent research has shifted towards exploring the multi-person motion prediction problem in 3D. Wang et al. [68] propose a multi-range Transformer to separately encode the local and global motion history. Guo et al. [22] design a cross-interaction-attention (XIA) module to model close interactions between pairs in duo dance scenarios. Vendrow et al. [63] introduce joint-aware attention and joint-wise query to predict the entire future sequence without recurrence. Peng et al. [48] present a social-aware motion attention mechanism that models both inter- and intra-individual motion relations. DuMMF [73] frames the problem as a dual-level generative task and instantiates it with various generative models. In contrast to previous work, we adopt a MARL perspective and model the strategic social interaction process using cognitive hierarchies.
2.4 Cognitive hierarchy theory
Cognitive hierarchy theory (CHT) is a model in behavioral economics and game theory that aims to describe human decision processes in strategic games. Several foundational works introduce the concept of CHT [10], and provide experimental evidence [45, 14, 23, 57, 56] supporting the existence of cognitive hierarchies in human decision-making processes. Taking it a step further, Wright et al. [71] develop a Poisson cognitive hierarchy model to predict human behavior in normal-form games. Li et al. [35] summarize how the information structure guided by the cognitive hierarchy supports belief generation and policy generation in game-theoretic multi-agent learning. In this work, we propose a cognitive hierarchy framework in conjunction with MARL, and demonstrate its effectiveness in a real-world social motion prediction scenario.
We present a multi-person 3D motion dataset with a special focus on strategic interactions called Wusi (Wusi Basketball Training Dataset). In the following, we first introduce the contents of the dataset, then outline the data collection pipeline, and finally provide statistical analysis and comparison with regard to existing datasets.
Our dataset captures the no-dribble-3-on-2 basketball drills performed by a team of professional basketball athletes. In each drill, three offensive players possess the ball, while another two players
Table 1: Dataset comparison. We compare our dataset with existing multi-person motion datasets employed by previous works on the multi-person motion prediction task.  ${}^{ \dagger  }$  denotes multi-person subset as utilized in previous works [2, 68].
Dataset	2D/3D	Frames	Duration (min)	No. of people	Interaction
PoseTrack [6]†	2D	8K	5.5	Multiple	Weak
CMU-Mocap [1]†	3D	34K	4.9	2	Weak
3DPW [65]†	3D	5K	6.1	2	Weak
MuPoTs-3D [41]	3D	8K	4.4	2-3	Weak
ExPI [22]	3D	30K	20.0	2	Cooperative
Wusi	3D	60K	40.3	5	Strategic
are on defense. The offensive team aims to accomplish as many successful passes as possible within the given time, while the defensive team strives to get steals, deflections, and slow down the offense. Since dribbling is prohibited, the drill requires the offensive players to make better passing decisions and foresee the defense before they pass; conversely, the defensive players need to anticipate the passing directions in order to steal the ball. Figure 1 showcases the diverse and dynamic human interactions present in our dataset. On the left, the offensive player successfully passes the ball to a teammate. In the middle, the defensive player accurately anticipates the passing trajectory and successfully steals the ball. On the right, the offensive player pretends to pass the ball to the right, then swiftly executes a pass to the left.
3.2 Data collection
Our motion capture system consists of 11 synchronized and calibrated wide-baseline cameras. Our setup ensures that the players are well-surrounded by the cameras from elevated shooting angles, enabling each body joint to be covered by at least four different camera views. We utilize a markerless multi-person 3D pose estimation algorithm [74] with JDETracker [69]. Subsequently, we apply a post-processing pipeline that includes automatic failure detection and temporal filtering [11]. A final manual check is conducted to guarantee the integrity of the data sequences.
3.3 Data analysis
In this section, we examine the statistics of Wusi dataset and compare it with existing datasets for multi-person motion prediction. As outlined in Table 1, current datasets exhibit two primary limitations for the multi-person motion prediction task. Firstly, high-quality 3D motion data is limited in terms of scale (duration and the number of people). Consequently, prior works [68, 48, 63, 73] resort to randomly mixing the motion sequences [1, 65, 41] in order to generate more multi-person motion sequences for training. However, such practice inevitably causes unnatural human interactions. Secondly, the interaction strength in existing datasets is restricted, mainly featuring simple actions such as walking together [6, 1, 65, 41]. More recently, the ExPI dataset [22] introduces cooperative 3D motions of two Lindy-hop dancers. In contrast, our dataset offers 2 to 9 times the video duration of existing datasets and includes strategic global interactions among 5 individuals.
Table 2: Comparison of pose diversity based on different thresholds.
Threshold	50mm	100mm
Human3.6M [28]	24%	12%
ExPI [22]	52%	23%
CMU-Mocap [1]	20%	9%
MuPoTs-3D [41]	37%	19%
Wusi	53%	27%
Figure 2: Comparison of motion intensity.
Additionally, we provide a quantitative comparison with existing datasets, emphasizing two critical aspects: pose diversity and motion intensity. Firstly, we examine the pose diversity within the datasets by calculating the ratio of unique poses to the total number of poses following [28, 22]. Table 2 shows that our dataset surpasses the previous datasets in terms of overall pose diversity. We further analyze the motion intensity by computing the average velocity for all the body joints. As illustrated in Table 2, our dataset exhibits more dynamic movements across all body joints. We suppose that these characteristics make Wusi not only challenging for motion prediction, but also potentially useful for other related tasks, e.g. motion generation [76].
The multi-person motion prediction problem can be formulated as follows: given the motion histories  $\{x_t^p\}_{1\leq t\leq T,1\leq p\leq P}$  of length  $T$  from  $P$  human subjects, predict their future motion  $\{x_{t}^{p}\}_{T\leq t\leq T + T^{\prime},1\leq p\leq P}$  of length  $T^{\prime}$ , where each  $x_{t}^{p}$  represents a 3D human pose.
We model the problem using a Markov Decision Process (MDP) [60] by dividing the prediction sequence into  $L$  steps with even step length  $m = \frac{T'}{L}$  following [66], and further extend the formulation to multiple agents. We use bold symbols to represent the collective variables of all agents following the convention in MARL. At each step  $i$ , The state  $\pmb{s}_i$  is defined as the aggregated motion history  $\{x_t^i\}_{1 \leq t \leq T + (i-1) \times m, 1 \leq p \leq P}$  for all the agents. The action for the  $p$ -th agent  $a^p$  is defined as a sequence of velocities  $\{v_t^p\}_{T + (i-1) \times m \leq t \leq T + i \times m-1}$ , where  $v_t^p = x_{t+1}^p - x_t^p$ . Therefore, the joint action of all  $P$  agents  $\pmb{a}_i$  deterministically transition the MDP into the new state  $\pmb{s}_{i+1}$  according to:
$$
x _ {t ^ {*}} ^ {p} = x _ {T + (i - 1) \times M} ^ {p} + \sum_ {t = T + (i - 1) \times M} ^ {t ^ {*}} v _ {t} ^ {p}, \quad T + (i - 1) \times m + \leq t ^ {*} \leq T + i \times m, 1 \leq p \leq P. \tag {1}
$$
The goal of this objective is to learn a policy  $\pi^p$  for each agent  $p$  that maps each possible state  $s$  to agent action  $a^p$  by  $a^p = \pi^p(s)$ . In the real world, people optimize their  $\pi^p$  to maximize an implicit team reward  $r$ . Since we have no access to  $r$  or the environment, we aim to learn  $\pi^p$  from expert demonstrations  $\mathcal{D}$  via imitation learning.
We parameterize  $\pi^p$  with a set of parameters  $\theta$ . In practice, we follow [68] to employ a global-range Transformer  $E_g$  to encode global state feature  $s_g = E_g(s)$ , and a local-range Transformer  $E_l$  to encode local state feature  $s_l^p = E_l(s)$  for each agent  $p$ . See Figure 3 for an overview. This approach helps to extract multi-level state features and enables all agents to share policy network parameters.
4.2 Behavioral Cloning
A straightforward method in imitation learning is Behavioral Cloning (BC) [7, 51], i.e., using expert demonstrations to minimize the difference between the actions produced by the policy and those taken by the experts:
$$
\theta^ {*} = \operatorname {a r g m i n} _ {\theta} \sum_ {\mathcal {D}} \sum_ {p = 1} ^ {P} \ell \left(\pi_ {\theta} ^ {p} (\boldsymbol {s}), \tilde {\pi} ^ {p} (\boldsymbol {s})\right), \tag {2}
$$
Figure 3: Framework overview. For the  $p$ -th agent, two state encoders first extract its local and global state features  $s_l^p$  and  $s_g$ , from which the level-0 policy network produces an initial action  $a_{(0)}^p$ . The level- $k$  policy network produces action  $a_{(k)}^p$  based on  $s_l^p$  and the joint actions of the previous level  $a_{(k-1)}$  ( $k \geq 1$ ).
where  $s$  is sampled from expert demonstrations  $\mathcal{D}$ ,  $\pi_{\theta}^{p}$  and  $\tilde{\pi}^p$  are policies of the model and the expert for agent  $p$ ,  $\ell$  is a distance function defined in the action space. While BC enjoys the benefit of being both computation-efficient and sample-efficient [53], it also faces several drawbacks. For instance, policies tend to be overfitted with respect to the state distribution encountered by the experts, leading to suboptimal generalization ability [66, 30]. To cope with this challenge, we further introduce GAIL and cognitive hierarchy reasoning.
4.3 Generative Adversarial Imitation Learning
Generative Adversarial Imitation Learning (GAIL) [26] is a model-free imitation learning approach. It leverages the adversarial training framework of Generative Adversarial Networks (GAN) [21] and the concept of inferring reward functions from Inverse Reinforcement Learning (IRL) [46]. Direct estimation of reward signals from expert demonstrations can be very difficult; therefore, GAIL transforms the IRL problem into its equivalent dual problem of occupancy measure matching. Specifically, the policy network is regularized to match its state-action pair distribution to that of the experts' policy through adversarial training. We implement GAIL using a global discriminator  $D$  parameterized by  $\omega$ , which differentiates the state-action pairs produced by agents' joint policy  $\pi$  and experts' joint policy  $\tilde{\pi}$ . The optimization objective can be formulated as:
$$
\min  _ {\theta} \max  _ {\omega} \mathbb {E} _ {\boldsymbol {\pi} _ {\theta}} \left[ \log D _ {\omega} (\boldsymbol {s}, \boldsymbol {\pi} _ {\theta} (\boldsymbol {s})) \right] + \mathbb {E} _ {\tilde {\boldsymbol {\pi}}} \left[ \log \left(1 - D _ {\omega} (\tilde {\boldsymbol {s}}, \tilde {\boldsymbol {\pi}} (\tilde {\boldsymbol {s}}))\right) \right], \tag {3}
$$
where  $s$  and  $\tilde{s}$  are sampled from expert demonstrations  $\mathcal{D}$  independently.
4.4 Cognitive hierarchies
Furthermore, we notice that in real-world situations, individuals would predict others' behaviors by forming beliefs about others' policies and then acting accordingly to maximize their own payoffs [20, 9]. This can be formulated as a hierarchical, game-theoretic decision-making process, in which lower-level agents adopt straightforward strategies, while higher-level players anticipate the strategies of lower-level agents and respond accordingly. We explicitly model this recursive reasoning process to learn more interpretable and robust agent policies. We adopt a specific type of cognitive hierarchy model [10] called level- $k$  thinking [57, 45, 15] to represent this process. Specifically, agents at each level (except for the lowest one) assume that others are reasoning at the previous level. As shown in Figure 3, a straightforward policy is to take actions based on local and global state features  $s_l^p$  and  $s_g$ , which we denote as level-0 policy  $\pi_{(0)}^p$ . Level-0 actions are thus defined as:
$$
a _ {(0)} ^ {p} = \pi_ {(0)} ^ {p} (\boldsymbol {s}) = \phi_ {(0)} \left(s _ {l} ^ {p}, \boldsymbol {s} _ {g}\right), \tag {4}
$$
where  $\phi_{(0)}$  is the level-0 policy network. Then for a level- $k$  agent ( $k \geq 1$ ), it takes actions based on the joint agent actions  $a_{(k-1)}$  at the previous level and its local state features  $s_l^p$ :
$$
a _ {(k)} ^ {p} = \pi_ {(k)} ^ {p} (\boldsymbol {s}) = \phi_ {(k)} \left(s _ {l} ^ {p}, \boldsymbol {a} _ {(k - 1)}\right), \quad 1 \leq k \leq K \tag {5}
$$
where  $\phi_{(k)}$  is the level-  $k$  policy network,  $K$  is the maximum strategic depth.
4.5 Training objectives
Finally, we introduce our training objectives which effectively integrate the aforementioned components. For level- $k$  ( $1 \leq k \leq K$ ) agents, we apply GAIL to regularize the distance between their joint policy and the experts' policy:
$$
\mathcal {L} _ {\mathrm {G A I L}} = \sum_ {k = 1} ^ {K} \mathbb {E} _ {\boldsymbol {\pi} _ {(k)}} \left[ \log D \left(\boldsymbol {s}, \boldsymbol {\pi} _ {(k)} (\boldsymbol {s})\right) \right] + \mathbb {E} _ {\tilde {\boldsymbol {\pi}}} \left[ \log \left(1 - D \left(\tilde {\boldsymbol {s}}, \tilde {\boldsymbol {\pi}} (\tilde {\boldsymbol {s}})\right)\right) \right]. \tag {6}
$$
In addition, we employ BC on level- $K$  agent policies:
$$
\mathcal {L} _ {\mathrm {B C}} = \sum_ {\mathcal {D}} \sum_ {p = 1} ^ {P} \ell \left(\pi_ {(K)} ^ {p} (\boldsymbol {s}), \tilde {\pi} ^ {p} (\boldsymbol {s})\right). \tag {7}
$$
We train the discriminator  $D$  and the policy networks  $\phi$  alternatively, where  $D$  aims to maximize  $\mathcal{L}_{\mathrm{GAIL}}$ , and  $\phi$  is optimized to minimize a linear combination of  $\mathcal{L}_{\mathrm{BC}}$  and  $\mathcal{L}_{\mathrm{GAIL}}$  using a policy gradient algorithm following [30]:
$$
\mathcal {L} _ {\phi} = \mathcal {L} _ {\mathrm {B C}} + \lambda \mathcal {L} _ {\mathrm {G A I L}}, \tag {8}
$$
where  $\lambda$  is a constant for balancing the loss terms.
Implementation details. We implement the presented framework to train and test on the proposed Wusi dataset. We employ Transformer encoder [62] for both the local and global state encoders, as well as Transformer decoders for the policy networks. Each Transformer consists of 3 layers with 8 attention heads. We share parameters for policy networks  $\phi_{(1)}\ldots \phi_{(K)}$ . We set the strategic reasoning depth  $K = 3$  unless otherwise stated. For all the methods, we provide 1s motion history and predict future 1s motion. For additional experimental details, please refer to the appendix.
Evaluation metrics. We compute the Mean Per Joint Position Error (MPJPE) between the predicted future motion and ground-truth. In order to disentangle the results of trajectory prediction and pose prediction, we further calculate the mean root position error and mean local pose error (MPJPE after root alignment). We report all the errors in millimeters.
5.2 Evaluation and comparison