PTP: Parallelized Tracking and Prediction with Graph Neural Networks and Diversity Sampling

TL;DR

The paper introduces PTP, a parallelized framework that jointly optimizes 3D multi-object tracking and trajectory forecasting to learn a shared representation of agent interaction. It leverages Graph Neural Networks for socially aware feature interaction and a diversity sampling mechanism based on Determinantal Point Processes to produce diverse, high-quality future trajectories via a CVAE. Empirically, PTP achieves state-of-the-art results on KITTI and nuScenes for both 3D MOT and trajectory forecasting, and ablations confirm the benefits of parallelization, GNNs, and DSF. The approach improves robustness and efficiency, offering a unified solution for perception systems in autonomous driving. The work also provides detailed architectural components and training losses that can inform future joint-tracking and prediction research.

Abstract

Multi-object tracking (MOT) and trajectory prediction are two critical components in modern 3D perception systems that require accurate modeling of multi-agent interaction. We hypothesize that it is beneficial to unify both tasks under one framework in order to learn a shared feature representation of agent interaction. Furthermore, instead of performing tracking and prediction sequentially which can propagate errors from tracking to prediction, we propose a parallelized framework to mitigate the issue. Also, our parallel track-forecast framework incorporates two additional novel computational units. First, we use a feature interaction technique by introducing Graph Neural Networks (GNNs) to capture the way in which agents interact with one another. The GNN is able to improve discriminative feature learning for MOT association and provide socially-aware contexts for trajectory prediction. Second, we use a diversity sampling function to improve the quality and diversity of our forecasted trajectories. The learned sampling function is trained to efficiently extract a variety of outcomes from a generative trajectory distribution and helps avoid the problem of generating duplicate trajectory samples. We evaluate on KITTI and nuScenes datasets showing that our method with socially-aware feature learning and diversity sampling achieves new state-of-the-art performance on 3D MOT and trajectory prediction. Project website is: https://www.xinshuoweng.com/projects/PTP

Figures (6)

• Figure 1: (Top) Previous work: 3D MOT and trajectory forecasting performed separately and connected as a sequential process. (Bottom) Ours: Joint and parallelized framework for MOT and forecasting. Two key innovations: (1) feature interaction using GNNs to obtain socially-aware features in the presence of multiple agents; (2) diversity sampling to improve sample efficiency and produce diverse and accurate trajectory samples.
• Figure 2: (Left) To leverage the location and motion cues, we extract the feature from object trajectories $\mathcal{O} = \{\boldsymbol{o}_1, \ldots, \boldsymbol{o}_M\}$ in the past using an LSTM model and extract the feature from detections $\mathcal{D} = \{\boldsymbol{d}_1, \ldots, \boldsymbol{d}_N\}$ in the current frame using a MLP. (Right) The GNN-based feature interaction mechanism is used to update object node feature $\mathcal{U}^l = \{\boldsymbol{u}_1^l, \ldots, \boldsymbol{u}_M^l\}$ and $\mathcal{V}^l = \{\boldsymbol{v}_1^l, \ldots, \boldsymbol{v}_N^l\}$ at GNN layer $l$ and iteratively through all GNN layers. At the final layer, we use the node features for tracked objects $\mathcal{U}^L$ for 3D MOT task (see Sec. \ref{['sec:mot']}), and use the edge features $\mathcal{E}^L$ (computed from $\mathcal{U}^L$ and $\mathcal{V}^L$) for trajectory forecasting task (see Sec. \ref{['sec:forecasting']} and \ref{['sec:diversity']}).
• Figure 3: 3D multi-object tracking head with an affinity loss.
• Figure 4: Trajectory Forecasting with Diversity Sampling. To produce diverse trajectory samples from our pretrained CVAE model, we learn a diversity sampling function (DSF) $\mathcal{S}_\gamma$ to map each object's node feature $\boldsymbol{u}_i^L$ to a set of latent codes, which can cover not only the major mode but also other modes in the CVAE latent space. Then, we decode those codes into diverse and accurate future trajectories for object $i$.
• Figure 5: 3D MOT results on sequence 0 (left) and 11 (right) of the KITTI test set. We show three frames (with an interval of 5 frames) for each sequence.