Sample-Efficient Expert Query Control in Active Imitation Learning via Conformal Prediction

TL;DR

This paper introduces CRSAIL, a novel active imitation learning method that reduces expert queries by selectively requesting expert actions only for under-represented states, using conformal prediction for robust, task-agnostic state novelty detection.

Contribution

CRSAIL is the first querying rule that leverages conformal prediction to adaptively control expert queries based on state novelty, improving efficiency and robustness in active imitation learning.

Findings

CRSAIL reduces expert queries by up to 96% compared to DAgger.

CRSAIL matches or exceeds expert-level reward in MuJoCo tasks.

CRSAIL is robust to parameter choices and applicable to unknown dynamics.

Abstract

Active imitation learning (AIL) combats covariate shift by querying an expert during training. However, expert action labeling often dominates the cost, especially in GPU-intensive simulators, human-in-the-loop settings, and robot fleets that revisit near-duplicate states. We present Conformalized Rejection Sampling for Active Imitation Learning (CRSAIL), a querying rule that requests an expert action only when the visited state is under-represented in the expert-labeled dataset. CRSAIL scores state novelty by the distance to the $K$-th nearest expert state and sets a single global threshold via conformal prediction. This threshold is the empirical $(1-\alpha)$ quantile of on-policy calibration scores, providing a distribution-free calibration rule that links $\alpha$ to the expected query rate and makes $\alpha$ a task-agnostic tuning knob. This state-space querying strategy is robust to outliers and, unlike safety-gate-based AIL, can be run without real-time expert takeovers: we roll out full trajectories (episodes) with the learner and only afterward query the expert on a subset of visited states. Evaluated on MuJoCo robotics tasks, CRSAIL matches or exceeds expert-level reward while reducing total expert queries by up to 96% vs. DAgger and up to 65% vs. prior AIL methods, with empirical robustness to $\alpha$ and $K$, easing deployment on novel systems with unknown dynamics.