用双层小批量加速多任务机器人扩散策略训练,节省95%时间
Mini Diffuser: Fast Multi-task Diffusion Policy Training Using Two-level Mini-batches
- 通过双层小批量配对视觉语言条件与多个噪声动作样本
- 在RLBench上达到95%性能,仅需5%训练时间和7%内存
- 适合需要高效训练多模态机器人策略的研究者
我们提出一种方法,将多任务视觉语言机器人扩散策略的训练时间与内存消耗降低一个数量级。该改进源于一个未被充分探索的差异:图像生成中目标维度高,而动作生成中目标维度相对较低,仅视觉条件维度高。为此,我们设计了微型扩散器(Mini Diffuser),采用两层小批量机制,将多个噪声动作样本与每个视觉语言条件配对,而非传统的单对单采样。为支持此方案,我们在扩散变压器中引入架构改进,防止样本间信息泄露,同时保持完整条件访问。在RLBench模拟环境中,Mini-Diffuser实现了顶尖多任务扩散策略95%的性能,训练时间仅为5%,内存使用仅7%。真实世界实验进一步验证,其保留了基于扩散策略的关键优势,包括建模多模态动作分布和根据多样化感知输入生成行为的能力。代码已公开于mini-diffuse-actor.github.io。
原文摘要 · Abstract (English)
We present a method that reduces, by an order of magnitude, the time and memory needed to train multi-task vision-language robotic diffusion policies. This improvement arises from a previously underexplored distinction between action diffusion and the image diffusion techniques that inspired it: In image generation, the target is high-dimensional. By contrast, in action generation, the dimensionality of the target is comparatively small, and only the image condition is high-dimensional. Our approach, \emph{Mini Diffuser}, exploits this asymmetry by introducing \emph{two-level minibatching}, which pairs multiple noised action samples with each vision-language condition, instead of the conventional one-to-one sampling strategy. To support this batching scheme, we introduce architectural adaptations to the diffusion transformer that prevent information leakage across samples while maintaining full conditioning access. In RLBench simulations, Mini-Diffuser achieves 95\% of the performance of state-of-the-art multi-task diffusion policies, while using only 5\% of the training time and 7\% of the memory. Real-world experiments further validate that Mini-Diffuser preserves the key strengths of diffusion-based policies, including the ability to model multimodal action distributions and produce behavior conditioned on diverse perceptual inputs. Code available at mini-diffuse-actor.github.io
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。