Year: 2021

Abstract:
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 influence 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.

heading: 1 INTRODUCTION
text: 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 difficult 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 significantly 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 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 tasks require models to predict the joint futures of multiple agents: models are expected to produce future predictions for all agents such that the agents futures are consistent 1 with one another.
1Marginal agent predictions may conflict with each other (have overlaps), while consistent joint predictions should have predictions where agents respect each other’s behaviors (avoid overlaps) within the same future.
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 reasonable, especially when agents have overlapping trajectories. We present a unified 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 efficiently 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-efficient, uniform, and scalable architecture.
We find that the resulting model, termed Scene Transformer, achieves superior performance on both independent (marginal) and interactive (joint) prediction benchmarks. We further show how we can structure the problem as 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 unified architecture and flexible 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 efficient and integrates the world state in a unified 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.

heading: 2 RELATED WORK
text: Motion prediction architectures. Motion prediction models have flourished in recent years, due to the rise in interest in self-driving applications and the release of related datasets and benchmarks (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, traffic 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 field 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 classification (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 flattened 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 unified 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 efficient model due to a single shared representation of world state in a common coordinate frame, but comes with the potential sacrifice 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 efficient reasoning while capturing arbitrarily distant interactions and high-fidelity 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). 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 & 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.

heading: 3 METHODS
text: 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 flexible, 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) (Zeng et al., 2019; Liu et al., 2021a) 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 specific 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.

heading: 3.1 SCENE-CENTRIC REPRESENTATION FOR AGENTS AND ROAD GRAPHS.
text: 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 first 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. traffic lights), also one feature vector per object. All three categories
2For WOMD, we center the scene with respect to the autonomous vehicle (AV). For Argoverse, we center the scene with respect to the agent that needs to be predicted. Both are centered around what would be the last visible time step in a motion prediction setup for all tasks.
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).

heading: 3.2 ENCODING TRANSFORMER
text: 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 briefly, but reserve details for the Appendix and Table 4. The majority of layers are a form of the Transformer layer (Vaswani et al., 2017) (Table 5). Attention layers are parameterized as matrices representing the query Q, key K, and value V , whose output y = softmax ( QKT )V√ dimk
. Each matrix is computed as a learned linear transformation of the underlying representation x, e.g. Q = Wq 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 final layer after the decoder is a 2-layer MLP that predicts 7 outputs. The first 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.
Efficient factorized self-attention. The bulk of the computation is performed with a Transformer (Vaswani et al., 2017) (Table 5). 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: two agents of the same type hidden at the same future location will have identical input representations. Thus, we design a factorized attention based on the time and agents axes (for related ideas, see Wang et al. (2020); Szegedy et al. (2016); Ho et al. (2019)).
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 specific 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 final 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 specific time step or agent. We highlight that the road graph representation is shared across all agents in the scene, whereas prior approaches have often used an agent-centric road graph representation.

heading: 3.3 PREDICTING PROBABILITIES FOR EACH FUTURES.
text: 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 first 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 artificial agent and time making our internal representation [A+ 1, T + 1, D] (Figure 2, left panel). This artificial 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 final layer, we slice out the artificial 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 classifier for a permutation equivariant estimate of probabilities for each futures.

heading: 3.4 JOINT AND MARGINAL LOSS FORMULATION.
text: 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 first 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 7). This allows the agents to select among different futures. One notable benefit of our approach is that we are able to produce predictions for all agent simultaneously.
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 L2 norm between a given trajectory and the ground truth (Alahi et al., 2016; Pellegrini et al., 2009). minADE reports the L2 norm of the trajectory with the minimal distance. minFDE likewise reports the L2 norm of the trajectory with the smallest distance only evaluated at the final 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).

heading: 4 RESULTS
text: 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.7 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.

heading: 4.1 MARGINAL MOTION PREDICTION
text: We first 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).
3Each dataset identifies a subset of agents to be predicted. We only include this subset in our loss calculation. For Argoverse, this is 1 agent; for WOMD this is 2-8 agents.
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 traffic light information.
Argoverse. We evaluate on the popular Argoverse (Chang et al., 2019) benchmark to demonstrate the efficacy 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 classification. 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).

heading: 4.2 JOINT MOTION PREDICTION
text: 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 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
4We exclude comparing to public leaderboard entries that have not been published since their details are not available, but note that our results are competitive on the leaderboard as of the submission date.
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 find 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 significantly 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.

heading: 4.3 FACTORIZED AGENTS SELF-ATTENTION
text: The factorized self-attention confers two benefits to our model: (a) it is more efficient 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 turned out to perform worse. On our Argoverse validation dataset, we obtained a minADE of 0.609 with the factored version, and a minADE of 0.639 with the non-factorized version.

heading: 4.4 ADVANTAGES OF A MASKED SEQUENCE MODELING STRATEGY
text: 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 final goals (goal-conditioned motion prediction)? This feature is 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).
5Note that the output of a joint model can be directly used in a marginal evaluation. However, converting the output of a marginal model into a joint evaluation is nuanced because there lacks an association of futures across agents (Figure 1). We employ a simple heuristic to convert the outputs of a marginal model for joint evaluation: we take the top 6 pairs of trajectories from the combination of both agent’s trajectories for 36 total pairs, and retain the top 6 pairs with the highest product of probabilities.
In the previous results, we used a MP-mask during training and inference, which corresponds to showing the first few timesteps and predicting the remaining timesteps for all agents. For this experiment, we further design two other 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 future trajectory, and "goal-conditioned motion prediction" (GCP) where we show the autonomous vehicles desired goal state (Figure 2). We train a multi-task model on each of the three masks (including MP masking) 1/3 of the time, and evaluate the model on each these tasks. We find that 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 10), indicating that the addition of tasks does not degrade standard motion prediction performance.
With the additional multi-task training, the resulting model may now condition on side information. As a first test, we measure in the CMP and GCP setting that the model reduces the overall minADE of the prediction (For CMP, 1.34→ 1.25 for non-AV; For GCP, 0.95→ 0.60 for AV). As a second test, 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 flexibly adapt to arbitrarily selected goal points for the AV (Figure 3). Although no quantitative benchmarks are available in the community for quantifying GCP predictions, we take these results as positive indication that the model is responding appropriately and save further exploration of counterfactual analysis for future work.

heading: 5 DISCUSSION
text: We propose a unified 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 field 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). 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).

heading: A ARCHITECTURE AND TRAINING DETAILS
text: Overview. Table 4 provides additional details about the Scene Transformer architecture and training. The bulk of the computation and parameters reside in the Transformer layers. Table 5 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 final position, headings and uncertainties of all agents (layers Z1 and Z2). 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 copied across F dimensions to generate [F,A, T,D]. We append on a one-hot encoding to the final dimension where a 1 indicates which of the F futures is the ground truth trajectory 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 specification. 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 specific 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 fidelity 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.
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 artificial 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 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 overfitting on both datasets. First, we use agent dropout, where we artificially remove non-predicted agents with some probability. We found a probability of 0.1 worked well for both datasets. We also
found it beneficial 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.
Argoverse classification label smoothing: Our best model on Argoverse uses D = 512 feature dimensions, but naively scaling the model this way leads to severe overfitting on the classification 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.2m) at the final 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 traffic 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 final 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 final feature per element.

heading: B ADDITIONAL MOTION PREDICTION RESULTS
text: B.1 UNDERSTANDING 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 10 "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 4, left), we find 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 difficult 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 find 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 find 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.

heading: C SLICING RESULTS

heading: D LOSS IMPLEMENTATION
text: 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 7). 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 final loss comprises of the regression loss on the trajectory and the classification loss of the best trajectory. A weighted linear combination of the loss terms is used to combine these two losses together. The classification 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)
Figure 7: Pseudo-code in TensorFlow (Abadi et al., 2015) demonstrating the joint versus marginal loss formulation.

Review summaries:
The paper proposes a new Transformer-based trajectory forecasting model that can predict multiple agents in a scene. It can be used for goal-directed trajectory forecasting as well. Experiments on Argoverse and Waymo with good results. Related work is sufficient. Technical novelty is low. Paper is clearly written.

The paper presents a unified architecture to predict future motion trajectories of all agents in the scene, while interactions between agents are captured. A Transformer based architecture is designed. It also uses a masking strategy as a query, enabling one to condition on hypothetical agent futures during the inference time, and achieves conditional motion prediction or goal conditioned prediction. The model was tested on two standard datasets, and achieves state of the art performance.

This paper proposes a transformer-style architecture, Scene Transformer, for joint trajectory prediction of multiple agents (including the autonomous vehicle). The key representation of Scene Transformer is a 3D tensor [A, T, D], with the dimensions being agent, time step, and feature dimension, and we mask out the entries that to be predicted. The output is [F, A, T, 7], where F is the number of future predictions. The model outputs 7 values for each waypoint, x/y position and uncertainty in 3D space and heading. The model uses factorized self-attention applied on the A axis and T axis separately to transform the [A, T, D] representation tensor, and uses a cross attention to fuse in the static map features. The model can be trained to predict either marginal predictions or joint predictions depending on the loss function. Also, based on what entries in the input tensor are masked out, Scene Transformer can perform motion prediction, conditioned motion prediction, and goal-conditioned motion prediction all in the same model.  The Scene Transformer model was evaluated on Argoverse and Waymo Open Motion Dataset (WOMD), and its marginal prediction achieves the state-of-the-art performance on both datasets. The evaluation also shows that the model is able to perform joint prediction for the WOMD Interaction Prediction task and achieves better performance on the scene-level metrics compared to the baselines. The evaluation also shows that the Scene Transformer model is able to do conditional motion prediction and goal-conditioned motion prediction.