# Provable Representation Learning for Imitation with Contrastive Fourier Features

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

## Key facts

- Cited by 34 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 28 of its 54 references that are in the map.
- First posted on arXiv 26 May 2021.

## Placed in

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

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

- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Efficient Estimation of Word Representations in Vector Space](https://mlsanity.com/paper/1301.3781/) (2013)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Constrained Policy Optimization](https://mlsanity.com/paper/1705.10528/) (2017)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Time-Contrastive Networks: Self-Supervised Learning from Video](https://mlsanity.com/paper/1704.06888/) (2017)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [FLAMBE: Structural Complexity and Representation Learning of Low Rank MDPs](https://mlsanity.com/paper/2006.10814/) (2020)
- [DeepMDP: Learning Continuous Latent Space Models for Representation Learning](https://mlsanity.com/paper/1906.02736/) (2019)
- [Reinforcement Learning in Feature Space: Matrix Bandit, Kernels, and Regret Bound](https://mlsanity.com/paper/1905.10389/) (2019)
- [Imitation Learning via Off-Policy Distribution Matching](https://mlsanity.com/paper/1912.05032/) (2019)
- [OPAL: Offline Primitive Discovery for Accelerating Offline Reinforcement Learning](https://mlsanity.com/paper/2010.13611/) (2020)
- [Scalable methods for computing state similarity in deterministic Markov Decision Processes](https://mlsanity.com/paper/1911.09291/) (2019)
- [State Representation Learning for Control: An Overview](https://mlsanity.com/paper/1802.04181/) (2018)
- [Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2101.05265/) (2021)
- [Dopamine: A Research Framework for Deep Reinforcement Learning](https://mlsanity.com/paper/1812.06110/) (2018)
- [Near-Optimal Representation Learning for Hierarchical Reinforcement Learning](https://mlsanity.com/paper/1810.01257/) (2018)
- [Playing hard exploration games by watching YouTube](https://mlsanity.com/paper/1805.11592/) (2018)
- [Representation Matters: Offline Pretraining for Sequential Decision Making](https://mlsanity.com/paper/2102.05815/) (2021)
- [Near Optimal Behavior via Approximate State Abstraction](https://mlsanity.com/paper/1701.04113/) (2017)
- [RL Unplugged: A Suite of Benchmarks for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.13888/) (2020)
- [Provable Representation Learning for Imitation Learning via Bi-level Optimization](https://mlsanity.com/paper/2002.10544/) (2020)
- [Scalable Kernel Methods via Doubly Stochastic Gradients](https://mlsanity.com/paper/1407.5599/) (2014)
- [A Separation Principle for Control in the Age of Deep Learning](https://mlsanity.com/paper/1711.03321/) (2017)

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

- [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021)
- [Implicit Behavioral Cloning](https://mlsanity.com/paper/2109.00137/) (2021)
- [Behavior Transformers: Cloning k modes with one stone](https://mlsanity.com/paper/2206.11251/) (2022)
- [Multi-Game Decision Transformers](https://mlsanity.com/paper/2205.15241/) (2022)
- [Foundation Models for Decision Making: Problems, Methods, and Opportunities](https://mlsanity.com/paper/2303.04129/) (2023)
- [Fundamental Limitations of Alignment in Large Language Models](https://mlsanity.com/paper/2304.11082/) (2023)
- [From Play to Policy: Conditional Behavior Generation from Uncurated Robot Data](https://mlsanity.com/paper/2210.10047/) (2022)
- [Making Linear MDPs Practical via Contrastive Representation Learning](https://mlsanity.com/paper/2207.07150/) (2022)
- [TRAIL: Near-Optimal Imitation Learning with Suboptimal Data](https://mlsanity.com/paper/2110.14770/) (2021)
- [SwitchTab: Switched Autoencoders Are Effective Tabular Learners](https://mlsanity.com/paper/2401.02013/) (2024)
- [Spectral Decomposition Representation for Reinforcement Learning](https://mlsanity.com/paper/2208.09515/) (2022)
- [Inverse Dynamics Pretraining Learns Good Representations for Multitask Imitation](https://mlsanity.com/paper/2305.16985/) (2023)
- [Chain of Thought Imitation with Procedure Cloning](https://mlsanity.com/paper/2205.10816/) (2022)
- [INFOrmation Prioritization through EmPOWERment in Visual Model-Based RL](https://mlsanity.com/paper/2204.08585/) (2022)
- [A Free Lunch from the Noise: Provable and Practical Exploration for Representation Learning](https://mlsanity.com/paper/2111.11485/) (2021)
- [Improving Zero-shot Generalization in Offline Reinforcement Learning using Generalized Similarity Functions](https://mlsanity.com/paper/2111.14629/) (2021)
- [Learning Bellman Complete Representations for Offline Policy Evaluation](https://mlsanity.com/paper/2207.05837/) (2022)
- [Agent-Controller Representations: Principled Offline RL with Rich Exogenous Information](https://mlsanity.com/paper/2211.00164/) (2022)
- [Contrastive Value Learning: Implicit Models for Simple Offline RL](https://mlsanity.com/paper/2211.02100/) (2022)
- [Latent Variable Representation for Reinforcement Learning](https://mlsanity.com/paper/2212.08765/) (2022)
- [Diffusion Spectral Representation for Reinforcement Learning](https://mlsanity.com/paper/2406.16121/) (2024)
- [Representation Learning in Deep RL via Discrete Information Bottleneck](https://mlsanity.com/paper/2212.13835/) (2022)
- [Bilinear Exponential Family of MDPs: Frequentist Regret Bound with Tractable Exploration and Planning](https://mlsanity.com/paper/2210.02087/) (2022)
- [Behavior Prior Representation learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2211.00863/) (2022)
- [PI-ARS: Accelerating Evolution-Learned Visual-Locomotion with Predictive Information Representations](https://mlsanity.com/paper/2207.13224/) (2022)
- [Understanding and Addressing the Pitfalls of Bisimulation-based Representations in Offline Reinforcement Learning](https://mlsanity.com/paper/2310.17139/) (2023)
- [Provable Representation with Efficient Planning for Partial Observable Reinforcement Learning](https://mlsanity.com/paper/2311.12244/) (2023)
- [State Sequences Prediction via Fourier Transform for Representation Learning](https://mlsanity.com/paper/2310.15888/) (2023)
- [PASTA: Pretrained Action-State Transformer Agents](https://mlsanity.com/paper/2307.10936/) (2023)
- [Is Inverse Reinforcement Learning Harder than Standard Reinforcement Learning? A Theoretical Perspective](https://mlsanity.com/paper/2312.00054/) (2023)

## Abstract (arXiv)

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.