Decentralized policy learning with partial observation and mechanical constraints for multiperson modeling

Here we consider a multi-agent problem that can be formulated as a decentralized POMDP [Bernstein et al.,, 2002, Kaelbling et al.,, 1998, Mao et al.,, 2019, Amato et al.,, 2019]. It is defined as a tuple $(K,S,A,\mathcal{T},R,O,Z,\gamma)$, where $K$ is the fixed number of agents,; $S$ is the set of states $s$; $A=[A_{1},\ldots,A_{K}]$ represents the set of joint action $\vec{a}\in A$ (for a variable number of agents), with $A_{k}$ being the set of local action $a_{k}$ that agent $k$ can take; $\mathcal{T}(s^{\prime}|s,\vec{a}):S\times A\times S\rightarrow[0,1]$ is the transition model; $R=[R_{1},\ldots,R_{K}]:S\times A\rightarrow\mathbb{R}^{K}$ is the joint reward function; $O=[O_{1},\ldots,O_{K}]$ is the set of joint observation $\vec{o}\in O$ controlled by the observation function $Z:S\times A\rightarrow O$; and $\gamma\in[0,1]$ is the discount factor. In on-policy reinforcement learning, the agent learns a policy $\pi_{k}:O_{k}\times A\rightarrow[0,1]$ that can maximize $\mathbb{E}[G_{k}]$, where $G_{k}=\sum_{t=1}^{T}\gamma^{t}R_{k}^{t}$ is the discounted return and $T$ is the time horizon. In a decentralized multi-agent system in complex real-world environments (e.g., team sports), transition and reward functions are difficult to design explicitly. Instead, if we can utilize the demonstrations of expert behaviors (e.g., trajectories of professional sports players), we can formulate and solve our problem as imitation learning (i.e., as a machine-learning problem). In other words, if the problem satisfies the two conditions, imitation learning is one of the better options. Here, for example, in team sports, the state includes all the position, velocity, and acceleration of all players. The actions include each player’s next velocity or acceleration, and the observations include information for a limited number of other players.

The goal in imitation learning is to learn a policy $\pi$ that imitates an expert policy $\pi_{E}$ given demonstrations from that expert [Schaal,, 1996, Ross et al.,, 2011]. In multi-agent imitation learning, we have $K$ agents to achieve a common goal. On the basis of the notation in [Ross et al.,, 2011], in the case of a fully centralized multi-agent policy for clarity, let $\vec{\pi}(\vec{o}):=\vec{a}$ denote the joint policy that maps the joint observation $\vec{o}=[o_{1},\ldots,o_{K}]$ into $K$ actions $\vec{a}=[a_{1},\ldots,a_{K}]$. Training data $\mathcal{D}$ consists of multiple demonstrations of $K$ agents. The decentralized setting decomposes the joint policy $\vec{\pi}=[\pi_{1},\ldots,\pi_{K}]$ into $K$ policies, sometimes tailored to each specific agent index or role. The loss function is then $\mathcal{L}_{imitation}=\sum_{k=1}^{K}\mathbb{E}_{s_{k}\sim d_{\pi_{k}}}\left[\ell(\pi_{k}(o_{k}))\right],$ where $d_{\pi_{k}}$ is the distribution of states experienced by joint policy $\pi_{k}$ ($o_{k}$ is determined by $s_{k}$ and $a_{k}$) and $\ell$ is the imitation loss defined over the demonstrations. This formulation assumes the assignment of appropriate roles for each agent such as solving a linear assignment problem [Papadimitriou and Steiglitz,, 1982] (see Appendix B).

Partial observation in the context of the POMDP typically indicates that the agent cannot directly observe the underlying state. However, in real multiperson systems, people sometimes actively ignore less informative others. Since these two cases cannot be distinguished from the data obtained, we formulate both cases as the former POMDP. Here we introduce several approaches for neural network–based partial observation using soft and hard attention mechanisms. The details and other approaches are described in Appendix A.

Attention is widely used in natural language processing [Bahdanau et al.,, 2014], computer vision [Wang et al.,, 2018], and multi-agent reinforcement learning (MARL) in virtual environments [Mao et al.,, 2019, Jiang and Lu,, 2018, Iqbal and Sha,, 2019]. Soft attention calculates an importance distribution of elements (e.g., agents). Soft attention is fully differentiable and thus can be trained with back-propagation. However, it usually assigns nonzero probabilities to unrelated elements; i.e., it cannot directly reduce the number of agents.

Hard attention focuses solely on an important element but is basically nondifferentiable because of the selection based on sampling. Among differentiable models with discrete variables, an approach using Gumbel-softmax reparameterization (see, e.g., Jang et al., [2017]) can be effectively implemented via a continuous relaxation of a discrete categorical distribution. However, the resulting one-hot vector is insufficient to represent a multi-hot observation. Our model enables a multi-hot representation as the observation in multi-agent systems for each agent. Other approaches, such as relational inference (see, e.g., Kipf et al., [2018]), can learn multi-agent interactions by learning graph structures (see also Appendix A). Since we aim to model decentralized long-term multi-agent behaviors, we adopt an approach to explicitly train independent dynamical models of each agent.

Agent trajectories or actions in the real world have been modeled recently as sequential generative models (see also Section 4). Here we consider single-agent modeling for simplicity. Let $a_{\leq T}=\{a_{1},\dots,a_{T}\}$ denote a sequence of actions of length $T$. The goal of sequential generative modeling is to learn the distribution over sequential data $\mathcal{D}$ consisting of multiple demonstrations. We assume that all sequences have the same length $T$, but in general, this does not need to be the case. A common approach is to factorize the joint distribution and then maximize the log-likelihood $\theta^{*}=\operatorname*{arg\,max}_{\theta}\sum_{a_{\leq T}\in\mathcal{D}}\log p_{\theta}(a_{\leq T})=\operatorname*{arg\,max}_{\theta}\sum_{a_{\leq T}\in\mathcal{D}}\sum_{t=1}^{T}\log p_{\theta}(a_{t}|a_{<t}),$ where $\theta$ denotes the model’s learnable parameters, such as recurrent neural networks (RNNs).

However, RNNs with simple output distributions often struggle to capture highly variable and structured sequential data (e.g., multimodal behaviors) [Zhan et al.,, 2019]. Recent work in sequential generative models addressed this issue by injecting stochastic latent variables into the model and optimization using amortized variational inference to learn the latent variables (see, e.g., Chung et al., [2015], Fraccaro et al., [2016], Goyal et al., [2017]); see Section 4. Among these methods, a variational RNN (VRNN) [Chung et al.,, 2015] has been widely used in base models for multi-agent trajectories [Yeh et al.,, 2019, Zhan et al.,, 2019] with unknown governing equations (see Section 3.1). Note that our VRNN-based approach is compatible with other sequential generative models. A VRNN is essentially a variational autoencoder (VAE) [Kingma and Welling,, 2014] conditioned on the hidden state of an RNN and is trained by maximization of the (sequential) evidence lower bound (ELBO):

where $z$ is a stochastic latent variable. The first term is the reconstruction term and $p_{\theta}(a_{t}\mid z_{\leq t},a_{<t})$ is the generative model. The second term is the Kullback-Leibler (KL) divergence between the approximate posterior or inference model $q_{\phi}(z_{t}\mid a_{\leq t},z_{<t})$ and the prior $p_{\theta}(z_{t}\mid a_{<t},z_{<t})$. Eq. (1) can be interpreted as the ELBO of a VAE summed over each timestep $t$ (see also Appendix C).

In robotics and computational neuroscience, the path planning in smooth and flexible biological motions is a classical problem (see, e.g., Flash and Hogan, [1985]). Most research has focused on individual multijoint motions. Since multi-agent datasets often include only a mass point for each agent, we can use the principles of smooth mass-point motions. A minimum-jerk principle [Flash and Hogan,, 1985] is a simple and well-known principle for smooth motor planning of voluntary motion. It minimizes the motor cost: $C=(1/2)\int_{0}^{T}||(d^{3}x/dt^{3})||^{2}_{2}dt$, where $x$ is a multidimensional position vector and $\|\cdot\|_{2}$ is the Euclidean norm. There is some research showing that this principle can be applied to human locomotion [Pham et al.,, 2007] and driving behavior [Dias et al.,, 2020]. We propose a biomechanical constraint for learning a multi-agent model inspired by this principle for the first time.

As an imitation-learning problem for decentralized multi-agent systems in a POMDP, we aim to learn policies $\vec{\pi}=[\pi_{1},\ldots,\pi_{K}]$ that imitate expert policies $\vec{\pi}_{E}=[\pi_{E_{1}},\ldots,\pi_{E_{K}}]$ given the expert action sequences of $K$ agents $a_{\leq T}=\{a_{\leq T,1},\dots,a_{\leq T,K}\}$ under biologically realistic constraints. As shown in Fig. 1A, the agent perceives the state $s_{t-1}$ and performs path planning for the action $a_{t,k}$. In this work, to interpret our model as a decentralized multi-agent policy, we train independent models for each agent.

Our model consists of a partial observation, local policy, macro goal, and mechanical constraints. According to the data flow, we first describe a perception and a binary vector $b_{t,k}$ of the partial observation model defined in Section 3.2. For example, we consider that, in team sports, the players perceive other players’ information and obtain an abstract representation (as an embedded vector) and selectively choose information for a limited number of players (as the partial observation). Next, the model also uses macro goals [Zhan et al.,, 2019] as auxiliary information (from a pretrained model) for long-term prediction modified for our decentralized and partially observable problem (see Appendix D). Note that human behaviors are generated by a policy function including tactical movements for which rules are partially unknown [Fujii et al.,, 2019] as shown in Fig. 1B. We therefore adopt a nonlinear transition model (see Section 4) for the local policy $\pi_{k}$, specifically based on a VRNN [Chung et al.,, 2015], which can generate a trajectory with multiple variations (see, e.g., Yeh et al., [2019], Zhan et al., [2019], Sun et al., [2019]). With the obtained observation as an input, the macro goal and local policy are learned. The macro-goal model is pretrained before the policy learning, and for the policy learning, we use the macro goal as weak supervision [Zhan et al.,, 2019]. For example, in team sports, players move, including the intention of where to move for a length of a few seconds (we model it as the macro goal) while adapting to the situation in the moment.

In addition, we consider mechanical constraints as a penalty of the objective function. In previous work (e.g., Zhan et al., [2019], Yeh et al., [2019]), the action generated by $\pi_{k}$ was usually the agent’s position. However, these studies reported difficulty in learning velocity and acceleration (e.g., frequent changes in direction of motion in basketball), which are critical to our problem because of the complex movements as shown in Fig. 1B. We aim to obtain policy models to generate biologically plausible actions in terms of position, velocity, acceleration, and their relations. Our solution incorporates mechanical constraints as a penalty function into policy learning. We describe the details for the policy model, macro goal, and mechanical constraints in Section 3.3.

Note that we entirely consider biologically plausible architectures and loss functions of the neural network, whereas we sometimes exploit existing modules based on previous work, which do not necessarily consider the biological plausibility and cannot often provide concrete examples such as in team sports.

To model partial observation in a multi-agent setting, we propose a binary partial observation model that can visualize whose information the agents utilize. In this work, this is represented by the binary vector $b_{t,k}$. As shown in Fig. 1A (perception and partial observation modules), the model is trained to map input state vector $s_{t}=[s_{t,1},\ldots,s_{t,K}]\in\mathbb{R}^{d_{s}\times K}$ to a latent binary vector $b_{t,k}$ of an agent $k$ at each time $t$, where $d_{s}$ is the dimension of the state for each agent. We consider a partially observable environment, where each agent $k$ receives a binary vector $b_{t,k}:=[e_{k,1}^{t},e_{k,2}^{t},\ldots,e_{k,K}^{t}]\in\{0,1\}^{K}$. That is, $b_{t,k}$ consists of an arbitrary number of 0’s and 1’s. In this article, $b_{t,k}$ for an agent $k$ is designed to represent the importance of all agents, used as the observation coefficients or interpretable partial observation described below, which can visualize whose information the agents utilize in a multi-agent system.

To obtain the binarized or multi-hot differentiable vector $b_{t,k}$, we linearly project a state vector $s_{t,k}$ into embedded matrices $f_{t,k}\in\mathbb{R}^{K\times d_{e}}$ (perception module with the full observation) and into dual channels $f^{\prime}_{t,k}\in\mathbb{R}^{K\times 2}$ (partial observation module), where $d_{e}$ is the embedding dimension. The former embedding can transform the state into a distributed representation, which can be directly weighted by $b_{t,k}$ as described below (e.g., Cartesian coordinates cannot be directly weighted in principle). The latter dual-channel projection inspired by [Liu et al.,, 2019] allows each dimension in $f^{\prime}_{t,k}$ (and finally in $b_{t,k}$) to represent multiple agents’ importance (not limited to the one-hot approach in Section 2.2). We perform a categorical reparameterization trick with Gumbel-softmax reparameterization [Jang et al.,, 2017] on the second dimension of the dual channel $f^{\prime}_{t,k}$ (rather than on the first dimension as used in the attention mechanisms). Since the two channels are linked together by Gumbel-softmax reparameterization, it is sufficient to simply pick one of them. That is, $e_{k,i}^{t}=[\operatorname{Gumbel-softmax}([f^{\prime}_{t,k}]_{i})]_{1}$, where $[\cdot]_{i}$ denotes $[\cdot]$’s $i$th element. This method is differentiable and allows direct back-propagation in our framework.

We then compute the observation vector $o_{t,k}$ as the input to the subsequent policy learning. We concatenate the elementwise product of binary vector $b_{t,k}$ (observation coefficients) and the embedded matrices $f_{t,k}$ for all agents: $o_{t,k}=[e_{k,1}^{t}[f_{t,k}]_{1},\ldots,e_{k,K}^{t}[f_{t,k}]_{K}]\in\mathbb{R}^{d_{e}K}$. This allows us to eliminate the information from unrelated agents, and to focus on only the important agents.

First, we perform VRNN modeling with conditional context information including the observation and macro goals, as illustrated in Fig. 2A. Here we describe three important components: decentralized macro goals, local policy, and mechanical constraints.

The objective function becomes a sequential ELBO of the hierarchical VRNN and the penalties of the mechanical constraints for the probabilistic model. The ELBO for agent $k$ is $\mathcal{L}_{vrnn}=$ $\mathbb{E}_{q_{\phi}(z_{\leq T}\mid o_{\leq t},g^{\prime}_{\leq t})}\sum_{t=1}^{T}[\log p_{\theta}(a_{t}\mid z_{\leq t},o_{<t},g^{\prime}_{\leq t})-D_{KL}(q_{\phi}(z_{t}\mid o_{\leq t},z_{<t},g^{\prime}_{\leq t})||p_{\theta}(z_{t}\mid o_{<t},z_{<t},g^{\prime}_{\leq t}))].$ For brevity, the symbol of the agent $k$ is not shown here. The penalties for mechanical constraints $\mathcal{L}_{body}$ are

where $m^{\prime}=\{vel,acc\}$ indicates velocity and acceleration as action outputs, depending on experimental conditions. The first term is the penalty for consistently learning the relationship between the directly and indirectly estimated dimensions. If we eliminate some dimensions (e.g., position), we can add NLLs as the penalty between the distribution of the indirect prediction and the ground truth: $-\log p_{\theta}(a_{m,t}\mid z_{\leq t},o_{<t},g^{\prime}_{<t})$. The second term is the biomechanical smoothing penalty, which minimizes the difference between the distribution of the directly or indirectly predicted acceleration ($\hat{a}_{acc,t}$ or $\tilde{a}_{acc,t}$) and the observed next acceleration $a_{acc,t+1}$. We jointly maximize $\mathcal{L}_{vrnn}+\lambda\mathcal{L}_{body}$ with respect to all the model parameters (for the specific regularization parameter $\lambda$, see Appendix E).

There has recently been increasing interest in deep generative models for sequential data, because of the flexibility of deep learning and (often probabilistic) generative models. In particular, many researchers intensively developed methods for physical systems with governing equations, such as a bouncing ball and a pendulum. These studies modeled latent linear state space dynamics (e.g., Karl et al., [2017], Fraccaro et al., [2017]) or ordinary differential equations (e.g., Chen et al., 2018b ) mainly without RNNs.

In biological systems with partially unknown governing equations (e.g., external forces and/or nontrivial interactions), RNN-based models are still used (see below). Among the RNN-based generative models (e.g., Chung et al., [2015], Fraccaro et al., [2016]), we incorporate a VRNN [Chung et al.,, 2015] with a stochastic latent state into the observation model and prior knowledge about mechanics.

In existing research on various biological multi-agent trajectories, pedestrian prediction problems including rule-based models (see, e.g., Helbing and Molnar, [1995]) have been widely investigated for a long time. Recent studies using RNN-based models (e.g., Alahi et al., [2016], Gupta et al., [2018]) effectively aggregated information across multiple persons using specialized pooling modules, but most of which predicted trajectories in a centralized manner. Vehicle trajectory prediction has been intensively researched, often with use of deep generative models based on RNNs [Bansal et al.,, 2018, Rhinehart et al.,, 2019, Tang and Salakhutdinov,, 2019], whereas little research has focused on animals [Eyjolfsdottir et al.,, 2017, Johnson et al.,, 2016, Tsutsui et al.,, 2021]. For sports multi-agent trajectory prediction such as in basketball and soccer, most methods have leveraged RNNs [Zheng et al.,, 2016, Le et al.,, 2017, Ivanovic et al.,, 2018, Teranishi et al.,, 2020], including VRNNs [Zhan et al.,, 2019, Yeh et al.,, 2019, Teranishi et al.,, 2022, Fujii et al.,, 2022], although some have used generative adversarial networks (e.g., Chen et al., 2018a , Hsieh et al., [2019]) without RNNs (see also the review in Fujii, [2021]). Most of these studies assumed full observation to achieve long-term prediction in a centralized manner (e.g., Zhan et al., [2019], Yeh et al., [2019]), except for an image-based study on partial observation [Sun et al.,, 2019]. In contrast, we aim to model decentralized biological multi-agent systems and visualize their partial observations for behavioral analyses in real-world agents.

In purely rule-based models, researchers investigating animals and vehicles conventionally proposed predefined observation rules based on specific distances and visual angles (e.g., Couzin et al., [2002]), the specific number of the nearest agents (e.g., Ballerini et al., [2008]), or other environments (e.g., Yoshihara et al., [2017]). However, it is difficult to define interactions between agents using predefined rules in general large-scale multi-agent systems. Thus, methods for learning observation of agents have been proposed for MARL in virtual environments [Mao et al.,, 2019, Jiang and Lu,, 2018, Iqbal and Sha,, 2019, Liu et al.,, 2020] and for real-world multi-agent systems [Hoshen,, 2017, Leurent and Mercat,, 2019, Guangyu et al.,, 2020, Fujii et al.,, 2020]. We describe the soft and hard attention approaches we used in Appendix A, whereas other approaches, such as relational or causal inference, can learn multi-agent interactions by learning graph structures (e.g., Kipf et al., [2018], Graber and Schwing, [2020], Löwe et al., [2020]) or sparse weights of the first layer (e.g., Tank et al., [2018], Khanna and Tan, [2019], Fujii et al., [2021]). These methods include a physically interpretable approach [Fujii and Kawahara,, 2019, Fujii et al.,, 2020] that can learn interactions especially in physical particles or oscillators. Instead, our approach aims to model decentralized smooth multi-agent systems that can also predict long-term behaviors.

We used the basketball dataset from the NBA 2015–2016 season (https://www.stats.com/data-science/), which contains tracking trajectories for professional basketball players and the ball. The data was preprocessed such that the offense team always moved toward the left side of the court. We chose 100 games so that the amount of data was similar to that for the subsequent soccer dataset. In total, the dataset contained 19968 training sequences, 2235 validation sequences, and 2608 test sequences.

To demonstrate the generality of our approach, we also used a soccer dataset [Le et al.,, 2017, Yeh et al.,, 2019], containing trajectories of soccer players and the ball from 45 professional soccer league games. We randomly split the dataset into 21504 training sequences, 2165 validation sequences, and 2452 test sequences. We did not model the goalkeepers since they tend to move very little. Table 1 (right) and Table 2 (right) indicate the effectiveness of our mechanical constraints and the similar trajectory prediction performance of our binary observation models and the baselines, results that were similar to those for the basketball dataset. For VRNN-macro-Bi-Mech, the averaged $b_{t,k}$ was $8.04\pm 1.54$ for each defender (maximum 23). The ratio of no relation for each defender in our method is $0\%$. Additional analysis to compare the distances in observation and an illustrative example are given in Appendices I and J, respectively.