<blockquote>
博雯 發自 凹非寺
量子位 報道 | 公眾號 QbitAI
</blockquote>
機器人要如何完成這樣一個動作？
<div><img src="https://imageproxy.pbkrs.com/https://p9.toutiaoimg.com/img/pgc-image/f623eeda24994690a05c9a9e326b3dcd~tplv-tt-shrink:640:0.image?x-oss-process=image/auto-orient,1/interlace,1/resize,w_1440,h_1440/quality,q_95/format,jpg" width="320" height="249" original-src="https://imageproxy.pbkrs.com/https://p9.toutiaoimg.com/img/pgc-image/f623eeda24994690a05c9a9e326b3dcd~tplv-tt-shrink:640:0.image"/></div>
我們一般會基於強化學習，在仿真環境中進行模擬訓練。
這時，如果在一台機器的 CPU 環境下進行模擬訓練，那麼需要幾個小時到幾天。
但現在，只需一個TPU/GPU，就能和數千個CPU 或 GPU 的計算集羣的速度一樣快，直接將所需時間縮短到幾分鐘！
相當於將強化學習的速度提升了1000 倍！
這就是來自谷歌的科學家們開發的物理模擬引擎Brax。
<div><img src="https://imageproxy.pbkrs.com/https://p9.toutiaoimg.com/img/pgc-image/1bfdb325c31549c39fd6b1290b898e44~tplv-tt-shrink:640:0.image?x-oss-process=image/auto-orient,1/interlace,1/resize,w_1440,h_1440/quality,q_95/format,jpg" width="640" height="67" original-src="https://imageproxy.pbkrs.com/https://p9.toutiaoimg.com/img/pgc-image/1bfdb325c31549c39fd6b1290b898e44~tplv-tt-shrink:640:0.image"/></div>
<h1>三種策略避免邏輯分支</h1>
現在大多數的物理模擬引擎都是怎麼設計的呢？
將重力、電機驅動、關節約束、物體碰撞等任務都整合在一個模擬器中，並行地進行多個模擬，以此來逼近現實中的運動系統。
<div>
<figure style="width: 640px" class="wp-caption alignnone"><img src="https://imageproxy.pbkrs.com/https://p9.toutiaoimg.com/img/pgc-image/8515ba54bb70412b993f5bf4c7bdf631~tplv-tt-shrink:640:0.image?x-oss-process=image/auto-orient,1/interlace,1/resize,w_1440,h_1440/quality,q_95/format,jpg" alt="" width="640" height="440" original-src="https://imageproxy.pbkrs.com/https://p9.toutiaoimg.com/img/pgc-image/8515ba54bb70412b993f5bf4c7bdf631~tplv-tt-shrink:640:0.image"/><figcaption>△對於每個模擬時間步長，力和力矩被整合在一起</figcaption></figure>
這種情況下，每個模擬器中的計算都不相同，且數據必須在數據中心內通過網絡傳輸。
</div>
這種並行佈局也就導致了較高的延遲時間——即學習者可能需要超過10000 納秒的等待時間，才能從模擬器中獲得經驗。
那麼怎樣才能縮短這種延遲時間呢？
Brax 選擇通過避免模擬中的分支來保證數千個並行環境中的計算完全統一，進而降低整個訓練架構的複雜度。
直到複雜度降低到可以在單一的 TPU 或 GPU 上執行，跨機器通信的計算開銷就隨之降低，延遲也就能被有效消除。
<div><img src="https://imageproxy.pbkrs.com/https://p26.toutiaoimg.com/img/pgc-image/a7851c941c8b4d1d857042ff6e59b982~tplv-tt-shrink:640:0.image?x-oss-process=image/auto-orient,1/interlace,1/resize,w_1440,h_1440/quality,q_95/format,jpg" width="638" height="326" original-src="https://imageproxy.pbkrs.com/https://p26.toutiaoimg.com/img/pgc-image/a7851c941c8b4d1d857042ff6e59b982~tplv-tt-shrink:640:0.image"/></div>
主要分為以下三個方法：
<ul>
<li>連續函數替換離散分支邏輯</li>
</ul>
比如，在計算一個小球與牆壁之間的接觸力時，就產生了一個分支：
如果球接觸牆壁，就執行模擬球從牆壁反彈的獨立代碼；
否則，就執行其他代碼；
這裏就可以通過符號距離函數來避免這種 if/else 的離散分支邏輯的產生。
<ul>
<li>使用 JAX 即時編譯中評估分支</li>
</ul>
在仿真時間之前評估基於環境靜態屬性的分支，例如兩個物體是否有可能發生碰撞。
<div><img src="https://imageproxy.pbkrs.com/https://p26.toutiaoimg.com/img/pgc-image/853cd8b206ca488d93337289de1235e6~tplv-tt-shrink:640:0.image?x-oss-process=image/auto-orient,1/interlace,1/resize,w_1440,h_1440/quality,q_95/format,jpg" original-src="https://imageproxy.pbkrs.com/https://p26.toutiaoimg.com/img/pgc-image/853cd8b206ca488d93337289de1235e6~tplv-tt-shrink:640:0.image"/></div>
<ul>
<li>在模擬中只選擇需要的分支結果</li>
</ul>
在使用了這三種策略之後，我們就得到了一個模擬由剛體、關節、執行器組成環境的物理引擎。
同時也是一種實現在這種環境中各類操作（如進化策略，直接軌跡優化等）的學習算法。
那麼 Brax 的性能究竟如何呢？
<h1>速度最高提升 1000 倍</h1>
Brax 測試所用的基準是 OpenAI Gym 中 Ant、HalfCheetah、Humanoid、Reacher 四種。
同時也增加了三個新環境：包括對物理的靈巧操作、通用運動（例如前往周圍任何一個放置了物體的地點）、以及工業機器人手臂的模擬：
<div><img src="https://imageproxy.pbkrs.com/https://p5.toutiaoimg.com/img/pgc-image/fadc33087f9242a792b8b6f81782cb88~tplv-tt-shrink:640:0.image?x-oss-process=image/auto-orient,1/interlace,1/resize,w_1440,h_1440/quality,q_95/format,jpg" width="160" height="115" original-src="https://imageproxy.pbkrs.com/https://p5.toutiaoimg.com/img/pgc-image/fadc33087f9242a792b8b6f81782cb88~tplv-tt-shrink:640:0.image"/><img src="https://imageproxy.pbkrs.com/https://p6.toutiaoimg.com/img/pgc-image/cbd7f8a19901407c8929466ba1169d51~tplv-tt-shrink:640:0.image?x-oss-process=image/auto-orient,1/interlace,1/resize,w_1440,h_1440/quality,q_95/format,jpg" width="160" height="115" original-src="https://imageproxy.pbkrs.com/https://p6.toutiaoimg.com/img/pgc-image/cbd7f8a19901407c8929466ba1169d51~tplv-tt-shrink:640:0.image"/><img src="https://imageproxy.pbkrs.com/https://p9.toutiaoimg.com/img/pgc-image/e44df815040840c3ab21a58b56fe5e34~tplv-tt-shrink:640:0.image?x-oss-process=image/auto-orient,1/interlace,1/resize,w_1440,h_1440/quality,q_95/format,jpg" width="160" height="115" original-src="https://imageproxy.pbkrs.com/https://p9.toutiaoimg.com/img/pgc-image/e44df815040840c3ab21a58b56fe5e34~tplv-tt-shrink:640:0.image"/></div>
研究人員首先測試了 Brax 在並行模擬越來越多的環境時，可以產生多少次物理步驟（也即對環境狀態的更新）。
測試結果中的 TPUv3 8×8 曲線顯示，Brax 可以在多個設備之間進行無縫擴展，每秒可達到數億個物理步驟：
<div><img src="https://imageproxy.pbkrs.com/https://p6.toutiaoimg.com/img/pgc-image/00d2575a6bd64613a62d8f0eb4cac628~tplv-tt-shrink:640:0.image?x-oss-process=image/auto-orient,1/interlace,1/resize,w_1440,h_1440/quality,q_95/format,jpg" width="640" height="528" original-src="https://imageproxy.pbkrs.com/https://p6.toutiaoimg.com/img/pgc-image/00d2575a6bd64613a62d8f0eb4cac628~tplv-tt-shrink:640:0.image"/></div>
而不僅是在 TPU 上，從 V100 和 P100 曲線也能看出，Brax 在高端GPU上同樣表現出色。
然後就是 Brax 在單個工作站（workstation）上運行一個強化學習實驗所需要的時間。
在這裏，研究人員將基於 Ant 基準環境訓練的 Brax 引擎與 MuJoCo 物理引擎做了對比：
<div><img src="https://imageproxy.pbkrs.com/https://p9.toutiaoimg.com/img/pgc-image/06c13150e27746b48acadee04d2197de~tplv-tt-shrink:640:0.image?x-oss-process=image/auto-orient,1/interlace,1/resize,w_1440,h_1440/quality,q_95/format,jpg" width="639" height="296" original-src="https://imageproxy.pbkrs.com/https://p9.toutiaoimg.com/img/pgc-image/06c13150e27746b48acadee04d2197de~tplv-tt-shrink:640:0.image"/></div>
可以看到，相對於 MuJoCo（藍線）所需的將近 3 小時時間，使用了 Brax 的加速器硬件最快只需要10 秒。
使用 Brax，不僅能夠提高單核訓練的效率，還可以擴展到大規模的並行模擬訓練。

論文地址： 
https://arxiv.org/abs/2106.13281
下載： 
https://github.com/google/brax
參考鏈接： 
https://ai.googleblog.com/2021/07/speeding-up-reinforcement-learning-with.html

谷歌-C

谷歌-A

機器人要如何完成這樣一個動作？

1 個 GPU 幾分鐘搞定強化學習訓練，谷歌新引擎讓深度學習提速 1000 倍