面向多任务学习的高效扩散Transformer策略与混合去噪专家
机器学习
2024-12-18 v1 机器人学
摘要
扩散策略在模仿学习中已得到广泛应用，具有生成多模态和不连续行为等吸引人的特性。随着模型变得越来越大以捕获更复杂的能力，其计算需求也随之增加，正如最近的缩放定律所示。因此，继续使用当前的架构将面临计算障碍。为了解决这一差距，我们提出了混合去噪专家（MoDE）作为一种新的模仿学习策略。MoDE超越了当前最先进的基于Transformer的扩散策略，同时通过稀疏专家和噪声条件路由实现了参数高效的扩展，通过专家缓存将活跃参数减少了40%，推理成本降低了90%。我们的架构将这种高效扩展与噪声条件自注意力机制相结合，使得在不同噪声水平下更有效地去噪。MoDE在四个已建立的模仿学习基准（CALVIN和LIBERO）的134个任务上实现了最先进的性能。值得注意的是，通过在多样化的机器人数据上预训练MoDE，我们在CALVIN ABC上达到了4.01，在LIBERO-90上达到了0.95。它在4个基准上平均超过基于CNN和Transformer的扩散策略57%，同时与默认的扩散Transformer架构相比，使用的FLOPs减少了90%，活跃参数更少。此外，我们对MoDE的组件进行了全面的消融研究，为设计高效且可扩展的扩散策略Transformer架构提供了见解。代码和演示可在https://mbreuss.github.io/MoDE_Diffusion_Policy/获取。
引用
@article{arxiv.2412.12953,
title = {Efficient Diffusion Transformer Policies with Mixture of Expert Denoisers for Multitask Learning},
author = {Moritz Reuss and Jyothish Pari and Pulkit Agrawal and Rudolf Lioutikov},
journal= {arXiv preprint arXiv:2412.12953},
year = {2024}
}