Scene Transformer: A unified architecture for predicting future trajectories of multiple agents

1 INTRODUCTION . 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 . 2 RELATED WORK . 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 .

Summary: 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.