[论文解读] Decoupling Representation Learning from Reinforcement Learning

Adam Stooke, Kimin Lee+2 more|arXiv (Cornell University)|Sep 14, 2020

Reinforcement Learning in Robotics参考文献 39被引用 65

一句话总结

本文介绍 Augmented Temporal Contrast (ATC)，一种将表征学习与策略学习在强化学习中解耦的无监督学习任务，并展示 ATC 训练的编码器在多种环境和任务中可匹配甚至超越端到端强化学习。

ABSTRACT

In an effort to overcome limitations of reward-driven feature learning in deep reinforcement learning (RL) from images, we propose decoupling representation learning from policy learning. To this end, we introduce a new unsupervised learning (UL) task, called Augmented Temporal Contrast (ATC), which trains a convolutional encoder to associate pairs of observations separated by a short time difference, under image augmentations and using a contrastive loss. In online RL experiments, we show that training the encoder exclusively using ATC matches or outperforms end-to-end RL in most environments. Additionally, we benchmark several leading UL algorithms by pre-training encoders on expert demonstrations and using them, with weights frozen, in RL agents; we find that agents using ATC-trained encoders outperform all others. We also train multi-task encoders on data from multiple environments and show generalization to different downstream RL tasks. Finally, we ablate components of ATC, and introduce a new data augmentation to enable replay of (compressed) latent images from pre-trained encoders when RL requires augmentation. Our experiments span visually diverse RL benchmarks in DeepMind Control, DeepMind Lab, and Atari, and our complete code is available at https://github.com/astooke/rlpyt/tree/master/rlpyt/ul.

研究动机与目标

• 激发为强化学习学习对奖励无关的视觉表征，以克服以奖励驱动的特征学习的局限性。
• 提出 ATC 作为无监督任务，从观测中学习鲁棒的编码器，与 RL 损失无关。
• 展示在多样化环境中使用 ATC 编码表示的在线 RL。
• 将 ATC 与其他无监督学习方法进行基准比较，并评估多任务泛化能力。
• 探索消融和数据增强效应，以理解 ATC 组件。

提出的方法

• 引入 Augmented Temporal Contrast (ATC)：将观测 o_t 与同一轨迹内近未来的 o_{t+k} 联系的对比任务。
• 通过共享的 CNN 编码器对增强后的观测进行编码；使用线性压缩器将其压缩为潜在代码；在潜在空间中使用残差预测器进行前向预测。
• 对正样本使用动量编码器，并对来自同一批次的负样本应用 InfoNCE 损失。
• 对观测应用随机数据增强（随机平移）；在对比之前使用预测层处理锚点编码。
• 离线使用 ATC 训练编码器（无监督），而策略学习仅在潜在表示上进行；在基线中可选择将编码器权重冻结在 RL 时。
• 可选地引入一种新的增强（亚像素随机平移）以实现潜在图像重放并降低计算量。

实验结果

研究问题

• RQ1 在不牺牲策略性能的前提下，表示学习是否可以从强化学习中解耦？
• RQ2 与现有的 UL 方法相比，ATC 是否是用于 RL 的编码器预训练的更优无监督任务？
• RQ3 通过 ATC 学习的多任务编码器能否迁移到新的环境或任务？
• RQ4 ATC 效果关键的消融和增强是什么？
• RQ5 ATC 在多样化基准（DMControl、DMLab、Atari）和强化学习算法上的表现如何？

主要发现

• ATC 训练的编码器在在线使用且与 RL 梯度分离时，在大多数 DMControl 和 DMLab 环境以及测试的 Atari 大约一半以上的游戏中，匹配或超越端到端 RL 编码器。
• 以冻结权重进行预训练的 ATC 编码器在 DMControl、DMLab 和 Atari 基准测试中通常超过其他领先的 UL 算法。
• 在多个 DMControl 环境上训练的单个 ATC 编码器可以在冻结权重的情况下推广到新的下游任务。
• ATC 消融结果显示数据增强对 DMControl 是必要的，亚像素随机平移使潜在图像重放在不增加额外计算的情况下成为可能。
• 跨环境的多任务预训练在某些任务上可以改善迁移，尽管 Atari 的结果显示跨游戏迁移受限，表明表示容量和领域差异影响显著。
• 将 ATC 作为辅助损失或用于权重初始化，在多款 Atari 游戏上进一步提升了性能。