如何借助新的物理模擬引擎來(lái)加速?gòu)?qiáng)化學(xué)習(xí)

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

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

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

通過(guò)回收模擬幀來(lái)提高訓(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）”中，我們提出了一種新的物理模擬引擎，該引擎可以在單個(gè) TPU 或 GPU 下可提供與大型計(jì)算集群相媲美的性能。

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

三種典型 RL 工作流。左邊顯示了一個(gè)標(biāo)準(zhǔn)的工作站流程：一臺(tái)配備 CPU 的機(jī)器，在這樣的環(huán)境中，訓(xùn)練需要數(shù)小時(shí)或數(shù)天。中間顯示了一個(gè)標(biāo)準(zhǔn)的分布式模擬流程：通過(guò)將模擬移交給數(shù)千臺(tái)機(jī)器，訓(xùn)練需要幾分鐘。右圖顯示了 Brax 流程：在單個(gè) CPU/GPU 芯片上并行推進(jìn)學(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è)計(jì)機(jī)會(huì)

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

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

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

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

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

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

Brax 設(shè)計(jì)

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

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

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

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

在模擬期間運(yùn)行分支的兩側(cè)，但只選擇所需的結(jié)果。 因?yàn)檫@將執(zhí)行一些最終沒(méi)有使用的代碼，所以相較于上文，它會(huì)浪費(fèi)操作。

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

Brax 環(huán)境

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

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

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

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

性能基準(zhǔn)

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

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

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

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

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

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

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

PPO

https://arxiv.org/abs/1707.06347

物理保真

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

SAC

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

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

我們還可以測(cè)量 Brax 保持線性動(dòng)量、角動(dòng)量和能量守恒的能力。

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

結(jié)論

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

我們的工作使快速、可擴(kuò)展的 RL 和機(jī)器人研究變得更加容易，以前只能通過(guò)大型計(jì)算集群才能實(shí)現(xiàn)的研究現(xiàn)在可以在工作站上運(yùn)行，或者托管在 Google Colaboratory 免費(fèi)運(yùn)行。我們的 Github 代碼庫(kù)不僅包括 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 代碼庫(kù)

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，感謝他們幫助孵化這個(gè)項(xiàng)目。

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

柔性物體

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

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

編輯：jq