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

arXiv: https://arxiv.org/abs/2503.01650 · HTML: https://mlsanity.com/paper/2503.01650/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Imitation-learning end-to-end planners, one of 9 directions in End-to-end autonomous driving.
- Builds on 37 of its 49 references that are in the map.
- First posted on arXiv 3 Mar 2025.

## Placed in

- End-to-end autonomous driving / [Imitation-learning end-to-end planners](https://mlsanity.com/driving/imitation-learning-end-to-end-planners/)
- Offline RL and imitation learning from data / [Imitation and offline RL for driving](https://mlsanity.com/offlinerl/imitation-and-offline-rl-for-driving/)

## Builds on (most cited first, 30 of 37)

- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [CARLA: An Open Urban Driving Simulator](https://mlsanity.com/paper/1711.03938/) (2017)
- [Prioritized Experience Replay](https://mlsanity.com/paper/1511.05952/) (2015)
- [Planning-oriented Autonomous Driving](https://mlsanity.com/paper/2212.10156/) (2022)
- [VAD: Vectorized Scene Representation for Efficient Autonomous Driving](https://mlsanity.com/paper/2303.12077/) (2023)
- [VectorNet: Encoding HD Maps and Agent Dynamics from Vectorized Representation](https://mlsanity.com/paper/2005.04259/) (2020)
- [Finite Scalar Quantization: VQ-VAE Made Simple](https://mlsanity.com/paper/2309.15505/) (2023)
- [ChauffeurNet: Learning to Drive by Imitating the Best and Synthesizing the Worst](https://mlsanity.com/paper/1812.03079/) (2018)
- [TNT: Target-driveN Trajectory Prediction](https://mlsanity.com/paper/2008.08294/) (2020)
- [Causal Confusion in Imitation Learning](https://mlsanity.com/paper/1905.11979/) (2019)
- [Bench2Drive: Towards Multi-Ability Benchmarking of Closed-Loop End-To-End Autonomous Driving](https://mlsanity.com/paper/2406.03877/) (2024)
- [Trajectory-guided Control Prediction for End-to-end Autonomous Driving: A Simple yet Strong Baseline](https://mlsanity.com/paper/2206.08129/) (2022)
- [Rethinking the Open-Loop Evaluation of End-to-End Autonomous Driving in nuScenes](https://mlsanity.com/paper/2305.10430/) (2023)
- [Think Twice before Driving: Towards Scalable Decoders for End-to-End Autonomous Driving](https://mlsanity.com/paper/2305.06242/) (2023)
- [DriveAdapter: Breaking the Coupling Barrier of Perception and Planning in End-to-End Autonomous Driving](https://mlsanity.com/paper/2308.00398/) (2023)
- [Imitation Is Not Enough: Robustifying Imitation with Reinforcement Learning for Challenging Driving Scenarios](https://mlsanity.com/paper/2212.11419/) (2022)
- [EnsembleDAgger: A Bayesian Approach to Safe Imitation Learning](https://mlsanity.com/paper/1807.08364/) (2018)
- [Input Perturbation Reduces Exposure Bias in Diffusion Models](https://mlsanity.com/paper/2301.11706/) (2023)
- [Query-Efficient Imitation Learning for End-to-End Autonomous Driving](https://mlsanity.com/paper/1605.06450/) (2016)
- [Hierarchical Model-Based Imitation Learning for Planning in Autonomous Driving](https://mlsanity.com/paper/2210.09539/) (2022)
- [Model-Based Runtime Monitoring with Interactive Imitation Learning](https://mlsanity.com/paper/2310.17552/) (2023)
- [Learning Realistic Traffic Agents in Closed-loop](https://mlsanity.com/paper/2311.01394/) (2023)
- [Combined Cleaning and Resampling Algorithm for Multi-Class Imbalanced Data with Label Noise](https://mlsanity.com/paper/2004.03406/) (2020)
- [Embedding Synthetic Off-Policy Experience for Autonomous Driving via Zero-Shot Curricula](https://mlsanity.com/paper/2212.01375/) (2022)
- [Blending Imitation and Reinforcement Learning for Robust Policy Improvement](https://mlsanity.com/paper/2310.01737/) (2023)
- [The Why, When, and How to Use Active Learning in Large-Data-Driven 3D Object Detection for Safe Autonomous Driving: An Empirical Exploration](https://mlsanity.com/paper/2401.16634/) (2024)
- [ActiveAD: Planning-Oriented Active Learning for End-to-End Autonomous Driving](https://mlsanity.com/paper/2403.02877/) (2024)
- [CIMRL: Combining IMitation and Reinforcement Learning for Safe Autonomous Driving](https://mlsanity.com/paper/2406.08878/) (2024)

## Abstract (arXiv)

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.