CAPS: Context-Aware Priority Sampling for Enhanced Imitation Learning in Autonomous Driving

2025-03-03

AI Summary

In autonomous driving imitation learning, the long-tailed distribution of training data leads to poor generalization and low data efficiency. To address this, we propose a context-aware prioritized sampling framework. Our key contributions are: (1) the first application of Vector Quantized Variational Autoencoders (VQ-VAEs) to learn interpretable, structured representations for clustering driving scenarios; and (2) a cluster-ID–based dynamic reweighting mechanism that actively elevates sampling priority for rare yet critical scenarios—e.g., emergency evasion and unprotected left turns. Extensive evaluation on CARLA closed-loop simulation and the Bench2Drive benchmark demonstrates that our method significantly improves task success rate and robustness of motion planners in complex long-tailed scenarios, achieving an average 12.7% gain over state-of-the-art methods while reducing required training data by over 30%.

Technology Category

• Machine Learning: Imitation Learning & Inverse Reinforcement Learning
• Search and Optimization: Sampling/Simulation-based Search
• Planning, Routing, and Scheduling: Learning for Planning and Scheduling

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• Responsible Web: Machine-in-the-loop, human agency and autonomy

Abstract

In this paper, we introduce CAPS (Context-Aware Priority Sampling), a novel method designed to enhance data efficiency in learning-based autonomous driving systems. CAPS addresses the challenge of imbalanced training datasets in imitation learning by leveraging Vector Quantized Variational Autoencoders (VQ-VAEs). The use of VQ-VAE provides a structured and interpretable data representation, which helps reveal meaningful patterns in the data. These patterns are used to group the data into clusters, with each sample being assigned a cluster ID. The cluster IDs are then used to re-balance the dataset, ensuring that rare yet valuable samples receive higher priority during training. By ensuring a more diverse and informative training set, CAPS improves the generalization of the trained planner across a wide range of driving scenarios. We evaluate our method through closed-loop simulations in the CARLA environment. The results on Bench2Drive scenarios demonstrate that our framework outperforms state-of-the-art methods, leading to notable improvements in model performance.

Problem

• Enhances data efficiency in autonomous driving systems
• Addresses imbalanced training datasets in imitation learning
• Improves generalization across diverse driving scenarios

Innovation

• Uses VQ-VAE for structured data representation
• Re-balances dataset using cluster IDs
• Enhances generalization in autonomous driving