THOMAS: Trajectory Heatmap Output with learned Multi-Agent Sampling

In this paper , we propose THOMAS , a joint multi-agent trajectory prediction framework allowing for an efficient and consistent prediction of multi-agent multimodal trajectories . We present a unified model architecture for simultaneous agent future heatmap estimation , in which we leverage hierarchical and sparse image generation for fast and memory-efficient inference . We propose a learnable trajectory recombination model that takes as input a set of predicted trajectories for each agent and outputs its consistent reordered recombination . This recombination module is able to realign the initially independent modalities so that they do no collide and are coherent with each other . We report our results on the Interaction multi-agent prediction challenge and rank 1st on the online test leaderboard . 1 INTRODUCTION . Motion forecasting is an essential step in the pipeline of an autonomous driving vehicle , transforming perception data into future prediction which are then leveraged to plan the future moves of the autonomous cars . The self-driving stacks needs to predict the future trajectories for all the neighbor agents , in a fast and coherent way . The interactivity between agents plays an important role for accurate trajectory prediction . Agents need to be aware of their neighbors in order to adapt their speed , yield right of way and merge in neighbor lanes . To do so , different interaction mechanisms have been developed , such as social pooling ( Alahi et al. , 2016 ; Lee et al. , 2017 ; Deo & Trivedi , 2018 ) , graphs ( Salzmann et al. , 2020 ; Zeng et al. , 2021 ) or attention ( Mercat et al. , 2020 ; Messaoud et al. , 2020 ; Luo et al. , 2020 ; Gao et al. , 2020 ; Liang et al. , 2020 ) , benefiting from the progress of powerful transformer architectures ( Li et al. , 2020 ; Yuan et al. , 2021 ; Girgis et al. , 2021 ; Ngiam et al. , 2021 ) . These mechanisms allow agents to look at and share features with neighbors and to take them into account in their own predictions . Multi-modality is another important aspect of the possible future trajectories . A car can indeed chose to turn right or left , or decide to realise a certain maneuver in various ways . Uncertainty modeled as variance of Gaussians is insufficient to model these multiple cases , as it can only represent a continuous spread and can not show multiple discrete possibilities . Therefore , current state-of-the-art produces not one but K possible trajectories for each agent predicted , and most recent benchmarks ( Caesar et al. , 2020 ; Chang et al. , 2019 ; Zhan et al. , 2019 ; Ettinger et al. , 2021 ) include multimodality in their metrics , taking only the minimum error over a predicted set of K trajectories . However , up until very recently and the opening of multi-agent joint interaction challenges ( Ettinger et al. , 2021 ; Zhan et al. , 2021 ) , no motion forecasting prediction datasets were taking into account the coherence of modalities between different agents predicted at the same time . As a result , the most probable predicted modality of a given agent could crash with the most probable modality of another agent . Our THOMAS model encodes the past trajectories of all the agents present in the scene , as well as the HD-Map lanelet graph , and then predicts for each agent a sparse heatmap representing the future probability distribution at a fixed timestep in the future , using hierarchical refinement for very efficient decoding . A deterministic sampling algorithm then iteratively selects the best K trajectory endpoints according to the heatmap for each agent , in order to represent a wide and diverse spec- trum of modalities . Given this wide spectrum of endpoints , a recombination module takes care of addressing consistency in the scene among agents . Our contributions are summarized as follow : • We propose a hierarchical heatmap decoder allowing for unconstrained heatmap generation with optimized computational costs , enabling efficient simultaneous multi-agent prediction . • We design a novel recombination model able to recombine the sampled endpoints to obtain scene-consistent trajectories across the agents . 2 RELATED WORK . Learning-based models have quickly overtaken physics-based methods for trajectory prediction for several reasons . First , the sequential nature of trajectories is a logical application for recurrent architectures ( Alahi et al. , 2016 ; Altché & de La Fortelle , 2017 ; Lee et al. , 2017 ; Mercat et al. , 2020 ; Khandelwal et al. , 2020 ) . Then , benefiting from the latest progresses in computer vision , convolutional layers can easily be applied to bird-view rasters of the map context ( Lee et al. , 2017 ; Tang & Salakhutdinov , 2019 ; Cui et al. , 2019 ; Hong et al. , 2019 ; Salzmann et al. , 2020 ; Chai et al. , 2020 ; Gilles et al. , 2021b ) . A more compact representation closer to the trajectory space can encode surrounding HD-Maps ( usually formalized as connected lanelets ) , using Graph Neural Networks ( Gao et al. , 2020 ; Liang et al. , 2020 ; Zeng et al. , 2021 ; Gilles et al. , 2021a ) . Finally , some pointbased approaches ( Ye et al. , 2021 ) can be applied in a broader way to trajectory prediction , as both lanes and trajectories can be considered as ordered set of points . Multi-modality in prediction can be obtained through a multiple prediction head in the model ( Cui et al. , 2019 ; Liang et al. , 2020 ; Ngiam et al. , 2021 ; Deo et al. , 2021 ) . However , some methods rather adopt a candidate-based approach where potential endpoints are obtained either from anchor trajectories through clustering ( Chai et al. , 2020 ; Phan-Minh et al. , 2020 ) or from a model-based generator ( Song et al. , 2021 ) . Other approaches use a broader set of candidates from the context graph ( Zhang et al. , 2020 ; Zhao et al. , 2020 ; Zeng et al. , 2021 ; Kim et al. , 2021 ) or a dense grid around the target agent ( Deo & Trivedi , 2020 ; Gu et al. , 2021 ; Gilles et al. , 2021b ; a ) . Another family of approaches use variational inference to generate diverse predictions through latent variables ( Lee et al. , 2017 ; Rhinehart et al. , 2018 ; Tang & Salakhutdinov , 2019 ; Casas et al. , 2020 ) or GAN ( Gupta et al. , 2018 ; Rhinehart et al. , 2018 ; Sadeghian et al. , 2019 ) but the sampling of these trajectories is stochastic and does not provide any probability value for each sample . While very little work has directly tackled multi-agent prediction and evaluation so far , multiple methods hint at the ability to predict multiple agents at the same time ( Liang et al. , 2020 ; Zeng et al. , 2021 ) even if they then focus on a more single-agent oriented framework . Other works ( Alahi et al. , 2016 ; Tang & Salakhutdinov , 2019 ; Rhinehart et al. , 2019 ; Girgis et al. , 2021 ) use autoregressive roll-outs to condition the future step of an agent on the previous steps of all the other agents . SceneTransformer ( Ngiam et al. , 2021 ) repeats each agent features across possible modalities , and performs self-attention operations inside each modality before using a loss computed jointly among agents to train a model and evaluate on the WOMD dataset ( Ettinger et al. , 2021 ) interaction track . ILVM ( Casas et al. , 2020 ) uses scene latent representations conditioned on all agents to generate scene-consistent samples , but its variational inference does not provide a confidence score for each modality , hence LookOut ( Cui et al. , 2021 ) proposes a scenario scoring function and a diverse sampler to improve sample efficiency . AIR2 ( Wu & Wu , 2021 ) extends Multipath ( Chai et al. , 2020 ) and produces a cross-distribution for two agents along all possible trajectory anchors , but it scales exponentially with the number of agents , making impractical for a real-time implementation that could encounter more than 10 agents at the same time . The approach most related to this paper is GOHOME ( Gilles et al. , 2021a ) , which uses a similar graph encoder and then leverage lane rasters to generate a probability heatmap in a sparse manner . However , this lane-generated heatmap remains constrained to the drivable area and to fixed lanewidths . Another close approach is the one of DenseTNT ( Gu et al. , 2021 ) , which also uses attention to a grid of points in order to obtain a dense prediction , but their grid also remains constrained to a neighborhood of the drivable area . Finally , none of these previous methods considers the problems of scene consistency for multi-agent prediction . 3 METHOD . Our goal is to predict the future T timesteps ofA agents using their past history made ofH timesteps and the HD-Map context . Similar to recent works ( Zhao et al. , 2020 ; Zeng et al. , 2021 ; Gu et al. , 2021 ) , we will divide the problem into goal-based prediction followed by full trajectory reconstruction . Our prediction pipeline is displayed in Fig . 1 . We first encode each agent trajectory and the HD-Map context graph into a common representation . We then decode a future probability heatmap for each agent in the scene , which we sample heuristically to maximize coverage . Finally , we recombine the sampled endpoints into scene-consistent modalities across agents and build the full trajectories for each agent . Our pipeline shares the same graph encoder , sampling algorithm and full trajectory generation as GOHOME ( Gilles et al. , 2021a ) , but uses a novel efficient hierarchical heatmap process that enables to scale to simultaneous multi-agent prediction . Furthermore , we add a novel scene-consistency module that recombines the marginal outputs into a joint prediction . 3.1 MODEL BACKBONE . 3.1.1 GRAPH ENCODER . We use the same encoder as the GOHOME model ( Gilles et al. , 2021a ) . The agent trajectories are encoded though TrajEncoder using a 1D CNN followed by a UGRU recurrent layer ( Rozenberg et al. , 2021 ) , and the HD-Map is encoded as a lanelet graph using a GNN GraphEncoder made of graph convolutions . We then run cross-attention Lanes2Agents to add context information to the agent features , followed by self-attention Agents2Agents to observe interaction between agents . The final result is an encoding Fa for each agent , where history , context and interactions have been summarized . This encoding Fa is used in the next decoder operations and is also stored to be potentially used in modality recombination described in Sec . 3.2 . The resulting architecture of these encoding operations is illustrated in the first half of Fig . 2 . 3.1.2 HIERARCHICAL GRID DECODER . Our aim here is to decode each agent encoding into a heatmap representing its future probability distribution at prediction horizon T . Since we create this heatmap for each agent in the scene , the decoding process has to be fast so that it can be applied to a great number of agents in parallel . We use hierarchical predictions at various levels of resolutions so that the decoder has the possibility of predicting over the full surroundings of the agent but learns to refine with more precision only in places where the agent will end up with high probability . This hierarchical process is illustrated in Fig . 3 . Starting from an initial full dense grid probability at low resolutionR0×R0 by pixels , we iteratively refine the resolution by a fixed factor f until we reach the desired final resolution Rfinal × Rfinal . At each iteration i , we select only the Ni highest ranking grid points from the previous iteration , and upsample only these points to the Ri × Ri = Ri−1/f × Ri−1/f . At each step , the grid points features are computed by a 2-layer MLP applied on the point coordinates , these features are then concatenated to the agent encoding and passed through a linear layer , and finally enriched by a 2- layer cross-attention on the graph lane features , before applying a linear layer with sigmoid to get the probability . For a given W output range , this hierarchical process allows the model to only operate on W/R0 × W/R0 + ∑ iNi × f2 grid points instead of the W/Rfinal × W/Rfinal available . In practice , for a final output range of 192 meters with desired Rfinal = 0.5m resolution , we start with an initial resolution of R0 = 8m and use two iterations of ( N1 , N2 ) = ( 16 , 64 ) points each and an upscaling factor f = 4 . This way , we compute only 1856 grid points from the 147 456 available , with no performance loss . The heatmap is trained on each resolution level using as pixel-wise focal loss as in Gilles et al . ( 2021b ) , detailed in Appendix A.3 , with the ground truth being a Gaussian centered at the target agent future position .

Summary: The paper proposes a joint trajectory prediction framework. This means that the model produces multi-modal future predictions at the scene level (i.e., for all actors jointly). There has only been very recent research in this area despite its importance to autonomy and motion planning in particular, so the relevance of the paper to the broader motion forecasting field is high. I identify two main contributions in the paper: (i) hierarchical and sparse heatmap output representation and (ii) a recombination method that is pluggable in any multi-modal marginal prediction model.