baller2vec++: A Look-Ahead Multi-Entity Transformer For Modeling Coordinated Agents

In many multi-agent spatiotemporal systems , agents operate under the influence of1 shared , unobserved variables ( e.g. , the play a team is executing in a game of bas-2 ketball ) . As a result , the trajectories of the agents are often statistically dependent3 at any given time step ; however , almost universally , multi-agent models implicitly4 assume the agents ’ trajectories are statistically independent at each time step . In5 this paper , we introduce baller2vec++1 , a multi-entity Transformer that can effec-6 tively model coordinated agents . Specifically , baller2vec++ applies a specially7 designed self-attention mask to a mixture of location and “ look-ahead ” trajectory8 sequences to learn the distributions of statistically dependent agent trajectories . We9 show that , unlike baller2vec ( baller2vec++ ’ s predecessor ) , baller2vec++10 can learn to emulate the behavior of perfectly coordinated agents in a simulated11 toy dataset . Additionally , when modeling the trajectories of professional basketball12 players , baller2vec++ outperforms baller2vec by a wide margin.13 1 Introduction and Related Work14 Whether it is a team executing a play in a game of basketball , a family navigating to an attraction in a15 theme park , or friends posting about a birthday party on a social media platform , humans frequently16 coordinate their behavior in response to shared information . When this coordinating information is17 unobserved ( which is often the case in many machine learning datasets ) , the individuals ’ observed18 behaviors become correlated , i.e. , the behavior of one individual at a specific moment contains19 information about the behavior of another individual at the same time . In the context of modeling20 agent trajectories in multi-agent spatiotemporal systems , this property translates to the trajectories21 being statistically dependent at each time step . However , nearly all multi-agent spatiotemporal models22 ( e.g. , [ 1–6 ] ) implicitly ( through their loss functions ) assume the trajectories of the agents at each time23 step are statistically independent given the agents ’ previous locations ( Figure 1 ) .24 Zhan et al . [ 7 ] explicitly focused on modeling coordinated multi-agent trajectories , using “ macro-25 intents ” [ 8 ] that are shared across agents to do so . The macro-intents are generated from a26 separately trained recurrent neural network ( RNN ) that learns to predict a future , coarse , “ sta-27 tionary ” location for each agent at each time step . The macro-intents for all of the agents at28 a specific time step are concatenated together to form a single , shared , macro-intent variable,29 which is then provided as input to the trajectories-generating model at that time step . However,30 similar to the previously mentioned multi-agent trajectory models , the macro-intents model im-31 1All data and code for the paper are available at : < anonymized > . Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . plicitly assumes the macro-intents for the agents at each time step are statistically independent,32 i.e. , the macro-intent for one agent does not depend on the macro-intents of the other agents.233 tially account for trajectory correla-44 tions through a global adversarial loss.45 Specifically , the trajectories for each46 agent are separately passed through47 an encoder RNN , and the final hid-48 den states for each agent are then pro-49 cessed with a graph attention network50 ( GAT ) [ 10 ] . The output of the GAT51 is then used as an input to a discrim-52 inator that classifies whether or not53 the input trajectories are real or gen-54 erated . This global adversarial loss,55 however , is only a single component56 of the full Social-BiGAT loss func-57 tion , and other components of the loss58 function do implicitly make the inde-59 pendence assumption . Further , inter-60 estingly , adding the global discriminator to a baseline model only improved the model ’ s performance61 for one out of six pedestrian datasets.62 In this paper , we describe a novel multi-agent spatiotemporal model that integrates information about63 concurrent actions of agents to predict statistically dependent distributions of trajectories . Specifically,64 we extend the recently introduced multi-entity Transformer baller2vec [ 6 ] by : ( 1 ) augmenting its65 input with a parallel sequence of “ look-ahead ” agent trajectories and ( 2 ) using a specially designed66 self-attention mask , which allows our model to exploit the chain rule of probability ( Section 3 ) . We67 find that:68 1. baller2vec++ is an effective learning algorithm for modeling coordinated agents . Unlike69 baller2vec , baller2vec++ can learn to emulate perfectly coordinated agents from a70 simulated toy dataset ( Section 5.1 ) . Further , baller2vec++ outperforms baller2vec by71 a wide margin ( 8.9 % ) when modeling the trajectories of professional basketball players72 ( Section 5.1 ) .73 2. baller2vec++ makes better predictions when conditioned on concurrent trajectory infor-74 mation from other agents , supporting our proposition that the commonly used independence75 assumption for agent trajectories is overly strong ( Section 5.2 ) .76 3 . Lastly , the joint probability assigned to a sequence by baller2vec++ is approximately77 permutation invariant with respect to the order of the agents , i.e. , baller2vec++ respects78 the properties of the chain rule ( Section 5.3 ) .79 2See the authors ’ implementation here : https : //github.com/ezhan94/ multiagent-programmatic-supervision/blob/a1d9152d4c8a287474953cba093c28fef2a05979/ models/macro_vrnn.py # L101 . 2 Background80 2.1 Multi-agent trajectory modelling81 Our problem description closely follows Alcorn and Nguyen [ 6 ] , whom we quote here:82 Let A = { 1 , 2 , . . . , B } be a set indexing B agents and P = { p1 , p2 , . . . , pK } ⊂83 A be the K agents involved in a particular sequence . [ Let ] Ct =84 { ( xt,1 , yt,1 ) , ( xt,2 , yt,2 ) , . . . , ( xt , K , yt , K ) } [ be ] an unordered set of K coordinate85 pairs such that ( xt , k , yt , k ) are the coordinates for agent pk at time step t. The86 ordered sequence of sets of coordinates C = ( C1 , C2 , . . . , CT ) , together with P ,87 thus defines the trajectories for the K agents over T time steps.88 In multi-agent trajectory modeling , the goal is to model a joint probability of the form:89 p ( ∆xt,1 , ∆xt,2 , . . . , ∆xt , K |x1 : t,1 , x1 : t,2 , . . . , x1 : t , K ) i.e. , the joint probability of the K agents ’ trajectories ∆xt , k at time step t given the agents ’ location90 histories x1 : t , k . We note here that the common practice of simultaneously predicting the trajectories91 for all of the agents at a specific time step is not required by theory . Using the chain rule of probability,92 the joint probability of the agents ’ trajectories can be factorized as , e.g . :93 p ( ∆xt,1 , ∆xt,2 , . . . , ∆xt , K ) = p ( ∆xt,1 ) p ( ∆xt,2|∆xt,1 ) . . . p ( ∆xt , K |∆xt,1 , ∆xt,2 , . . . , ∆xt , K−1 ) where we omit the conditional historical trajectories for brevity . As a result , it is perfectly acceptable to94 generate trajectories agent-wise , using the previously generated trajectories as additional conditioning95 information when generating the trajectories for later agents ( see Figure 2 ) .96 2.2 baller2vec is a ( conditional ) generative model.97 baller2vec is a recently described multi-entity Transformer that can model sequences of sets98 ( the underlying data structure for multi-agent spatiotemporal systems ) , as opposed to sequences99 of individual inputs ( like words in a sentence ) . When used to model the game of basketball , the100 input at each time step for baller2vec is a set of feature vectors where each feature vector contains101 information about the identity and location of a player on the court . baller2vec maps each input102 feature vector to an output feature vector , which is then used to “ classify ” the binned trajectory for103 that specific player at that specific time step.104 Here , we provide a probabilistic interpretation of baller2vec , which establishes the theoret-105 ical grounds for using the chain rule to generate trajectories agent-wise at each time step in106 baller2vec++ . Without loss of generality , we only consider one-dimensional trajectories for107 a single agent here . To briefly summarize , the outputs of the softmax function over the n binned108 trajectories in baller2vec can be interpreted as mixture proportions for a mixture of uniform distri-109 butions with predetermined bounds that partition the Euclidean trajectory space . Further , because110 xt+1 = xt + ∆xt , baller2vec is in fact a conditional generative model that assigns a probability111 to a sequence of trajectories given the initial position of the agent , i.e. , p ( ∆x1 , ∆x2 , . . . , ∆xT |x1 ) .112 Using the chain rule , we decompose the joint probability of the trajectories as:113 p ( ∆x1 , ∆x2 , . . . , ∆xT ) = p ( ∆x1 ) p ( ∆x2|∆x1 ) . . . p ( ∆xT |∆x1 , ∆x2 , . . . , ∆xT−1 ) to reflect their temporal structure ( we omit the conditional initial position term for brevity ) . Therefore,114 new trajectories can be generated from baller2vec with the following procedure ( see Figure 3 ) :115 1 . First , sample one of the n different mixture components using the mixture proportions116 output from the classifier f ( i.e. , baller2vec ) conditioned on the agent ’ s current position,117 i.e. , i ∼ Categorical ( π1 , π2 , . . . , πn ) where [ π1 , π2 , . . . , πn ] = f ( xt ) .118 2 . Next , sample a trajectory from the uniform distribution associated with the sampled compo-119 nent , i.e. , ∆xt ∼ U ( ai , bi ) .120 3 . Finally , add the sampled trajectory to the agent ’ s input position to generate the agent ’ s121 position at the start of the next time step , i.e. , xt+1 = xt + ∆xt.122 Let [ ∆xmin , ∆xmax ) be an interval on the real line such that any trajectory ∆x < 123 ∆xmin or ∆x ≥ ∆xmax has zero density ( i.e. , such trajectories are humanly impossible ) .124 Let { [ ai , bi ) } ni=1 be a set of n intervals that par-125 tition the interval [ ∆xmin , ∆xmax ) into n bins,126 i.e. , ∪ni=1 [ ai , bi ) = [ ∆xmin , ∆xmax ) and i 6=127 j =⇒ [ ai , bi ) ∩ [ aj , bj ) = ∅ . Recall that the128 probability density function ( PDF ) for a uniform129 distribution with bounds −∞ < a < b < ∞ is:130 p ( ∆x ) = { 1 b−a for ∆x ∈ [ a , b ) 0 otherwise Letting ci = 1bi−ai , the PDF for a mixture of131 uniforms with these bounds is thus:132 p ( ∆x ) = n∑ i=1 πiU ( ∆x ; ai , bi ) = n∑ i=1 πici ( 1 ) where p ( ∆x ) is the density assigned to ∆x133 by the mixture , πi is the mixture proportion134 for the mixture component indexed by i ( i.e.,135 0 ≤ πi ≤ 1 and ∑ πi = 1 ) , and U ( ∆x ; ai , bi ) 136 is the density assigned to ∆x by the uniform dis-137 tribution with bounds −∞ < ai < bi < ∞.138 Because the bounds of the uniform distribu-139 tions partition [ ∆xmin , ∆xmax ) , Equation ( 1 ) 140 reduces to:141 p ( ∆x ) = πi′ci′ where ∆x ∈ [ ai′ , bi′ ) ( because the other uni-142 form distributions will assign a density of zero to143 ∆x ) . The likelihood for data D ( with |D| = N ) 144 is then:145 L ( D ) = N∏ j=1 p ( ∆xj ) = N∏ j=1 πj , i′cj , i′ where πj , i′ is the mixture proportion assigned to the component with ∆xj ∈ [ ai′ , bi′ ) and cj , i′ is the146 associated density . Taking the negative logarithm of the likelihood gives:147 − ln ( L ( D ) ) = − N∑ j=1 ln ( πj , i′ ) − N∑ j=1 ln ( cj , i′ ) ( 2 ) Because the bounds are fixed , the second summation is a constant , and Equation ( 2 ) becomes:148 − ln ( L ( D ) ) = − N∑ j=1 ln ( πj , i′ ) + C where C = − ∑N j=1 ln ( cj , i′ ) . Therefore , minimizing the loss of baller2vec as a classifier of149 binned trajectories is equivalent to minimizing the loss of the model when assuming the trajectories150 are generated from a mixture of uniform distributions as specified in Equation ( 1 ) .151 3 Model Architecture152 We motivate our baller2vec++ architecture by first highlighting an issue that arises in baller2vec153 when trying to model agent trajectories using the chain rule . The baller2vec self-attention mask154 can be adapted so that baller2vec “ looks ahead ” at the future positions of agents whose trajectories155 are generated prior to the agent being processed in the current time step ( Figure 4 ) . However , this156 look-ahead self-attention mask can only be used with the final layer of the Transformer ; otherwise,157 the model needs to see the future ( Figure 4 ) . As a result , baller2vec is severely limited in the158 conditional distribution functions it can learn.159 baller2vec++ ( Figure 5 ) overcomes this limitation by : ( 1 ) augmenting the baller2vec input with160 two other sets of feature vectors and ( 2 ) using a specially designed self-attention mask . The three sets161 of feature vectors in baller2vec++ take the following forms:162 1. zt , k = gz ( [ e ( pk ) , xt , k , yt , k , ht , k ] ) ( current location information ) 163 2. ut , k = gu ( [ e ( pk ) , xt+1 , k , yt+1 , k , ht , k , ∆xt , k , ∆yt , k ] ) ( “ look-ahead ” information ) 164 3. rk = gr ( [ e ( pk ) , x1 , k , y1 , k , h1 , k ] ) ( initial location information ) 165 where gz , gu , and gr are multilayer perceptrons ( MLPs ) , e is an agent embedding layer , and ht , k is a166 vector of optional contextual features for agent pk at time step t. zt , k is the same location feature167 vector used in baller2vec and contains information about a specific agent ’ s identity and the agent ’ s168 location at time step t. ut , k is a “ look-ahead ” trajectory feature vector that contains information about169 a specific agent ’ s identity , the agent ’ s location at the next time step t+ 1 , and the agent ’ s trajectory170 at time step t , i.e. , ( xt+1 , k − xt , k , yt+1 , k − yt , k ) . Lastly , rk is a starting location feature vector that171 contains information about a specific agent ’ s identity and the agent ’ s location at time step t = 1.172 The rk feature vectors are necessary so that baller2vec++ can “ see ” the initial locations of all the173 agents when processing the agents agent-wise in the first time step.174 These three sets of feature vectors are combined to form a ( K + 2TK ) × F matrix Z such that175 the first K rows consist of the K rk feature vectors , and the remaining 2TK rows consist of the176 TK zt , k and TK ut , k feature vectors interleaved with one another , i.e. , each zt , k is followed by its177 corresponding ut , k in the matrix . This matrix is passed into the Transformer along with the specially178 designed self-attention mask , which encodes the following dependencies ( see Figure 5 ) :179 1 . When processing rk1 , baller2vec++ is exclusively allowed to “ look ” at each rk2 ( i.e.,180 baller2vec++ can not look at any location or look-ahead feature vectors when processing181 rk1 ) .182 2 . When processing zt2 , k2 , baller2vec++ is allowed to “ look ” at : ( i ) each rk1 , ( ii ) any zt1 , k1183 where ( a ) t1 < t2 or ( b ) t1 = t2 and k1 ≤ k2 , and ( iii ) any ut1 , k1 where ( a ) t1 < t2 or ( b ) 184 t1 = t2 and k1 < k2.185 3 . When processing ut2 , k2 , baller2vec++ is allowed to “ look ” at : ( i ) each rk1 , ( ii ) any zt1 , k1186 where ( a ) t1 < t2 or ( b ) t1 = t2 and k1 ≤ k2 , and ( iii ) any ut1 , k1 where ( a ) t1 < t2 or ( b ) 187 t1 = t2 and k1 ≤ k2.188 Each processed zt , k feature vector is then passed through a linear layer that is followed by a softmax,189 which gives a probability distribution over the trajectory bins for agent pk at time step t. Similar to190 baller2vec , the loss for each sample is:191 L = T∑ t=1 K∑ k=1 − ln ( f ( Z ) t,2k−1 [ vt , k ] ) ( 3 ) where f ( Z ) t,2k−1 [ vt , k ] is the probability assigned to the trajectory bin vt , k ( where vt , k =192 Bin ( ∆xt , k , ∆yt , k ) is an integer from one to n2 ) by f , i.e. , Equation ( 3 ) is the negative log-likelihood193 ( NLL ) of the data according to the model.194 Because any ordering of a chain rule decomposition of a joint probability produces the same value,195 e.g . :196 p ( ∆xt,1 ) p ( ∆xt,2|∆xt,1 ) p ( ∆xt,3|∆xt,1∆xt,2 ) = p ( ∆xt,3 ) p ( ∆xt,2|∆xt,3 ) p ( ∆xt,1|∆xt,3∆xt,2 ) like [ 11 ] , we shuffled the order of the agents in each training sequence to encourage the model to197 learn joint probabilities of the agent trajectories that are approximately permutation invariant with198 respect to the ordering of the agents.199 4 Experiments200 We tested baller2vec++ on two different datasets . To highlight the pathological behavior of models201 that assume agent trajectories are statistically independent at each time step , we trained scaled down202 versions of baller2vec++ and baller2vec on a toy dataset consisting of simulated trajectories203 for two perfectly coordinated agents . Additionally , to demonstrate the efficacy of baller2vec++204 in real world settings , we trained baller2vec++ and baller2vec on a dataset of trajectories for205 professional basketball players.206 4.1 Toy dataset207 Each training sample was initialized with the agents starting at ( −1 , 0 ) and ( 1 , 0 ) on a grid in random208 order ( i.e. , the first agent could be placed to either the left or the right of the origin ) . At each time209 step , one of nine actions ( corresponding to the 3×3 grid surrounding the agent ) was sampled from a210 uniform distribution , and each of the agents was translated along this trajectory . This process was211 repeated for 20 time steps ( see Figure 6 ( a ) for a sample ) .212 4.2 Basketball dataset213 We used the same National Basketball Association ( NBA ) dataset3 employed by Alcorn and Nguyen214 [ 6 ] , whom we paraphrase here:215 The NBA dataset consists of trajectories from 631 games from the 2015-2016216 season , which were split into 569/30/32 training/validation/test games , respectively.217 During training , each sequence was sampled using the following procedure : ( 1 ) 218 randomly select a training game , ( 2 ) randomly select a starting time from the game,219 ( 3 ) take the following four seconds of data and downsample it to 5 Hz from the220 original 25 Hz , and then ( 4 ) randomly ( with a probability of 0.5 ) rotate the court221 180◦ . This sampling procedure gave us access to on the order of ∼82 million222 different ( albeit overlapping ) training sequences . For both the validation and test223 sets , ∼1,000 different , non-overlapping sequences were selected for evaluation by224 dividing each game into d 1,000N e non-overlapping chunks ( where N is the number225 of games ) , and using the starting four seconds from each chunk as the evaluation226 sequence.227 4.3 Model228 Our baller2vec++ and baller2vec models for the basketball dataset closely followed [ 6 ] , and so229 largely resemble the original Transformer architecture [ 12 ] . Specifically , the Transformer settings230 were : dmodel = 512 ( the dimension of the input and output of each Transformer layer ) , eight attention231 heads , dff = 2048 ( the dimension of the inner feedforward layers ) , six layers , no dropout , and no232 positional encoding . Each MLP ( i.e. , gz , gu , and gr ) had 128 , 256 , and 512 nodes in its three layers,233 respectively , and a ReLU nonlinearity following each of the first two layers . The player embeddings234 [ 13 ] had 20 dimensions , and ht , k was a binary variable indicating the side of the frontcourt for235 player pk ( i.e. , the direction of his team ’ s hoop ) at time step t. Lastly , the 11 ft× 11 ft 2D Euclidean236 trajectory space was binned into 121 1 ft× 1 ft squares.237 We used the Adam optimizer [ 14 ] with an initial learning rate of 10−6 , β1 = 0.9 , β2 = 0.999 , and238 = 10−9 to update the model parameters , of which there were ∼19 million . The learning rate was239 3https : //github.com/linouk23/NBA-Player-Movements reduced to 10−7 after 20 epochs of the validation loss not improving . Models were implemented in240 PyTorch and trained on a single NVIDIA GTX 1080 Ti GPU for ∼650 epochs ( seven days ) where241 each epoch consisted of 20,000 training samples , and the validation set was used for early stopping.242 For the toy dataset , we used scaled down versions of the basketball models with dmodel = 128 , four243 attention heads , dff = 512 , and two layers in the Transformer . Additionally , each MLP had two layers244 with 64 and 128 nodes , respectively . The models were trained for 50 epochs of 500 samples per245 epoch ( ∼10.5 minutes ) using a single learning rate of 10−5.246 5 Results247 5.1 baller2vec++ can effectively model coordinated agents in both simulated and real248 settings249 For the toy dataset , the training loss for baller2vec converged to ∼2.2 ≈ − ln ( 19 ) , i.e. , the model250 was simply independently guessing the trajectories for both agents at every time step . In contrast,251 the training loss for baller2vec++ converged to ∼1.1 ≈ − ln ( 19 ) ÷ 2 , which is the expected loss252 for a model that perfectly learns the deterministic relationship between the agents ’ trajectories253 ( because the prediction for the second agent will always contribute − ln ( 1.0 ) = 0 to the loss ) .254 When generating trajecto-255 ries with baller2vec , the256 agents are completely unco-257 ordinated , with each agent258 following an independent259 random walk around the260 grid ( Figure 6 ( b ) ) . In con-261 trast , trajectories generated262 by baller2vec++ display263 the same coordinated agent264 behavior as the training data265 ( Figure 6 ( c ) ) .266 For the basketball dataset,267 baller2vec++ achieved268 an average NLL of 0.472269 on the test set , 8.9 % better270 than the average NLL for271 baller2vec ( 0.518 ) ( see Figure S1 for trajectories generated by baller2vec++ and baller2vec ) .272 As was observed in [ 6 ] , the trajectory bin distributions for baller2vec become much more certain273 after observing a portion of the sequence ( Figure 7 ) , which suggests baller2vec may be inferring274 some of the shared hidden variables ( e.g. , plays ) influencing the players . If that hypothesis was true,275 the performance gap between baller2vec++ and baller2vec should be largest at the beginning276 of the sequence ( before any shared hidden variables can be inferred by baller2vec ) . Indeed , the277 average NLL for baller2vec++ in the first time step of each test set sequence ( 1.567 ) is 16.1 % 278 better than the average NLL for baller2vec ( 1.869 ) , while the average NLL for baller2vec++ in279 the last time step of each test set sequence ( 0.420 ) is only 9.7 % better than the average NLL for280 baller2vec ( 0.465 ) ( see Figure 7 ) .281 5.2 baller2vec++ makes better predictions when conditioned on concurrent trajectory282 information from other agents283 Implicit in much of our discussion has been the intuition that providing a model with additional284 ( relevant ) information will improve its performance . To empirically test this conjecture , we compared285 the performance of baller2vec++ when predicting the trajectory of a specific basketball player286 placed in the first position of the player order ( i.e. , when k = 1 ) vs. predicting the trajectory for that287 same player placed in the last position ( i.e. , when k = 10 ) . Specifically , for each player in each test288 sequence , we calculated the NLL of the player ’ s trajectory in the first time step4 with the player in the289 4Because , as previously discussed , the benefits of baller2vec++ were most pronounced in the first time step . first position of the player order . Next , we moved the player to the last position of the player order,290 and then randomly shuffled the remaining nine players 10 times , calculating the NLL for the player291 in the last position each time . Finally , we calculated the average percent change in the last position292 NLLs relative to their corresponding first position NLLs . On average , moving a player from the first293 to the last position improved the NLL for the player ’ s trajectory by 14.6 % .294 5.3 The joint probability assigned to a sequence by baller2vec++ is approximately295 permutation invariant with respect to the order of the agents296 To determine whether or not baller2vec++ respects the fact that any ordering of a chain rule297 decomposition of a joint probability produces the same value , we measured how much the average298 NLL for each test sequence in the basketball dataset varied when the order of the agents changed.299 Specifically , for each test set sequence , we shuffled the order of the agents 10 times . Then , for each300 permuted sequence , we calculated the percent error5 in the average NLL relative to the original,301 unshuffled sequence . Across all test sequences , the average percent error was only ±1.5 % . Further,302 as can be seen in Figure 7 , the shuffled average NLLs are highly correlated with their corresponding303 unshuffled average NLLs ( Pearson correlation coefficient = 0.997 ) , i.e. , the joint probability assigned304 to a sequence by baller2vec++ is approximately permutation invariant with respect to the order of305 the agents.306 6 Conclusion and Future Work307 In this paper , we have shown how the commonly used independence assumption of many multi-agent308 spatiotemporal models can severely limit their ability to learn to emulate coordinated agents . By309 relaxing this independence assumption in baller2vec , baller2vec++ was able to more accurately310 model the trajectories of professional basketball players . Models for other multi-agent spatiotemporal311 environments , such as pedestrian traffic ( see [ 15 ] for a survey ) and vehicle traffic ( e.g. , [ 16–19 ] ) , may312 also benefit from the look-ahead approach used by baller2vec++ . Additionally , the interleaved313 input design of baller2vec++ could be useful when modeling other systems involving many entities314 interacting through time , such as social media platforms ( e.g. , [ 20 , 21 ] ) . However , confronting the315 quadratic complexity of the Transformer attention mechanism as the number of entities grows large316 in these datasets is an open problem , but recent work in sparse Transformers ( e.g. , [ 22–27 ] ) shows317 encouraging progress.318 5See : https : //en.wikipedia.org/wiki/Approximation_error # Formal_Definition . References319 [ 1 ] Panna Felsen , Patrick Lucey , and Sujoy Ganguly . Where will they go ? predicting fine-grained320 adversarial multi-agent motion using conditional variational autoencoders . In Proceedings of321 the European Conference on Computer Vision ( ECCV ) , pages 732–747 , 2018.322 [ 2 ] Agrim Gupta , Justin Johnson , Li Fei-Fei , Silvio Savarese , and Alexandre Alahi . Social gan:323 Socially acceptable trajectories with generative adversarial networks . In Proceedings of the324 IEEE Conference on Computer Vision and Pattern Recognition , pages 2255–2264 , 2018.325 [ 3 ] Amir Sadeghian , Vineet Kosaraju , Ali Sadeghian , Noriaki Hirose , Hamid Rezatofighi , and326 Silvio Savarese . Sophie : An attentive gan for predicting paths compliant to social and physical327 constraints . In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern328 Recognition , pages 1349–1358 , 2019.329 [ 4 ] Raymond A Yeh , Alexander G Schwing , Jonathan Huang , and Kevin Murphy . Diverse genera-330 tion for multi-agent sports games . In Proceedings of the IEEE Conference on Computer Vision331 and Pattern Recognition , pages 4610–4619 , 2019.332 [ 5 ] Cunjun Yu , Xiao Ma , Jiawei Ren , Haiyu Zhao , and Shuai Yi . Spatio-temporal graph transformer333 networks for pedestrian trajectory prediction . In Proceedings of the European Conference on334 Computer Vision ( ECCV ) , August 2020.335 [ 6 ] Michael A. Alcorn and Anh Nguyen . baller2vec : A multi-entity transformer for multi-agent336 spatiotemporal modeling . arXiv preprint arXiv:2102.03291 , 2021.337 [ 7 ] Eric Zhan , Stephan Zheng , Yisong Yue , Long Sha , and Patrick Lucey . Generating multi-agent338 trajectories using programmatic weak supervision . In International Conference on Learning339 Representations , 2019 . URL https : //openreview.net/forum ? id=rkxw-hAcFQ.340 [ 8 ] Stephan Zheng , Yisong Yue , and Jennifer Hobbs . Generating long-term trajectories using deep341 hierarchical networks . Advances in Neural Information Processing Systems , 29:1543–1551,342 2016.343 [ 9 ] Vineet Kosaraju , Amir Sadeghian , Roberto Martín-Martín , Ian Reid , Hamid Rezatofighi , and344 Silvio Savarese . Social-bigat : Multimodal trajectory forecasting using bicycle-gan and graph345 attention networks . In H. Wallach , H. Larochelle , A. Beygelzimer , F. d'Alché-Buc , E. Fox,346 and R. Garnett , editors , Advances in Neural Information Processing Systems , volume 32.347 Curran Associates , Inc. , 2019 . URL https : //proceedings.neurips.cc/paper/2019/348 file/d09bf41544a3365a46c9077ebb5e35c3-Paper.pdf.349 [ 10 ] Petar Veličković , Guillem Cucurull , Arantxa Casanova , Adriana Romero , Pietro Liò , and Yoshua350 Bengio . Graph Attention Networks . International Conference on Learning Representations,351 2018 . URL https : //openreview.net/forum ? id=rJXMpikCZ . accepted as poster.352 [ 11 ] Zhilin Yang , Zihang Dai , Yiming Yang , Jaime Carbonell , Ruslan Salakhutdinov , and Quoc V353 Le . Xlnet : Generalized autoregressive pretraining for language understanding . arXiv preprint354 arXiv:1906.08237 , 2019.355 [ 12 ] Ashish Vaswani , Noam Shazeer , Niki Parmar , Jakob Uszkoreit , Llion Jones , Aidan N Gomez,356 Łukasz Kaiser , and Illia Polosukhin . Attention is all you need . In Advances in Neural Informa-357 tion Processing Systems , pages 5998–6008 , 2017.358 [ 13 ] Michael A Alcorn . ( batter|pitcher ) 2vec : Statistic-free talent modeling with neural player359 embeddings . In MIT Sloan Sports Analytics Conference , 2018.360 [ 14 ] Diederik P Kingma and Jimmy Ba . Adam : A method for stochastic optimization . In Interna-361 tional Conference on Learning Representations , 2015.362 [ 15 ] Andrey Rudenko , Luigi Palmieri , Michael Herman , Kris M Kitani , Dariu M Gavrila , and Kai O363 Arras . Human motion trajectory prediction : A survey . The International Journal of Robotics364 Research , 39 ( 8 ) :895–935 , 2020.365 [ 16 ] Nachiket Deo and Mohan M Trivedi . Convolutional social pooling for vehicle trajectory366 prediction . In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition367 Workshops , pages 1468–1476 , 2018.368 [ 17 ] Ming-Fang Chang , John Lambert , Patsorn Sangkloy , Jagjeet Singh , Slawomir Bak , Andrew369 Hartnett , De Wang , Peter Carr , Simon Lucey , Deva Ramanan , and James Hays . Argoverse:370 3d tracking and forecasting with rich maps . In Proceedings of the IEEE/CVF Conference on371 Computer Vision and Pattern Recognition ( CVPR ) , June 2019.372 [ 18 ] Tianyang Zhao , Yifei Xu , Mathew Monfort , Wongun Choi , Chris Baker , Yibiao Zhao , Yizhou373 Wang , and Ying Nian Wu . Multi-agent tensor fusion for contextual trajectory prediction . In374 Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition ( CVPR ) ,375 June 2019.376 [ 19 ] Rohan Chandra , Uttaran Bhattacharya , Aniket Bera , and Dinesh Manocha . Traphic : Trajectory377 prediction in dense and heterogeneous traffic using weighted interactions . In Proceedings of the378 IEEE/CVF Conference on Computer Vision and Pattern Recognition ( CVPR ) , June 2019.379 [ 20 ] Srijan Kumar , Xikun Zhang , and Jure Leskovec . Predicting dynamic embedding trajectory380 in temporal interaction networks . In Proceedings of the 25th ACM SIGKDD International381 Conference on Knowledge Discovery & Data Mining , pages 1269–1278 , 2019.382 [ 21 ] Emanuele Rossi , Ben Chamberlain , Fabrizio Frasca , Davide Eynard , Federico Monti , and383 Michael Bronstein . Temporal graph networks for deep learning on dynamic graphs . In ICML384 2020 Workshop on Graph Representation Learning , 2020.385 [ 22 ] Rewon Child , Scott Gray , Alec Radford , and Ilya Sutskever . Generating long sequences with386 sparse transformers . arXiv preprint arXiv:1904.10509 , 2019.387 [ 23 ] Manzil Zaheer , Guru Guruganesh , Kumar Avinava Dubey , Joshua Ainslie , Chris Alberti,388 Santiago Ontanon , Philip Pham , Anirudh Ravula , Qifan Wang , Li Yang , and Amr Ahmed . Big389 bird : Transformers for longer sequences . In H. Larochelle , M. Ranzato , R. Hadsell , M. F.390 Balcan , and H. Lin , editors , Advances in Neural Information Processing Systems , volume 33,391 pages 17283–17297 . Curran Associates , Inc. , 2020 . URL https : //proceedings.neurips.392 cc/paper/2020/file/c8512d142a2d849725f31a9a7a361ab9-Paper.pdf.393 [ 24 ] Iz Beltagy , Matthew E Peters , and Arman Cohan . Longformer : The long-document transformer.394 arXiv preprint arXiv:2004.05150 , 2020.395 [ 25 ] Nikita Kitaev , Lukasz Kaiser , and Anselm Levskaya . Reformer : The efficient transformer . In396 International Conference on Learning Representations , 2020 . URL https : //openreview.397 net/forum ? id=rkgNKkHtvB.398 [ 26 ] Jonathan Ho , Nal Kalchbrenner , Dirk Weissenborn , and Tim Salimans . Axial attention in399 multidimensional transformers . arXiv preprint arXiv:1912.12180 , 2019.400 [ 27 ] Gedas Bertasius , Heng Wang , and Lorenzo Torresani . Is space-time attention all you need for401 video understanding ? arXiv preprint arXiv:2102.05095 , 2021.402 Checklist403 1 . For all authors ... 404 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s405 contributions and scope ? [ Yes ] 406 ( b ) Did you describe the limitations of your work ? [ Yes ] See Section 6.407 ( c ) Did you discuss any potential negative societal impacts of your work ? [ No ] Our work408 does not introduce new ethical challenges.409 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to410 them ? [ Yes ] 411 2 . If you are including theoretical results ... 412 ( a ) Did you state the full set of assumptions of all theoretical results ? [ Yes ] 413 ( b ) Did you include complete proofs of all theoretical results ? [ Yes ] 414 3 . If you ran experiments ... 415 ( a ) Did you include the code , data , and instructions needed to reproduce the main experi-416 mental results ( either in the supplemental material or as a URL ) ? [ Yes ] 417 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they418 were chosen ) ? [ Yes ] 419 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running experi-420 ments multiple times ) ? [ N/A ] 421 ( d ) Did you include the total amount of compute and the type of resources used ( e.g. , type422 of GPUs , internal cluster , or cloud provider ) ? [ Yes ] 423 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 424 ( a ) If your work uses existing assets , did you cite the creators ? [ Yes ] 425 ( b ) Did you mention the license of the assets ? [ No ] We link directly to the dataset.426 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ Yes ] 427 ( d ) Did you discuss whether and how consent was obtained from people whose data you ’ re428 using/curating ? [ N/A ] 429 ( e ) Did you discuss whether the data you are using/curating contains personally identifiable430 information or offensive content ? [ N/A ] 431 5 . If you used crowdsourcing or conducted research with human subjects ... 432 ( a ) Did you include the full text of instructions given to participants and screenshots , if433 applicable ? [ N/A ] 434 ( b ) Did you describe any potential participant risks , with links to Institutional Review435 Board ( IRB ) approvals , if applicable ? [ N/A ] 436 ( c ) Did you include the estimated hourly wage paid to participants and the total amount437 spent on participant compensation ? [ N/A ] 438

Summary: This work presents an extension to the existing baller2vec transformer architecture. In particular, a new attention mask is used internally that is lower-triangular (offset by one block diagonal) rather than block-lower-triangular. This enables baller2vec++ to attend over entities at both previous and one-step future timesteps (in particular by attending over one-step predictions). The method was demonstrated in two experimental scenarios, an illustrative toy example and a dataset of NBA basketball trajectories.