Generative Active Learning for Long-tail Trajectory Prediction via Controllable Diffusion Model

Abstract:

Predicting future trajectories of dynamic traffic agents is crucial in autonomous systems. While data-driven methods enable large-scale training, they often underperform on rarely observed tail samples, yielding a long-tail problem. Prior works have tackled this by modifying model architectures, such as using a hypernetwork. In contrast, we propose refining the training procedure to unlock each model’s potential without altering its structure. To this end, we introduce the Generative Active Learning for Trajectory prediction (GALTraj), which iteratively identifies tail samples and augments them via a controllable generative diffusion model. By incorporating the augmented samples in each iteration, we directly mitigate dataset imbalance. To ensure effective augmentation, we design a new tail-aware generation method that categorizes agents (tail, head, relevant) and applies tailored guidance of the diffusion model. It enables producing diverse and realistic trajectories that preserve tail characteristics while respecting traffic constraints. Unlike prior traffic simulation methods focused on producing diverse scenarios, ours is the first to show how simulator-driven augmentation can benefit long-tail learning for trajectory prediction. Experiments on multiple trajectory datasets (WOMD, Argoverse2) with popular backbones (QCNet, MTR) confirm that our method significantly boosts performance on tail samples and also enhances accuracy on head samples.