Provable Representation Learning for Imitation with Contrastive Fourier Features

By Ofir Nachum, Mengjiao Yang. 2021

40 citations recorded by the source corpus.

In imitation learning, it is common to learn a behavior policy to match an unknown target policy via max-likelihood training on a collected set of target demonstrations. In this work, we consider using offline experience datasets - potentially far from the target distribution - to learn low-dimensional state representations that provably accelerate the sample-efficiency of downstream imitation learning. A central challenge in this setting is that the unknown target policy itself may not exhibit low-dimensional behavior, and so there is a potential for the representation learning objective to alias states in which the target policy acts differently. Circumventing this challenge, we derive a representation learning objective that provides an upper bound on the performance difference between the target policy and a lowdimensional policy trained with max-likelihood, and this bound is tight regardless of whether the target policy itself exhibits low-dimensional structure. Moving to the practicality of our method, we show that our objective can be implemented as contrastive learning, in which the transition dynamics are approximated by either an implicit energy-based model or, in some special cases, an implicit linear model with representations given by random Fourier features. Experiments on both tabular environments and high-dimensional Atari games provide quantitative evidence for the practical benefits of our proposed objective.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Statistical and Nonlinear Physics, Representation & Self-Supervised Learning

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Statistical and Nonlinear Physics

In Representation & Self-Supervised Learning

Citations in the corpus

What Matters in Learning from Offline Human Demonstrations for Robot Manipulation (2021, background): “The choice of observation space for imitation merits careful consideration – other work has also conﬁrmed the importance of feature representations used for ofﬂine policy learning [52, 53].”
Implicit Behavioral Cloning (2021, background): “, value functions) for 30 various reinforcement learning (RL) methods [13, 14, 15, 16], our work presents a distinct yet simple method 31 for robot learning: do BC with implicit models.”
Behavior Transformers: Cloning k modes with one stone (2022, background): “Finally, a class of algorithms, most notably [50, 25, 43, 58] do not directly learn a generative model but instead learn energy based models.”