Modular Diffusion Policy Training: Decoupling and Recombining Guidance and Diffusion for Offline RL

📄 arXiv: 2506.03154v1 📥 PDF

作者: Zhaoyang Chen, Cody Fleming

分类: cs.LG

发布日期: 2025-05-19


💡 一句话要点

提出模块化扩散策略训练,解耦引导与扩散模型,提升离线强化学习性能。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱四:生成式动作 (Generative Motion)

关键词: 离线强化学习 扩散模型 引导模块 模块化训练 策略优化

📋 核心要点

  1. 现有基于扩散模型的离线强化学习方法依赖于引导模块和扩散模型的联合训练,早期不准确的引导会产生噪声。
  2. 本文提出模块化训练方法,解耦引导模块和扩散模型,先独立训练引导模块,再用其引导扩散模型。
  3. 实验表明,该方法降低了分数方差,提高了样本效率和最终性能,并验证了跨模块的可迁移性。

📝 摘要(中文)

本文提出了一种模块化的扩散策略训练方法,旨在解耦引导模块和扩散模型,用于离线强化学习。现有方法通常联合训练引导模块和扩散模型,但在早期阶段,由于引导不准确,可能产生噪声学习信号。本文基于三个关键发现,提出了模块化训练方法:引导必要性,即引导的有效性随训练阶段和算法选择而变化;引导优先的扩散训练,即首先独立训练引导模块作为价值估计器,然后冻结该模块以引导扩散模型;跨模块可迁移性,即独立训练的引导模型可以在不同算法间复用。实验结果表明,该方法降低了分数方差,提高了样本效率和最终性能,为离线强化学习提供了一种模块化、可复用和可组合的训练流程。

🔬 方法详解

问题定义:现有基于扩散模型的离线强化学习方法通常采用联合训练的方式训练引导模块(用于提供奖励引导)和扩散模型。这种联合训练的缺点在于,在训练初期,引导模块的性能较差,提供的引导信号可能带有噪声,从而影响扩散模型的训练效果。此外,联合训练计算开销大,效率较低。

核心思路:本文的核心思路是将引导模块和扩散模型的训练解耦,采用模块化的训练方式。具体来说,首先独立训练引导模块,使其能够准确地估计状态-动作价值。然后,将训练好的引导模块固定,利用其提供的引导信号来训练扩散模型。这样可以避免早期不准确的引导信号对扩散模型的影响,并提高训练效率。

技术框架:该方法主要包含两个阶段:1) 引导模块训练阶段:使用离线数据集,采用现有的离线强化学习算法(如IDQL、DQL等)训练引导模块,使其能够准确地估计状态-动作价值。引导模块可以是一个神经网络,输入状态和动作,输出对应的价值估计。2) 扩散模型训练阶段:固定训练好的引导模块,利用其提供的奖励引导信号,训练扩散模型。扩散模型的目标是生成能够最大化累积奖励的动作序列。

关键创新:该方法最重要的创新点在于将引导模块和扩散模型的训练解耦,实现了模块化的训练方式。这种解耦使得可以独立地优化引导模块和扩散模型,避免了联合训练带来的问题。此外,该方法还探索了引导模块的跨模块可迁移性,即训练好的引导模块可以在不同的扩散模型或不同的离线强化学习算法之间复用。

关键设计:在引导模块训练阶段,可以使用不同的离线强化学习算法,如IDQL、DQL等。损失函数通常采用时序差分误差的均方误差。在扩散模型训练阶段,可以使用classifier-free guidance,将奖励引导信号融入到扩散模型的训练过程中。具体的网络结构可以根据具体任务进行选择,例如可以使用Transformer网络作为扩散模型的骨干网络。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,该方法在Bullet D4RL benchmark上取得了显著的性能提升。例如,通过使用独立训练的引导模型,可以将归一化分数方差降低86%。此外,实验还验证了引导模块的跨模块可迁移性,即使用IDQL训练的引导模块可以直接用于DQL,并达到基线水平的性能,无需额外训练。

🎯 应用场景

该研究成果可应用于各种需要离线强化学习的场景,例如机器人控制、自动驾驶、推荐系统等。通过模块化训练,可以更高效地训练策略,并实现策略的迁移和复用,降低开发成本,加速应用落地。此外,该方法为离线强化学习提供了一种新的训练范式,有望推动离线强化学习领域的发展。

📄 摘要(原文)

Classifier free guidance has shown strong potential in diffusion-based reinforcement learning. However, existing methods rely on joint training of the guidance module and the diffusion model, which can be suboptimal during the early stages when the guidance is inaccurate and provides noisy learning signals. In offline RL, guidance depends solely on offline data: observations, actions, and rewards, and is independent of the policy module's behavior, suggesting that joint training is not required. This paper proposes modular training methods that decouple the guidance module from the diffusion model, based on three key findings: Guidance Necessity: We explore how the effectiveness of guidance varies with the training stage and algorithm choice, uncovering the roles of guidance and diffusion. A lack of good guidance in the early stage presents an opportunity for optimization. Guidance-First Diffusion Training: We introduce a method where the guidance module is first trained independently as a value estimator, then frozen to guide the diffusion model using classifier-free reward guidance. This modularization reduces memory usage, improves computational efficiency, and enhances both sample efficiency and final performance. Cross-Module Transferability: Applying two independently trained guidance models, one during training and the other during inference, can significantly reduce normalized score variance (e.g., reducing IQR by 86%). We show that guidance modules trained with one algorithm (e.g., IDQL) can be directly reused with another (e.g., DQL), with no additional training required, demonstrating baseline-level performance as well as strong modularity and transferability. We provide theoretical justification and empirical validation on bullet D4RL benchmarks. Our findings suggest a new paradigm for offline RL: modular, reusable, and composable training pipelines.