[论文解读] TACO: Temporal Latent Action-Driven Contrastive Loss for Visual Reinforcement Learning

Ruijie Zheng, Xiyao Wang+6 more|arXiv (Cornell University)|Jun 22, 2023

Reinforcement Learning in Robotics被引用 5

一句话总结

TACO 提出了一种时间对比学习框架，通过最大化当前状态与动作序列之间及其对应未来状态之间的互信息，联合学习状态和动作表征。在经历一百万次环境交互步骤后，TACO 在视觉连续控制任务上相较于最先进方法实现了 40% 的性能提升，并在不同质量的数据集上均改善了离线强化学习的性能。

ABSTRACT

Despite recent progress in reinforcement learning (RL) from raw pixel data, sample inefficiency continues to present a substantial obstacle. Prior works have attempted to address this challenge by creating self-supervised auxiliary tasks, aiming to enrich the agent's learned representations with control-relevant information for future state prediction. However, these objectives are often insufficient to learn representations that can represent the optimal policy or value function, and they often consider tasks with small, abstract discrete action spaces and thus overlook the importance of action representation learning in continuous control. In this paper, we introduce TACO: Temporal Action-driven Contrastive Learning, a simple yet powerful temporal contrastive learning approach that facilitates the concurrent acquisition of latent state and action representations for agents. TACO simultaneously learns a state and an action representation by optimizing the mutual information between representations of current states paired with action sequences and representations of the corresponding future states. Theoretically, TACO can be shown to learn state and action representations that encompass sufficient information for control, thereby improving sample efficiency. For online RL, TACO achieves 40% performance boost after one million environment interaction steps on average across nine challenging visual continuous control tasks from Deepmind Control Suite. In addition, we show that TACO can also serve as a plug-and-play module adding to existing offline visual RL methods to establish the new state-of-the-art performance for offline visual RL across offline datasets with varying quality.

研究动机与目标

• 解决从原始像素输入的视觉强化学习中的样本效率低下问题。
• 学习联合状态和动作表征，以捕捉最优策略和价值函数所需的控制相关信息。
• 克服先前对比方法依赖于特定策略的正样本对或在连续控制中忽略动作表征的局限性。
• 提供一个即插即用的模块，兼容现有的在线和离线强化学习算法。
• 通过理论基础坚实的互信息最大化，提升高维连续控制任务中的稳定性和泛化能力。

提出的方法

• TACO 使用基于 InfoNCE 损失的对比学习目标，以最大化当前状态表征与动作序列表征之间，以及对应未来状态表征之间的互信息。
• 它联合训练状态编码器和动作编码器，以生成捕捉时序动态和控制相关性信息的潜在表征。
• 该方法将动作序列及其对应未来状态视为正样本对，而其他序列则作为负样本。
• 对比目标被设计为与策略无关，避免了由策略依赖的正样本对带来的不稳定性。
• TACO 兼容任何现成的强化学习算法，可无缝集成到在线和离线强化学习训练流程中。
• 该方法具有理论基础，表明所学习的表征包含实现最优控制所需的充分信息。

实验结果

研究问题

• RQ1 联合状态与动作表征学习是否能提升视觉强化学习中的样本效率？
• RQ2 一种不依赖策略的对比目标，若结合动作序列，是否能带来更稳定和泛化能力更强的表征？
• RQ3 TACO 是否能在在线连续控制基准测试中超越最先进模型无关和模型相关的视觉强化学习算法？
• RQ4 TACO 在不同质量数据集的离线视觉强化学习中，能在多大程度上提升性能？
• RQ5 与直接建模潜在动态相比，TACO 的互信息最大化在稳定性和表征质量方面表现如何？

主要发现

• 在九个 DeepMind Control Suite 任务中，TACO 在经历一百万次环境交互步骤后，平均性能相比最先进模型无关视觉强化学习算法提升了 40%。
• 当与 DrQ-v2 结合时，TACO 在在线视觉强化学习设置中超越了最强的模型基线方法 Dreamer-v3。
• TACO 通过与 TD3+BC 和 CQL 集成，在离线视觉强化学习中实现了性能提升，在不同数据质量的离线数据集上均创下新的 SOTA 记录。
• 由于其与策略无关的对比目标，TACO 展现出稳定的训练过程，避免了由策略依赖正样本对带来的不稳定性。
• TACO 的联合状态-动作表征学习提升了连续控制任务中的泛化能力和长期规划性能。
• 尽管效果显著，TACO 仍需要较大的批量大小，这影响了计算效率，是未来优化的一个方向。