【论文阅读】Steering Your Diffusion Policy with Latent Space Reinforcement Learning

1 团队与发表时间

• 团队：主要由 加州大学伯克利分校（UC Berkeley） 的研究人员组成，包括 Mitsuhiko Nakamoto, Andrew Wagenmaker, Sergey Levine 等，此外还有来自华盛顿大学和亚马逊（Amazon）的合作者。
• 发表时间：该论文的 arXiv 版本更新于 2025 年 6 月 25 日。

2. 问题背景与核心思路

• 问题背景：虽然基于行为克隆（BC）的扩散策略（Diffusion Policy）在机器人领域表现出色，但当预训练模型性能不足时，通常需要昂贵的人工演示来改进。传统的强化学习（RL）虽然能自主改进，但在高维动作空间中样本效率极低，且容易破坏扩散模型学到的先验分布，甚至造成训练不稳定。

  • 多步去噪导致的梯度爆炸/消失：扩散模型生成动作不是“一步到位”的，而是经过几十步（如 50 步）去噪。如果你想通过动作a aa的奖励来反向传播更新模型权重，梯度必须链式传导 50 层网络。这就像训练一个极深的循环神经网络（RNN），梯度在传导过程中极易失控，导致训练瞬间崩盘。

• 核心思路：提出 DSRL。其核心想法是不在原始动作空间运行 RL，而是在扩散模型的“潜噪声空间（Latent-noise Space）”运行 RL。通过调整每一步去噪时加入的噪声（即“转向”），在保留原始行为分布的同时，引导策略向高奖励区域偏移。

3. 具体做法

3.1 模型设计

• 潜空间引导（Latent Steering）：将预训练好的扩散策略视为一个“黑盒”。扩散策略通过多步去噪（从高斯噪声ϵ \epsilonϵ迭代到动作a aa）生成动作。

  • 做法：

    • 原本的扩散模型（比如π 0 \pi_0π0​或其他 Diffusion Policy）被当作一个固定的“基座”。
    • 训练对象：DSRL 额外训练的是一个轻量级的 Actor 网络（通常是一个简单的多层感知机 MLP）。这个 Actor 的任务是根据当前状态s ss，预测一个噪声位移（Noise Shift）Δ ϵ \Delta \epsilonΔϵ。
    • 运行逻辑：在去噪过程中，原本输入给模型的纯高斯噪声ϵ \epsilonϵ会被替换为ϵ + Δ ϵ \epsilon + \Delta \epsilonϵ+Δϵ。扩散模型依然按照原来的权重运行，但因为输入的“种子”变了，最终生成的动作就会向高奖励区域偏移。

  • 架构：使用了基于 Soft Actor-Critic (SAC) 框架的结构，包含 Actor 网络（预测噪声位移）和 Critic 网络（评估状态-噪声对的价值）。

  • 优势：

    • 防止先验崩溃（Maintaining Priors）：扩散模型里蕴含了大量从人类演示中学到的平滑动作先验。如果直接用 RL 微调权重，非常容易出现“灾难性遗忘”，导致机器人动作变得抖动、不自然。冻结权重能确保机器人永远在“人类可能的动作分布”附近进行微调。
    • 极高的样本效率：更新一个 8B 模型的参数需要海量数据和算力。而 DSRL 只训练一个极小的 Actor 网络（用来预测噪声偏移），这使得它在 20-40 次真机尝试内就能收敛，这在全参数微调中几乎是不可能的。
    • 先天的保守优势：无论你给这个扩散模型输入什么样的初始噪声w ww（哪怕是奇形怪状的w ww），这个被冻结的“黑盒”模型最终吐出来的动作a aa几乎总是符合人类演示风格的、在分布内（In-distribution）的动作。因此噪声w ww的探索是无穷的，而策略并不会出分布。
      

3.2 Loss 设计

• Actor Loss：基于 SAC 的目标函数，包含最大化预期奖励和熵正则项。通过最小化Q ( s , Δ ϵ ) Q(s, \Delta \epsilon)Q(s,Δϵ)的负值来优化，同时约束Δ ϵ \Delta \epsilonΔϵ的大小，以防偏离原始分布太远。
• Critic Loss：使用标准的时间差分（TD）误差来训练Q QQ函数（Critic），评估在当前状态下采用特定噪声转向后的长期价值。
• 软约束：通过 KL 散度或熵正则化，确保微调后的策略不会完全丧失预训练模型中蕴含的人类演示先验。

3.3 数据设计

• 离线预训练数据：利用已有的行为克隆（BC）数据集（如人类演示）来初始化扩散策略。
• 在线交互数据：在强化学习阶段，机器人通过执行带有噪声转向的动作与环境交互，收集新的三元组( s , Δ ϵ , r ) (s, \Delta \epsilon, r)(s,Δϵ,r)用于更新 Actor 和 Critic。
• 样本效率：由于是在低维且分布简单的噪声空间优化，DSRL 相比于直接在连续动作空间搜索，所需的数据量大幅减少。

4 实验效果

• 模拟环境（Libero, Aloha）：在多个标准机器人基准测试中，DSRL 在极少的样本下（通常只需几十个或几百个 episodes）就能显著提升预训练策略的成功率。
• 真机实验：在“打开烤箱”和“放勺子到盘子里”等任务中，通过与π 0 \pi_0π0​（Physical Intelligence 的基础模型）结合，DSRL 仅需 20-40 次真机交互 就能将成功率从接近 0% 提升到 80% 以上。实验证明其样本效率比直接在动作空间运行 SAC 高出数倍。

5 结论

DSRL 证明了在扩散模型的潜噪声空间进行强化学习是实现快速、自主策略改进的有效途径。它解决了扩散策略难以进行样本高效微调的难题，实现了“黑盒”式微调（无需深入修改扩散模型内部架构）。

6 扩散模型进行RL的方案

6.1 纯离线设置 (Purely Offline Setting)

在离线 RL 中，我们只有一份历史数据集（包含状态、动作和奖励），目标是从中学习一个能获得高奖励的策略。

• 按价值加权的 BC (Weighted BC by Value):

  • 方法：这种方法不改变扩散模型的基本训练方式，但在损失函数上加权。给那些高奖励（高Q QQ值或V VV值）的样本更大的权重，给表现差的样本小权重。
  • 本质：它是行为克隆（BC）的变体。模型依然在模仿，但它被告知“多模仿好学生，少模仿坏学生”。

• 直接最大化奖励 (Directly maximizing a reward):

  • 方法：不只是模仿数据，而是直接将R RR（奖励）或Q QQ（价值）作为目标函数。通过对扩散策略的参数求导，让模型生成的动作序列朝着奖励更高的方向偏移。
  • 难点：扩散模型是多步去噪，求导链条非常长，计算开销大且容易梯度爆炸。

• 拒绝采样 (Rejection Sampling):

  • 方法：先用预训练好的扩散策略生成一大堆候选动作（比如生成 10 个），然后用一个学好的Q QQ函数（评分器）对这 10 个动作打分，最后只执行分最高的那一个。
  • 本质：这是一种“后处理”策略，不需要重新训练模型，但在推理（Inference）时计算量翻倍。

6.2 在线设置 (Online Setting)

在线 RL 允许机器人与环境实时交互，根据反馈不断调整策略。

• 基于 PPO 微调前几步去噪 (Finetune first few steps with PPO):
  • 方法：扩散模型去噪通常有几十步。有些方法只用经典的 PPO 算法去更新前几步（高噪声阶段）的神经网络参数。
  • 理由：早期步决定了动作的大致轮廓，后期步只是修饰细节。这种方法试图通过只动“大方向”来降低 RL 的学习难度。
• 匹配 Q 函数的 Score (Matching the score of the Q-function):
  • 方法：数学上，扩散模型学习的是概率分布的“分值（Score）”。如果能学到一个奖励函数的分值∇ log ⁡ R \nabla \log R∇logR，并将其加到扩散模型的梯度中，就能把动作推向高奖励区。
  • 本质：通过数学推导，将“寻找高奖励动作”转化成“沿着奖励梯度的方向去噪”。
• 迭代式 BC (Iterative BC / Best-of-N):
  • 方法：机器人先去跑，选出表现好的轨迹（即那些Q QQ值大的），把这些轨迹放回训练集，再次进行行为克隆（BC）。
  • 本质：自我进化。不断重复“尝试 -> 筛选好样本 -> 重新模仿好样本”的过程。

6.3 残差策略 (Residual Policy)

• 方法：保持预训练好的扩散策略（a d i f f a_{diff}adiff​）完全不动，另外训练一个轻量级的神经网络（残差网络）输出一个修正值Δ a \Delta aΔa。最终执行的动作是a = a d i f f + Δ a a = a_{diff} + \Delta aa=adiff​+Δa。
• 本质：给扩散模型打一个“补丁”，用 RL 专门学习这个补丁。