L ATENT V ARIABLE S EQUENTIAL S ET T RANSFORMERS FOR J OINT M ULTI -A GENT M OTION P REDICTION
Anonymous authors
Paper under double-blind review
A BSTRACT
Robust multi-agent trajectory prediction is essential for the safe control of robotic systems.
A major challenge is to efﬁciently learn a representation that approximates the true joint distribution of contextual, social, and temporal information to enable planning.
We propose Latent Variable Sequential Set Transformers which are encoder-decoder architectures that generate scene-consistent multi-agent trajectories.
We refer to these architectures as “AutoBots”.
The encoder is a stack of interleaved temporal and social multi-head self-attention modules which alternately perform equivariant processing across the temporal and social dimensions.
The decoder can produce either the trajectory of one ego-agent or a distribution over the future trajectories for all agents in the scene.
For the single-agent prediction case, our model achieves top results on the global nuScenes vehicle motion prediction leaderboard, and produces strong results on the Argoverse vehicle prediction challenge.
In the multi-agent setting, we evaluate on the synthetic partition of TrajNet ++ dataset to showcase the model’s socially-consistent predictions.
We also demonstrate our model on general sequences of sets and provide illustrative experiments modelling the sequential structure of the multiple strokes that make up symbols in the Omniglot data.
A distinguishing feature of our method is that all models are trainable on a single desktop GPU (1080 Ti) in under 48h.
I NTRODUCTION
Many problems require processing complicated hierarchical compositions of sequences and sets.
For example, multiple choice questions are commonly presented as a query (e.g. natural language) and an unordered set of options (Kembhavi et al., 2017; Richardson et al., 2013).
Machine learning models that perform well in this setting should be insensitive to the order in which the options are presented.
Similarly, for motion prediction tasks (Biktairov et al., 2020; Sadeghian et al., 2018) where the inputs are agent trajectories evolving over time and outputs are future agent trajectories, models should be insensitive to agent ordering.
In this work, we focus on the generative modelling of sequences and sets and demonstrate this method on a variety of different motion forecasting tasks.
Suppose we have a set of M ∈ N sequences X
= { ( x 1 , . . .
, x K ) 1 , . . .
, ( x 1 , . . .
, x K ) M } each with K ∈ N elements.
Allowing also that X evolves across some time horizon T , we denote this sequence of sets as X =
( X 1 , . . .
, X T ) .
This problem setting often requires the model to capture high-order interactions and diverse futures.
Generally, solutions are built with auto-regressive sequence models which include t steps of context.
In our approach here, to better model the multi-modal distribution over possible futures, we also allow to include a discrete latent variable Z , to create sequential models of sets of the form
[Equation]
In this work, we propose this type of parameterized conditional mixture model over sequences of sets, and call this model a Latent Variable Sequential Set Transformer (affectionately referred to as “AutoBot”).
We evaluate AutoBot on nuScenes (Caesar et al., 2020) and Argoverse (Chang et al., 2019), two autonomous driving trajectory prediction benchmarks, on the synthetic partition of the TrajNet ++ (Kothari et al., 2021) dataset, a pedestrian trajectory forecasting benchmark, and Omniglot (Lake et al., 2015), a dataset of hand-drawn characters that we use for predicting strokes.
All of these provide empirical validation of our model’s ability to perform robust multi-modal prediction.
Speciﬁcally we make the following contributions:
• We propose a novel approach for modelling sequences of set structured continuous variables, and extend this to a latent-variable formulation to capture multi-modal distributions.
• We provide a theoretical analysis highlighting our model’s permutation sensitivity with respect to different components of the input data structures.
• We validate our method with strong empirical results on diverse tasks, including the largescale nuScenes and Argoverse datasets for autonomous driving (Caesar et al., 2020;
Chang et al., 2019), the multi-lingual Omniglot dataset of handwritten characters (Lake et al., 2015), and the synthetic partition of the TrajNet ++ pedestrian trajectory prediction task (Sadeghian et al., 2018).
With this submission we included an implementation of our method in a Jupyter notebook for easy reproduction and to demonstrate that the architecture is performant enough to train on a desktop
GPU.
B ACKGROUND
We now review several components of the Transformer (Vaswani et al., 2017) and Set Transformer(Lee et al., 2019) architectures, mostly following the notation found in their manuscripts.
For additional background, see their works and
Appendix A.
Multi-Head Self Attention (MHSA) can be thought of as an information retrieval system, where a query is executed against a key-value database, returning values where the key matches the query best.
While (Vaswani et al., 2017) deﬁnes MHSA on three tensors, for convenience, we input a single set-valued argument.
Internally, MHSA then performs intra-set attention by casting the input set X to query, key and value matrices and adding a residual connection, MHSA ( X ) = X + MHSA ( X , X , X ) .
Multi-Head Attention Blocks (MAB) resemble the encoder proposed in Vaswani et al.
(2017), but lack the position encoding and dropout (Lee et al., 2019).
Speciﬁcally, they consist of the MHSA operation described in Eq.
8 (appendix) followed by a row-wise feed-forward neural network (rFFN), with residual connections and layer normalization (LN)
(Ba et al., 2016) after each block.
Given an input set X containing n x elements of dimension d , and some conditioning variables C containing n c elements of dimension d , the MAB can be described by the forward computation,
[Equation]
where H = LN ( X + MHSA ( X , X , X
[Equation]
[Equation]
Multi-head Attention Block Decoders (MABD) were also introduced in Vaswani et al.
and were used to produce decoded sequences.
Given input matrices X and Y representing sequences, the decoder block performs the following computations:
[Equation]
Our model, described throughout the following section, makes extensive use of these functions.
L ATENT V ARIABLE S EQUENTIAL S ET T RANSFORMERS
Latent Variable Sequential Set Transformers is a class of encoder-decoder architectures that process sequences of sets (see Fig 2).
In this section, we describe the encoding and decoding procedures, the latent variable mechanism, and the training objective.
Additionally, we provide proof of the permutation equivariance of our model in Appendix B and implementation details in Appendix C as well as a discussion of the special case of AutoBot-Ego which is similar to AutoBots but predicts a future sequence for only one element in a scene.
3. E NCODER : I NPUT S EQUENCE R EPRESENTATION
AutoBot takes as an input a sequence of sets, X 1: t =
( X 1 , . . .
, X t ) , which in the motion prediction setting may be viewed as the state of a scene evolving over t timesteps.
Let us describe each set as having M elements (or agents) with K attributes (e.g. x/y position).
To process the social and temporal information the encoder applies the following two transformations.
Appendix B.2 provides a proof of equivariance with respect to the agent dimension for this encoder.
First, the AutoBots encoder injects temporal information into the sequence of sets using a sinusoidal positional encoding function PE ( . )
(Vaswani et al., 2017).
For this step, we view the data as a collection of matrices, { X 0 , . . .
, X M } , that describe the temporal evolution of agents.
The encoder processes temporal relationships between sets using a MAB, performing the following operation on each set where m ∈ { 1 , . .
, M } :
[Equation]
where rFFN ( . ) is an embedding layer projecting each element in the set to the size of the hidden state d K .
Next we process temporal slices of S , retrieving sets of agent states S τ at some timestep τ and is processed using S τ = MAB ( S τ ) .
These two operations are repeated L enc times to produce a context tensor C ∈ R ( d K ,M,t ) summarizing the entire input scene, where t is the number of timesteps in the input scene.
3. D ECODER : M ULTIMODAL S EQUENCE G ENERATION
The goal of the decoder is to generate temporally and socially consistent predictions in the context of multi-modal data distributions.
To generate c ∈ N different predictions for the same input scene, the AutoBot decoder employs c matrices of learnable seed parameters Q i ∈ R ( d K ,T ) where T is the prediction horizon, where i ∈ { 1 , ..., c } .
Intuitively, each matrix of learnable seed parameters corresponds to a setting of the discrete latent variable in AutoBot.
Each learnable matrix Q i is then repeated across the agent dimension M times to produce the input tensor Q (cid:48)
i ∈ R ( d K ,M,T ) Additional contextual information (such as a rasterized image of the environment) is encoded using a convolutional neural network to produce a vector of features m i where i ∈ { 1 , ..., c } .
In order to provide the contextual information to all future timesteps and to all set elements, we copy this vector along the M and T dimensions, producing the tensor M i ∈ R ( d K ,M,T ) .
Each tensor Q (cid:48) i is concatenated with M i along the d K dimension, as shown on the left side of Figure 2 (bottom).
This tensor is then processed using a rFFN ( . ) to produce the tensor H ∈ R ( d K ,M,T ) .
We begin to decode by processing the time dimension, conditioning on the encoder’s output C , and the encoded seed parameters and environmental information in H .
The decoder processes each of m agents in H separately with an MABD layer:
[Equation]
where H (cid:48) is the output tensor that encodes the future time evolution of each element in the set independently.
In order to ensure that the future scene is socially consistent between set elements, we process each temporal slice of H (cid:48) , retrieving sets of agent states H (cid:48)
τ at some future timestep τ where each element h (cid:48)
τ ∈ H τ is processed by a MAB layer.
Essentially, the MAB layer here performs a per-timestep attention between all set elements, similar to the Set Transformer (Lee et al., 2019).
These two operations are repeated L dec times to produce a ﬁnal output tensor for mode i .
The decoding process is repeated c times with different learnable seed parameters Q i and additional contextual information m i .
The output of the decoder is a tensor O ∈ R ( d K ,M,T,c ) which can then be processed element-wise using a neural network φ ( . ) to produce the desired output representation.
In several of our experiments, we generate trajectories in ( x, y ) space, and as such, φ produces the parameters of a bivariate Gaussian distribution. Note on the learnable seed parameters.
One of the main contributions that make AutoBot’s inference and training time faster than prior work is the use of the decoder seed parameters Q .
These parameters serve a dual purpose.
The ﬁrst is that they account for the diversity in future prediction, where each matrix Q i ∈ R ( d k ,T ) (where i ∈ { 1 , ..., c } ) corresponds to one setting of the discrete latent variable.
Secondly, they contribute to the speed of AutoBot by allowing it to perform inference on the entire scene with a single forward pass through the decoder without sequential sampling. Prior work (e.g., Tang and Salakhutdinov (2019)) uses auto-regressive sampling of the future scene with interlaced social attention, which allows for socially consistent multi-agent predictions. Other work (e.g., Messaoud et al. (2020)) resort to having a single output model generate the future of each agent independently .
Using the attention mechanisms of transformers, AutoBot combines the advantages of both approaches by deterministically decoding the future scene starting from each seed parameter matrix.
3. T RAINING O BJECTIVE
Given a dataset D = (cid:8) ( X 1 , . . .
, X T ) (cid:9) N i
=
1 consisting of N sequences, each having T sets
, our goal is to maximize the likelihood of the future trajectory of all elements in the set given the input sequence of sets, i.e., max p ( X t +1: T | X 1: t ) .
In order to simplify notation, we will henceforth refer to the ego-agent’s future trajectory by Y =
X t +1: T .
As discussed above, AutoBot employs discrete latent variables, which allows us to compute the log-likelihood exactly.
The gradient of our objective function can then be expressed as follows:
[Equation]
As previously discussed in Tang and Salakhutdinov (2019), computing the posterior likelihood p θ ( Z |Y , X 1: t ) is difﬁcult in general and varies with θ .
As discussed in Bishop (2006), one can introduce a distribution over the latent variables, q ( Z ) , that is convenient to compute.
With this distribution, we can now decompose the log-likelihood as follows:
[Equation]
where D KL ( . || . ) is the Kullback-Leibler divergence between the approximating posterior and the actual posterior.
A natural choice for this approximating distribution is q ( Z ) =
p θ old ( Z |Y , X 1: t ) , where θ old corresponds to AutoBot’s parameters before performing the parameter update.
With this choice of q ( Z ) , the objective function can be re-written as
[Equation]
where the computation of p θ ( Z | X 1: t ) is described in Appendix C.2.
Note that the posterior p θ old ( Z |Y , X 1: t ) can be computed exactly in the model we explore here since AutoBot operates with discrete latent variables.
Therefore, our ﬁnal objective function becomes to maximize Q ( θ, θ old ) and minimize D KL ( p θ old ( Z |Y , X 1: t )
||
p
θ ( Z | X 1: t )) .
As mentioned, for our experiments we use a function φ to produce a bivariate Gaussian distribution for each agent at every timestep.
We found that to ensure that each mode’s output sequence does not have high variance, we introduce a mode entropy (ME) regularization term to penalize large entropy values of the output distributions,
[Equation]
As we can see, this entropy regularizer penalizes only the mode with the maximum entropy.
Our results in Section D.2 show the importance of this term in learning to cover the modes of the data.
E XPERIMENTS
We now demonstrate the capability of AutoBots to model sequences of set-structured data across several domains.
In subsection 4.1 we show strong performance on the competitive motion-prediction dataset nuScenes (Caesar et al., 2020), a dataset recorded from a self-driving car, capturing other vehicles and the intersection geometry.
In subsection 4.2, we demonstrate our model on the Argoverse autonomous driving challenge.
In subsection 4.3, we show results on the synthetic partition of the TrajNet ++ (Sadeghian et al., 2018) dataset which contains multi-agent scenes with a high degree of interaction between agents.
Finally in subsection 4.4, we show that our model can be applied to generate characters conditioning across strokes, and verify that AutoBot-Ego is permutation-invariant across an input set.
See Appendix D.2, for an illustrative toy experiment that highlights our model’s ability to capture discrete latent modes.
4. N U S CENES , A R EAL -W ORLD D RIVING D ATASET
In this section, we present AutoBot-Ego’s results on the nuScenes dataset.
The goal of this benchmark dataset is to predict the ego-agent’s future trajectory ( 6 seconds at 2 hz) given the past trajectory (up to 2 seconds at 2 hz).
Additionally, the benchmark provides access to all neighbouring agents’ past trajectories and a birds-eye-view RGB image of the road network in the vicinity of the ego-agent.
Figure 3 shows three example predictions produced by AutoBot-Ego trained with c = 10 .
We observe that the model learns to produce trajectories that are in agreement with the road network and covers most possible futures directions.
We also note that for each direction, AutoBot-Ego assigns different modes to different speeds to more efﬁciently cover the possibilities.
We compare our results on this benchmark to the state-of-the-art in Table 1.
Min ADE (5) and (10) are the average of pointwise L2 distances between the predicted trajectory and ground truth over 5 and 10 most likely predictions respectively.
A prediction is classiﬁed as a miss if the maximum pointwise L2 distance between the prediction and ground truth is greater than 2 meters.
Miss Rate
Top-5 (2m) and Top-10 (2m) are the proportion of misses over all agents, where for each agent, 5 and 10 most likely predictions respectively are evaluated to check if they’re misses.
Min FDE (1) is the L2 distance between the ﬁnal points of the prediction and ground truth of the most likely prediction averaged over all agents.
Off Road Rate is the fraction of predicted trajectories that are not entirely contained in the drivable area of the map
[-Our model was trained on a single Nvidia GTX 
 Ti for 3 hours.-][Table]
We can see that AutoBot-Ego obtains the best performance on the Min ADE (10) metric and on Off
Road Rate metric, and strong performance on the other metrics. Furthermore, using an ensemble of three Autobot-Ego models slightly improves the performance.
This demonstrates our model’s ability to capture multi-modal distributions over future trajectories due to the latent variables.
AutoBot-Ego excelled in maintaining predictions on the road, showing that the model has effectively used its capacity to attend to the map information.
Finally, we note that like all other models, AutoBot-Ego struggles to predict the probability of the correct mode, which results in poor performance of the Min FDE (1).
The most interesting part of AutoBot-Ego is its computational requirements.
We highlight that AutoBot-Ego was trained on a single Nvidia GTX 1080 Ti for 3 hours.
This is in large part due to the use of the decoder seed parameters which allow us to predict the future trajectory of the ego-agent in a single forward pass.
Note, that in the appendix C, we present the computational complexity of AutoBot-Egocompared to WIMP (Khandelwal et al., 2020), which shows a 2x inference speedup .
Ablation studies on various components of AutoBot-Ego are presented in Appendix D.1.
4. A RGOVERSE R ESULTS
Table 2 shows the results of AutoBot-Ego (with c = 6 ) on the Argoverse dataset test set.
While our method does not achieve top scores, it uses much less computation, less data augmentation and no ensembling compared to alternatives.
Further, we compare quite well to the 2020 Argoverse competition winner, a model called “Jean”.
In addition, we note that the Argoverse challenge does not restrict the use of additional training data and that many of the top entries on the leaderboard have no publicly available papers or codebases.
Many methods are therefore optimized for Argoverse by means of various data augmentations.
For example, Scene Transformer (Ngiam et al., 2021) used a data augmentation scheme that includes agent dropout and scene rotation which we have not, and the current top method and winner of the 2021 challenge (QCraft) makes use of ensembling.
We believe that a fair comparison of the performance between neural network methods should ensure that each architecture is parameter-matched and the training is FLOP-matched.
While the
Argoverse leaderboard does not require entrants to report their parameter or FLOP count, we report available public statistics to enable comparison here.
First, Scene Transformer (Ngiam et al., 2021) reports that their model trains in 3 days on TPUs and TPU-v3s have a peak performance of
TeraFLOPS, which means their model could have consumed up to 3 days × 24 hours × 420 TFLOPS, or 108,000,000 TFLOPS.
On the other hand, AutoBot-Ego was trained on a Nvidia 1080ti forhours, consuming roughly 396,000 TFLOPs, or 0.4% as much.
Jean’s method (Mercat et al., 2020) consumes roughly 8,600,000 TFLOPs (training our method uses 5% as much compute).
While the method currently topping the leaderboard has not revealed their model details, the fact that they used ensembles very likely requires signiﬁcantly more compute than our method.
[Table]
4. T RAJ N ET ++ S YNTHETIC D ATASET
In this section, we demonstrate the utility of the having social and temporal attention in the encoder and decoder of AutoBot. We chose to evaluate on the synthetic partition of the TrajNet ++ dataset since this was speciﬁcally designed to have a high-level on interaction between scene agents Kothari et al. In this task, we are provided with the state of all agents for the past 9 timesteps and are tasked with predicting the next 12 timesteps for all agents. The scene is initially normalized around one (ego) agent such that the agent’s ﬁnal position in the input sequence is at (0 , 0) and its heading is aligned with the positive y direction.
There are a total of 54 , 513 unique scenes in the dataset, which we split into 49 , 062 training scenes and 5 , 451 validation scenes.
For this experiment, we apply our general architecture (AutoBot) to a multi-agent scene forecasting, and ablate the social attention components in the decoder only, and decoder and encoder.
[Table]
[Table]
Table 3 presents three variants of the model [Table]
and baseline extrapolation model which shows that the prediction problem cannot be trivially solved.
The ﬁrst, AutoBot-AntiSocial, corresponds to a model without the intra-set attention functions in the encoder and decoder (see Figure 2).
The second variant, AutoBot-Ego, only sees the past social context and has no social attention during the trajectory generation process.
The third model corresponds to the general architecture, AutoBot.
All models were trained with c = 6 and we report scene-level error metrics as deﬁned by Casas et al.
Using the social attention in the decoder signiﬁcantly reduces the number of collisions between predicted agents’ futures and signiﬁcantly outperforms its ablated counterparts on scene-level minADE and minFDE.
We expect this at a theoretical level since the ego-centric formulation makes the independence assumption that the scene’s future evolution can be decomposed as the product of the individual agent’s future motion in isolation, as remarked in (Casas et al., 2020).
AutoBot does not make this assumption as we condition between agents across all future timesteps.
We refer the reader to Appendix D.3 for additional ablation experiments and qualitative results on this dataset.
4. O MNIGLOT , S TROKE - BASED I MAGE G ENERATION
To demonstrate the general efﬁcacy of our approach to sequence generation, we demonstrate AutoBots on a character completion, a diverse alternative to automotive and pedestrian trajectory prediction.
We set up two different tasks on the Omniglot dataset (Lake et al., 2015): (a) stroke-completion task in which the model is presented with the ﬁrst half of each stroke and has to complete all of them in parallel and (b) character-completion task in which the model is given several full strokes and has to draw the last stroke of the character.
In task (a), the model is trained on the entire Omniglot dataset (training set), in which strokes are encoded as sequential sets of points, and we are comparing the performance of AutoBot qualitatively against that of an LSTM baseline with social attention and discrete latents (similar to in AutoBot) on a held-out test set.
With the ﬁrst task, we illustrate that our model is ﬂexible enough to learn to legibly complete characters from a variety of different languages.
In task (b), we train AutoBot-Ego only on characters “F” and “H” to demonstrate that our discrete latent variable captures plausible character completions given the context of the other strokes, e.g. a character “F” can be completed as an “H” when the top stroke is missing.
The results of these experiments can be found in Fig.
Implementation details and additional results can be found in
Appendix section D.5.
R ELATED W ORK
Processing set structured data.
Models of set-valued data (i.e. unordered collections) should be permutation-invariant (see Appendix B for formal deﬁnitions), and be capable of processing sets with arbitrary cardinality.
Canonical forms of feed forward neural networks and recurrent neural network models do not have these properties (Lee et al., 2019).
One type of proposed solution is to integrate pooling mechanisms with feed forward networks to process set data (Zaheer et al., 2017; Su et al., 2015; Hartford et al., 2016).
This is referred as social pooling in the context of motion prediction (Alahi et al., 2016; Deo and Trivedi, 2018; Lee et al., 2017).
Attention-based models (Lee et al., 2019; Casas et al., 2020; Yuan et al., 2021; Ngiam et al., 2021; Zhao et al., 2020b) or graph-based models (Zhao et al., 2020a; Messaoud et al., 2020; Li et al., 2020b; Gao et al., 2020) have performed well in set processing tasks that require modelling high-order element interactions.
In contrast, we provide a joint single pass attention-based approach for modelling more complex data-structures composed by sets and sequence, while also addressing how to generate diverse and heterogeneous data-structures.
Although those properties exist in concurrent motion prediction papers (Tang and
Salakhutdinov, 2019; Casas et al., 2020; Yuan et al., 2021; Ngiam et al., 2021; Li et al., 2020a), our proposed method is targeting a more general set of problems while still getting competitive results in motion prediction benchmarks at a fraction of other models’ compute.
Diverse set generation.
Generation of diverse samples using neural sequence models is an important part of our model and a well-studied problem.
For example, in Bowman et al.
(2015) a variational (CVAE) approach is used to construct models that condition on attributes for dialog generation.
The
CVAE approach has also been integrated with transformer networks (Lin et al., 2020; Wang and Wan, 2019) for diverse response generation.
For motion prediction (Lee et al., 2017; Deo and Trivedi, 2018) use a variational mechanism where each agent has their own intentions encoded. (Li et al., 2020a) output Gaussian mixtures at every timestep and sample from it to generate multiple futures.
The most similar strategy to ours was developed by (Tang and Salakhutdinov, 2019).
Other recent works (Suo et al., 2021; Casas et al., 2020) use continuous latent variables to model the multiple actors’ intentions jointly.
We built on this method by using one joint transformer conditioned on the latent variables and added a maximum entropy likelihood loss that improves the variability of the generated outcomes.
Sequential state modelling.
We are interested in data that can be better represented as a sequence.
The main example of this is language modelling (Vaswani et al., 2017; Devlin et al., 2018), but also includes sound (Kong et al., 2020), video (Becker et al., 2018), and others (Kant et al., 2020;
Yang et al., 2021).
For motion prediction, prior work has largely focused on employing RNNs (LSTMs (Hochreiter and Schmidhuber, 1997) or GRUs (Cho et al., 2014)) to model the input and/or output sequence (Alahi et al., 2016; Lee et al., 2017; Deo and Trivedi, 2018; Tang and Salakhutdinov, 2019; Messaoud et al., 2020;
Casas et al., 2020; Li et al., 2020b;a).
With its recent success in natural language processing, Transformers (Vaswani et al., 2017) have been adopted by recent approaches for motion forecasting. Yu et al. (2020) uses a graph encoding strategy for pedestrian motion prediction. Giuliari et al.
(2021) and Liu et al.
(2021) propose a transformer based encoding to predict the future trajectories of agents separately .
In contrast, AutoBots produces representations that jointly encode both the time and social dimensions by treating this structure as a sequence of sets.
This provides reuse of the model capacity and a lightweight approach.
Concurrent to our method, Yuan et al.
(2021), also jointly encodes social and time dimensions but ﬂattens both of those axis into a single dimension.
C ONCLUSION
In this paper, we propose the Latent Variable Sequential Set Transformer to model time-evolution of sequential sets using discrete latent variables.
We make use of the multi-head attention block to efﬁciently perform intra-set attention, to model the time-dependence between the input and output sequence, and to predict the prior probability distribution over the latent modes.
We validate our
AutoBot-Ego by achieving competitive results on the nuScenes benchmark for ego-centric motion forecasting in the domain of autonomous vehicles.
We further demonstrate that AutoBot can model diverse sequences of sets that adhere to social conventions.
We validate that since our model can attend to the hidden state of all agents during the generative process, it produces scene-consistent predictions on the TrajNet ++ dataset when compared with ego-centric forecasting methods.
Our approach is limited to operate in a setting where a perceptual pipeline is providing structured sets representing the other entities under consideration.
If our approach is used in a robot or autonomous car, any errors or unknown entities not captured by sets used as input to the algorithm would not be modelled by our approach.
Memory limitations may also limit the number of sets that can be processed at any time.
In future work, we plan to extend this line of work to broader and more interpretable scene-consistency.
Further, we would like to investigate the combination of a powerful sequence model like AutoBot with a reinforcement learning algorithm to create a data-efﬁcient model-based RL agent for multi-agent settings.
A B ACKGROUND
We now review several components of the Transformer (Vaswani et al., 2017) and Set Transformer (Lee et al., 2019) architectures, mostly following the notation found in their manuscripts.
Multi-Head Self Attention (MHSA) as proposed in (Vaswani et al., 2017) is a function deﬁned on n q query vectors Q ∈ R n q × d q , and key-value pairs ( K ∈ R n k × d q and V ∈ R n k × d v ).
With a single attention head, the function resolves the queries Q by performing the computation:
[Equation]
where d q is the dimension of the query and key vectors and d v is the dimension of the value vector.
MHSA consists of h attention heads performing the operation shown in Eq.
7 with h linear projections of keys, queries, and values.
The ﬁnal output is a linear projection of the concatenated output of each attention head.
These computations can be expressed as
[Equation]
where
[Equation]
and where W Qi , W Vi and W Ki along with the output projection matrix W O are the learnable projection matrices of each attention head in the MHSA.
Note that W Qi , W Vi and W Ki project the initial set of inputs to smaller dimensionality.
For example, W Qi ∈ R d q × d Mq projects the original queries with dimension d q to d Mq which is typically chosen to be d q /h .
This operation has useful applications in learning a representation of an input set where all its elements interact with each other.
An interesting application of MHSA is to perform self-attention on an input set (Lee et al., 2019).
Given an input set X , one can perform intra-set attention by using X as the queries, keys and values and having a residual connection, i.e., MHSA ( X ) = X + MHSA ( X , X , X ) .
This operation has useful applications in learning a representation of an input set where all its elements interact with each other.
Multi-head Attention Block Decoders (MABD) were also introduced in Vaswani et al.
and were used to produce decoded sequences.
Given an input matrix X and an output matrix Y representing sequences, the decoder block performs the following computations:
[Equation]
During training, MABD can be trained efﬁciently by shifting Y backward by one (with a start token at the beginning) and computing the output with one forward pass.
During testing where one does not have access to Y , the model would then generate the future autoregressively.
In order to avoid conditioning on the future during training, the MHSA ( Y ) operation in MABD employs time-masking.
This would prevent the model from accessing privileged future information during training.
B P ERMUTATION E QUIVARIANCE OF THE M ODEL
In this section, we prove the permutation equivariance of the Latent Variable Sequential Set Transformer with respect to the dimension M .
To do so, we analyze the model’s response to permutation along this dimension of the input tensor.
This tensor, X , is of dimension ( K, M, t ) , where
K, M, t ∈ N .
The dimensions t and K are ordered, i.e. if we index into the tensor along dimension
M (we denote this operation as X m where m ∈ { 1 , . . .
, M } ) then we retrieve a ( K, t ) -dimensional matrix.
However, the M dimension is unordered.
This implies that when we index by another dimension, for example t , we retrieve an invertible multiset of vectors, denoted X τ = { x 1 , . . .
, x M } where τ ∈ { 1 , . . .
, t } and x m ∈
X τ are K -dimensional.
Our proof that our model is permutation equivariant on M demonstrates a clear difference with the properties of similar models.
For example, the model proposed in Set Transformers (Lee et al., 2019) features a permutation equivariant encoder combined with a permutation invariant decoder.
And while Deep Sets (Zaheer et al., 2017) provides a theoretical framework for designing neural network models that operating on sets, they do not directly address the processing of hetereogeneous data structures as we do here.
The rest of the section will proceed in the following manner:
Provide deﬁnitions and lemmas of the mathematical structures we will use in our proofs.
Show the permutation equivariance of the AutoBot encoder on M . 3.
Show the permutation equivariance of the AutoBot decoder on M .
For clarity, we re-state the motivation for proving permutation equivariance in this particular model.
In tasks like motion prediction, we may not be able to uniquely identify moving objects (perhaps due to occlusion or similar visual features).
Therefore, we may wish to represent these objects as a collection of states evolving over time.
Readers may beneﬁt from this perspective on X as a time-evolving ( t ) collection of M objects with K attributes.
We described X τ as a multiset for generality, as it permits repeated elements.
However, in motion prediction tasks it might be safe to assume that no two distinct objects will share the same attribute vector.
By proving that our model is equivariant with respect to M , we ensure that it will not waste capacity on spurious order-speciﬁc features.
The inductive bias for this type of task differs from other settings like machine translation where word order is meaningful.
B. D EFINITIONS AND L EMMAS
A permutation equivariant function is one that achieves the same result if you permute the set then apply the function to the permuted set, as you would achieve if you had applied the function to the set then applied the same permutation to the output set.
We provide a formal description in Def B.1.
Deﬁnition B.1 (Permutation Equivariance) .
Let X be the set of n vectors { x 1 , . . .
, x n } , and S n the group of permutations of n elements π : X → X .
Suppose we have a function f : X → X , we call f
Permutation Equivariant iff f
( π ( X )) =
π ( f ( X )) .
One of the central components of the Sequential Set Transformer is the Multi-Head Self Attention function (MHSA : X → X ) originally introduced in (Vaswani et al., 2017).
MHSA computes an updated representation for one input vector in a set of input vectors by processing its relation with a collection of other vectors.
Without the explicit addition of positional information to the content of these vectors, this operation is invariant to input order.
MHSA is a content-based attention mechanism, and updates the contents of a single input vector based on the contents of the others.
However, if we deﬁne f MHSA as the function applying MHSA to update each input vector, then we have created a permutation equivariant function, described in Lemma B.1.
Lemma B.1 ( f MAB is Permutation Equivariant on M ) .
Let X be a set of D -dimensional vectors { x 1 , . . .
, x M } , and f MAB : R D × M → R D
× M be a function that applies MAB (as deﬁned in
Equation 1 of the main text) to each x m ∈ X .
Then f MAB ( X τ ) is permutation equivariant because ∀ π ∈ S m , f MAB ( π ( X ))
=
π ( f MAB ( X )) .
Having discussed the preliminaries, we now show the permutation equivariance of AutoBot’s encoder.
B. A UTO B OT E NCODER P ROPERTIES
Theorem B.2 (AutoBot Encoder is Permutation Equivariant on M ) .
Let X be a tensor of dimension ( K, M, t ) where K, M, t ∈ N , and where selecting on t retrieves an invertible multiset of sequences, denoted X τ = { x 1 , . . .
, x M } , where τ ∈ { 1 , . .
. , t } and x m ∈
X τ are K -dimensional.
The AutoBot encoder, as deﬁned in Section 3.1, is permutation equivariant on M .
Proof.
What follows is a direct proof of the permutation equivariance of the AutoBot encoder.
The encoder is a composition of three functions, rFFN: R K × t → R d K × t , f MAB : R d k × M → R d k × M , and f MAB : R d K × t × M → R d K × t × M .
First, we show that the embedding function rFFN is equivariant on M .
Let rFFN represent the application of a function e : R K → R d K to each element of a set X τ where τ ∈ { 1 , . . .
, t } .
Then rFFN is equivariant with respect to M by Deﬁnition B.1.
Speciﬁcally, we can see that the function rFFN satisﬁes the equation ∀ π ∈ S M , rFFN ( π ( X τ ))
=
π ( rFFN ( X τ )) .
In the AutoBot encoder, the
Next, we show that the function f MAB , which applies a Multi-Head Attention Block (MAB) to each set, ( X 1 , . . .
, X t ) , is equivariant on M .
f MAB is permutation equivariant with respect to the M dimension of X because MAB is permutation equivariant on each set X τ by Lemma B.1.
Finally, we deﬁne a function f MAB which applies a positional encoding along the time dimension of X then a Multi-head Attention Block (MAB) to each matrix in X , { X 1 , . . .
, X M } .
f MAB is permutation equivariant with respect to the M dimension of X because the collection of matrices { X 1 , . . .
, X M } is an unordered set, and the uniform application of a function to each element of this set satisﬁes permutation equvariance.
A composition of equivariant functions is equivariant, so the encoder is equivariant on M .
B. A UTO B OT D ECODER P ROPERTIES
We now provide a direct proof the the permutation equivariance of the AutoBot decoder on M .
The decoder is an auto-regressive function initialized with a seed set X t containing M vectors of dimension d k , each concatenated with a vector of latent variables and other context of dimension 2 d k .
For τ ∈ { 1 , . . .
, T − t } iterations, we concatenate the output of the decoder with the seed, the previous decoder output, and the latents to produce a tensor of dimension (3 d K , M, τ ) .
In particular, the decoder is a composition of three functions, rFFN dec : R 3 d K × τ → R d K × τ , f MAB : R d k × M → R d k × M , and f MABD : R d K × M × τ → R d K × M × τ .
Theorem B.3 (AutoBot Decoder is Permutation Equivariant on M ) .
Given an invertible multiset X t with M vector-valued elements, and an auto-regressive generator iterated for τ ∈ { 1 , . . .
, T − t } , the AutoBot decoder, formed by a composition of three functions, rFFN
dec : R 3 d K × M → R d K × M , f MAB : R d k × M → R d k × M , and f MABD :
R d K × τ × M → R d K × τ × M , is equivariant on M .
Proof.
First, we must establish that the function rFFN dec : R 3 d K × M
→ R d K × M is equivariant with respect to M .
The AutoBot decoder applies rFFN dec to each set X τ where τ ∈ { 1 , . .
, T − t } , transforming the M vectors of dimension 3 d K → d k .
Because rFFN dec represents a uniform application of e dec to each element of the set, we can see that it satisﬁes the deﬁnition of permutation equivariance, speciﬁcally that ∀ π ∈ S M , rFFN dec ( π ( X τ ))
=
π ( rFFN dec ( X τ )) .
Next, we take the function f MABD , which ﬁrst applies a function PE : R d k ,M,τ → R d k ,M,τ to X adding position information along the temporal dimension, then applies f MABD :
R d K × τ → R d K × τ to each matrix { X 1 , . . .
, X M } .
f MABD is permutation equivariant with respect to the M dimension of
X because the collection of matrices { X 1 , . . .
, X M } is an unordered set, and the uniform application of a function to transform each element independently does not impose an order on this set.
Similar to the ﬁnal step of the previous proof, we see that a function f MAB applying a multi-head selfattention block (MAB) to each set { X t +1 , . . .
, X τ } is equivariant on M because MAB is permutation equivariant on each set (see Lemma B.1).
The composition of equivariant functions is equivariant, so the decoder is equivariant.
C M ODEL D ETAILS C. E GO -C ENTRIC F ORECASTING
AutoBot can be used to predict the future trajectories of all agents in a scene.
However, in some settings we may wish to forecast the future trajectory of only a single agent in the set.
We call this variant “AutoBot-Ego”.
Referring back to the encoder presented in Section 3.1, instead of propagating the encoding of all elements in the set, we can instead select the state of that element produced by the intra-set attention, s ∈ S τ .
With this, our encoder proceeds with the rest of the computation identically to produce the tensor C m 1: t , which is an encoding of the ego element’s history conditioned on the past sequence of all elements.
AutoBot-Ego’s decoder proceeds as presented in section 3.2, with one exception.
Since AutoBot-Ego only generates the future trajectory of an ego element, we do not perform intra-set attention in the
Example of the birds-eye-view road network provided by the Nuscenes dataset.
Our CNN encoder produces a volume which we partition equally into K modes.
This allows each generated trajectory to be conditioned on a different representation of the input image.
decoder. As a consequence of only dealing with a single ego-agent in the future generation, we do not repeat the seed parameter matrices Q i across the M dimension, where i ∈ { 1 , . . , c } .
The objective function is updated to compute the likelihood of one element’s future given the entire input sequence.
That is, in AutoBot-Ego, Y = X mt
+1: T .
C. E XPECTATION M AXIMIZATION D ETAILS
Our model also computes the distribution P ( Z | X 1: t ) of discrete random variables.
To achieve this, we employ c learnable vectors (one for each mode) concatenated into the matrix P ∈ R ( d K ,c ) , which behave like seeds to compute distribution over modes prior to observing the future (these seed parameters should not be confused with the decoder seed parameters Q ).
The distribution is generated by performing the computations p
[Equation]
[Equation]
where F = MABD ( P 1: c , C ) , where rLin is a row-wise linear projection layer to a vector of size c .
C. C ONTEXT E NCODING D ETAILS
In the nuScenes dataset, we are provided with a birds-eye-view 128 × 128 RGB image of the road network, as shown in Figure 5.
In the TrajNet dataset, we are also provided with a birds-eye-view of the scene, as shown in Figure 6.
We process this additional context using a 4 layer convolutional neural network (CNN) which encodes the map information into a volume of size 7 × 7 × (7 ∗ c ) where c is the number of modes.
We apply a 2 D dropout layer with a rate of 0 .
1 on this output volume before processing it further.
As our model employs discrete latent variables, we found it helpful to divide this volume equally among all c modes, where each mode receives a ﬂattened version of the 7 × 7 ×volume.
As described in Section 3.2, this context is copied across the M and T dimensions, and concatenated with the decoder seed parameters during the sequence generation process.
Intuitively, each mode’s generated trajectory is conditioned on a different representation of the context.
C. I MPLEMENTATION AND T RAINING DETAILS
We implemented our model using the Pytorch open-source framework.
Table 4 shows the values of parameters used in our experiments across the different datasets.
The MAB encoder and decoder blocks in all parts of AutoBots use dropout.
We train the model using the Adam optimizer with an initial learning rate.
For the Nuscenes and TrajNet ++ datasets, we anneal the learning rate every 10 epochs for the ﬁrst 20 epochs by a factor of 2 , followed by annealing it by a factor of 1 .
33 every 10 epochs for the next 30 epochs.
For the Argoverse dataset, we anneal the learning rate everyepochs by a factor 2 for the ﬁrst 20 epochs.
In all datasets, we found it helpful to clip the gradients [Table] to a maximum magnitude of 5 . 0 .
For the Nuscenes experiments, our model takes approximately 80 epochs to converge, which corresponds to approximately 3 hours of compute time on a single Nvidia Geforce GTX 1080Ti GPU, using approximately 2 GB of VRAM.
We used mirroring of all trajectories and the map information as a data augmentation method since the dataset’s scenes originate from Singapore and Boston. For Argoverse experiments, our model takes approximatelyepochs to converge, which corresponds to approximately 10 hours on the same resources. In order to improve the performance, we use other agents in the certain scenes as an ego-agent, effectively doubling the size of the dataset. We chose agents that move by at least 10 meters.
This, in-turn, increases the training time to approximately 24 hours. For the TrajNet, TrajNet ++ and Omniglot datasets, our model takes approximately 100 epochs which corresponds to a maximum 2 hours of compute time on the same resources. Table 5 highlights AutoBot-Ego’s inference speed compared to WIMP, highlighting its potential to be employed in real-world applications.
[Table]
Table 6 shows the results of ablating various components from AutoBot-Ego on the Nuscenes dataset.
As we can see, the model trained without the map information performs signiﬁcantly worse across all metrics. This shows that indeed AutoBot-Ego learns to attend to the map information effectively. The ‘No Latent Variable’ version of our model is a model without the latent variable, and therefore, without the decoder seed parameters. This model then becomes auto-regressive and samples from the output bivariate Gaussian distribution at every timestep.
We note that this version of the model is signiﬁcantly slower to train (taking approximately three times as long) due to its autoregressive nature.
In order to generate ten trajectories, the generation process is repeated c = 10 times.
This model uses the same transformer encoder and decoder blocks as described in Section 3.
As we can see from its performance on the ADE metrics, without latent variables, the model is not capable of generating diverse enough trajectories.
This is expected since the output model φ is now tasked with accounting for both the aleatoric and epistemic uncertainties in the process.
Finally, the ‘No Data
Augmentation’ model shows that although it is helpful to augment the data with mirroring, even without it the model is still capable of producing strong performance.
[Table]
Study. . We study the performance beneﬁts of AutoBot-Egoover various ablated counterparts. D. A NALYSIS OF E NTROPY R EGULARIZATION
This dataset contains the trajectories of a single particle having identical pasts but diverse futures, with different turning rates and speeds, as shown in Figure 7 (top row).
In the second row of Figure 7, we show that our model trained without entropy regularization, λ e = 0 . 0 , can cover all modes.
However, the short trajectory variants are all represented using only one of the modes (2nd row, 3rd column), as shown by the growth of the uncertainty ellipses as the trajectory moves forward.
This shows that AutoBots learn to increase the variance of the bivariate Gaussian distribution in order to cover the space of trajectories.
The third row of Figure 7, shows the learned trajectories for this variant trained with λ e = 3 . 0 .
As we can see, the entropy regularization penalizes the model if the uncertainty of bivariate Gaussians is too large.
By pushing the variance of the output distribution to a low magnitude, we restrict the model such that the only way it can achieve high likelihood is to minimize the distance between the mean of the bivariate distributions and ground-truth trajectories.
Ablation Study - Importance of Entropy Loss Term .
In Table 7, we show an ablation study to evaluate the importance of the entropy component on training the model for nuScenes.
Again, we observe that for multi-modal predictions, enforcing low entropy on the predicted bivariate distributions has a great positive impact on the results, with an increased λ e resulting in better performance.
This was observed on all evaluated metrics with the exception of the min FDE, which was poor regardless of the model.
Furthermore, we can see from Table 1 that, in general, all prior models have a difﬁculty obtaining strong performance on the min FDE (1) metric, with the best value having an overall error of 8 . 49 meters.
This may be due to the difﬁculty of predicting the correct mode when one only has access to two seconds of the past sequence.
[Table]
D. T RAJ N ET ++ A BLATION AND Q UALITATIVE RESULTS
Table 8 shows the effect of increasing the number of social/temporal encoder and decoder layers in the general AutoBot architecture.
We can see that increasing the number of layers in the decoder has a positive impact on the performance of the joint metrics.
This is expected since as we increase the number of social layers, the model can use high-order interactions to model the scene.
Figureshows an example scene generated by each of the model variants.
[Table]
) ( ↓ )
L enc = 1 , L dec =0.
L enc = 2 , L dec =0.
L enc = 3 , L dec =0.
L enc = 1 , L dec =0.
L enc = 1 , L dec =0.
L enc = 2 , L dec =0.
L enc = 3 , L dec =0.
D. T RAJ N ET , A P EDESTRIAN M OTION P REDICTION D ATASET
We further validate AutoBots on a pedestrian trajectory prediction dataset TrajNet (Sadeghian et al., 2018) where we are provided with the state of all agents for the past eight timesteps and are tasked with predicting the next 12 timesteps.
In this experiment, we apply our general architecture (AutoBot) to a multi-agent scene forecast situation, and ablate various the social components in the decoder only, and decoder and encoder.
[Table]
In Table 9, we present three variants of the model.
The ﬁrst, AutoBot-AntiSocial corresponds to a model without the intra-set attention functions in the encoder and decoder (see Figure 2).
The second variant, AutoBot-Ego, only sees the past social context and has no social attention during the trajectory generation process.
The third model corresponds to the general architecture, AutoBot.
All models were trained with c = 5 and we report scene-level error metrics as deﬁned by (Casas et al., 2020).
Using the social data with the multi agent modelling in the decoder greatly reduces the number of collisions between predicted agent futures.
Furthermore, we can see that AutoBot signiﬁcantly outperforms its ablated counterparts on generating scenes that match ground-truth, as observed by the reduced scene-level minADE and minFDE.
We expect this theoretically since the ego-centric formulation makes the independence assumption that the scene’s future evolution can be decomposed as the product of the individual agent’s future motion in isolation, as remarked in (Casas et al., 2020).
AutoBot does not make this assumption as we condition between agents at every timestep.
Figure 9 shows an example of multi-agent trajectory generation using the three different approaches for predicting the scene evolution.
We present four of the ﬁve modes generated by the model for this example input scene.
One can observe that the only model achieving scene-consistent predictions across all modes is AutoBot .
Interestingly, we also observe that the different modes correspond to different directions of motion, highlighting the utility of using modes.
In fact, the modes in AutoBot condition the entire scene future evolution, producing alternative realistic multi-agent futures.
We refer the reader to Figures 10, 11, 12 and 13 where we show some additional qualitative results of AutoBots’ variants compared to AutoBot on the TrajNet dataset.
These results highlight the effectiveness of the sequential set attention mechanism in the decoding process of our method and how important it is to predict all agents in parallel compared to a 1-by-1 setup.
D. O MNIGLOT A DDITIONAL D ETAILS AND R ESULTS
The LSTM baseline used in our Omniglot experiments have a hidden size of 128 and use c =latent modes.
The input is ﬁrst projected into this space using an embedding layer, which is identical to the one used in AutoBot.
All input sequences of sets are encoded independently using a shared LSTM encoder.
During the decoding phase, we autoregressively generate the next stroke step using an LSTM decoder.
In order to ensure consistency between the predicted future strokes, inspired by
Social LSTM (Alahi et al., 2016), the LSTM baseline employs an MHSA layer at every timestep These trajectories are plotted over the birds-eye-view image of the scene where we zoom into interesting trajectories can see that only AutoBot produces trajectories that are realistic in a group setting across all operating on the hidden state of the different strokes making up the character.
We concatenate a transformation of the one-hot vector representing the latent mode with the socially encoded hidden state at every timestep.
The output model is identical to the one used AutoBot, generating a sequence of bivariate Gaussian distributions for each stroke.
We performed a hyperparameter search on the learning rate and found the optimal learning rate to be 5 e − 4 .
Furthermore, as with all our other experiments, we found it helpful to employ gradient clipping during training.
We provide additional results on the two tasks deﬁned in Section 4.4.
Figure 14 shows additional successful AutoBot results on task 1 (completing multiple strokes) compared to an LSTM baseline equipped with social attention.
These results highlight the effectiveness of sequential set transformers for generating consistent diverse futures given an input sequence of sets.
Figure 15 shows examples where both models fail to correctly complete the character in task 1.
Figure 16 compares AutoBot-Ego with the LSTM baseline on predicting a new stroke given the two ﬁrst strokes of a character (task 2).
These results highlight that although not all modes predict the correct character across all modes, all generated predictions are consistent with realistic characters (i.e., confusing “F" with “ Π " or “H" with “ Π ").
Again, we see that AutoBot produces trajectories consistent with the all modes (e g., not crashing into the bushes) and maintains the social distance between the walking agents