# Generalization Capability for Imitation Learning

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Robot manipulation imitation learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 15 of its 16 references that are in the map.
- First posted on arXiv 25 Apr 2025.

## Placed in

- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

## Builds on (most cited first, 15 of 15)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Octo: An Open-Source Generalist Robot Policy](https://mlsanity.com/paper/2405.12213/) (2024)
- [GR00T N1: An Open Foundation Model for Generalist Humanoid Robots](https://mlsanity.com/paper/2503.14734/) (2025)
- [DROID: A Large-Scale In-The-Wild Robot Manipulation Dataset](https://mlsanity.com/paper/2403.12945/) (2024)
- [Where are we in the search for an Artificial Visual Cortex for Embodied Intelligence?](https://mlsanity.com/paper/2303.18240/) (2023)
- [Emergence of Invariance and Disentanglement in Deep Representations](https://mlsanity.com/paper/1706.01350/) (2017)
- [Data Scaling Laws in Imitation Learning for Robotic Manipulation](https://mlsanity.com/paper/2410.18647/) (2024)
- [Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks](https://mlsanity.com/paper/1710.11029/) (2017)
- [How Does Information Bottleneck Help Deep Learning?](https://mlsanity.com/paper/2305.18887/) (2023)
- [InfoBot: Transfer and Exploration via the Information Bottleneck](https://mlsanity.com/paper/1901.10902/) (2019)
- [OTTER: A Vision-Language-Action Model with Text-Aware Visual Feature Extraction](https://mlsanity.com/paper/2503.03734/) (2025)
- [Dynamics Generalization via Information Bottleneck in Deep Reinforcement Learning](https://mlsanity.com/paper/2008.00614/) (2020)
- [Conditional Graph Information Bottleneck for Molecular Relational Learning](https://mlsanity.com/paper/2305.01520/) (2023)
- [Correlation Information Bottleneck: Towards Adapting Pretrained Multimodal Models for Robust Visual Question Answering](https://mlsanity.com/paper/2209.06954/) (2022)

## Built on by (most cited first, 2 of 2)

- [Robust Finetuning of Vision-Language-Action Robot Policies via Parameter Merging](https://mlsanity.com/paper/2512.08333/) (2025)
- [Bidirectional Tutoring for Developmental Motor Learning in Robots: Co-Developed Interaction Dynamics Support Stable Learning](https://mlsanity.com/paper/2606.19728/) (2026)

## Abstract (arXiv)

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.