Steering Your Diffusion Policy with Latent Space Reinforcement Learning
作者: Andrew Wagenmaker, Mitsuhiko Nakamoto, Yunchu Zhang, Seohong Park, Waleed Yagoub, Anusha Nagabandi, Abhishek Gupta, Sergey Levine · CoRL 2025
许可协议
- 开放获取
- 论文: CC BY 4.0
以下内容均依据论文全文,按其 CC BY 4.0 许可协议使用;其中的图已注明作者。
一句话总结
DSRL 通过训练一个选择输入噪声的小型 RL 策略来改进冻结的扩散或流策略,在数十到数百个回合内完成对真实机器人和 π0 的适应。
为何重要
团队可以就地改进一个已能工作的行为克隆策略,只需几十个真实回合和稀疏的成功奖励,而无需触碰其权重,这也使黑盒或仅提供 API 的模型变得可适应。在 π0 实验中,DSRL 在一块 8 GB 的 RTX 3070 上完成训练,而 π0 代码库说明 LoRA 微调至少需要 22.5 GB。
问题
行为克隆策略(如今通常是扩散策略)无法从自身的部署经验中学习,因此要修正较弱的性能,通常意味着要收集更多人类示范。RL 可以自主改进它们,但用 RL 微调扩散策略需要通过其多步去噪链进行反向传播,计算量大且数值不稳定,而标准 RL 所需的样本对真实机器人来说太多。
提出的方法
DSRL 将预训练的扩散策略(使用确定性的 DDIM 或流采样)视为环境的一部分,并学习一个单独的潜在噪声策略,由它选择初始噪声,而不是从 N(0, I) 中采样。这将适应过程转化为在一个“潜在动作 MDP”上的普通 RL,只需要基础策略的前向传播,而不需要其权重。噪声混叠变体 DSRL-NA 在真实动作上训练 critic,并将其蒸馏为噪声空间的 critic,因此可以利用离线数据,并利用映射到同一动作的众多噪声值。
关键见解
- 在确定性的 DDIM 或流采样下,输入噪声完全决定动作,因此选择噪声即可在不改变策略的情况下引导它。
- 由于基础策略只输出分布内的动作,在其噪声上进行优化无需显式惩罚项即可获得离线 RL 的保守性。
- 许多噪声值映射到几乎相同的动作;将动作空间的 critic 蒸馏到噪声空间的 critic 中可以利用这种混叠,所需样本约为在噪声空间上直接使用 SAC 的一半。
- 对于 π0 的 1,600 维分块动作空间,在整个 chunk 上重复同一个 32 维噪声向量,其表达能力已足以引导该策略。
- 只要基础策略表现出一定的解决任务的行为,引导就对不同规模、训练时长和数据质量的基础策略都有效。
主要结果
- 在 Robomimic 和 OpenAI Gym 上的在线适应达到了接近最优的性能,所需样本比 DPPO、IDQL、DQL、DIPO 和 QSM 至少少 5-10 倍。
- 在 10 个 OGBench 离线任务上,DSRL 在约一半任务上达到最先进水平,并较未引导的流策略有大幅提升。
- 真实 Franka 单任务(将方块放入碗中):在 3,500 个在线步(约 40 个回合)内成功率从 20% 提升到接近 100%;RLPD 无论有无人工干预都未能完成该任务。
- 真实 WidowX 250 搭配在 Bridge V2 上训练的多任务策略:在 100-150 个回合内,抓取放置、关抽屉和堆积木任务均有显著提升,而 RLPD 未能学会。
- π0 在 LIBERO 上:经过约 10,000 个在线样本后,成功率从约 20% 提升到约 100%;ResiP 和 V-GPS 未能显著改进它。
- 真实 Franka 搭配公开的 π0-DROID 检查点:在 80 和 65 个回合(约 11,000 和 10,000 步)内改进了烤面包机任务和把勺子放到盘子上的任务,作者称这是首次在真实世界中对 π0 进行 RL 微调。
局限
- 探索受基础策略的限制:无法保证每个扩散策略都可被引导,而非常狭窄的动作分布可能留下的选择太少(作者)。
- 目前还无法预先量化某个给定基础策略可以被改进多少(作者)。
- 与任何 RL 方法一样,它需要奖励信号、在线变体所需的在线 rollout 以及真实世界中的重置,这些都超出了模仿学习的要求(作者)。
- 真实世界任务使用人工定义的 0-1 成功奖励,且评估规模较小(WidowX 上每个任务 10-20 次试验)。
对领域的影响
入围 CoRL 2025 最佳论文 / 最佳学生论文奖决选,目前已被引用 242 次。论文提供了项目网站链接(diffusion-steering.github.io),并基于公开的 π0 检查点构建;获得许可的 PMLR 版本未注明代码仓库。
摘要
从人类示范中学习的机器人控制策略已在许多真实世界应用中取得了令人瞩目的成果。然而,当初始性能不令人满意时(这在新的开放世界场景中很常见),这类通过行为克隆(BC)学到的策略通常需要收集更多人类示范来进一步改进其行为——这是一个昂贵且耗时的过程。相比之下,强化学习(RL)有望实现自主的在线策略改进,但由于通常需要大量样本,往往难以实现这一点。在这项工作中,我们朝着通过高效的真实世界 RL 实现 BC 训练策略的快速自主适应迈出了几步。我们特别关注扩散策略——一种最先进的 BC 方法——并提出基于强化学习的扩散引导(DSRL):通过在 BC 策略的潜在噪声空间上运行 RL 来适应该策略。我们表明,DSRL 具有很高的样本效率,只需要对 BC 策略的黑盒访问,并能实现有效的真实世界自主策略改进。此外,DSRL 避免了与微调扩散策略相关的许多挑战,完全无需修改基础策略的权重。我们在仿真基准、真实世界机器人任务以及预训练通用策略的适应上展示了 DSRL,说明了它的样本效率以及在真实世界策略改进中的有效性能。
链接
- arXiv 2506.15799: https://arxiv.org/abs/2506.15799
- DOI 10.48550/arXiv.2506.15799: https://doi.org/10.48550/arXiv.2506.15799
- PDF: https://raw.githubusercontent.com/mlresearch/v305/main/assets/wagenmaker25a/wagenmaker25a.pdf
- 项目主页: https://diffusion-steering.github.io/