如何借助新的物理模擬引擎來加速強化學習

強化學習 （Reinforcement Learning） 是一種指導機器人在現(xiàn)實世界完成導航和執(zhí)行操作的熱門方法，其本身可以簡化并表示為剛性物體 ［1］（即受外力作用時不會變形的固體物理對象）之間的相互作用。

RL 通常會利用模擬，在有效的時間段內(nèi)收集訓練數(shù)據(jù)。在模擬過程中，我們會對任意數(shù)量的復雜對象做近似處理，將其視作由關(guān)節(jié)組成、且由致動器提供動力的許多剛體。但挑戰(zhàn)也應(yīng)運而生：即便是行走、使用工具或組裝積木等簡單的任務(wù)，RL 智能體也通常需要數(shù)百萬甚至數(shù)十億的模擬幀才能熟練地完成。

雖然通過回收模擬幀來提高訓練效率的方法已經(jīng)取得進展，但一些 RL 工具需仍通過模擬幀的生成分布到許多模擬器來回避這個問題。這些分布式模擬平臺已經(jīng)取得了亮眼表現(xiàn)，訓練速度也非?？?。但它們必須在擁有數(shù)千個 CPU 或 GPU 的計算集群上運行，往往大多數(shù)研究人員無法訪問這些 CPU 或 GPU。

通過回收模擬幀來提高訓練效率

https://ai.googleblog.com/2020/08/tackling-open-challenges-in-offline.html

在“Brax：用于大規(guī)模剛體模擬的微分物理引擎 （Brax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation）”中，我們提出了一種新的物理模擬引擎，該引擎可以在單個 TPU 或 GPU 下可提供與大型計算集群相媲美的性能。

該引擎旨在保障在單個加速器上高效并行運行數(shù)千個物理模擬和機器學習 （ML） 算法，以及在相互連接的加速器 pod 中無縫擴展數(shù)百萬個模擬。我們開源了引擎、參考 RL 算法和模擬環(huán)境，這些都可以通過 Colab 進行訪問。通過使用這個新的平臺，我們展現(xiàn)了比傳統(tǒng)工作站安裝快 100-1000 倍的訓練速度。

三種典型 RL 工作流。左邊顯示了一個標準的工作站流程：一臺配備 CPU 的機器，在這樣的環(huán)境中，訓練需要數(shù)小時或數(shù)天。中間顯示了一個標準的分布式模擬流程：通過將模擬移交給數(shù)千臺機器，訓練需要幾分鐘。右圖顯示了 Brax 流程：在單個 CPU/GPU 芯片上并行推進學習和大批量模擬

Brax：用于大規(guī)模剛體模擬的微分物理引擎

https://arxiv.org/abs/2106.13281

引擎、參考 RL 算法和模擬環(huán)境

https://github.com/google/brax

Colab

https://colab.sandbox.google.com/github/google/brax/blob/main/notebooks/training.ipynb

物理模擬引擎設(shè)計機會

剛體物理在電子游戲、機器人、分子動力學、生物力學、圖形和動畫等領(lǐng)域均有應(yīng)用。為了對這類系統(tǒng)進行準確建模，模擬器集成了來自重力、電機驅(qū)動、關(guān)節(jié)約束、物體碰撞等方面的力用于模擬物理系統(tǒng)根據(jù)時間所做出的動作。

模擬情形：三個球體、一個壁、兩個關(guān)節(jié)和一個致動器。系統(tǒng)會整合每個模擬時步中的力和力矩，以更新每個物體的位置、旋轉(zhuǎn)和速度

在仔細研究當今大多數(shù)物理模擬引擎的設(shè)計方式過后，不難看出，目前存在一些提高效率的潛在機會。正如我們上面所提到的，典型的機器人學習流水線會將單個學習器置于已有的、許多且并行模擬的緊密反饋中，但在分析這個架構(gòu)后，人們發(fā)現(xiàn)：

1. 這種架構(gòu)帶來了巨大的延遲瓶頸。由于數(shù)據(jù)必須通過數(shù)據(jù)中心內(nèi)的網(wǎng)絡(luò)進行傳輸，因此學習器必須等待 10，000 多納秒才能從模擬器中提取經(jīng)驗。如果這種經(jīng)驗在與學習器神經(jīng)網(wǎng)絡(luò)相同的設(shè)備上運行，延遲時間則會下降到不足 1 納秒。

2. 訓練智能體（一模擬步后隨之更新智能體的神經(jīng)網(wǎng)絡(luò)）所需的計算量遠不及打包數(shù)據(jù)（即在引擎內(nèi)編組數(shù)據(jù)，轉(zhuǎn)換為諸如 protobuf 之類的有線格式，然后進入 TCP 緩沖區(qū)，在學習器端撤消所有這些步驟）所需的計算量。

3. 每個模擬器中發(fā)生的計算非常相似，但并不完全相同。

Brax 設(shè)計

針對這些觀察結(jié)果，Brax 旨在使物理計算在其數(shù)千個并行環(huán)境中的每一個環(huán)境中保持完全相同，方法是確保模擬沒有分支（即模擬會因環(huán)境狀態(tài)而發(fā)生相應(yīng)變化的 “if” 邏輯）。

物理引擎的分支示例之一是在球和墻壁之間施加接觸力：系統(tǒng)將根據(jù)球是否接觸墻，執(zhí)行不同的代碼路徑。也就是說，如果 （if） 球接觸墻壁，則將執(zhí)行用于模擬球從墻壁反彈的單獨代碼。Brax 混合使用以下三種策略來避免分支：

用連續(xù)函數(shù)代替離散分支邏輯，例如使用符號距離函數(shù) （Signed distance function） 模擬球壁之間的接觸力。這種方法可以最大程度地提高效率。

在 JAX 的即時編譯期間評估分支。 許多基于環(huán)境靜態(tài)屬性的分支，比如兩個對象是否可能發(fā)生碰撞，系統(tǒng)可能會在模擬時間之前對其進行評估。

在模擬期間運行分支的兩側(cè)，但只選擇所需的結(jié)果。 因為這將執(zhí)行一些最終沒有使用的代碼，所以相較于上文，它會浪費操作。

在保證計算一致完全后，整個訓練體系結(jié)構(gòu)的復雜度就會降低，且更便于在單個 TPU 或 GPU 上執(zhí)行。這樣即可消除跨機器通信的計算開銷和延遲。在實踐中，對于相同的工作負載，這些更改將訓練成本降低了 100 -1000 倍。

Brax 環(huán)境

環(huán)境 （Environments） 是微小的封裝世界，定義了 RL 智能體要學習的任務(wù)。環(huán)境不僅包含模擬世界的方法，還包含如何觀察這個世界以及如何在該世界中定義目標等功能。

近年來出現(xiàn)了一些用于測試新的 RL 算法，并使用研究員熟悉的指標來評估這些算法影響的基準環(huán)境。Brax 內(nèi)含熱門 OpenAI Gym 中的 4 個即用型環(huán)境：Ant、HalfCheetah、Humanoid，以及 Reacher。

從左到右依次為：Ant、HalfCheetah、Humanoid 和 Reacher（均為適合 RL 環(huán)境的熱門基準環(huán)境）

Brax 還內(nèi)含三個新穎的環(huán)境：靈巧操作對象（機器人領(lǐng)域中的熱門挑戰(zhàn)），廣義運動（智能體前往放置在其周圍任意位置的目標），和工業(yè)機器人手臂的模擬。

性能基準

分析 Brax 性能的第一步是測試它模擬大批量環(huán)境的速度，這是學習器需要克服的關(guān)鍵瓶頸。只有這樣才能獲取足夠的經(jīng)驗來開展快速學習。

下面的這兩個圖顯示了 Brax 可以產(chǎn)生的實際步數(shù)（隨環(huán)境狀態(tài)變化而更新），它的任務(wù)是并行模擬越來越多的環(huán)境。左圖顯示，在 Brax 中，每秒步數(shù)與并行環(huán)境的數(shù)量成線性比例，僅在 10，000 個環(huán)境時遇到內(nèi)存帶寬瓶頸，這不僅足以訓練單個智能體，但適合訓練整個群體智能體。

第一，Brax 不僅在 TPU 上表現(xiàn)良好，且在高端 GPU 上也表現(xiàn)良好（參見 V100 和 P100 曲線），第二，通過利用 JAX 的設(shè)備并行原語，Brax 可在多個設(shè)備之間無縫擴展，實現(xiàn)每秒數(shù)億個物理步數(shù)（參見 TPUv3 8x8 曲線，即 64 個 TPUv3 芯片通過高速互連直接相互連接）。

分析 Brax 性能的另一種方法是評估其對在單個工作站上運行強化學習實驗所需時間的影響。在這里，我們對比了 Brax 在 Ant 與 OpenAI 基準測試環(huán)境中訓練，后者由 MuJoCo 物理引擎提供支持。

Ant

https://github.com/google/brax/blob/main/brax/envs/ant.py

OpenAI

https://github.com/openai/gym/blob/master/gym/envs/mujoco/ant.py

在下圖中，藍線代表已安裝的標準工作站，其中學習器在 GPU 上運行，模擬器在 CPU 上運行。我們看到通過訓練 Ant 按照合理的熟練度（y 軸上的分數(shù)為 4000）來運行，所需的時間從藍線的大約 3 小時下降到在加速器硬件上使用 Brax 時的大約 10 秒。

有趣的是，即使僅使用 CPU（灰線），Brax 的執(zhí)行速度也加快了一個數(shù)量級。這得益于處于同一進程中的學習器和模擬器。

Brax 的優(yōu)化 PPO 與標準 GPU 支持的 PPO 在學習 MuJoCo-Ant-v2 環(huán)境時，其執(zhí)行 1 千萬步時的表現(xiàn)對比情況。請注意，x 軸是執(zhí)行作業(yè)所需的實際時間（以秒為單位）。陰影區(qū)域表示超過 5 個副本的最低和最高性能種子，實線表示平均值

PPO

https://arxiv.org/abs/1707.06347

物理保真

設(shè)計與現(xiàn)實世界行為相匹配的模擬器是一個尚未解決的已知難題。盡管如此，將 Brax 與參考模擬器進行比較，至少確保其產(chǎn)生同樣有效的輸出。在這種情況下，我們再次比較 Brax 與 MuJoCo，后者因其模擬質(zhì)量而廣受好評。我們希望看到的是，在其他條件相同的情況下，無論是在 MuJoCo 還是 Brax 中訓練，策略都具有相似的獎勵軌跡。

SAC

https://arxiv.org/pdf/1812.05905.pdf

這些曲線表明，隨著兩個模擬器的獎勵以大致相同的速度增長，這兩個引擎計算物理的復雜性和難度都相當。由于兩條曲線都以大致相同的獎勵達到頂峰，所以我們相信，相同的常規(guī)物理限制適用于在任一模擬中盡其所能操作的智能體。

我們還可以測量 Brax 保持線性動量、角動量和能量守恒的能力。

這種物理模擬質(zhì)量的衡量標準由 MuJoCo 的作者率先提出，其目的是了解模擬在計算越來越大的時間步長時的偏航情況。在這里，Brax 的表現(xiàn)與其近鄰曲線相似。

結(jié)論

我們邀請研究人員通過在 Brax Training Colab 中訓練他們自己的策略，對 Brax 的物理保真度進行更定性的測量。學習到的軌跡與在 OpenAI Gym 中看到的軌跡非常相似。

我們的工作使快速、可擴展的 RL 和機器人研究變得更加容易，以前只能通過大型計算集群才能實現(xiàn)的研究現(xiàn)在可以在工作站上運行，或者托管在 Google Colaboratory 免費運行。我們的 Github 代碼庫不僅包括 Brax 模擬引擎，還包括大量用于快速訓練的參考 RL 算法。我們迫不及待地想看看 Brax 能夠?qū)崿F(xiàn)什么樣的新研究。

托管在 Google Colaboratory

https://colab.research.google.com/github/google/brax/blob/main/notebooks/training.ipynb

Github 代碼庫

https://github.com/google/brax

致謝

我們要感謝我們的論文共同作者：Anton Raichuk、Sertan Girgin、Igor Mordatch 和 Olivier Bachem。我們還要感謝 Erwin Coumans 關(guān)于構(gòu)建物理引擎的建議、Blake Hechtman 和 James Bradbury 為 JAX 和 XLA 提供的優(yōu)化幫助，以及 Luke Metz 和 Shane Gu 的建議。我們還要感謝 Vijay Sundaram、Wright Bagwell、Matt Leffler、Gavin Dodd、Brad Mckee 和 Logan Olson，感謝他們幫助孵化這個項目。

［1］ 由于現(xiàn)實世界的復雜性，我們?nèi)匀辉陂_展探索柔性物體的物理研究。

柔性物體

https://ai.googleblog.com/2021/05/learning-to-manipulate-deformable.html

發(fā)布人：Google Research 高級軟件工程師 C. Daniel Freeman 和資深軟件工程師 Erik Frey

編輯：jq