機器人要如何完成這樣一個動作?我們一般會基於強化學習,在仿真環境中進行模擬訓練。這時,如果在一台機器的CPU環境下進行模擬訓練,那麼需要幾個小時到幾天。但現在,只需 一個TPU/GPU,就能和數千個CPU或GPU的計算集群的速度一樣快,直接將所需時間縮短到幾分鐘!
博雯 發自 凹非寺
量子位 報道 | 公眾號 QbitAI
相當於將強化學習的速度提升了 1000倍!
這就是來自Google的科學家們開發的物理模擬引擎 Brax。
三種策略避免邏輯分支
現在大多數的物理模擬引擎都是怎麼設計的呢?
將重力、電機驅動、關節約束、物體碰撞等任務都整合在一個模擬器中, 并行地進行多個模擬,以此來逼近現實中的運動系統。
△對於每個模擬時間步長,力和力矩被整合在一起
這種情況下,每個模擬器中的計算都不相同,且數據必須在數據中心內通過網絡傳輸。
這種并行布局也就導致了較高的延遲時間——即學習者可能需要超過 10000納秒的等待時間,才能從模擬器中獲得經驗。
那麼怎樣才能縮短這種延遲時間呢?
Brax選擇通過 避免模擬中的分支來保證數千個并行環境中的計算完全統一,進而降低整個訓練架構的複雜度。
直到複雜度降低到可以在單一的TPU或GPU上執行,跨機器通信的計算開銷就隨之降低,延遲也就能被有效消除。
主要分為以下三個方法:
-
連續函數替換離散分支邏輯
比如,在計算一個小球與牆壁之間的接觸力時,就產生了一個分支:
如果球接觸牆壁,就執行模擬球從牆壁反彈的獨立代碼;
否則,就執行其他代碼;
這裡就可以通過符號距離函數來避免這種if/else的離散分支邏輯的產生。
-
使用JAX即時編譯中評估分支
在仿真時間之前評估基於環境靜態屬性的分支,例如兩個物體是否有可能發生碰撞。
-
在模擬中只選擇需要的分支結果
在使用了這三種策略之後,我們就得到了一個模擬由剛體、關節、執行器組成環境的物理引擎。
同時也是一種實現在這種環境中各類操作(如進化策略,直接軌跡優化等)的學習算法。
那麼Brax的性能究竟如何呢?
速度最高提升1000倍
Brax測試所用的基準是OpenAI Gym中Ant、HalfCheetah、Humanoid、Reacher四種。
同時也增加了三個新環境:包括對物理的靈巧操作、通用運動(例如前往周圍任何一個放置了物體的地點)、以及工業機器人手臂的模擬:
研究人員首先測試了Brax在并行模擬越來越多的環境時,可以產生多少次物理步驟(也即對環境狀態的更新)。
測試結果中的TPUv3 8×8曲線顯示,Brax可以在多個設備之間進行無縫擴展,每秒可達到 數億個物理步驟:
而不僅是在TPU上,從V100和P100曲線也能看出,Brax在高端 GPU上同樣表現出色。
然後就是Brax在單個工作站 (workstation)上運行一個強化學習實驗所需要的時間。
在這裡,研究人員將基於Ant基準環境訓練的Brax引擎與MuJoCo物理引擎做了對比:
可以看到,相對於MuJoCo(藍線)所需的將近3小時時間,使用了Brax的加速器硬件最快只需要 10秒。
使用Brax,不僅能夠提高單核訓練的效率,還可以擴展到大規模的并行模擬訓練。
論文地址:
https://arxiv.org/abs/2106.13281
下載:
https://github.com/google/brax
參考鏈接:
https://ai.googleblog.com/2021/07/speeding-up-reinforcement-learning-with.html