一句话定位

Sparse Diffusion Policy(SDP)把 Mixture-of-Experts(MoE)层塞进 transformer-based Diffusion Policy 的 FFN 位置,将专家(expert)当作可复用的”技能”、路由器(router)当作”技能规划器”,用任务专属路由 + 互信息正则(而非常规负载均衡损失)实现稀疏、可复用、可扩展的机器人多任务策略:2D/3D 仿真与真实机器人上,多任务成功率优于同参数量密集基线(仅多约 1% 激活参数),持续学习新任务时激活参数量恒定不变(不像 LoRA 随任务数线性增长)且不遗忘旧任务,任务迁移仅微调 <0.5% 参数量的路由器即可超过从零训练。CoRL 2024 收录。

背景与定位

机器人多任务/终身学习长期依赖单一稠密策略网络(如 Diffusion Policy、Perceiver-Actor 等),即便是”推一下积木”这类简单任务也要激活全部参数;面对新任务往往需要代价高昂的全量微调,带来灾难性遗忘风险,而 LoRA 这类任务专属适配器虽然缓解遗忘,但推理时需要不断叠加激活参数。另一条路线是技能发现(skill discovery)与技能链(chain of skills),如 Lotus、SkillDiffuser 等,但这些方法依赖视觉特征或语言提示等人工设计的知识引导来切分技能,技能抽象模块本身不可扩展,网络结构也未针对稀疏计算设计。

MoE 已在 NLP、视觉、多模态大模型中证明有效(如 Mod-Squad、Uni-MoE、Branch-Train-Mix 等工作),其”路由选择子集专家”的机制天然具备稀疏性、跨任务/跨时间的专家复用能力,以及不破坏已有专家前提下扩容的灵活性。SDP 的问题设定是:仅靠 MoE 结构本身(不引入额外人工设计的技能知识)能否同时解决多任务效率、持续学习遗忘、任务迁移这三个问题。

SDP 的策略骨干直接沿用 transformer-based Diffusion Policy(Chi et al. 2023)并将其 FFN 替换为 MoE 层;2D 仿真基准使用 MimicGen(论文称是首个在 MimicGen 上做多任务训练的工作),3D 仿真则把同样的 MoE 改造套用到 3D Diffusion Policy 的 FFN block 上。同期另一条把 MoE 用进 diffusion policy 的工作是 MoDE(Efficient Diffusion Transformer Policies with Mixture of Expert Denoisers),但 MoDE 的出发点是用噪声条件路由的去噪专家换取推理加速与更强的单策略多任务基准表现,不像 SDP 强调专家即”技能”、路由器即”规划器”这一可解释框架,也未针对持续学习/任务迁移做设计。同年处理 diffusion policy 效率问题的还有 consistency-policy(一致性蒸馏加速推理)、DPPO(策略梯度微调)、AdaFlow(自适应 ODE 求解器),但它们的目标轴是”单任务推理延迟”,而非本文的”多任务参数复用与持续学习”。

模型架构

  • 策略骨干:transformer-based Diffusion Policy(Chi et al. 2023),不使用任何预训练 VLM;输入历史状态(图像或点云观测),输出未来动作序列(diffusion 去噪生成,连续动作,非离散 token/flow-matching)。
  • MoE 层设计:第 n 层 MoE 含 L 个专家 {E_l^n}(每个为 MLP)与一个路由器 R^n;路由器将输入 x∈ℝ^{1×M} 与专家嵌入 W^n∈ℝ^{M×L} 比较打分,取 Softmax 后做 Top-K 门控:y = Σ_l R^n(x,l)·E_l^n(x),R^n(x,l)=Top-K(Softmax(xW^n), l)(未入选的 l 权重为 0)。用 MoE 层替换 transformer 每个 block 中的 FFN。
  • 任务专属路由(多任务):每个任务 j 拥有独立路由器 R_j^n,基于同一套历史状态但不同任务路由到不同专家组合;同一专家可以在同一任务的不同时间步、或跨不同任务被重复调用,体现”专家=技能,跨任务/跨时间可复用”。
  • 持续学习扩展:面对新任务时冻结已学专家、已学路由器及其余(非 MoE)模块,仅新增一小批可训练专家 + 一个新任务专属路由器并只训练这部分新增参数;旧任务的推理路径完全不受影响,因此计算量不随任务数增加。
  • 任务迁移:冻结全部已学专家和旧路由器,仅新训练一个轻量路由器(配合解冻的视觉编码器)来为未见的新任务选择已有专家组合,不新增专家。
  • 专家组合数:作者给出组合上界公式 (L!/(L-K)!)^N(N 为 MoE 层数),示例 L=4, N=2, K=1 时,每层新增 1 个专家(共 2 个)即可新产生 9 种技能组合,用以论证少量专家扩容即可覆盖大量技能。
  • 关键配置数字(Appendix A.1):
    • 2D 仿真(MimicGen 多任务):12 层 transformer block,512 维 embedding,batch size 64,Adam 优化器,lr 1e-4;观测步数 2、规划 8 个动作步(执行首步);SDP 每层 8 个专家、激活 top-2;训练 300 epoch,每 50 epoch 评一次。
    • 3D 仿真(叠加在 3D Diffusion Policy 上):2 层 transformer block,256 维 embedding,Adam,lr 1e-4;激活专家的参数量被设定为与原始(稠密)网络相当;训练 6000 epoch(因每任务仅 20-40 条演示),每 200 epoch 评一次。
    • 持续学习设置:8 层 transformer block,256 维 embedding,Adam,lr 1e-4;每个新任务新增 8 个专家、激活 top-2;训练 500 epoch,每 50 epoch 评一次。
    • 真实机器人:FANUC LRMate 200iD/7L 机械臂 + SMC 夹爪,导纳控制(admittance control)实现顺应运动以保证操作安全;hang 任务用可学习的傅里叶嵌入(Fourier embedding)区分两个不同的悬挂目标点;训练 2000 epoch。
  • 参数量对比:2D 多任务 SDP 总激活参数 53.3M,仅比密集基线(TH/TT/TCD,52.6-52.7M)多约 1%(对应摘要中”negligible increase”的表述);持续学习中 MoE(Ours) 的激活参数在 3 个阶段恒为 9.2M,而 LoRA 从 9.0M(Stage1)增至 12.0M(Stage2)、14.9M(Stage3)。

数据

  • 2D 仿真(MimicGen 基准):Square、Stack、Coffee、Hammer(Cleanup)、Mug(Cleanup)、Nut(Assembly)、Stack Three、Thread(ing) 共 8 个任务,每任务 1K-10K 条人类演示(含较宽的初始状态分布),用于多任务泛化评测;论文称这是首个在 MimicGen 上做多任务训练的工作。
  • 3D 仿真(DexArt 基准):Toilet、Faucet、Laptop 3 个任务,点云输入,每任务仅 20-40 条人类演示,演示数量少是把训练轮数拉到 6000 epoch 的原因。
  • 真实机器人(FANUC 机械臂):Pull(拉圆环至桌面中心红色区域,20 条演示)、Pick-and-place(拾取杯子放到盘子上,20 条演示)、Hang(悬挂杯子,两个不同悬挂目标点各 20 条,共 40 条演示,模型用可学习傅里叶嵌入区分两个目标点);三任务联合多任务训练。数据采集/执行均通过导纳控制实现顺应机械臂运动,论文未进一步说明具体遥操作方式。
  • 持续学习(robomimic):Can(拾放罐子)→Lift(举起方块)→Square(方形轴插入孔位)三阶段序列任务;论文未重述各任务的具体演示条数(仅引用 robomimic 数据集本身)。
  • 任务迁移(MimicGen):基础任务为 Coffee + Mug Cleanup(预训练专家池),迁移目标为 Coffee Preparation——一个长时程、与基础任务部分不同的组合任务(Coffee Preparation 要求从抽屉里取咖啡,而基础任务是从桌面取咖啡放入抽屉),用以检验专家是否能被复用而非死记硬背。
  • 数据混合/配比:多任务训练中各任务损失按等权重求和(Eq. 1 的 L_bc=Σ_j L_bc^j),互信息正则项假设各任务先验概率相等 p(T_j)=1/J;论文未披露除此之外的跨任务采样权重或课程设计,也未使用仿真到真实(sim-to-real)迁移——真实世界数据全部来自实机演示。

训练方法

  • 目标函数:L = L_bc − γ·Σ_n I(T, E^n),即行为克隆损失(多任务时对各任务损失求和,持续学习/任务迁移时仅对当前新任务损失)减去一个乘以系数 γ 的任务-专家互信息项 I(T,E^n)=Σ_j Σ_l p(T_j,E_l^n)·log[p(T_j,E_l^n)/(p(T_j)p(E_l^n))](对每个 MoE 层求和)。
  • 为何不用常规负载均衡损失:NLP 领域 MoE 常用负载均衡损失防止个别专家过载,但作者观察到机器人任务中某些技能(如 pick-and-place)本就应该被频繁复用,专家过载是常态而非问题;真正需要避免的是”任务过载”——多个任务因路由未充分区分而退化成共用同一套专家、丧失任务特异性。因此互信息项转而鼓励专家对任务做专门化(specialization),而不是让专家使用频次均衡。
  • 持续学习训练流程:冻结旧专家、旧路由器及除 MoE 外的其余模块;每个 MoE 层新增固定数量的可训练专家(8 个)与一个新任务路由器;仅用新任务的行为克隆损失(+互信息正则)训练这些新增参数。消融(Ours-MI_loss,Fig. 5)显示:若去掉互信息损失,由于新增专家初期输出接近随机动作,路由器会倾向偏向旧的、已训练好的冻结专家,从而影响新任务学习效果,互信息损失能纠正这种偏向、强迫路由器选择任务专属专家(论文以描述性结论呈现,Fig. 5 为柱状图,正文未重复给出具体数值)。
  • 任务迁移训练流程:先在 Coffee、Mug Cleanup 两个基础任务上预训练完整 SDP;迁移到 Coffee Preparation 时只微调 MoE 路由器并解冻视觉编码器,训练 100 epoch;对照组是在 Coffee Preparation 上从零训练 100 epoch 的完整策略(25.9M 可训练参数)。
  • 持续学习消融:视觉编码器是否全量微调(Appendix A.4, Table 6):若把视觉编码器也做全量微调(FFT),无论策略解码器用 LoRA 还是 MoE(Ours),Stage 3 时对 Can、Lift 两个旧任务的成功率都归零(即视觉编码器全量微调本身就会诱发遗忘,与策略解码器选择无关);仅在当前新任务 Square 上 MoE(Ours) 仍优于 LoRA(0.75 对 0.57)。
  • 优化器与学习率:全部实验统一用 Adam,学习率 1e-4(2D 仿真、3D 仿真、持续学习设置均相同,见 Appendix A.1)。
  • 不涉及:论文全程为模仿学习(行为克隆),未使用强化学习,也未对模型本身做知识蒸馏。

Infra(训练 / 推理工程)

  • 训练硬件:2D 仿真实验(300 epoch)在单张 NVIDIA A6000 GPU 上训练 130-150 小时;3D 仿真、持续学习、真实机器人实验的 GPU 型号/数量/总训练 GPU 时均未披露。
  • 并行策略:未披露(仅提及单卡 A6000,无分布式训练描述)。
  • 训练精度:未披露(论文未提及混合精度/fp16/bf16)。
  • 推理 FPS / 控制频率 / 端侧延迟:未披露。论文未给出 SDP 推理耗时、动作生成频率或与稠密基线的延迟对比数字;真实机器人执行依赖导纳控制实现顺应运动,但未给出对应的控制 Hz。

评测 benchmark

Table 1(2D 仿真,MimicGen 8 任务,取最优 3 个 checkpoint 的平均成功率与激活参数量):

MethodActive ParamsSquareStackCoffeeHammerMugNutStack threeThreadAvg.
TH52.6M0.760.980.720.970.630.520.730.550.73
TT w/ 3Layer52.6M0.730.950.760.990.650.490.680.590.73
TCD52.7M0.630.950.770.920.530.440.620.560.68
SDP(Ours)53.3M0.740.990.830.980.700.420.760.650.76

SDP 以几乎相同的激活参数量(仅多约 1%)在 8 个任务中的 6 个上超过全部密集基线,整体平均成功率 0.76 对基线最高的 0.73/0.73。

Table 2(3D 仿真,DexArt 基准,点云输入):

MethodToiletFaucetLaptopAvg.
TT w/ 1Layer0.730.350.850.64
TCD0.720.330.800.62
SDP(Ours)0.750.430.820.67

Table 3(真实机器人,FANUC 机械臂):

MethodPullPick PlaceHangAvg.
TCD1.00.00.350.45
SDP(Ours)1.01.00.950.98

TCD 在 Pick Place 任务上完全失败(0.0),论文将其归因于 Appendix A.3 描述的”model collapse”现象:稠密基线无法区分不同任务间的多模态动作分布,执行了训练演示中从未出现过的动作(Fig. 8 可视化)。

Table 4(持续学习,robomimic Can→Lift→Square 三阶段,AP=策略网络激活参数量):

Policy DecoderStage1 CanAPStage2 CanLiftAPStage3 CanLiftSquareAP
FFT0.979.0M0.001.009.0M0.000.000.899.0M
LoRA0.949.0M0.941.0012.0M0.941.000.7314.9M
MOE(Ours)0.969.2M0.941.009.2M0.941.000.759.2M

FFT(全量微调)在学完 Lift 后 Can 成功率归零(0.97→0.00),典型的灾难性遗忘;LoRA 与 MoE(Ours) 都能保住旧任务成功率,但 LoRA 的激活参数量随任务数线性增长(9.0M→12.0M→14.9M),MoE(Ours) 恒定为 9.2M,且在最难的 Square 任务上(0.75)优于 LoRA(0.73)。

Table 5(任务迁移,Coffee Preparation):

MethodTrainable ParamsCoffee Preparation
Scratch25.9M0.70
Router only(Ours)0.1M0.80

仅微调路由器(约 0.1M 参数,占策略网络参数量不到 0.5%)即超过从零训练完整策略(25.9M 参数)的成功率,验证预训练专家池的可复用性。

Table 6(Appendix A.4,视觉编码器全量微调下的持续学习对照):

Vision EncoderPolicy DecoderStage1 CanStage2 CanLiftStage3 CanLiftSquare
FFTLoRA0.970.001.000.000.000.57
FFTMOE(Ours)0.930.001.000.000.000.75

视觉编码器一旦全量微调,无论策略解码器是 LoRA 还是 MoE,旧任务 Can/Lift 在 Stage3 都归零,但 MoE(Ours) 在当前任务 Square 上仍明显优于 LoRA(0.75 对 0.57)。

消融实验(Fig. 5,Square Stage 3,仅描述性结论,具体数值原文未重复给出):完整方法(Ours)成功率最高;去掉互信息损失(Ours-MI_loss)性能下降;仅保留 Policy MoE(PoMoE,不含其余 MoE 层)与仅保留 Vision MoE + Policy MoE(PoViMoE)均低于完整方法,说明需要在多个模块(不止策略解码器)引入 MoE 结构、并配合互信息损失才能取得最佳持续学习效果。

创新点与影响

  • 首次把 MoE 结构系统性地整合进 transformer-based diffusion policy,用于机器人多任务/持续学习/任务迁移三合一场景:把专家视为可复用”技能”、路由器视为”技能规划器”,无需额外的人工技能发现或语言/视觉引导即可获得稀疏、可复用、可扩展的策略结构。
  • 提出任务-专家互信息正则替代常规 MoE 负载均衡损失:论证机器人多任务场景下专家过载(同一技能被多任务复用)是合理现象,真正要避免的是任务层面路由未充分区分(任务过载),因而用最大化任务-专家互信息鼓励专家做任务特异性选择。
  • 三个场景的量化收益:多任务学习仅增加约 1% 激活参数即在 2D/3D 仿真与真实机器人上普遍超过稠密基线;持续学习中激活参数量恒定(不像 LoRA 随任务数线性增长)且不发生灾难性遗忘;任务迁移仅需微调 <0.5% 参数量的轻量路由器即超过从零训练完整策略。
  • 发现并命名真实机器人多任务场景下的”model collapse”现象:稠密基线(TCD)在无法区分任务间多模态动作分布时会退化输出从未在演示中出现的动作,而 SDP 凭借任务专属稀疏路由规避了这一问题。
  • 作者自陈局限:(1) 若网络内共享知识过于有限,即便结构上是 MoE,不同任务路由器仍可能激活相同专家组合,导致方法失效;(2) 路由器是任务专属的(需要显式的任务标识来选路由),这限制了模型做”通用任务执行”的能力(无法仅凭感知/语言自主判断当前是什么任务);作者提出未来结合大语言模型、让 SDP 接受语言条件以支持更广泛场景是可能的改进方向。

原始链接

一手源存档(sources/)