CAPS: Context-Aware Priority Sampling for Enhanced Imitation Learning in Autonomous Driving
Hamidreza Mirkhani, Behzad Khamidehi, Ehsan Ahmadi, mohammed elmahgiubi, Weize Zhang, Fazel Arasteh, Umar Rajguru, Kasra Rezaee, Dongfeng Bai
AI summary
Problem
Imitation learning for autonomous driving struggles with imbalanced datasets dominated by trivial scenarios, causing policies to overfit and fail on rare, critical edge cases. Existing balancing methods rely on costly manual labeling or ignore crucial environmental context.
Approach
CAPS employs a VQ-VAE to encode scene context and trajectories into discrete cluster IDs, then re-weights training samples by inverse cluster frequency to prioritize rare scenarios during planner training.
Key results
- Introduces a two-stage CAPS framework for context-aware trajectory clustering and priority sampling
- Achieves ~10% improvement in Driving Score and Success Rate on Bench2Drive
- Outperforms endpoint and anchor-based clustering baselines under similar compute budgets
- Delivers robust closed-loop CARLA performance without extra expert data or deployment overhead
Why it matters
Enables scalable, context-aware data balancing that significantly improves the safety and generalization of learning-based autonomous driving planners without manual annotation costs.
Abstract
In this paper, we introduce Context-Aware Pri- ority Sampling (CAPS), a novel method designed to enhance data efficiency in learning-based autonomous driving systems. CAPS addresses the challenge of imbalanced datasets in im- itation learning by leveraging Vector Quantized Variational Autoencoders (VQ-VAEs). In this way, we can get structured and interpretable data representations, which help to 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. We evaluate our method through closed-loop experiments in the CARLA simulator. The results on Bench2Drive scenarios demonstrate the effectiveness of CAPS in enhancing model generalization, with substantial improvements in both driving score and success rate.