7 月 22 日消息，機器人要如何完成這樣一個動作?我們一般會基于強化學習，在仿真環境中進行模擬訓練。

這時，如果在一臺機器的 CPU 環境下進行模擬訓練，那么需要幾個小時到幾天。

但現在，只需一個 TPU/GPU，就能和數千個 CPU 或 GPU 的計算集群的速度一樣快，直接將所需時間縮短到幾分鐘，相當于將強化學習的速度提升了 1000 倍!

這就是來自谷歌的科學家們開發的物理模擬引擎 Brax。

三種策略避免邏輯分支

現在大多數的物理模擬引擎都是怎么設計的呢?

將重力、電機驅動、關節約束、物體碰撞等任務都整合在一個模擬器中，并行地進行多個模擬，以此來逼近現實中的運動系統。

這種情況下，每個模擬器中的計算都不相同，且數據必須在數據中心內通過網絡傳輸。

這種并行布局也就導致了較高的延遲時間 —— 即學習者可能需要超過 10000 納秒的等待時間，才能從模擬器中獲得經驗。

那么怎樣才能縮短這種延遲時間呢?

Brax 選擇通過避免模擬中的分支來保證數千個并行環境中的計算完全統一，進而降低整個訓練架構的復雜度。

直到復雜度降低到可以在單一的 TPU 或 GPU 上執行，跨機器通信的計算開銷就隨之降低，延遲也就能被有效消除。

主要分為以下三個方法：

連續函數替換離散分支邏輯

比如，在計算一個小球與墻壁之間的接觸力時，就產生了一個分支：

如果球接觸墻壁，就執行模擬球從墻壁反彈的獨立代碼;

否則，就執行其他代碼;

這里就可以通過符號距離函數來避免這種 if/else 的離散分支邏輯的產生。

使用 JAX 即時編譯中評估分支

在仿真時間之前評估基于環境靜態屬性的分支，例如兩個物體是否有可能發生碰撞。

在模擬中只選擇需要的分支結果

在使用了這三種策略之后，我們就得到了一個模擬由剛體、關節、執行器組成環境的物理引擎。

同時也是一種實現在這種環境中各類操作(如進化策略，直接軌跡優化等)的學習算法。

那么 Brax 的性能究竟如何呢?

速度最高提升 1000 倍

Brax 測試所用的基準是 OpenAI Gym 中 Ant、HalfCheetah、Humanoid、Reacher 四種。

同時也增加了三個新環境：包括對物理的靈巧操作、通用運動(例如前往周圍任何一個放置了物體的地點)、以及工業機器人手臂的模擬：

研究人員首先測試了 Brax 在并行模擬越來越多的環境時，可以產生多少次物理步驟(也即對環境狀態的更新)。

測試結果中的 TPUv3 8x8 曲線顯示，Brax 可以在多個設備之間進行無縫擴展，每秒可達到數億個物理步驟：

而不僅是在 TPU 上，從 V100 和 P100 曲線也能看出，Brax 在高端 GPU 上同樣表現出色。

然后就是 Brax 在單個工作站(workstation)上運行一個強化學習實驗所需要的時間。

在這里，研究人員將基于 Ant 基準環境訓練的 Brax 引擎與 MuJoCo 物理引擎做了對比：

可以看到，相對于 MuJoCo(藍線)所需的將近 3 小時時間，使用了 Brax 的加速器硬件最快只需要 10 秒。

使用 Brax，不僅能夠提高單核訓練的效率，還可以擴展到大規模的并行模擬訓練。

關鍵詞： 谷歌