[论文解读] CoBERL: Contrastive BERT for Reinforcement Learning

Andrea Banino, Adrià Puidomenech Badia+4 more|arXiv (Cornell University)|Jul 12, 2021

Reinforcement Learning in Robotics被引用 8

一句话总结

CoBERL 提出了一种对比性 BERT 启发的表示学习目标，结合门控 LSTM-Transformer 架构，以提升强化学习中的数据效率。通过利用掩码预测和无需手工设计增强的时间域对比损失，CoBERL 在 Atari、控制套件和 3D 环境中均实现了更优的样本效率和性能，始终优于 GTrXL 等基线模型。

ABSTRACT

Many reinforcement learning (RL) agents require a large amount of experience to solve tasks. We propose Contrastive BERT for RL (CoBERL), an agent that combines a new contrastive loss and a hybrid LSTM-transformer architecture to tackle the challenge of improving data efficiency. CoBERL enables efficient, robust learning from pixels across a wide range of domains. We use bidirectional masked prediction in combination with a generalization of recent contrastive methods to learn better representations for transformers in RL, without the need of hand engineered data augmentations. We find that CoBERL consistently improves performance across the full Atari suite, a set of control tasks and a challenging 3D environment.

研究动机与目标

• 通过从原始像素中学习更优表示，提升深度强化学习的数据效率。
• 通过大型密集网络在强化学习训练中应对噪声梯度和强时间相关性的问题。
• 开发一种无需手工设计数据增强的自监督表示学习方法。
• 通过可学习门控机制整合 LSTM 与 Transformer 的优势，实现动态知识利用。
• 在包括 2D 和 3D 控制任务在内的多样化强化学习环境中，实现稳健且可泛化的表示学习。

提出的方法

• 提出一种新颖的对比损失，结合双向掩码预测（受 BERT 启发）与 ReLIC 的时间域泛化形式，用于自监督表示学习。
• 使用因果掩码的门控 Transformer-XL（GTrXL）处理观测并预测被掩码的输入，确保自回归一致性。
• 引入可学习门控机制，融合 GTrXL 和 LSTM 的输出，实现根据任务需求动态选择表示路径。
• 通过强化学习策略损失训练门控，实现整个架构的端到端优化。
• 使用预测掩码输入的归一化嵌入与原始观测之间的对比损失，基于实例判别和不变性惩罚构建对比目标。
• 使用残差卷积神经网络编码器对原始像素观测进行嵌入，再输入到 Transformer 和 LSTM 组件中。

实验结果

研究问题

• RQ1 受 BERT 启发的对比表示学习目标是否能在无需手工数据增强的情况下提升强化学习的数据效率？
• RQ2 通过可学习门控将 LSTM 与 Transformer 结合，对多样化强化学习环境中性能与样本效率的影响如何？
• RQ3 在时间域中整合掩码预测与对比学习，是否能带来更一致且更具泛化能力的表示？
• RQ4 CoBERL 在多个基准环境中的表现相较于现有架构（如 GTrXL）的优越程度如何？
• RQ5 掩码比例对辅助预训练目标及最终强化学习性能的影响是什么？

主要发现

• CoBERL 在全部 57 款 Atari 游戏中均实现数据效率的持续提升，样本效率方面取得显著性能增益。
• 在 DeepMind Control Suite 中，CoBERL 展现出更高的样本效率和最终得分，优于 GTrXL 及其他基线模型。
• 在 3D DMLab-30 环境中，CoBERL 展现出更优的记忆与推理能力，AUC 指标优于 GTrXL，t(60) = 2.616，p = 0.0011。
• 消融实验证实，对比损失与门控架构均独立贡献于性能提升，其中对比目标在表示学习中尤为有效。
• 较高的掩码比例（如 50%）会降低性能，表明过度掩码会减少辅助任务所需的时间上下文信息。
• 该方法在不同领域间泛化良好，在 2D 平台游戏和 3D 第一人称/第三人称控制任务中均表现强劲。