Uncertainty-Guided Robotic Manipulation Through Variational Information Bottleneck in Imitation Learning

Abstract

Traditional imitation learning methods typically rely on high-quality expert demonstrations and exhibit poor generalization when deployed in unfamiliar environments. A key limitation is their inability to effectively quantify and utilize epistemic uncertainty in the decision-making process. To address these limitations, this paper introduces a novel imitation learning framework that explicitly incorporates epistemic uncertainty estimation into policy learning. We leverage the Variational Information Bottleneck (VIB) to learn a compact and robust representation of the input data while simultaneously quantifying the uncertainty associated with each decision. Our method enables the model to generalize better to unseen scenarios and to make safer and more reliable decisions by reasoning about its own confidence in the predictions. Experimental results on various robotic manipulation tasks show that our method significantly improves performance compared to standard imitation learning methods, achieving better stability and adaptability. The code is available at <uri xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">https://github.com/Dear-Chen/VIB-ILCtrl</uri>.