S CENE T RANSFORMER : A UNIFIED ARCHITECTURE FOR PREDICTING MULTIPLE AGENT TRAJECTORIES
A BSTRACT
Predicting the motion of multiple agents is necessary for planning in dynamic environments.
This task is challenging for autonomous driving since agents (e.g., vehicles and pedestrians) and their associated behaviors may be diverse and inﬂuence one another.
Most prior work have focused on predicting independent futures for each agent based on all past motion, and planning against these independent predictions.
However, planning against independent predictions can make it challenging to represent the future interaction possibilities between different agents, leading to sub-optimal planning.
In this work, we formulate a model for predicting the behavior of all agents jointly, producing consistent futures that account for interactions between agents.
Inspired by recent language modeling approaches, we use a masking strategy as the query to our model, enabling one to invoke a single model to predict agent behavior in many ways, such as potentially conditioned on the goal or full future trajectory of the autonomous vehicle or the behavior of other agents in the environment.
Our model architecture employs attention to combine features across road elements, agent interactions, and time steps.
We evaluate our approach on autonomous driving datasets for both marginal and joint motion prediction, and achieve state of the art performance across two popular datasets.
Through combining a scene-centric approach, agent permutation equivariant model, and a sequence masking strategy, we show that our model can unify a variety of motion prediction tasks from joint motion predictions to conditioned prediction.
I NTRODUCTION
Motion planning in a dense real-world urban environment is a mission-critical problem for deploying autonomous driving technology.
Autonomous driving is traditionally considered too difﬁcult for a single end-to-end learned system (Thrun et al., 2006).
Thus, researchers have opted to split the task into sequential sub-tasks (Zeng et al., 2019): (i) perception, (ii) motion prediction, and (iii) planning.
Perception is the task of detecting and tracking objects in the scene from sensors such as LiDARs and cameras.
Motion prediction involves predicting the futures actions of other agents in the scene.
Finally, planning involves creating a motion plan that navigates through dynamic environments.
Dividing the larger problem into sub-tasks achieves optimal performance when each sub-task is truly independent.
However, such a strategy breaks down when the assumption of independence does not hold.
For instance, the sub-tasks of motion prediction and planning are not truly independent—the autonomous vehicle’s actions may signiﬁcantly impact the behaviors of other agents.
Similarly, the behaviors of other agents may dramatically change what is a good plan.
The goal of this work is to take a step in the direction of unifying motion prediction and planning by developing a model that can exploit varying forms of conditioning information, such as the AV’s goal, and produce joint consistent predictions about the future for all agents simultaneously.
While the motion prediction task has traditionally been formulated around per-agent independent predictions, recent datasets (Ettinger et al., 2021; Zhan et al., 2019) have introduced interaction prediction tasks that enable us to study joint future prediction (Figure 1).
These interaction prediction
[Figure]
tasks require models to predict the multiple to produce future predictions for all agents such that the agents futures are consistent 1 with one another.
A naive approach to producing joint futures is to consider the exponential number of combinations of marginal agent predictions.
Many of the combinations are not consistent, especially when agents have overlapping trajectories.
We present a uniﬁed model that naturally captures the interactions between agents, and can be trained as a joint model to produce scene-level consistent predictions across all agents (Figure 1, right).
Our model uses a scene-centric representation for all agents (Lee et al., 2017;
Hong et al., 2019; Casas et al., 2020a; Salzmann et al., 2020) to allow scaling to large numbers of agents in dense environments.
We employ a simple variant of self-attention (Vaswani et al., 2017) in which the attention mechanism is efﬁciently factorized across the agent-time axes.
The resulting architecture simply alternates attention between dimensions representing time and agents across the scene, resulting in a computationally-efﬁcient, uniform, and scalable architecture.
We ﬁnd that the resulting model, termed Scene Transformer , achieves superior performance on both independent (marginal) and interactive (joint) prediction benchmarks.
Moreover, we demonstrate a novel formulation of the task using a masked sequence model, inspired by recent advances in language modeling (Brown et al., 2020; Devlin et al., 2019), to allow conditioning of the model on the autonomous vehicle (AV) goal state or full trajectory.
In this reformulation, a single model can naturally perform tasks such as motion prediction, conditional motion prediction, and goal-conditioned prediction simply by changing which data is visible at inference time.
We hope that our uniﬁed architecture and ﬂexible problem formulation opens up new research directions for further combining motion prediction and planning.
In summary, our key contributions in this work are:
• A novel, scene-centric approach that allows us to gracefully switch training the model to produce either marginal (independent) and joint agent predictions in a single feed-forward pass.
Our model achieves state-of-the-art on both marginal and joint prediction tasks on both the Argoverse and the Waymo Open Motion Dataset.
• A permutation equivariant Transformer-based architecture factored over agents, time, and road graph elements that exploits the inherent symmetries of the problem.
The resulting architecture is efﬁcient and integrates the world state in a uniﬁed way.
• A masked sequence modeling approach that enables us to condition on hypothetical agent futures at inference time, enabling conditional motion prediction or goal conditioned prediction.
R ELATED W ORK
Motion prediction architectures.
Motion prediction models have ﬂourished in recent years, due to the rise in interest in self-driving applications and the release of related datasets and bench- marks (Kesten et al., 2019;
Chang et al., 2019; Caesar et al., 2020; Ettinger et al., 2021).
Successful models must take into account the history of agent motion, and the elements of the road graph (e.g., lanes, stop lines, trafﬁc light dynamic state).
Furthermore, such models must learn the relationships between these agents in the context of the road graph environment.
One class of models draws heavily upon the computer vision literature, rendering inputs as a multichannel rasterized top-down image (Cui et al., 2019; Chai et al., 2019; Lee et al., 2017; Hong et al., 2019;
Casas et al., 2020a; Salzmann et al., 2020; Zhao et al., 2019).
In this approach, relationships between scene elements are captured via convolutional deep architectures.
However, the localized structure of the receptive ﬁeld makes capturing spatially-distant interactions challenging.
A popular alternative is to use an entity-centric approach.
With this approach, agent state history is typically encoded via sequence modeling techniques like RNNs (Mercat et al., 2020; Khandelwal et al., 2020;
Lee et al., 2017; Alahi et al., 2016;
Rhinehart et al., 2019) or temporal convolutions (Liang et al., 2020).
Road elements are approximated with basic primitives (e.g. piecewise-linear segments) which encode pose information and semantic type.
Modeling relationships between entities is often presented as an information aggregation process, and models employ pooling (Zhao et al., 2020;
Gao et al., 2020; Lee et al., 2017; Alahi et al., 2016; Gupta et al., 2018), soft-attention (Mercat et al., 2020;
Zhao et al., 2020; Salzmann et al., 2020) as well as graph neural networks (Casas et al., 2020a; Liang et al., 2020; Khandelwal et al., 2020).
Like our proposed method, several recent models use Transformers (Vaswani et al., 2017), composed of multi-headed attention layers.
Transformers are a popular state-of-the-art choice for sequence modeling in natural language processing (Brown et al., 2020; Devlin et al., 2019), and have recently shown promise in core computer vision tasks such as detection (Bello et al., 2019; Carion et al., 2020;
Srinivas et al., 2021), tracking (Hung et al., 2020) and classiﬁcation (Ramachandran et al., 2019;
Vaswani et al., 2021; Dosovitskiy et al., 2021; Bello, 2013; Bello et al., 2019).
For motion modeling, recent work has employed variations of self-attention and Transformers for modeling different axes: temporal trajectory encoding and decoding (Yu et al., 2020; Giuliari et al., 2020; Yuan et al., 2021), encoding relationships between agents (Li et al., 2020; Park et al., 2020; Yuan et al., 2021; Yu et al., 2020; Mercat et al., 2020;
Bhat et al., 2020), and encoding relationships with road elements.
When applying self-attention over multiple axes, past work used independent self-attention for each axis (Yu et al., 2020), or ﬂattened two axes together into one joint self-attention layer (Yuan et al., 2021) – by comparison, our method proposes axis-factored attention to model relationships between time steps, agents, and road graph elements in a uniﬁed way.
Scene-centric versus agent-centric representations.
Another key design choice is the frame of reference in which the representation is encoded.
Some models do a majority of modeling in a global, scene-level coordinate frame, such as work that employs a rasterized top-down image (Cui et al., 2019;
Chai et al., 2019; Lee et al., 2017; Hong et al., 2019; Casas et al., 2020a; Salzmann et al., 2020).
This can lead to a more efﬁcient model due to a single shared representation of world state in a common coordinate frame, but comes with the potential sacriﬁce of pose-invariance.
On the other hand, models that reason in the agent-coordinate frame (Mercat et al., 2020;
Zhao et al., 2020; Khandelwal et al., 2020) are intrinsically pose-invariant, but scale linearly with the number of agents, or quadratically with the number of pairwise interactions between agents.
Many works employ a mix of a top-down raster representation for road representation fused with a per-agent representations (Rhinehart et al., 2019; Tang & Salakhutdinov, 2019; Lee et al., 2017).
Similar to our own work, LaneGCN (Liang et al., 2020) is agent-centric yet representations are in a global frame – to the best of our knowledge, this is the only other work to do so.
This enables efﬁcient reasoning while capturing arbitrarily distant interactions and high-ﬁdelity state representations without rasterization.
Representing multi-agent futures.
A common way to represent agent futures is via a weighted set of trajectories per agent (Alahi et al., 2016;
Biktairov et al., 2020; Buhet et al., 2020; Casas et al., 2020a;a; Chai et al., 2019;
Cui et al., 2019; Gao et al., 2020; Hong et al., 2019; Lee et al., 2017; Marchetti et al., 2020; Mercat et al., 2020; Salzmann et al., 2020; Zhao et al., 2020; Chandra et al., 2018).
This representation is encouraged by benchmarks which primarily focus on per-agent distance error metrics (Caesar et al., 2020; Chang et al., 2019; Zhan et al., 2019).
We argue in this work that modeling joint futures in a multi-agent environment (Figure 1, right) is an important concept that has been minimally explored in prior work.
Some prior work consider a factorized pairwise joint distribution, where a subset of agent futures are conditioned on other agents – informally, modeling
P ( X ) and P ( Y | X ) for agents X and Y (Khandelwal et al., 2020; Tolstaya et al., 2021; Salzmann et al., 2020).
To generalize joint prediction to arbitrary multi-agent settings, other work (Tang [Figure] & Salakhutdinov, 2019; Rhinehart et al., 2019; Casas et al., 2020b; Suo et al., 2021; Yeh et al., 2019) iteratively roll out samples per-agent, where each agent is conditioned on previously sampled trajectory steps.
In contrast, our model directly decodes a set of k distinct joint futures with associated likelihoods.
M ETHODS
The Scene Transformer model has three stages: (i) Embed the agents and the road graph into a high dimensional space, (ii) Employ an attention-based network to encode the interactions between agents and the road graph, (iii) Decode multiple futures using an attention-based network.
The model takes as input a feature for every agent at every time step, and also predicts an output for every agent at every time step.
We employ an associated mask, where every agent time step has an associated indicator of 1 (hidden) or 0 (visible), indicating whether the input feature is hidden (i.e. removed) from the model.
This approach mirrors the approach of masked-language models such as BERT (Devlin et al., 2019).
The approach is ﬂexible, enabling us to simultaneously train a single model for motion prediction (MP) (Cui et al., 2019;
Chai et al., 2019; Lee et al., 2017; Hong et al., 2019; Casas et al., 2020a; Salzmann et al., 2020; Casas et al., 2020a; Liang et al., 2020; Khandelwal et al., 2020), conditional motion prediction (CMP) (Khandelwal et al., 2020; Tolstaya et al., 2021; Salzmann et al., 2020) and goal-conditioned prediction (GCP) (Deo & Trivedi, 2020) simply by changing what data is shown to the model (Figure 2, left).
We summarize the key contributions below, and reserve details for the Appendix.
Multi-task representation.
The key representation in the model is a 3-dimensional tensor of A agents with D feature dimensions across T time steps.
At every layer within the architecture, we aim to maintain a representation of shape [ A, T, D ] , or when decoding, [ F, A, T, D ] across F potential futures.
Each task (MP, CMP, GCP) can be formulated as a query with a speciﬁc masking strategy by setting the indicator mask to 0, thus providing that data to the model (Figure 2, left).
The goal of the model is to impute the features for each shaded region corresponding to subsets of time and agents in the scenario that are masked.
3. S CENE - CENTRIC REPRESENTATION FOR AGENTS AND ROAD GRAPHS .
We use a scene-centric embedding where we use an agent of interest’s position as the origin 2 , and encode all roadgraph and agents with respect to it.
This is contrast to approaches which use an agent-centric representation, where the representations are computed separately for each agent, treating each agent in turn as the origin.
In detail, we ﬁrst generate a feature for every agent time step if that time step is visible.
Second, we generate a set of features for the static road graph, road elements static in space and time, learning one feature vector per polyline (with signs being polylines of length 1) using a PointNet (Qi et al., 2017).
Last, we generate a set of features for the dynamic road graph, which are road elements static in space but dynamic in time (e.g. trafﬁc lights), also one feature vector per object.
All three categories have xyz position information, which we preprocess to center and rotate around the agent of interest and then encode with sinusoidal position embeddings (Vaswani et al., 2017).
3. E NCODING TRANSFORMER
We focus on a simple encoder-decoder attention-based architecture which maintains a representation of [ A, T, D ] throughout (Figure 2, right).
We summarize the architecture brieﬂy, but reserve details for the Appendix and Table 5.
The majority of layers are a form of the Transformer layer (Vaswani et al., 2017)
(Table 6).
Attention layers are parameterized as matrices representing the query Q , key
K , and value V , whose output y = softmax ( QK T ) V √ dim k .
Each matrix is computed as a learned linear transformation of the underlying representation x , e.g. Q = W q x .
Each attention layer is followed by a feed-forward layer of the same hidden dimension, and a skip connection addition of the result with the input to the whole Transformer layer.
All layers of the encoder and decoder employ a D feature dimension.
The ﬁnal layer after the decoder is a 2-layer MLP that predicts 7 outputs.
The ﬁrst 6 outputs correspond to the 3-dimensional position of an agent at a given time step in absolute coordinates (e.g. meters) with respect to the agent of interest, and the corresponding uncertainty parameterized by a Laplace distribution (Meyer & Thakurdesai, 2020).
The remaining dimension predicts the heading.
Efﬁcient factorized self-attention.
The bulk of the computation is performed with a Transformer (Vaswani et al., 2017)
(Table 6).
One naive approach to use the Transformer would be to perform attention directly on the entire set of agent and time step features (i.e., attention across AT dimensions).
However, this approach is computationally expensive, and also suffers from an identity symmetry challenge: since we do not add any speciﬁc agent identity indicator, two agents of the same type with the same masked future time-step will have the same input representation to the transformer, resulting in the same output.
Thus, we design a factorized attention based on the time and agents axes (for related ideas, see Wang et al.
(2020); Szegedy et al.
; Ho et al.
Applying attention only across time allows the model to learn smooth trajectories independent of the identity of the agent.
Likewise, applying attention only across agents allows the model to learn multi-agent interactions independent of the speciﬁc time step.
Finally, in order to capture both time and agent dependencies, the model simply alternates attention across agents and time in subsequent layers (Figure 2, right panel).
The model is also permutation equivariant to the ordering of the agents at input, since the attention operation is permutation equivariant.
Cross-attention.
In order to exploit side information, which in our case is a road graph, we use cross-attention to enable the agent features to be updated by attending to the road graph.
Concretely, we calculate the queries from the agents, but the keys and values come from the embeddings of the road graph.
The road graph embeddings are ﬁnal after the per-polyline PointNet, and therefore not updated during these attention layers.
This requires that the model learn interactions between the road structure and agents that are independent of the speciﬁc time step or agent.
We highlight that the road graph representation is also permutation-equivariant and shared across all agents in the scene, whereas prior approaches have often used a per-agent road graph representation.
3. P REDICTING PROBABILITIES FOR EACH FUTURES .
Our model also needs to predict a probability score for each future (in the joint model) or trajectory (in the marginal model).
In order to do so, we need a feature representation that summarizes the scene and each agent.
After the ﬁrst set of factorized self-attention layers, we compute the mean of the agent features tensor across the agent and time dimension separately, and add these as an additional artiﬁcial agent and time making our internal representation [ A + 1 , T + 1 , D ] (Figure 2, left panel).
This artiﬁcial agent and time step propagates through the network, and provides the model with extra capacity for representing each agent, that is not tied to any timestep.
At the ﬁnal layer, we slice out the artiﬁcial agent and time step to obtain summary features for each agent (the additional time per agent), and for the scene (the ‘corner’ feature that is both additional time and agent).
This feature is then processed by a 2-layer MLP producing a single logit value that we use with a softmax classiﬁer for a permutation equivariant estimate of probabilities for each futures.
3. J OINT AND MARGINAL LOSS FORMULATION .
The output of our model is a tensor of shape [ F, A, T, 7] representing the location and heading of each agent at the given time step.
Because the model uses a scene-centric representation for the locations through positional embeddings, the model is able to predict all agents simultaneously in a single feed-forward pass.
This design also makes it possible to have a straight-forward switch between joint future predictions and marginal future predictions.
To perform joint future prediction, we treat each future (in the ﬁrst dimension) to be coherent futures across all agents.
Thus, we aggregate the displacement loss across all agents 3 and time steps to build a loss tensor of shape [ F ] .
We only back-propagate the loss through the individual future that most closely matches the ground-truth in terms of displacement loss (Gupta et al., 2018; Yeh et al., 2019).
For marginal future predictions, each agent is treated independently.
After computing the displacement loss of shape [ F, A ] , we do not aggregate across agents.
Instead, we select the future with minimum loss for each agent separately, and back-propagate the error correspondingly (Appendix, Figure 8).
Evaluation metrics for motion prediction.
We evaluate the quality of k weighted trajectory hypotheses using the standard evaluation metrics: minADE, minFDE, miss rate, and mAP.
Each evaluation metric attempts to measure how close the top k trajectories are to ground truth observation.
A simple and common distance-based metric is to measure the L 2 norm between a given trajectory and the ground truth (Alahi et al., 2016; Pellegrini et al., 2009).
minADE reports the L 2 norm of the trajectory with the minimal distance.
minFDE likewise reports the L 2 norm of the trajectory with the smallest distance only evaluated at the ﬁnal location of the trajectory.
We additionally report the miss rate (MR) and mean average precision (mAP) to capture how well a model predicts all of the future trajectories of agents probabilistically (Yeh et al., 2019;
Chang et al., 2019; Ettinger et al., 2021).
For joint future evaluation settings, we measure the scene-level equivalents (minSADE, minSFDE, and
SMR) that evaluate the prediction of the best single consistent future (Casas et al., 2020b).
R ESULTS
We evaluate the Scene Transformer on motion prediction tasks from the Argoverse dataset (Chang et al., 2019) and Waymo Open Motion Dataset (WOMD)
(Ettinger et al., 2021).
The Argoverse dataset consists of 324,000 run segments (each 5 seconds in length) from 290 km of roadway containing 11.
million agent tracks in total, and focuses on single agent (marginal) motion prediction.
The WOMD dataset consists of 104,000 run segments (each 20 seconds in length) from 1,750 km of roadway containing 7.64 million unique agent tracks.
Importantly, the WOMD has two tasks, each with their own set of evaluation metrics: a marginal motion prediction challenge that evaluates the quality of the motion predictions independently for each agent (up to 8 per scene), and a joint motion prediction challenge that evaluates the quality of a model’s joint predictions of exactly 2 agents per scene.
We train each model on a Cloud TPU (Jouppi et al., 2017) See Appendix for all training details.
First, in Section 4.1, we focus on the marginal prediction task and show that Scene Transformer achieves competitive results on both Argoverse (Chang et al., 2019) and the WOMD (Ettinger et al., 2021).
In Section 4.2, we focus on the joint prediction task and train the Scene Transformer with our joint loss formulation (see Section 3.4).
We show that with a single switch of the loss formula, we can achieve superior joint motion prediction performance.
In Section 4.3 we discuss factorized versus non-factorized attention.
Finally, in Section 4.4 we show how our masked sequence model formulation allows us to train a multi-task model capable of motion prediction, conditional motion prediction, and goal-conditioned prediction.
In Appendix B.1 we discuss the trade-off between marginal and joint models. [Table]
[Table]
Results presented on the standard splits of the validation and test datasets (Ettinger et al., 2021) evaluated with traditional marginal metrics for t = 8 seconds.
minADE, minFDE reported for k =predictions (Alahi et al., 2016; Pellegrini et al., 2009); Miss Rate (MR) (Chang et al., 2019) withinmeters of the target. See Appendix for t = 3 or 5 seconds. We include the challenge winner results in this table (Waymo, 2021). 4. M ARGINAL MOTION PREDICTION
We ﬁrst evaluate the performance of Scene Transformer trained and evaluated as a traditional marginal, per-agent motion prediction model.
This is analogous to the problem illustrated in Figure 1 (left).
For all results until Section 4.4, we use a masking strategy that provides the model with all agents as input, but with their futures hidden.
We also mask out future trafﬁc light information.
Argoverse.
We evaluate on the popular Argoverse (Chang et al., 2019) benchmark to demonstrate the efﬁcacy of our architecture.
During training and evaluation, the model is only required to predict the future of the single agent of interest.
Our best Argoverse model uses D = 512 feature dimensions and label smoothing for trajectory classiﬁcation.
Our model achieves state-of-the-art results compared to published, prior work 4 in terms of minADE and minFDE (Table 1).
Waymo Open Motion Dataset (WOMD).
We next evaluate the performance of our model with
D = 256 on the recently released WOMD (Ettinger et al., 2021) for the marginal motion prediction task.
This task is a standard motion prediction task where up to 8 agents per scene are selected to have their top 6 motion predictions evaluated independently.
Our model trained with the marginal loss achieves state-of-the-art results on the minADE, minFDE, and miss rate metrics (Table 2).
4. J OINT MOTION PREDICTION
To evaluate the effectiveness of Scene Transformer when trained with a joint loss formulation (Section 3.4), we evaluate our model on the Interaction Prediction challenge in WOMD (Ettinger et al., 2021).
This task measures the performance of the model at predicting the joint future trajectories of two [Table] interacting agents (Figure 1, right), and employs joint variants of the common minADE, minFDE, and Miss Rate (MR) metrics denoted as minSADE, minSFDE, SMR.
Note that the “ S ” indicates “scene-level” metrics.
These metrics aim to measure the quality and consistency of the two agents joint prediction - for example, the joint variant of Miss Rate (SMR) only records a "hit" if both interacting agent’s predicted trajectories are within the threshold of their respective ground truths.
We ﬁnd that for the Interaction Prediction challenge our joint model’s joint predictions easily outperforms the WOMD provided baseline as well as a marginal version of the model converted into a joint prediction 5 into joint predictions.
(Table 3).
This shows that beyond the strength of our overall architecture and approach, that explicitly training a model as a joint model signiﬁcantly improves joint performance on joint metrics.
A notable observation is that even though the Interaction Prediction task only requires predicting the joint trajectories of two agents, our method is fully general and predicts joint consistent futures of all agents.
4. F ACTORIZED AGENTS SELF - ATTENTION
Factorized self-attention confers two beneﬁts to our model: (a) it is more efﬁcient since the attention is over a smaller set, and (b) it provides an implicit identity to each agent during the attention across time.
We ran an experiment where we replaced each axis-factorized attention layer (each pair of time and agent factorized layer) with a non-axis factored attention layer.
This increased the computational cost of the model and performed worse on the Argoverse validation dataset: the factorized version achieved a minADE of 0.609 with the factored version, and 0.639 with the non-factorized version.
4. A DVANTAGES OF A MASKED SEQUENCE MODELING STRATEGY
Our model is formulated as a masked sequence model (Devlin et al., 2019), where at training and inference time we specify which agent timesteps to mask from the model.
This formulation allows us to select which information about any agent at any timestep to supply to the model, and measure how the model exploits or responds to this additional information.
We can express several motion prediction tasks at inference time in mask space (Figure 2), in effect providing a multi-task model.
We can use this unique capability to query the models for counterfactuals.
What would the model predict given a subset of agents full trajectories (conditional motion prediction), or given a subset of agents ﬁnal goals (goal-conditioned motion prediction)?
This feature could be particularly useful for autonomous vehicle planning to predict what various future rollouts of the scene would look like given a desired goal for the autonomous vehicle (Figure 3).
[Figure]
In the previous results, we used a preﬁx-mask during training and inference, which shows the ﬁrst few steps and predicts the remaining timesteps for all agents.
For this experiment, we employ two masking strategies to test on the interactive split of the WOMD, namely "conditional motion prediction" (CMP) where we show one of the two interacting agents’ full trajectory, and "goalconditioned motion prediction" (GCP) where we show the autonomous vehicle’s desired goal state (Figure 2).
We train a model using each strategy (including MP masking) 1 / 3 of the time, and evaluate the model on each all tasks.
We ﬁnd that across the 44k evaluation segments, the multi-task model matches the performance of our MP-only trained joint model on both joint (Table 3) and marginal metrics (see also Appendix, Table 11); multi-tasking training does not signiﬁcantly degrade standard motion prediction performance. We qualitatively examine the performance of the multi-task model in a GCP setting, and observe that the joint motion predictions for the AV and non-AV agents ﬂexibly adapt to selected goal points for the AV (Figure 3).
We note that dataset average metrics like min(S)ADE do not capture forecasting subtleties (Ivanovic & Pavone, 2021) that are critical for measuring progress in the ﬁeld. Unfortunately, the Waymo
Motion Open Dataset Interactive Prediction set does not contain pre-deﬁned metadata necessary to evaluate speciﬁc interesting slices using different models. In Appendix C we provide some initial analysis of slices that we could compute on the dataset, which illustrate how joint and marginal model predictions behavior differently as the scene a) scales with more agents in complicated scenes, b) differs in the average speed of agents. We also show that joint prediction models demonstrate joint consistent futures via lower inter-prediction overlap compared to marginal models.
D ISCUSSION
We propose a uniﬁed architecture for autonomous driving that is able to model the complex interactions of agents in the environment.
Our approach enables a single model to perform motion prediction, conditional motion prediction, and goal-conditioned prediction.
In the ﬁeld of autonomous driving, elaborations of this problem formulation may result in learning models for planning systems that quantitatively improve upon existing systems (Buehler et al., 2009; Montemerlo et al., 2008; Ziegler et al., 2014; Zeng et al., 2019; Liu et al., 2021a).
Likewise, such modeling efforts may be used to directly go after issues of identifying interacting agents in an environment, and potentially provide an important task for identifying causal relationships (Arjovsky et al., 2019; Schölkopf et al., 2021).
Appendix A A RCHITECTURE AND T RAINING D ETAILS
Overview.
Table 5 provides additional details about the Scene Transformer architecture and training.
The bulk of the computation and parameters reside in the Transformer layers.
Table 6 lists all of the operations and learned parameters in our implementation of a Transformer.
MLPs are employed to embed the original data (e.g, layers A , B , C ), build F futures from the encoding (layer T ), and predict the ﬁnal position, headings and uncertainties of all agents (layers Z 1 and Z 2 ).
The resulting attention-based architecture using D = 256 feature dimensions contains 15,296,136 parameters.
Nearly all settings used for both WOMD and Argoverse are identical, except that for Argoverse we use D = 512 and label smoothing to get our best marginal results.
The resulting model is trained on a TPU custom hardware accelerator (Jouppi et al., 2017) and converges in about 3 days of training.
Decoding multiple futures.
In order to allow the decoder to output F distinct futures, we perform the following operations.
The decoder receives as input a tensor of shape [ A, T, D ] corresponding to
A agents across T time steps and D feature dimensions.
The following series of simple operations restructures the representation to predict F futures.
First, the representation is tiled F times to generate [ F, A, T, D ] .
We append a one-hot encoding to the ﬁnal dimension where a 1 indicates which the identity of each future, resulting in a tensor of shape [ F, A, T, D + F ] .
The one-hot encoding allows the network to learn an embedding for each of the F futures.
For computational simplicity, the resulting representation is propagated through a small MLP to produce a return a tensor of the shape [ F, A, T, D ] .
Padding and Hidden Masks.
Padding and hidden masks are important to get right in the implementation of such a model.
In particular, we need to ensure that the masks do not convey additional future information to the model (e.g., if the model knows which timesteps an agent is visible or occluded based on the padding in the data, it may take advantage of this and not be able to generalize).
We use the concept of padding to indicate the positions where the input is absent or provided.
This is distinct from the hidden mask of shape [ A, T ] that is used for task speciﬁcation .
The hidden mask is used to query the model to inform it on which locations to predict, while the padding tells us which locations have inputs and ground-truth to compute a loss over.
All padded positions are set to be hidden during preprocessing, regardless of the masking scheme, so the model tries to predict their values.
All layers in our implementation are padding aware, including normalization layers (like batch normalization used in our encoding MLPs) and attention operations.
Our attention operations set the attention weights to 0 for padded elements.
If after providing our task speciﬁc hidden mask, no non-padded (valid) time steps exist for an agent, the whole agent is set to be padded.
This prevents agent slots with no valid data from being used in the model.
Predicting uncertainties.
For each predicted variate such as position or heading, we additionally predict a corresponding uncertainty.
In early experiments we found that predicting a paired uncertainty improves the predictive performance of the original variate and provides a useful signal for interpreting the ﬁdelity of the predicted value.
We employ a loss function that predicts a parameterization of the uncertainty corresponding to a Laplace distribution (Meyer & Thakurdesai, 2020).
Loss in lateral and longitudinal coordinates.
We use a Laplace distribution parameterization for the loss, with a diagonal covariance such that each axis is independent.
To enable the model to learn meaningful uncertainties, we rotate the scene per box prediction such that each prediction’s associated ground-truth is axis aligned.
This formulation results in uncertainty estimates that correspond to the agent’s lateral and longitudinal coordinates.
Heading representation: Our output format is 7 dimensional, where one of the dimensions, heading, is represented in radians.
To supervise the loss on the heading dimension, we employ a now standard method of "wrapping" the angle difference between the predicted and groundtruth heading for each agent; see Figure 4.
Marginal predictions from a joint model.
The model has the added feature of being easily adapted for making marginal predictions for each agent.
To produce per-trajectory scores for each agents predictions, we attach an artiﬁcial extra time step to the end of the [ A, T, D ] agent feature matrix to give us a matrix of shape [ A, T + 1 , D ] .
This additional feature provides the model with extra def WrapAngleRadians(angles_rad, min_val=-np.pi, max_val= np.pi): max_min_diff = max_val - min_val return min_val + tf.math.floormod(angles_rad + max_val, max_min_diff) heading_diff =
pred_heading - gt_heading heading_error =
WrapAngleRadians(heading_diff, min_val=-PI, max_val=PI) heading_loss
=
HuberLoss(heading_error)
[Figure]
capacity for representing each agent, that is not tied to any timestep.
This additional feature can be used to predict which of the k trajectories most closely matches the ground truth for that agent.
Data augmentation.
We use two methods for data augmentation to aid generalization and combat overﬁtting on both datasets.
First, we use agent dropout, where we artiﬁcially remove non-predicted agents with some probability.
We found a probability of 0 .
1 worked well for both datasets.
We also found it beneﬁcial to randomly rotate the entire scene between [ − π 2 , π 2 ] after centering to the agent of interest.
On Argoverse, this agent of interest is the agent the task designates to be predicted, where as on the WOMD we always center around the autonomous vehicle (AV), even though for some tasks it is not one of the predicted agents.
Lastly, each Argoverse scene contains many agents that the model is not required to predict; we employ these contextual agents as additional target training data if the contextual agents moved by at least 6m. On WOMD, Table 4 shows data augmentation on the standard validation set does improve minADE modestly, though does not account for the majority of the beneﬁt of the model.
Argoverse classiﬁcation label smoothing: Our best model on Argoverse uses D = 512 feature dimensions, but naively scaling the model this way leads to severe overﬁtting on the classiﬁcation subtask.
During training of our best model we employ label smoothing (0.1 + 1/6 for negatives and 0.9 + 1/6 for positive target).
WOMD redundant trajectory combination: The AP evaluation metric on the WOMD expects that exactly one of the trajectories is given a high probability and the other trajectories a low probability.
For example, the future prediction for a stationary agent is expected to have only one future with a high probability score, with a low score for the rest – even though the trajectories may all be the same, in this case, stationary.
Our best model combines redundant trajectories together if they are close in spatial location (less than 3 . 2 m) at the ﬁnal timestep.
Embedding of agents and road graph.
To generate input features, we use sinusoidal positional embeddings (Vaswani et al., 2017) to embed the time (for agents and dynamic roadgraph) and xyz -coordinates separately into a D dimensional features per dimension.
We encode the type of each object using a one-hot encoding (e.g. object type, lane type, etc), and concatenate any other features provided in the data set such as yaw, width, length, height, and velocity.
Dynamic road graphs have a second one-hot encoding indicating state, like the trafﬁc light state.
Lastly, for agents and the dynamic road graph, we add a binary indicator on whether the agent is hidden at that time step.
If the agent or dynamic road graph element is hidden, all input features (e.g. position, type, velocity, state, etc) are set to 0 before encoding except the time embedding, which are linearly spaced values at the dataset’s update rate starting at 0, and the hidden indicator.
For agents and the dynamic road graph, we use a 2 layer MLP with a hidden and output dimension of
D to produce a ﬁnal feature per agent or object and per time step.
For the static road graph, we must reduce a point cloud of up to 20 , 000 road graph points, each belonging to a smaller set of polylines, to a single vector per polyline.
Because some lanes can be extremely long, we break up any polyline longer than 20 points into a new set of smaller polylines.
Then, we apply a small PointNet (Qi et al., 2017) architecture with a 2 layer MLP with a hidden and output dimension of D to each point, and use max pooling per polyline to get a ﬁnal feature per element.
Inference latency. Scene Transformer produces a prediction for every agent in the scene in a single pass. Its inference speed depends on the number of agents in the scene. Our preliminary proﬁling of inference speed ranged from 52 ms (32 agents) to 175 ms (128 agents) on the Waymo Open Motion
Dataset (128 agents, 91 timesteps for each agent, up to 1400 roadgraph elements). This was measured on an Nvidia V100 using a standard TensorFlow inference graph in ﬂoat32 without optimization tricks, optimization tuning, etc. This is in line with the expected linear scaling of the factorized attention modules.
[Table]
[Figure]
[Table] [Figure] [Table] [Figure]
[Table]
[Figure]
None
[Table]
[Figure]
[Table]
B A DDITIONAL M OTION P REDICTION R ESULTS
[Table]
[Table]
[Table]
[Table] [Table] for t = 8 seconds . Please see Table 2 for details.
We additionally include "joint-as-marginal" results for our standard MP masked joint model
[Table]
and our multi-task joint model. These are joint models evaluated as if they were marginal models (with no changes in the outputs)
[Table]
cyc
[Table]
cyc veh ped cyc veh ped cyc
[Table]
B. U NDERSTANDING THE TRADE - OFFS BETWEEN THE JOINT AND MARGINAL MODELS
While the joint mp-only model performs better on the interactive prediction task than the marginal model when evaluated with joint metrics , it performs worse on the motion prediction task when evaluated with marginal metrics (see Table 11 "ours (joint)").
This is because the joint model is trained to produce predictions that are internally consistent between agents, while the marginal model does not, which is a strictly harder task.
In this section, we examine this quality difference and the internal consistency of the predictions from both models.
In the WOMD dataset, each example has a different number of agents to be predicted.
By slicing the marginal minADE results based on the number of agents to be predicted (Figure 5, left), we ﬁnd that the joint model performs worse as there are more predicted agents, while the marginal model performs the same.
This is expected: when there are more agents, the joint model has a more difﬁcult task since it needs to produce internally consistent predictions.
One may expect a joint model to need an exponential number of trajectory combinations to perform competitively with the marginal model, but we are encouraged to ﬁnd that this is not actually the case.
We believe this is due to the fact that when many agents are interacting, the number of realistic scenarios is actually heavily constrained by interactions - most combinations of marginal predictions don’t actually make sense together.
However, when more agents are to be predicted, the possibility of interactions are higher and we would expect that the joint model is able to capture these interactions through internally consistent predictions.
We measure the models’ internal consistency by selecting the best trajectory and measuring the inter-prediction overlap rate (details below).
We ﬁnd that the joint model has a consistently lower inter-prediction overlap rate, showing that it is able to capture the interactions between agents.
The ability to model interactions enables the model to be suitable for conditional motion prediction, and goal conditioned prediction tasks, which we discuss in Section 4.4.
Measuring agent inter-prediction overlap.
We measure predicted agent overlap on the best set of trajectory produced by the model.
For the joint model, this corresponds to the joint prediction that has the higher probability score.
For the marginal model, we take the top scoring trajectory for each agent to obtain the best set.
For every predicted agent in the trajectory, we determine if it has an overlap with any other predicted agent by comparing the rotated upright 3D bounding boxes using our predicted xyz and heading.
The inter-prediction overlap rate is the number of predicted agents that are involved in some overlap, divided by the total number of predicted agents.
We count two agents as overlapping if the intersection over the union (IoU) exceeds 0.01.
C S LICING RESULTS
D L OSS I MPLEMENTATION
The Scene Transformer model is both agent permutation equivariant and scene-centric.
These properties allow us to easily switch between a marginal vs joint loss formulation (Figure 8).
In the marginal formulation, we reduce the loss of the best future for every agent separately; while in the joint formulation, we reduce the loss of the best future for all agents jointly.
In practice, this determines when the reduce_min operation is performed.
Our ﬁnal loss composes the regression loss on the trajectory and the classiﬁcation loss of the best trajectory.
A weighted linear combination of the loss terms is used to combine these two losses together.
The classiﬁcation loss weight was set to be 0.1, while the regression losses have weight 1.0.
The weights were determined using the hold-out validation set.
# agent_predictions and agent_gt are # [F, A, T, 7] Tensors with [x, y, z], uncertainty terms, and yaw.
# Get the KL Divergence of the predictions vs ground truth.
# Use LaplaceKL method from (Meyer & Thakurdesai, 2020) loss = LaplaceKL(agent_predictions, agent_gt)
# Now reduce across all timesteps and values to produce # a tensor of shape [F, A] loss = tf.reduce_sum(loss, axis=[2, 3])
# The marginal loss, we only apply the loss to the best trajectory # per agent (so min across the future dimension).
marginal_loss = tf.reduce_min(loss, axis=0) # Then sum over the agent dimension marginal_loss = tf.reduce_sum(marginal_loss)
# The joint loss, we sum over all agents to get a loss value # per future.
joint_loss = tf.reduce_sum(loss, axis=1)
# Then only apply the loss to the best future prediction joint_loss = tf.reduce_min(joint_loss)