[Paper Review] THOMAS: Trajectory Heatmap Output with learned Multi-Agent Sampling

Thomas Gilles, Stefano Sabatini+3 more|arXiv (Cornell University)|Oct 13, 2021

Autonomous Vehicle Technology and SafetyEngineering64 references59 citations

TL;DR

THOMAS is a joint multi-agent trajectory prediction framework that uses hierarchical heatmap decoding and a learned recombination module to produce scene-consistent multi-modal trajectories for all agents, achieving state-of-the-art results on the Interaction Interpret track.

ABSTRACT

In this paper, we propose THOMAS, a joint multi-agent trajectory prediction framework allowing for an efficient and consistent prediction of multi-agent multi-modal 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 $1^{st}$ on the online test leaderboard.

Motivation & Objective

• Predict future multi-agent trajectories with coherence across agents in a multi-modal setting.
• Encode past trajectories and HD-Map context to produce per-agent future heatmaps.
• Generate multiple endpoints per agent and reconstruct full trajectories.
• Recombine sampled endpoints to ensure collision-free, scene-consistent predictions.

Proposed method

• Use a shared graph encoder for agent histories and lanelet-based HD-Map context with cross- and self-attention to produce per-agent encodings.
• Decode per-agent future probability heatmaps via a hierarchical grid refinement scheme that grows from a low-resolution grid to final resolution, selectively refining high-probability regions.
• Sample K endpoints per agent from the heatmaps using a MissRate-based decoding strategy to cover diverse modalities.
• Generate full trajectories for each endpoint with an MLP conditioned on history and endpoint.
• Introduce a scene-consistency recombination module that learns to align and combine agent modalities into coherent scene modalities using cross-attention and winner-takes-all loss.

Experimental results

Research questions

• RQ1 How can multi-agent trajectory predictions be made jointly coherent across agents while preserving multi-modality?
• RQ2 Can a hierarchical heatmap decoder scale to many agents and maintain fast inference?
• RQ3 Does a learned recombination module improve scene consistency and reduce collisions without sacrificing marginal accuracy?
• RQ4 How does THOMAS compare to state-of-the-art joint-prediction methods on interpretability benchmarks?

Key findings

• THOMAS achieves state-of-the-art joint metrics on the Interaction Interpret track compared to ILVM and SceneTransformer.
• Joint SMR improves by about 25% and SCR by ~30% relative to baselines, with cSMR also reduced by >25%.
• Hierarchical heatmap decoding computes only a fraction of grid points (e.g., 1856 of 147,456) with no performance loss for a 192 m range at 0.5 m resolution.
• The recombination module reduces collisions and yields more coherent scene modalities by reusing and aligning existing agent modalities.
• Compared to GOHOME, THOMAS offers faster inference (32-128 agents) and improved joint consistency metrics.