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

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

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

通過回收模擬幀來提高訓(xùn)練效率

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

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

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è)計機會

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

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

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

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

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

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í)行一些最終沒有使用的代碼，所以相較于上文，它會浪費操作。

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

Brax 環(huán)境

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

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

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

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

性能基準(zhǔn)

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

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

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

分析 Brax 性能的另一種方法是評估其對在單個工作站上運行強化學(xué)習(xí)實驗所需時間的影響。在這里，我們對比了 Brax 在 Ant 與 OpenAI 基準(zhǔn)測試環(huán)境中訓(xù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

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

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

Brax 的優(yōu)化 PPO 與標(biāo)準(zhǔn) GPU 支持的 PPO 在學(xué)習(xí) 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 中訓(xùn)練，策略都具有相似的獎勵軌跡。

SAC

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

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

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

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

結(jié)論

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

我們的工作使快速、可擴展的 RL 和機器人研究變得更加容易，以前只能通過大型計算集群才能實現(xiàn)的研究現(xiàn)在可以在工作站上運行，或者托管在 Google Colaboratory 免費運行。我們的 Github 代碼庫不僅包括 Brax 模擬引擎，還包括大量用于快速訓(xùn)練的參考 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)實世界的復(fù)雜性，我們?nèi)匀辉陂_展探索柔性物體的物理研究。

柔性物體

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

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

編輯：jq