Provable Representation Learning for Imitation Learning via Bi-level Optimization
This work addresses sample efficiency in imitation learning for AI systems, but it is incremental as it builds on existing frameworks with theoretical analysis.
The paper tackles the problem of representation learning for imitation learning in Markov decision processes by formulating it as a bi-level optimization problem, showing theoretically that this approach provides sample complexity benefits for behavior cloning and observation-alone settings, with proof-of-concept experiments verifying the theory.
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.