Mini Diffuser: Fast Multi-task Diffusion Policy Training Using Two-level Mini-batches
作者: Yutong Hu, Pinhao Song, Kehan Wen, Renaud Detry
分类: cs.RO, cs.LG
发布日期: 2025-05-14 (更新: 2025-06-05)
💡 一句话要点
Mini Diffuser：利用双层Mini-batch加速多任务扩散策略训练，显著降低时间和内存需求。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 扩散模型 机器人学习 多任务学习 视觉语言 策略训练
📋 核心要点
- 现有扩散策略训练计算成本高昂，尤其是在多任务机器人学习中，限制了其应用。
- Mini Diffuser利用动作空间低维特性，提出双层minibatch策略，显著减少计算量。
- 实验表明，Mini Diffuser在性能损失很小的情况下，大幅降低了训练时间和内存占用。
📝 摘要（中文）
本文提出了一种名为Mini Diffuser的方法，该方法能够显著降低训练多任务视觉-语言机器人扩散策略所需的时间和内存，降幅可达一个数量级。这一改进源于对动作扩散和启发它的图像扩散技术之间一个先前未被充分探索的区别：在图像生成中，目标是高维的。相比之下，在动作生成中，目标的维度相对较小，只有图像条件是高维的。Mini Diffuser通过引入双层minibatch来利用这种不对称性，它将多个带噪声的动作样本与每个视觉-语言条件配对，而不是传统的one-to-one采样策略。为了支持这种批处理方案，我们对扩散Transformer进行了架构调整，防止样本间的信息泄露，同时保持完全的条件访问。在RLBench模拟中，Mini-Diffuser实现了最先进的多任务扩散策略95%的性能，同时仅使用5%的训练时间和7%的内存。真实世界的实验进一步验证了Mini-Diffuser保留了基于扩散策略的关键优势，包括建模多模态动作分布和产生以各种感知输入为条件的行为的能力。
🔬 方法详解
问题定义：现有基于扩散模型的机器人策略学习方法，尤其是在多任务场景下，需要大量的计算资源和内存，训练时间长，难以应用到实际场景中。主要痛点在于，传统方法对每个视觉-语言条件，只采样一个动作样本，计算效率较低。
核心思路：论文的核心思路是利用动作空间维度远低于图像空间的特性，通过对每个视觉-语言条件，采样多个带噪声的动作样本，形成双层minibatch。这样可以在不显著增加计算量的前提下，提高训练效率，减少内存占用。
技术框架：Mini Diffuser的整体框架仍然是基于扩散模型的策略学习，主要包括以下几个阶段：1）前向扩散过程：将动作样本逐步加入噪声，直到完全变为噪声；2）反向去噪过程：利用扩散Transformer网络，逐步去除噪声，恢复原始动作；3）条件输入：视觉-语言信息作为条件输入到扩散Transformer网络中，指导动作生成。
关键创新：最重要的技术创新点是双层minibatch策略。传统方法是one-to-one的采样，而Mini Diffuser是one-to-many的采样。此外，为了支持双层minibatch，论文还对扩散Transformer进行了架构调整，防止不同动作样本之间的信息泄露，同时保证所有样本都能充分利用条件信息。
关键设计：为了防止信息泄露，论文修改了Transformer的注意力机制，使得每个动作样本只能关注到自身的条件信息，而不能关注到同一batch中其他动作样本的信息。具体的实现方式是修改了注意力权重计算公式，引入了一个mask矩阵，将不同样本之间的注意力权重设置为0。此外，论文还对损失函数进行了调整，使得每个动作样本的损失都能够被正确计算和反向传播。
🖼️ 关键图片
📊 实验亮点
实验结果表明，Mini Diffuser在RLBench模拟环境中，仅使用5%的训练时间和7%的内存，就能够达到最先进的多任务扩散策略95%的性能。在真实世界的实验中，Mini Diffuser也表现出了良好的性能，验证了其在实际应用中的可行性。这些结果表明，Mini Diffuser是一种高效、实用的多任务扩散策略训练方法。
🎯 应用场景
Mini Diffuser具有广泛的应用前景，尤其是在需要快速训练和部署机器人策略的场景中。例如，它可以用于工业机器人、服务机器人、自动驾驶等领域，实现多任务、高效率的机器人学习。该方法降低了训练成本，使得在资源受限的平台上训练复杂的机器人策略成为可能，加速了机器人技术在实际场景中的应用。
📄 摘要（原文）
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