Aligning Diffusion Behaviors with Q-functions for Efficient Continuous Control

📄 arXiv: 2407.09024v2 📥 PDF

作者: Huayu Chen, Kaiwen Zheng, Hang Su, Jun Zhu

分类: cs.LG

发布日期: 2024-07-12 (更新: 2024-10-30)

备注: Accepted by NeurIPS 2024

💡 一句话要点

提出EDA算法，通过与Q函数对齐的扩散模型实现高效连续控制

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 离线强化学习 扩散模型 Q函数 策略对齐 连续控制

📋 核心要点

  1. 现有离线强化学习方法难以有效利用大规模无奖励行为数据，泛化能力受限，且对标注数据需求高。
  2. EDA算法通过预训练扩散模型生成行为策略，并利用Q函数进行对齐微调，从而提升泛化性和效率。
  3. 实验表明，EDA在D4RL基准测试中超越现有方法，即使仅使用少量Q标注数据也能保持高性能。

📝 摘要（中文）

本文借鉴语言模型对齐的最新进展，将离线强化学习建模为一个两阶段优化问题：首先在无奖励的行为数据集上预训练表达能力强的生成策略，然后微调这些策略以对齐任务特定的标注（如Q值）。这种策略允许我们利用丰富多样的行为数据来增强泛化能力，并使用最少的标注快速适应下游任务。特别地，我们提出了用于解决连续控制问题的Efficient Diffusion Alignment (EDA)。EDA利用扩散模型进行行为建模。与以往方法不同，我们将扩散策略表示为标量神经网络关于动作输入的导数。这种表示至关重要，因为它能够直接计算扩散模型的密度，使其与现有的LLM对齐理论兼容。在策略微调期间，我们扩展了基于偏好的对齐方法（如Direct Preference Optimization (DPO)）以将扩散行为与连续Q函数对齐。在D4RL基准上的评估表明，EDA在整体性能上超过了所有基线方法。值得注意的是，即使在微调期间仅使用1%的Q标注数据，EDA仍能保持约95%的性能，并且仍然优于多个基线。

🔬 方法详解

问题定义：离线强化学习旨在利用预先收集的数据训练策略，而无需与环境交互。然而，现有方法通常难以有效利用大规模、多样化的无奖励行为数据，导致泛化能力不足。此外，许多方法需要大量的任务特定标注（例如奖励或Q值）才能进行有效的策略学习，限制了其在实际应用中的可行性。

核心思路：EDA的核心思路是将离线强化学习问题分解为两个阶段：首先，利用扩散模型在无奖励行为数据集上预训练一个表达能力强的生成策略。然后，通过与任务特定的Q函数进行对齐，对预训练的策略进行微调。这种两阶段方法允许模型从大量无标注数据中学习通用的行为模式，并仅使用少量标注数据快速适应特定任务。

技术框架：EDA的整体框架包括两个主要阶段：1) 行为建模阶段：使用扩散模型对行为数据进行建模，学习行为的概率分布。关键在于将扩散策略表示为标量神经网络关于动作输入的导数，从而实现直接的密度计算。2) 策略对齐阶段：利用扩展的偏好对齐方法（类似于DPO），将扩散模型的行为与Q函数进行对齐。具体来说，通过优化一个目标函数，使得策略生成的动作的Q值尽可能高。

关键创新：EDA的关键创新在于将扩散模型与Q函数对齐的思想，以及将扩散策略表示为标量神经网络导数的方法。这种表示方式使得扩散模型能够直接计算密度，从而与现有的LLM对齐理论兼容，并能够利用偏好对齐方法进行策略微调。此外，EDA能够有效利用大规模无奖励数据，并仅需少量标注数据即可实现高性能。

关键设计：在行为建模阶段，扩散模型采用基于能量的表示，其中能量函数由一个神经网络表示。扩散过程通过逐步添加噪声将行为数据转换为噪声数据，然后训练模型学习逆向过程，即从噪声数据恢复原始行为数据。在策略对齐阶段，使用类似于DPO的损失函数，该损失函数基于Q函数的差异来调整策略的参数。具体来说，损失函数鼓励策略生成具有更高Q值的动作，并惩罚生成具有较低Q值的动作。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，EDA在D4RL基准测试中取得了显著的性能提升，超过了所有基线方法。特别是在halfcheetah-medium-replay数据集上，EDA的平均得分达到了68.5，显著高于其他方法。更重要的是，即使仅使用1%的Q标注数据进行微调，EDA仍然能够保持约95%的性能，并优于多个基线，证明了其高效性和数据利用率。

🎯 应用场景

EDA算法具有广泛的应用前景，例如机器人控制、自动驾驶、游戏AI等领域。通过利用大量离线数据和少量标注数据，EDA可以快速训练出高性能的控制策略，从而降低开发成本和时间。此外，EDA的泛化能力使其能够适应不同的环境和任务，具有很强的实用价值。

📄 摘要（原文）

Drawing upon recent advances in language model alignment, we formulate offline Reinforcement Learning as a two-stage optimization problem: First pretraining expressive generative policies on reward-free behavior datasets, then fine-tuning these policies to align with task-specific annotations like Q-values. This strategy allows us to leverage abundant and diverse behavior data to enhance generalization and enable rapid adaptation to downstream tasks using minimal annotations. In particular, we introduce Efficient Diffusion Alignment (EDA) for solving continuous control problems. EDA utilizes diffusion models for behavior modeling. However, unlike previous approaches, we represent diffusion policies as the derivative of a scalar neural network with respect to action inputs. This representation is critical because it enables direct density calculation for diffusion models, making them compatible with existing LLM alignment theories. During policy fine-tuning, we extend preference-based alignment methods like Direct Preference Optimization (DPO) to align diffusion behaviors with continuous Q-functions. Our evaluation on the D4RL benchmark shows that EDA exceeds all baseline methods in overall performance. Notably, EDA maintains about 95\% of performance and still outperforms several baselines given only 1\% of Q-labelled data during fine-tuning.