[论文解读] Provable Representation Learning for Imitation with Contrastive Fourier Features

Ofir Nachum, Mengjiao Yang|arXiv (Cornell University)|May 26, 2021

Reinforcement Learning in Robotics参考文献 41被引用 7

一句话总结

该论文提出了一种基于对比傅里叶特征的可证明有效的表示学习方法，用于增强模仿学习的样本效率。该方法推导出一个在目标策略为低维时仍保持紧致的策略性能差异上界，并通过对比学习结合随机傅里叶特征或能量模型实现该目标，在Atari和表格环境上均表现出显著性能提升。

ABSTRACT

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.

研究动机与目标

• 通过利用大规模离线数据集，解决行为克隆在示范数据稀缺时的样本效率低下问题。
• 克服表示学习中的混叠问题，即不同目标策略行为的状态被压缩为同一表示。
• 提供一个理论基础坚实的表示学习目标，确保下游行为克隆能够匹配任意内在维度的目标策略。
• 弥合实际对比学习方法与离线模仿学习中理论保证之间的差距。
• 证明所提方法在表格环境和高维环境中均优于现有的潜在空间模型（如DeepMDP和DBC）。

提出的方法

• 推导出通过最大似然训练的低维策略与目标策略之间性能差异的上界，该上界在目标策略为低维或非低维时均成立。
• 制定一个表示学习目标，通过离线数据最小化该上界，当可用时整合动力学和奖励信息。
• 将该目标实现为对比学习过程，其中正样本对基于隐式能量模型近似的状态转移动力学形成。
• 使用随机傅里叶特征参数化表示，以实现线性动力学近似的高效且可扩展的学习。
• 将表示学习阶段与下游行为克隆分离，在BC微调过程中固定表示，以确保理论保证。
• 通过在表示目标中使用线性参数化动力学和奖励，将方法扩展至对数线性策略，同时保持理论边界。

实验结果

研究问题

• RQ1 我们能否设计一种模仿学习的表示学习目标，可证明地减少目标策略与低维策略之间的性能差距，且不依赖于目标策略的内在维度？
• RQ2 当目标策略不表现出低维结构时，如何避免表示学习中的混叠问题，尤其是在远离目标分布的离线数据存在的情况下？
• RQ3 使用随机傅里叶特征的对比学习能否以可扩展且实用的方式有效近似线性动力学模型，用于模仿学习？
• RQ4 在样本效率和Atari游戏性能方面，所提方法与现有潜在空间模型（如DeepMDP和DBC）相比表现如何？
• RQ5 理论边界在实际中是否成立？能否被有效利用以在多样化环境中实现一致的性能提升？

主要发现

• 所提的对比傅里叶特征方法在60款Atari 2600游戏中超过一半的游戏性能提升超过40%，显著优于基线行为克隆。
• 能量模型参数化的对比目标也带来显著性能增益，优于基线行为克隆，并与端到端微调的潜在空间模型性能相当或更优。
• 尽管DeepMDP和DBC是先进的潜在空间模型，但作为独立损失使用时，其性能改进微乎其微，甚至在端到端训练下仍逊于对比学习方法。
• 无论目标策略是否为低维，策略性能差异的理论边界始终保持紧致，验证了方法的鲁棒性。
• 在行为克隆训练过程中固定表示，相比端到端微调，能实现更好的泛化与性能，尤其在通过对比傅里叶特征学习表示时更为显著。
• 该方法能有效利用大规模离线数据集（如来自DQN重放池的5000万条过渡数据），显著提升模仿学习的样本效率，即使离线数据远离目标策略分布亦然。