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