SSTP: Efficient Sample Selection for Trajectory Prediction
This work addresses the challenge of efficient and reliable trajectory prediction for autonomous driving by mitigating dataset imbalance, though it is incremental as it builds on existing methods with a novel selection strategy.
The paper tackles the problem of training trajectory prediction models efficiently by addressing dataset imbalance, where high-density traffic scenes are underrepresented, leading to poor performance in safety-critical cases. The proposed SSTP framework constructs a compact, density-balanced dataset that achieves comparable performance to full-dataset training using only half the data, with substantial improvements in high-density scenes and reduced training time.
Trajectory prediction is a core task in autonomous driving. However, training advanced trajectory prediction models on existing large-scale datasets is both time-consuming and computationally expensive. More critically, these datasets are highly imbalanced in scenario density, with normal driving scenes (low-moderate traffic) overwhelmingly dominating the datasets, while high-density and safety-critical cases are underrepresented. As a result, models tend to overfit low/moderate-density scenarios and perform poorly in high-density scenarios. To address these challenges, we propose the SSTP framework, which constructs a compact yet density-balanced dataset tailored to trajectory prediction. SSTP consists of two main stages: (1)Extraction, where a baseline model is pretrained for a few epochs to obtain stable gradient estimates, and the dataset is partitioned by scenario density. (2)Selection, where gradient-based scores and a submodular objective select representative samples within each density category, while biased sampling emphasizes rare high-density interactions to avoid dominance by low-density cases. This approach significantly reduces the dataset size and mitigates scenario imbalance, without sacrificing prediction accuracy. Experiments on the Argoverse 1 and Argoverse 2 datasets with recent state-of-the-art models show that SSTP achieves comparable performance to full-dataset training using only half the data while delivering substantial improvements in high-density traffic scenes and significantly reducing training time. Robust trajectory prediction depends not only on data scale but also on balancing scene density to ensure reliable performance under complex multi agent interactions.