[论文解读] Switch Trajectory Transformer with Distributional Value Approximation for Multi-Task Reinforcement Learning

Qinjie Lin, Han Liu+1 more|arXiv (Cornell University)|Mar 14, 2022

Reinforcement Learning in Robotics被引用 4

一句话总结

SwitchTT 提出了一种基于稀疏专家的 Transformer 架构（Switch Transformer）和分布值估计的多任务离线强化学习方法，以提升样本效率并减少计算成本。在 10 个稀疏奖励任务上，其性能比 Trajectory Transformer 高出 10%，训练速度最快提升 90%，展示了在多任务设置下卓越的可扩展性和值估计能力。

ABSTRACT

We propose SwitchTT, a multi-task extension to Trajectory Transformer but enhanced with two striking features: (i) exploiting a sparsely activated model to reduce computation cost in multi-task offline model learning and (ii) adopting a distributional trajectory value estimator that improves policy performance, especially in sparse reward settings. These two enhancements make SwitchTT suitable for solving multi-task offline reinforcement learning problems, where model capacity is critical for absorbing the vast quantities of knowledge available in the multi-task dataset. More specifically, SwitchTT exploits switch transformer model architecture for multi-task policy learning, allowing us to improve model capacity without proportional computation cost. Also, SwitchTT approximates the distribution rather than the expectation of trajectory value, mitigating the effects of the Monte-Carlo Value estimator suffering from poor sample complexity, especially in the sparse-reward setting. We evaluate our method using the suite of ten sparse-reward tasks from the gym-mini-grid environment.We show an improvement of 10% over Trajectory Transformer across 10-task learning and obtain up to 90% increase in offline model training speed. Our results also demonstrate the advantage of the switch transformer model for absorbing expert knowledge and the importance of value distribution in evaluating the trajectory.

研究动机与目标

• 解决在多样化、异构数据集下多任务离线强化学习中计算成本高且样本效率差的挑战。
• 克服在多任务设置中共享策略参数导致性能下降的局限性。
• 改善稀疏奖励环境中值估计的性能，因为蒙特卡洛估计器在样本复杂度方面表现不佳。
• 在不带来推理成本成比例增加的前提下，实现从大规模多任务数据集中吸收高容量知识。

提出的方法

• 采用带有专家混合（MoE）层的 Switch Transformer 架构，其中门控网络将输入路由到特定任务的专家网络，实现稀疏参数激活并减少计算量。
• 集成三个 Transformer 头：动作预测、动态建模和返回到目标（RTG）预测，形成统一的序列建模框架以生成轨迹。
• 将标准的期望值估计替换为分布值估计器，该估计器对轨迹回报的类别分布进行建模，从而提升不确定性感知的评估能力。
• 在规划过程中，基于分布值估计结果对生成的轨迹进行束搜索（beam search），以选择高回报的动作序列。
• 在离线多任务数据集上端到端训练模型，将完整的状态、动作、回报和任务标识符序列作为输入 token。
• 利用 Switch Transformer 的高容量能力，在保持计算效率的同时吸收来自多个任务的多样化专家演示。

实验结果

研究问题

• RQ1 稀疏激活的 Transformer 架构是否能在不损失模型容量或性能的前提下，降低多任务离线强化学习中的计算成本？
• RQ2 与仅估计期望回报相比，对轨迹回报分布的建模是否能在稀疏奖励场景中带来更鲁棒、更准确的值估计？
• RQ3 在多任务、稀疏奖励环境中，SwitchTT 与标准 Trajectory Transformer 相比，在性能和训练效率方面表现如何？
• RQ4 Switch Transformer 架构在离线强化学习中，能在多大程度上有效学习并共享多样化、异构任务之间的表示？

主要发现

• 在 gym-mini-grid 环境的 10 个稀疏奖励任务中，SwitchTT 的性能比 Trajectory Transformer 提高了 10%。
• 由于专家计算的稀疏性，该方法将离线模型训练时间最多减少了 90%，相比标准 Trajectory Transformer。
• 分布值估计器显著提升了低样本、稀疏奖励场景下的值估计准确性，缓解了蒙特卡洛估计器样本复杂度差的问题。
• Switch Transformer 架构促进了任务间的更好参数共享和知识迁移，尤其在任务分布不同的情况下效果更明显。
• 基于分布值估计的束搜索能带来更可靠的轨迹规划和更高的累积奖励。
• 实证结果证实，将值建模为分布而非期望值，能提升多任务离线强化学习中策略的性能和稳定性。