TrACT: A Training Dynamics Aware Contrastive Learning Framework for Long-tail Trajectory Prediction

TL;DR

TrACT tackles the long-tail problem in autonomous driving trajectory prediction by leveraging training dynamics and rich contextual features. It constructs a dataset map from per-sample training errors and variances, clusters samples into four difficulty-based groups, and derives cluster prototypes to drive a prototypical contrastive learning stage. The method achieves state-of-the-art results on nuScenes and ETH-UCY, with notable gains in both accuracy metrics (ADE/FDE) on challenging tails and improvements in scene compliance (HOR/SOR). The approach also demonstrates that a dataset-map-based strategy can reduce training bias without sacrificing performance on the head of the distribution, offering practical benefits for safety-critical deployment.

Abstract

As a safety critical task, autonomous driving requires accurate predictions of road users' future trajectories for safe motion planning, particularly under challenging conditions. Yet, many recent deep learning methods suffer from a degraded performance on the challenging scenarios, mainly because these scenarios appear less frequently in the training data. To address such a long-tail issue, existing methods force challenging scenarios closer together in the feature space during training to trigger information sharing among them for more robust learning. These methods, however, primarily rely on the motion patterns to characterize scenarios, omitting more informative contextual information, such as interactions and scene layout. We argue that exploiting such information not only improves prediction accuracy but also scene compliance of the generated trajectories. In this paper, we propose to incorporate richer training dynamics information into a prototypical contrastive learning framework. More specifically, we propose a two-stage process. First, we generate rich contextual features using a baseline encoder-decoder framework. These features are split into clusters based on the model's output errors, using the training dynamics information, and a prototype is computed within each cluster. Second, we retrain the model using the prototypes in a contrastive learning framework. We conduct empirical evaluations of our approach using two large-scale naturalistic datasets and show that our method achieves state-of-the-art performance by improving accuracy and scene compliance on the long-tail samples. Furthermore, we perform experiments on a subset of the clusters to highlight the additional benefit of our approach in reducing training bias.

Figures (5)

• Figure 1: A dataset map constructed using the training dynamics information of the baseline model on nuScenes. $variance$ measures the variance of the samples' errors across all training epochs, and $error$ measures the samples' errors on the last epoch. The dataset is divided into four clusters, namely $easy$, $hard$, $confusing$, or $trained$, based on a $variance$ threshold and an $error$ threshold. Each sample receives one cluster assignment based on their training behaviors. Only $20\%$ of the data is plotted on the log-scale axis for better visibility.
• Figure 2: Error evolution during training for one sample from each one of the $hard$, $confusing$, $easy$, and $trained$ clusters on nuScenes.
• Figure 3: An overview of the training dynamics-aware contrastive learning framework. The orange module represents the first training phase to construct the dataset map and performs the cluster assignment of each sample for prototype computation. The green module represents the second training phase with the prototypical contrastive loss.
• Figure 4: Qualitative results on the nuScenes dataset: left shows a complicated scene layout in which the baseline produces an off-road prediction, whereas TrACT does not; Right shows TrACT producing lane compliance predictions compared to the baseline.
• Figure 5: Parameter sensitivity of $\lambda$ on minFDE, plotted for the top 1-5$\%$ challenging subsets and the rest 95$\%$ of the data.