arXiv:2409.14411cs.RO2024-09ICRA被引 54

将扩散策略模型扩展至10亿参数,显著提升机器人操作性能。

Scaling Diffusion Policy in Transformer to 1 Billion Parameters for Robotic Manipulation

  • 分解观测特征嵌入并引入非因果注意力,改善训练稳定性。
  • 在MetaWorld上平均提升21.6%,真实机器人任务提升最高75%。
  • 适合追求大规模视觉-运动学习的机器人研究者参考。

扩散策略是学习端到端视觉-运动机器人控制的强大工具,理论上具备可扩展性,即模型越大性能越好。然而我们发现,基于Transformer架构的扩散策略( extit{DP})难以有效扩展,甚至增加少量层数也会导致训练效果下降。为此,我们提出可扩展的扩散变压器策略( extbf{ extit{ScaleDP}}),引入两个模块:首先,将观测特征嵌入分解为多个仿射层并整合进Transformer块,缓解梯度过大问题;其次,采用非因果注意力机制,使策略网络在预测时能“看到”未来动作,减少误差累积。实验表明, extit{ScaleDP}成功将扩散策略从1000万参数扩展至10亿参数。在MetaWorld的50个任务中,其最大模型平均优于原始 extit{DP} 21.6%;在7个真实机器人任务中,单臂任务平均提升36.25%,双臂任务提升达75%。本工作为视觉-运动学习模型的规模化提供了可行路径。

原文摘要 · Abstract (English)

Diffusion Policy is a powerful technique tool for learning end-to-end visuomotor robot control. It is expected that Diffusion Policy possesses scalability, a key attribute for deep neural networks, typically suggesting that increasing model size would lead to enhanced performance. However, our observations indicate that Diffusion Policy in transformer architecture (\DP) struggles to scale effectively; even minor additions of layers can deteriorate training outcomes. To address this issue, we introduce Scalable Diffusion Transformer Policy for visuomotor learning. Our proposed method, namely \textbf{\methodname}, introduces two modules that improve the training dynamic of Diffusion Policy and allow the network to better handle multimodal action distribution. First, we identify that \DP~suffers from large gradient issues, making the optimization of Diffusion Policy unstable. To resolve this issue, we factorize the feature embedding of observation into multiple affine layers, and integrate it into the transformer blocks. Additionally, our utilize non-causal attention which allows the policy network to \enquote{see} future actions during prediction, helping to reduce compounding errors. We demonstrate that our proposed method successfully scales the Diffusion Policy from 10 million to 1 billion parameters. This new model, named \methodname, can effectively scale up the model size with improved performance and generalization. We benchmark \methodname~across 50 different tasks from MetaWorld and find that our largest \methodname~outperforms \DP~with an average improvement of 21.6\%. Across 7 real-world robot tasks, our ScaleDP demonstrates an average improvement of 36.25\% over DP-T on four single-arm tasks and 75\% on three bimanual tasks. We believe our work paves the way for scaling up models for visuomotor learning. The project page is available at scaling-diffusion-policy.github.io.

扩散模型机器人控制可扩展性Transformer

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。