RL訓練太慢?UC伯克利重磅開源FastTD3:并行+大批次+分布Critic,人形機器人3小時搞定

論文鏈接:??https://arxiv.org/pdf/2505.22642??
項目鏈接:?https://younggyo.me/fast_td3/??

亮點直擊
- 提出了FastTD3,一種簡單、快速且高效的強化學習算法,能夠解決多種運動和操作任務,而之前的RL算法需要數十小時才能完成或無法解決。
- 可以通過一個極其簡單的方案實現:使用大批次更新、并行仿真、分布式強化學習以及精心調優的超參數來訓練TD3智能體。展示了各種設計選擇的有效性。
- 發布了一個易于使用的開源FastTD3實現,以加速機器人領域的強化學習研究。該實現支持流行的測試平臺,如HumanoidBench、IsaacLab和MuJoCo Playground。
總結速覽
解決的問題
- 訓練速度慢:現有RL算法(如PPO、PQL等)在復雜機器人任務(如HumanoidBench、IsaacLab)中訓練時間過長(如48小時未收斂)。
- 實現復雜性高:部分高效算法(如PQL)依賴異步并行等復雜設計,難以廣泛應用。
- 樣本效率與穩定性的權衡:PPO等on-policy算法樣本效率低,而off-policy算法(如TD3)雖高效但訓練不穩定或速度不足。
效果展示
FastTD3解決了MuJoCo游樂場的一系列任務,實現了與PPO相似或更快的掛鐘時間:

FastTD3在單個GPU上在不到3小時的時間內解決了HumanoidBench中的各種任務,而之前的算法即使經過數十小時的訓練也很困難。然而可以看到,由于獎勵設計的局限性,行為非常不自然:

證明了FastTD3可用于模擬到真實的傳輸。使用非常快速和方便的MuJoCo Playground在模擬中訓練Booster T1策略,并使用Booster Gym將其轉移到真實的Booster T1機器人上:

不同的 RL 算法可能需要不同的獎勵函數:

提出的方案
- 并行仿真(Parallel Simulation):加速數據收集。
- 大批次更新(Large-Batch Updates):提升訓練穩定性。
- 分布化Critic(Distributional Critic):改進價值估計。
- 精細調參:針對機器人任務優化的超參數。
應用的技術
- 算法基礎:TD3(雙延遲確定性策略梯度)的off-policy框架。
- 加速技術:GPU并行化仿真(單A100實現)、大批次梯度更新。
- 性能優化:分布化Critic、簡化異步流程。
達到的效果
- 訓練速度:在HumanoidBench等任務中,3小時內完成訓練(對比PPO更快,尤其在復雜地形任務中)。
- 穩定性:訓練過程穩定,無需復雜調參。
- 易用性:提供輕量級PyTorch實現,支持快速部署和擴展(兼容HumanoidBench/IsaacLab/MuJoCo Playground)。
- 兼容性:與前沿RL研究(如SR-SAC、TDMPC2等)正交,可靈活集成新技術。
FastTD3:面向人形機器人控制的簡單、快速、高效強化學習算法
FastTD3是基于雙延遲深度確定性策略梯度算法的高性能變體,專為復雜機器人任務優化。這些優化借鑒了Li等人的研究發現:并行仿真、大批次訓練和分布化評論家對離線策略強化學習算法的性能提升至關重要。
設計選擇
并行環境
與Li等人的觀察一致,發現大規模并行環境能顯著加速TD3訓練。假設:確定性策略梯度算法與并行仿真的結合具有特殊優勢,因為并行環境產生的隨機性增加了數據分布的多樣性。這使得TD3既能充分發揮其價值函數高效利用的優勢,又能緩解探索能力不足的缺陷。
大批次訓練
實驗表明,使用32,768的超大批次規模訓練FastTD3智能體效果顯著。推測:在并行環境下,大批次更新通過確保每次梯度更新的數據多樣性,為評論家提供了更穩定的學習信號。否則,除非保持高更新數據比(UTD),否則大量數據可能永遠不被智能體利用。雖然增大批次規模會增加單次更新的耗時,但由于訓練效率提升,總體訓練時間往往反而縮短。
分布化強化學習

截斷雙Q學習(CDQ)
Nauman等人指出,當結合層歸一化時,使用Q值均值比CDQ采用的Q值最小值效果更好。實驗顯示,在沒有層歸一化時趨勢相反——CDQ仍是關鍵設計選擇,采用最小值通常表現更優。這表明CDQ仍是需要按任務調節的重要超參數。
網絡架構
采用1024-512-256單元數的多層感知機(MLP),行動者采用512-256-128單元數。實驗發現,更小的模型會降低時間效率和樣本效率。雖然嘗試過類似BRO或Simba的殘差連接與層歸一化,但這些設計反而會拖慢訓練且無顯著收益。推測:并行仿真和大批次訓練提供的數據多樣性降低了更新的"離線策略性",從而緩解了自助法、函數逼近和離線策略學習"致命三角"帶來的不穩定性。因此即使沒有殘差連接或層歸一化等穩定結構,訓練過程仍能保持穩定。
探索噪聲調度

更新數據比(UTD)
已有研究表明:提高UTD(即每環境步長的梯度更新次數)通常需要額外技術或架構調整。但FastTD3在標準3層MLP(無歸一化)下,樣本效率隨UTD提升而改善,僅需付出更長的訓練耗時。這歸因于FastTD3工作在極低UTD(通常每128-4096個并行環境步長僅2/4/8次更新),降低了高UTD導致的早期過擬合風險。
經驗回放池大小

實現細節
并行環境
對于IsaacLab和MuJoCo Playground,直接使用它們原生的并行仿真支持。然而,HumanoidBench不支持基于 GPU 的并行化,因此使用Stable Baselines3的 ??SubprocVecEnv??。HumanoidBench 的默認配置會為每個仿真啟動一個基于 GPU 的渲染器,這使得運行超過 100 個環境變得困難。已向 HumanoidBench 提交了一個Pull Request,添加了禁用默認 GPU 渲染器的支持,并已合并到主分支。

環境封裝器為了構建一個支持不同測試套件(各自采用不同配置)的易用代碼庫,為每個套件開發或使用了封裝器:
- MuJoCo Playground:使用原生?
?RSLRLBraxWrapper???,將 JAX 張量轉換為 PyTorch 張量,并遵循RSL-RL的 API。由于該封裝器不支持在重置環境前保存最終觀測值,在一個獨立的分支中實現了這一功能,并計劃將其合并到主倉庫。 - IsaacLab:由于原生支持 PyTorch,實現了一個符合RSL-RL API的簡單封裝器。目前,本文的實現不支持訓練期間的渲染,因為IsaacLab不允許并發運行多個仿真。
- HumanoidBench:開發了一個遵循RSL-RL API的封裝器,并將 NumPy 數組轉換為 PyTorch 張量。
非對稱 Actor-Critic
對于IsaacLab和MuJoCo Playground(它們通常為評論家網絡提供特權狀態),實現了非對稱 Actor-Critic的支持。
AMP 與 ??torch.compile??
盡管基于JAX的 RL 實現在近年來因速度優勢而流行,但仍選擇PyTorch作為基礎,因其簡單性和靈活性。在單塊A100 GPU上,使用AMP(自動混合精度)和bfloat16可將訓練速度提升高達40%,且未觀察到不穩定性。還基于LeanRL 支持 ??torch.compile???,實驗顯示其可帶來35%的加速。當同時使用AMP和 ??torch.compile?? 時,訓練速度最高可提升70%。
經驗回放池
沒有固定全局回放池大小,而是設定為。這種方式能更好地解耦回放池大小與并行環境數量的影響。例如:
- 如果全局回放池固定為100萬,任務片段長度為1000,而用戶將并行環境數從1000增加到2000,則回放池只能保存每個環境軌跡的一半(在500 步后就會開始丟棄早期樣本),這可能影響性能。
- 但如果為每個環境設定 **N=1000,則回放池能完整保存軌跡,不受并行環境數量的影響。
由于專注于非視覺領域,所有數據均存儲在GPU上,以避免CPU-GPU 數據傳輸的開銷。
實驗
實驗設置
對于DreamerV3、SAC和TDMPC2基線算法,使用 HumanoidBench 代碼庫中提供的三次運行學習曲線。由于每次訓練耗時48 小時,通過插值計算實際時間戳來繪制曲線。對于SimbaV2,使用官方代碼庫在涉及靈巧手的任務上進行實驗。SimbaV2 采用單次運行結果,未來計劃增加更多實驗。所有 FastTD3 結果均為三次運行的平均值。實驗在配備單塊 NVIDIA A100 GPU和16 CPU 核心的云實例上完成。
實驗結果
在下圖3中提供了每個套件的所有任務的匯總結果,在下圖4中提供了選定任務集的單獨結果。在下圖5和圖6中提供了廣泛的實驗結果,研究了各種設計選擇的影響。




不同RL算法需要不同的獎勵函數
在 MuJoCo Playground 訓練人形機器人運動策略時,發現PPO和FastTD3即使使用相同獎勵函數也會產生顯著不同的步態(見下圖7a-b)。推測:現有獎勵函數通常針對 PPO 調優,不同算法可能需要不同的獎勵結構來產生理想行為。為此專門為 FastTD3 調整了獎勵函數——增加懲罰項強度。得益于 FastTD3 的快速訓練,該調參過程非常高效。如圖7c 所示,調整后的獎勵使 FastTD3 學習到比圖7a 更穩定美觀的步態。而將 FastTD3 調優的獎勵用于 PPO 訓練時,卻會導致步態緩慢不可用(圖7d)。這表明標準指標——片段回報——可能無法有效衡量策略的實際可用性。

FastSAC 實驗
為驗證方法通用性,本文將 FastTD3 方案遷移至 SAC開發出FastSAC。實驗顯示 FastSAC 顯著快于原始 SAC(下圖8),但訓練穩定性較差,推測源于高維動作空間中動作熵最大化的固有難度。鑒于 SimbaV2在主實驗中明顯快于原始 SAC,未來可將此類最新進展融入 FastTD3/FastSAC。

FastTD3 的仿真到現實遷移
下圖1的仿真到現實實驗使用Booster Gym ,該平臺支持12自由度 Booster T1 人形機器人控制(固定手臂/腰部/朝向)。為簡化流程,將 Booster Gym 的機器人配置和獎勵函數移植到原本支持23自由度控制的 MuJoCo Playground。相比基于 IsaacGym 的 Booster Gym,在 MuJoCo Playground 中訓練 FastTD3 能大幅簡化和加速迭代周期。

討論
本文提出了FastTD3——一種簡潔、高效且性能優異的強化學習算法,該算法成功解決了HumanoidBench、IsaacLab和MuJoCo Playground中的各類運動控制與機械操作任務。研究表明,在未引入新架構或訓練技術的前提下,僅通過精心調優的超參數結合簡潔算法,即可為復雜機器人任務構建出人意料的強基線。隨本報告同步開源的FastTD3實現具有輕量化、易用化特點,提供用戶友好功能與預設超參數配置。
需要特別說明的是,本研究目標并非宣稱算法創新性或優越性。本文的方法直接建立在學界已有洞見之上:并行Q學習(PQL)揭示了如何通過大規模并行仿真擴展離策略強化學習,后續的并行Q網絡(PQN)在離散控制中得出了相似結論。Raffin與Shukla同樣證明SAC算法通過并行仿真與精細超參數調優可實現高效擴展。本工作旨在將這些洞見提煉為簡潔算法,提供全面的設計選擇實驗分析,并發布開箱即用的實現方案。
本文對多個未來方向充滿期待:首先,FastTD3與近期強化學習進展具有正交性,這些改進可輕松集成以推動技術前沿,預計整合過程將高效順暢;其次,FastTD3作為離策略算法,非常適用于人形機器人控制的示范驅動強化學習場景,以及通過真實環境交互微調仿真訓練策略;最后,FastTD3的快速迭代特性可賦能基于語言模型生成獎勵的逆向強化學習框架,為長期困擾人形控制領域的獎勵設計難題提供新思路。
本文轉自AI生成未來 ,作者:AI生成未來

















