[论文解读] Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning

Rishabh Agarwal, Marlos C. Machado+2 more|arXiv (Cornell University)|Jan 13, 2021

Reinforcement Learning in Robotics参考文献 64被引用 27

一句话总结

本论文提出了一种策略相似性度量（PSM）和对比学习流程，以创建策略相似性嵌入（PSEs），从而提升强化学习策略在未见环境中的泛化能力。

ABSTRACT

Reinforcement learning methods trained on few environments rarely learn policies that generalize to unseen environments. To improve generalization, we incorporate the inherent sequential structure in reinforcement learning into the representation learning process. This approach is orthogonal to recent approaches, which rarely exploit this structure explicitly. Specifically, we introduce a theoretically motivated policy similarity metric (PSM) for measuring behavioral similarity between states. PSM assigns high similarity to states for which the optimal policies in those states as well as in future states are similar. We also present a contrastive representation learning procedure to embed any state similarity metric, which we instantiate with PSM to obtain policy similarity embeddings (PSEs). We demonstrate that PSEs improve generalization on diverse benchmarks, including LQR with spurious correlations, a jumping task from pixels, and Distracting DM Control Suite.

研究动机与目标

• 将RL的序列结构融入表示学习，以提高跨环境的泛化。
• 定义一个与奖励无关的策略相似性度量（PSM），捕捉跨状态/环境的长期最优行为相似性。
• 学习嵌入（PSEs），使用基于PSM的对比目标将行为相似的状态聚类。
• 在多样化基准上展示泛化改进，包括像素级跳跃任务、带虚假相关性的LQR，以及分散干扰的DM 控制套件。

提出的方法

• 使用最优策略之间的距离（Dist）以及对下一个状态分布的1-Wasserstein项，得到递归方程 d*(x,y)。
• 使用 Dist（对于离散动作为 TV，或对连续动作采用平均动作的 L1）和 Wasserstein 项来界定跨环境策略转移误差（定理1）。
• 将 d 转换为通过高斯核的相似度 Gamma，并构建对比学习目标（CME）以学习满足跨环境状态相似性的嵌入 z_theta(t)。
• 以 PSM 实例化 CME，产生将具有相似长期最优行为的状态聚类的策略相似性嵌入（PSEs），将 RL 模仿/训练损失与 CME 辅助损失结合。
• 提供一个端到端架构，其中编码器、投影头和策略头共同训练（SimCLR 风格的对比损失 + RL 模仿损失）。
• 给出一个实用算法（算法1），用 Gamma 下的最近邻正样本以及软对比损失，从训练 MDP 对中计算 CME。

实验结果

研究问题

• RQ1 奖励无关的策略相似性度量（PSM）能否在 RL 中提供鲁棒的跨环境泛化？
• RQ2 编码 PSM 基于状态相似性的对比学习嵌入（PSEs）是否能在普通正则化或数据增强之外提升泛化？
• RQ3 与现有方法相比，PSEs 在多种泛化基准（像素跳跃任务、带干扰项的 LQR、Distracting DM Control Suite）上的表现如何？

主要发现

数据增强	方法	宽	窄	随机
✗ Dropout 和 ℓ2 正则。	双仿真传输	17.8 (2.2)	10.2 (4.6)	9.3 (5.4)
✓ RandConv	PSEs	33.6 (10.0)	9.3 (5.3)	37.7 (10.4)
✗ RandConv + π*-双仿真	RandConv + π*-双仿真	41.4 (17.6)	17.4 (6.7)	33.4 (15.6)
✓ RandConv + PSEs	RandConv + PSEs	87.0 (10.1)	52.4 (5.8)	83.4 (10.1)

• PSEs 在多项基准上实现零-shot 泛化，在若干情形下优于正则化和双仿真传输基线。
• PSEs 与 RandConv 数据增强结合时显著提升泛化，在某些配置下甚至优于仅使用增强。
• 结合 PSM 使用 CME 的情况下，PSEs 的消融结果优于如 l2-嵌入或 π*-双仿真等替代方案，凸显学习到的不变性的重要性。
• PSEs 对中等策略次优（epsilon-suboptimal 策略）保持鲁棒，当无法获得精确策略时，可以利用近似 π*。
• 在颜色修改版本的 Jumping Task 中，PSEs 捕捉到与任务相关的不变性，并保持优于仅使用 RandConv 的泛化，在该方法下增强会模糊任务特定的区别。
• 在带干扰项的 LQR 和 Distracting DM Control Suite 中，基于 PSM 的表示在不访问真实最优策略的情况下也能实现强泛化。