# CAPS: Context-Aware Priority Sampling for Enhanced Imitation Learning in Autonomous Driving Source: [CAPS: Context-Aware Priority Sampling for Enhanced Imitation Learning in Autonomous Driving]() Collection index: [llms.txt]() Data updated: 2026-03-10T05:00:12.119Z Authors: Hamidreza Mirkhani, Behzad Khamidehi, Ehsan Ahmadi, Mohammed Elmahgiubi, Weize Zhang, Fazel Arasteh, Umar Rajguru, Kasra Rezaee, Dongfeng Bai Published: 2025-03-03T15:27:11.000Z ## Abstract In this paper, we introduce Context-Aware Priority Sampling (CAPS), a novel method designed to enhance data efficiency in learning-based autonomous driving systems. CAPS addresses the challenge of imbalanced datasets in imitation 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. ## Recorded benchmark results A research result applies to its stated dataset, split, metric and evaluation configuration. Compare like-for-like settings, preserve the metric's unit and direction, and distinguish paper-reported results from independent reproduction. The best result recorded here is not a claim of exhaustive, current global coverage. - [Autonomous Driving on Bench2Drive 220 short segment scenarios](): Autonomous Driving; Bench2Drive 220 short segment scenarios; Driving Score: 68.91. ## Original sources and implementation [Original paper]() [PDF]()