Provable Representation Learning for Imitation Learning via Bi-level Optimization

Sanjeev Arora, Simon S. Du, Sham Kakade, Yuping Luo, Nikunj Saunshi

TL;DR: Using a bi-level optimization framework, we learn representations by leveraging multiple imitation learning tasks to provably reduce the sample complexity of learning a policy for a new task

Abstract: A common strategy in modern learning systems is to learn a representation which is useful for many tasks, a.k.a, representation learning. We study this strategy in the imitation learning setting 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 cases where the imitation setting being behavior cloning and observation alone. Theoretically, we provably show using our framework that representation learning can reduce the sample complexity of imitation learning in both settings. We also provide proof-of-concept experiments to verify our theoretical findings.

Keywords: imitation learning, representation learning, multitask learning, theory, behavioral cloning, imitation from observations alone, reinforcement learning