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/