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

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.

@article{mirkhani2025_2503.01650,
  title={ CAPS: Context-Aware Priority Sampling for Enhanced Imitation Learning in Autonomous Driving },
  author={ Hamidreza Mirkhani and Behzad Khamidehi and Ehsan Ahmadi and Fazel Arasteh and Mohammed Elmahgiubi and Weize Zhang and Umar Rajguru and Kasra Rezaee },
  journal={arXiv preprint arXiv:2503.01650},
  year={ 2025 }
}