[论文解读] Multi-task Batch Reinforcement Learning with Metric Learning

Jiachen Li, Quan Vuong+6 more|arXiv (Cornell University)|Sep 25, 2019

Reinforcement Learning in Robotics参考文献 77被引用 10

一句话总结

本文提出MBML，一种多任务批量强化学习方法，通过三元组损失进行度量学习，以在离线数据集间存在分布偏移的情况下提升任务推理的鲁棒性。通过重新标记奖励以生成困难负样本，该方法迫使策略关注奖励而非虚假的状态-动作相关性，在微调未见任务时实现高达80%的收敛速度提升。

ABSTRACT

We tackle the Multi-task Batch Reinforcement Learning problem. Given multiple datasets collected from different tasks, we train a multi-task policy to perform well in unseen tasks sampled from the same distribution. The task identities of the unseen tasks are not provided. To perform well, the policy must infer the task identity from collected transitions by modelling its dependency on states, actions and rewards. Because the different datasets may have state-action distributions with large divergence, the task inference module can learn to ignore the rewards and spuriously correlate $ extit{only}$ state-action pairs to the task identity, leading to poor test time performance. To robustify task inference, we propose a novel application of the triplet loss. To mine hard negative examples, we relabel the transitions from the training tasks by approximating their reward functions. When we allow further training on the unseen tasks, using the trained policy as an initialization leads to significantly faster convergence compared to randomly initialized policies (up to $80\%$ improvement and across 5 different Mujoco task distributions). We name our method $ extbf{MBML}$ ($ extbf{M} ext{ulti-task}$ $ extbf{B} ext{atch}$ RL with $ extbf{M} ext{etric}$ $ extbf{L} ext{earning}$).

研究动机与目标

• 解决当离线数据集存在较大的状态-动作分布偏移时，多任务批量强化学习中泛化能力差的挑战。
• 缓解任务推理模块因数据分布缺乏重叠而学习到状态-动作对与任务身份之间虚假相关性的风险。
• 通过使策略能够从上下文转移中推断任务身份（而无需访问真实任务标签），提升在未见任务上的零样本性能。
• 通过使用训练好的多任务策略作为有效初始化，加速在新任务上的微调。
• 证明通过奖励重标记进行困难负样本挖掘的度量学习，相比标准方法能实现更好的泛化性能。

提出的方法

• 该方法引入三元组损失目标，以训练一个任务嵌入网络，使其能够区分正样本和负样本上下文转移。
• 通过近似每个任务的奖励函数，并将其他任务的转移重新标记为这些近似奖励，生成困难负样本。
• 策略网络通过同时使用状态-动作对和奖励来推断任务身份，三元组损失促使不同任务在嵌入空间中实现分离。
• 该方法采用具有共享表示和任务特定头的多任务Q网络架构，通过来自多样化任务的离线数据进行训练。
• 该方法与SAC等异策略算法集成，使用多任务策略作为策略微调的初始化。
• 采用SAC的变体评估微调性能，消融研究显示奖励重标记和三元组损失均必不可少。

实验结果

研究问题

• RQ1 在分布偏移下，通过困难负样本挖掘的度量学习是否能提升多任务批量强化学习中的任务推理鲁棒性？
• RQ2 将其他任务的奖励重新标记是否有助于策略避免学习到状态-动作对与任务身份之间的虚假相关性？
• RQ3 与标准离线强化学习基线相比，该方法在零样本泛化到未见任务时表现如何？
• RQ4 多任务策略在多大程度上可作为新任务微调的有效初始化？
• RQ5 三元组损失、奖励重标记和多任务训练各组件对整体性能的贡献如何？

主要发现

• 所提出的MBML方法在未见任务的微调过程中，样本效率相比随机初始化策略最高提升80%。
• 消融研究显示，若移除三元组损失或奖励重标记，性能均显著下降，证实了二者的重要性。
• MBML在所有评估的任务分布（包括AntDir、AntGoal、HDir、UGoal、HalfCheetahVel和WalkerParam）中均优于PEARL和上下文BCQ等强基线模型。
• 即使在推理阶段无法访问真实任务身份或奖励函数，该方法也能在未见任务上实现良好泛化。
• 训练完整MBML模型的运行时间在3.9至26.6小时之间不等，具体取决于任务分布，且使用MBML初始化的SAC微调显著更快。