Abstract
Imitation learning holds the promise of equipping robots with versatile skills by learning from expert demonstrations. However, policies trained on finite datasets often struggle to generalize beyond the training distribution. In this work, we present a unified perspective on the generalization capability of imitation learning, grounded in both information theorey and data distribution property. We first show that the generalization gap can be upper bounded by (i) the conditional information bottleneck on intermediate representations and (ii) the mutual information between the model parameters and the training dataset. This characterization provides theoretical guidance for designing effective training strategies in imitation learning, particularly in determining whether to freeze, fine-tune, or train large pretrained encoders (e.g., vision-language models or vision foundation models) from scratch to achieve better generalization. Furthermore, we demonstrate that high conditional entropy from input to output induces a flatter likelihood landscape, thereby reducing the upper bound on the generalization gap. In addition, it shortens the stochastic gradient descent (SGD) escape time from sharp local minima, which may increase the likelihood of reaching global optima under fixed optimization budgets. These insights explain why imitation learning often exhibits limited generalization and underscore the importance of not only scaling the diversity of input data but also enriching the variability of output labels conditioned on the same input.
Abstract (translated)
模仿学习通过从专家演示中学习，承诺赋予机器人多种技能。然而，基于有限数据集训练的策略往往难以超越训练分布进行泛化。在这项工作中，我们提出了一种统一的观点来探讨模仿学习的泛化能力，该观点既以信息理论为依据，也考虑了数据分布特性。首先，我们表明泛化差距可以被上界限定为（i）中间表示上的条件信息瓶颈和（ii）模型参数与训练数据集之间的互信息。这一特征提供了设计有效训练策略的理论指导，在模仿学习中特别有助于确定是否冻结、微调或从头开始训练大规模预训练编码器（例如，视觉-语言模型或视觉基础模型），以实现更好的泛化能力。 此外，我们还证明了输入到输出方向上的高条件熵会诱导出更平坦的概率景观，从而减少泛化差距的上界。同时，它缩短了随机梯度下降（SGD）从尖锐局部最小值逃逸的时间，在固定优化预算下可能增加达到全局最优解的可能性。 这些见解解释了为什么模仿学习通常表现出有限的泛化能力，并强调不仅需要扩大输入数据的多样性，而且还需要丰富给定相同输入时输出标签的变化性。
URL
https://arxiv.org/abs/2504.18538
