In-Context Adaptation for Generalizable Imitation Learning

Junlin Xie, Xu Luo, Hao Wu, Ji Zhang, Youguang Xing, Lianli Gao, Jingkuan Song

Keywords: In-Context Adaptation, Imitation Learning, Zero-Shot Generalization

Abstract: While imitation learning on large-scale robot data produces robot policies with impressive task performance, these policies are typically \emph{reactive} and lack the ability to adapt to novel conditions at test time. This limitation stands in stark contrast to Large Language Models (LLMs), which excel at in-context learning and adaptation. In this work, we take the first steps toward bridging this gap, exploring how imitation learning can instill in-context adaptation into robot policies. We specifically address the challenge of varying action dynamics, a scenario requiring online inference and adjustment. Our experiments with Diffusion Policy reveal that enabling such adaptation hinges on two critical components: conditioning the policy on histories of both observations and actions, and training on a diverse sampling of action dynamics. The resulting method successfully generalizes to unseen, out-of-distribution dynamics, representing a key advancement toward behavioral generalization in imitation learning.

Submission Number: 10