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

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

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

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

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

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

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

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

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

https://arxiv.org/abs/2106.13281

引擎、參考 RL 算法和模擬環境

https://github.com/google/brax

Colab

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

物理模擬引擎設計機會

剛體物理在電子游戲、機器人、分子動力學、生物力學、圖形和動畫等領域均有應用。為了對這類系統進行準確建模，模擬器集成了來自重力、電機驅動、關節約束、物體碰撞等方面的力用于模擬物理系統根據時間所做出的動作。

模擬情形：三個球體、一個壁、兩個關節和一個致動器。系統會整合每個模擬時步中的力和力矩，以更新每個物體的位置、旋轉和速度

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

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

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

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

Brax 設計

針對這些觀察結果，Brax 旨在使物理計算在其數千個并行環境中的每一個環境中保持完全相同，方法是確保模擬沒有分支（即模擬會因環境狀態而發生相應變化的 “if” 邏輯）。

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

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

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

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

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

Brax 環境

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

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

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

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

性能基準

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

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

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

分析 Brax 性能的另一種方法是評估其對在單個工作站上運行強化學習實驗所需時間的影響。在這里，我們對比了 Brax 在 Ant 與 OpenAI 基準測試環境中訓練，后者由 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 軸上的分數為 4000）來運行，所需的時間從藍線的大約 3 小時下降到在加速器硬件上使用 Brax 時的大約 10 秒。

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

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

PPO

https://arxiv.org/abs/1707.06347

物理保真

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

SAC

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

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

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

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

結論

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

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

托管在 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 關于構建物理引擎的建議、Blake Hechtman 和 James Bradbury 為 JAX 和 XLA 提供的優化幫助，以及 Luke Metz 和 Shane Gu 的建議。我們還要感謝 Vijay Sundaram、Wright Bagwell、Matt Leffler、Gavin Dodd、Brad Mckee 和 Logan Olson，感謝他們幫助孵化這個項目。

［1］ 由于現實世界的復雜性，我們仍然在開展探索柔性物體的物理研究。

柔性物體

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

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

編輯：jq