AMD: Adaptive Momentum and Decoupled Contrastive Learning Framework for Robust Long-Tail Trajectory Prediction

TL;DR

The paper introduces AMD, a novel framework combining contrastive learning and data augmentation to improve long-tail trajectory prediction in autonomous driving, effectively handling rare and complex traffic scenarios.

Contribution

It proposes an adaptive momentum and decoupled contrastive learning framework with dynamic pseudo-labeling for better long-tail trajectory recognition.

Findings

AMD achieves state-of-the-art long-tail trajectory prediction performance.

The framework improves overall accuracy on nuScenes and ETH/UCY datasets.

Contrastive learning enhances recognition of rare trajectory patterns.

Abstract

Accurately predicting the future trajectories of traffic agents is essential in autonomous driving. However, due to the inherent imbalance in trajectory distributions, tail data in natural datasets often represents more complex and hazardous scenarios. Existing studies typically rely solely on a base model's prediction error, without considering the diversity and uncertainty of long-tail trajectory patterns. We propose an adaptive momentum and decoupled contrastive learning framework (AMD), which integrates unsupervised and supervised contrastive learning strategies. By leveraging an improved momentum contrast learning (MoCo-DT) and decoupled contrastive learning (DCL) module, our framework enhances the model's ability to recognize rare and complex trajectories. Additionally, we design four types of trajectory random augmentation methods and introduce an online iterative clustering strategy, allowing the model to dynamically update pseudo-labels and better adapt to the distributional shifts in long-tail data. We propose three different criteria to define long-tail trajectories and conduct extensive comparative experiments on the nuScenes and ETH$/$UCY datasets. The results show that AMD not only achieves optimal performance in long-tail trajectory prediction but also demonstrates outstanding overall prediction accuracy.