# Provable Representation Learning for Imitation Learning via Bi-level Optimization

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

## Key facts

- Cited by 41 papers in the map (data through 14 Oct 2026).
- Placed in Imitation learning methods, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 23 of its 43 references that are in the map.
- First posted on arXiv 24 Feb 2020.

## Placed in

- Offline RL and imitation learning from data / [Imitation learning methods](https://mlsanity.com/offlinerl/imitation-learning-methods/)

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [A Theoretical Analysis of Contrastive Unsupervised Representation Learning](https://mlsanity.com/paper/1902.09229/) (2019)
- [Meta-Learning with Implicit Gradients](https://mlsanity.com/paper/1909.04630/) (2019)
- [Behavioral Cloning from Observation](https://mlsanity.com/paper/1805.01954/) (2018)
- [Rapid Learning or Feature Reuse? Towards Understanding the Effectiveness of MAML](https://mlsanity.com/paper/1909.09157/) (2019)
- [Meta-Learning: A Survey](https://mlsanity.com/paper/1810.03548/) (2018)
- [Batch Policy Learning under Constraints](https://mlsanity.com/paper/1903.08738/) (2019)
- [One-Shot Imitation Learning](https://mlsanity.com/paper/1703.07326/) (2017)
- [One-Shot Visual Imitation Learning via Meta-Learning](https://mlsanity.com/paper/1709.04905/) (2017)
- [The Benefit of Multitask Representation Learning](https://mlsanity.com/paper/1505.06279/) (2015)
- [Adaptive Gradient-Based Meta-Learning Methods](https://mlsanity.com/paper/1906.02717/) (2019)
- [Reinforcement and Imitation Learning via Interactive No-Regret Learning](https://mlsanity.com/paper/1406.5979/) (2014)
- [Imitating Latent Policies from Observation](https://mlsanity.com/paper/1805.07914/) (2018)
- [Provably Efficient Imitation Learning from Observation Alone](https://mlsanity.com/paper/1905.10948/) (2019)
- [Task-Embedded Control Networks for Few-Shot Imitation Learning](https://mlsanity.com/paper/1810.03237/) (2018)
- [Learning to Search Better Than Your Teacher](https://mlsanity.com/paper/1502.02206/) (2015)
- [Complete Dictionary Recovery over the Sphere I: Overview and the Geometric Picture](https://mlsanity.com/paper/1511.03607/) (2015)
- [Learning-to-Learn Stochastic Gradient Descent with Biased Regularization](https://mlsanity.com/paper/1903.10399/) (2019)
- [Truncated Horizon Policy Search: Combining Reinforcement Learning & Imitation Learning](https://mlsanity.com/paper/1805.11240/) (2018)
- [Incremental Learning-to-Learn with Statistical Guarantees](https://mlsanity.com/paper/1803.08089/) (2018)

## Built on by (most cited first, 30 of 41)

- [Investigating Bi-Level Optimization for Learning and Vision from a Unified Perspective: A Survey and Beyond](https://mlsanity.com/paper/2101.11517/) (2021)
- [Toward the Fundamental Limits of Imitation Learning](https://mlsanity.com/paper/2009.05990/) (2020)
- [A Survey of Deep RL and IL for Autonomous Driving Policy Learning](https://mlsanity.com/paper/2101.01993/) (2021)
- [FedNest: Federated Bilevel, Minimax, and Compositional Optimization](https://mlsanity.com/paper/2205.02215/) (2022)
- [Lower Bounds and Accelerated Algorithms for Bilevel Optimization](https://mlsanity.com/paper/2102.03926/) (2021)
- [Near-optimal Representation Learning for Linear Bandits and Linear RL](https://mlsanity.com/paper/2102.04132/) (2021)
- [TRAIL: Near-Optimal Imitation Learning with Suboptimal Data](https://mlsanity.com/paper/2110.14770/) (2021)
- [Provable Benefits of Representational Transfer in Reinforcement Learning](https://mlsanity.com/paper/2205.14571/) (2022)
- [Provable Representation Learning for Imitation with Contrastive Fourier Features](https://mlsanity.com/paper/2105.12272/) (2021)
- [Decentralized Gossip-Based Stochastic Bilevel Optimization over Communication Networks](https://mlsanity.com/paper/2206.10870/) (2022)
- [Inverse Dynamics Pretraining Learns Good Representations for Multitask Imitation](https://mlsanity.com/paper/2305.16985/) (2023)
- [On the Power of Multitask Representation Learning in Linear MDP](https://mlsanity.com/paper/2106.08053/) (2021)
- [Chain of Thought Imitation with Procedure Cloning](https://mlsanity.com/paper/2205.10816/) (2022)
- [Provable Benefit of Multitask Representation Learning in Reinforcement Learning](https://mlsanity.com/paper/2206.05900/) (2022)
- [Multi-Task Imitation Learning for Linear Dynamical Systems](https://mlsanity.com/paper/2212.00186/) (2022)
- [Language Models Can Learn from Verbal Feedback Without Scalar Rewards](https://mlsanity.com/paper/2509.22638/) (2025)
- [Nearly Minimax Algorithms for Linear Bandits with Shared Representation](https://mlsanity.com/paper/2203.15664/) (2022)
- [Provably Breaking the Quadratic Error Compounding Barrier in Imitation Learning, Optimally](https://mlsanity.com/paper/2102.12948/) (2021)
- [Provable General Function Class Representation Learning in Multitask Bandits and MDPs](https://mlsanity.com/paper/2205.15701/) (2022)
- [Behavior Prior Representation learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2211.00863/) (2022)
- [Aligning Robot and Human Representations](https://mlsanity.com/paper/2302.01928/) (2023)
- [Understanding and Addressing the Pitfalls of Bisimulation-based Representations in Offline Reinforcement Learning](https://mlsanity.com/paper/2310.17139/) (2023)
- [Learning an Explicit Hyperparameter Prediction Function Conditioned on Tasks](https://mlsanity.com/paper/2107.02378/) (2021)
- [Min-Max Bilevel Multi-objective Optimization with Applications in Machine Learning](https://mlsanity.com/paper/2203.01924/) (2022)
- [Fast and Sample Efficient Multi-Task Representation Learning in Stochastic Contextual Bandits](https://mlsanity.com/paper/2410.02068/) (2024)
- [A Generalized Alternating Method for Bilevel Learning under the Polyak-Lojasiewicz Condition](https://mlsanity.com/paper/2306.02422/) (2023)
- [IL-SOAR : Imitation Learning with Soft Optimistic Actor cRitic](https://mlsanity.com/paper/2502.19859/) (2025)
- [Learning distinct features helps, provably](https://mlsanity.com/paper/2106.06012/) (2021)
- [Adaptive Clustering and Personalization in Multi-Agent Stochastic Linear Bandits](https://mlsanity.com/paper/2106.08902/) (2021)
- [Investigating Pre-Training Objectives for Generalization in Vision-Based Reinforcement Learning](https://mlsanity.com/paper/2406.06037/) (2024)

## Abstract (arXiv)

A common strategy in modern learning systems is to learn a representation that is useful for many tasks, a.k.a. representation learning. We study this strategy in the imitation learning setting for Markov decision processes (MDPs) where multiple experts' trajectories are available. We formulate representation learning as a bi-level optimization problem where the "outer" optimization tries to learn the joint representation and the "inner" optimization encodes the imitation learning setup and tries to learn task-specific parameters. We instantiate this framework for the imitation learning settings of behavior cloning and observation-alone. Theoretically, we show using our framework that representation learning can provide sample complexity benefits for imitation learning in both settings. We also provide proof-of-concept experiments to verify our theory.