一塊GPU頂數(shù)千個CPU內(nèi)核，英偉達(dá)的這個強(qiáng)化學(xué)習(xí)利器技術(shù)細(xì)節(jié)終于公開了

2021-08-29 10:02

來源：澎湃新聞·澎湃號·湃客

機(jī)器之心報道

機(jī)器之心編輯部

很多機(jī)器人強(qiáng)化學(xué)習(xí)任務(wù)都面臨計算需求和仿真速度的瓶頸，而英偉達(dá)這個仿真環(huán)境可以將過去需要數(shù)千個 CPU 核參與訓(xùn)練的任務(wù)移植到單個 GPU 上完成訓(xùn)練。

強(qiáng)化學(xué)習(xí)已經(jīng)成為機(jī)器學(xué)習(xí)中最有前途的研究領(lǐng)域之一，在解決復(fù)雜問題方面展現(xiàn)出了巨大的潛力?；趶?qiáng)化學(xué)習(xí)的系統(tǒng)在很多具有挑戰(zhàn)性的任務(wù)中展現(xiàn)出了超越人類的性能，比如圍棋、國際象棋、「星際爭霸」、DOTA 等等。

此外，基于強(qiáng)化學(xué)習(xí)的方法也為機(jī)器人應(yīng)用帶來了希望，比如解魔方、通過模仿動物來學(xué)習(xí)運(yùn)動等。

但直到現(xiàn)在，多數(shù)強(qiáng)化學(xué)習(xí)機(jī)器人研究者都不得不使用 CPU 和 GPU 的組合來運(yùn)行強(qiáng)化學(xué)習(xí)系統(tǒng)。二者各司其職：CPU 用于模擬環(huán)境物理、計算獎勵和運(yùn)行環(huán)境，而 GPU 用于在訓(xùn)練和推理期間加速神經(jīng)網(wǎng)絡(luò)模型，以及在需要時進(jìn)行渲染。

然而，在為序列任務(wù)優(yōu)化的 CPU 內(nèi)核和提供大規(guī)模并行運(yùn)算的 GPU 之間來回切換本質(zhì)上是低效的，需要于訓(xùn)練過程中在系統(tǒng)的不同部分之間進(jìn)行多點(diǎn)數(shù)據(jù)傳輸。因此，機(jī)器人深度強(qiáng)化學(xué)習(xí)的可擴(kuò)展性面臨兩個關(guān)鍵瓶頸：1）龐大的計算需求；2）有限的仿真速度。在擁有多個自由度的機(jī)器人學(xué)習(xí)長視野行為時，這個問題尤其具有挑戰(zhàn)性。

流行的物理引擎，如 MuJoCo、PyBullet、DART、Drake、V-Rep 等，需要大型 CPU 集群來解決具有挑戰(zhàn)性的強(qiáng)化學(xué)習(xí)任務(wù)，自然會面臨這些瓶頸。例如，OpenAI 2019 年推出了一個會玩魔方的機(jī)械手，為了訓(xùn)練這個機(jī)械手，OpenAI 動用了大約 30,000 個 CPU 內(nèi)核（920 臺 32 核計算機(jī)）。

在一個類似的手持立方體定位任務(wù)中，OpenAI 也用到了 6144 個 CPU 內(nèi)核組成的 384 個系統(tǒng)集群，再加上 8 個 Volta V100 GPU，而且需要近 30 個小時的訓(xùn)練才能達(dá)到最佳效果。這種手持立方體定位任務(wù)非常復(fù)雜，背后有著復(fù)雜的物理學(xué)、動力學(xué)原理和高維連續(xù)控制空間。

加速模擬和訓(xùn)練的一種方法是使用硬件加速器。GPU 已經(jīng)在計算機(jī)圖形學(xué)中取得了巨大的成功，自然也適用于高度并行的模擬。之前已經(jīng)有研究采用了這種方法，并在 GPU 上運(yùn)行仿真顯示了非常有前景的結(jié)果，證明使用 RL 可以極大地減少訓(xùn)練時間，并解決龐大的計算資源需求。然而，一些瓶頸仍然沒有解決，即模擬是在 GPU 上，但物理狀態(tài)被復(fù)制回 CPU。在這種模式中，觀察和獎勵使用優(yōu)化的 c++ 代碼計算，然后復(fù)制回 GPU。此外，只有簡化的基于物理的場景得到了訓(xùn)練，而不是典型的機(jī)器人環(huán)境。

為了克服這些瓶頸，去年，英偉達(dá)發(fā)布了用于強(qiáng)化學(xué)習(xí)的物理模擬環(huán)境 Isaac Gym 預(yù)覽版。它可以借助 GPU 的并行計算能力，將過去需要數(shù)千個 CPU 核參與訓(xùn)練的任務(wù)移植到單個 GPU 上完成訓(xùn)練。

為什么 Isaac Gym 可以如此高效？昨天，英偉達(dá)機(jī)器人研究科學(xué)家（前 OpenAI 研究科學(xué)家）Ankur Handa 在推特上公布，他們已經(jīng)發(fā)布了 Isaac Gym 的技術(shù)報告，里面詳細(xì)解釋了 Isaac Gym 的構(gòu)建細(xì)節(jié)。

論文鏈接：https://arxiv.org/pdf/2108.10470.pdf

Isaac Gym 是一個端到端的高性能機(jī)器人仿真平臺，它運(yùn)行一個端到端的 GPU 加速訓(xùn)練 pipeline，使研究人員能夠克服上述限制，在連續(xù)控制任務(wù)中實(shí)現(xiàn) 2-3 個數(shù)量級的訓(xùn)練加速。

具體來說，Isaac Gym 利用 NVIDIA PhysX 提供了一個 GPU 加速的模擬后端，允許它以只有使用高度并行才能達(dá)到的速度收集機(jī)器人 RL 所需的經(jīng)驗(yàn)數(shù)據(jù)。它提供了一個基于 PyTorch 張量的 API，可以在 GPU 上本地訪問物理模擬的結(jié)果。觀測張量可以作為策略網(wǎng)絡(luò)的輸入，而產(chǎn)生的動作張量可以直接反饋到物理系統(tǒng)中。

借助端到端的方法，觀察、獎勵和動作的緩沖可以在整個學(xué)習(xí)過程中保留在 GPU 上，因此不需要從 CPU 讀取數(shù)據(jù)。這種設(shè)置使得一個 GPU 上可以同時處理數(shù)以萬計的環(huán)境，研究者可以輕松地在本地運(yùn)行實(shí)驗(yàn)。

基于這一平臺，研究者復(fù)現(xiàn)了 OpenAI 的手持立方體定位任務(wù)。實(shí)驗(yàn)表明，在 A100 上，Isaac Gym 可以實(shí)現(xiàn)與 OpenAI 相似的性能結(jié)果，即使用前饋連續(xù)成功 20 次，使用 LSTM 網(wǎng)絡(luò)連續(xù)成功 37 次，但分別只用了 1 小時和 6 小時左右。相比之下，之前 OpenAI 的研究分別需要 30 個小時和 17 個小時。

此外，研究者還在其他多個機(jī)器人任務(wù)上進(jìn)行了實(shí)驗(yàn)，Isaac Gym 都實(shí)現(xiàn)了明顯的加速。

以下是報告細(xì)節(jié)節(jié)選。

端到端 GPU 強(qiáng)化學(xué)習(xí)

Isaac Gym 通過利用英偉達(dá)的 PhysX GPU 加速模擬引擎實(shí)現(xiàn)了這些結(jié)果，使其能夠收集機(jī)器人 RL 所需的經(jīng)驗(yàn)數(shù)據(jù)。除了快速物理模擬以外，Isaac Gym 還支持在 GPU 上進(jìn)行觀察和獎勵計算，從而避免嚴(yán)重的性能瓶頸，特別是消除了 GPU 與 CPU 之間昂貴的數(shù)據(jù)傳輸。通過這種方式實(shí)施，Isaac Gym 實(shí)現(xiàn)了完整的端到端 GPU 強(qiáng)化學(xué)習(xí) pipeline。

Isaac Gym 提供了一個簡單的 API，用于創(chuàng)建包括機(jī)器人和物體場景，支持從常見的 URDF（Unified Robot Description Format）和 MJCF 文件格式加載數(shù)據(jù)。每個環(huán)境都可以根據(jù)需要進(jìn)行多次復(fù)制，同時保留復(fù)制之間的變化能力（例如通過域隨機(jī)化）。該環(huán)境可以同時并行模擬，而不需要與其他環(huán)境交互。

使用一個完全由 GPU 加速的仿真和訓(xùn)練 pipeline 可以幫助降低研究阻礙，使得以前只能在大型 CPU 集群完成的任務(wù)，現(xiàn)在僅用一個 GPU 即可解決。Isaac Gym 還包括一個基本的近似策略優(yōu)化 (PPO) 實(shí)現(xiàn)和一個簡單的 RL 任務(wù)系統(tǒng)，用戶可以根據(jù)需要替換為其他任務(wù)系統(tǒng)或 RL 算法。當(dāng)示例使用的是 PyTorch 時，用戶也能夠通過定制與 TensorFlow 訓(xùn)練庫集成。

圖 2：系統(tǒng)概述。Tensor API 提供了一個 Python 代碼接口，可以直接在 GPU 上實(shí)現(xiàn) PhysX 后端以及 get 和 set 模擬器狀態(tài)，在整個 RL 訓(xùn)練 pipeline 中能實(shí)現(xiàn) 100 倍到 1000 倍的速度提升，同時提供高保真度的仿真，并且能夠與現(xiàn)有機(jī)器人模型連接。

這項(xiàng)研究的主要貢獻(xiàn)包括:

開發(fā)了用于機(jī)器人學(xué)習(xí)任務(wù)的高保真 GPU 加速的機(jī)器人模擬器；

Python 中的 Tensor API 將物理緩沖區(qū)封裝到 PyTorch 張量中，從而提供對物理緩沖區(qū)的直接訪問，而無需受到任何 CPU 限制；

實(shí)施多個高度復(fù)雜的機(jī)器人操作環(huán)境，這些環(huán)境可以在一個 GPU 上以每秒數(shù)十萬 step 的速度模擬；

在充滿挑戰(zhàn)的機(jī)器人環(huán)境中，使用 Isaac Gym 和深度強(qiáng)化學(xué)習(xí)進(jìn)行高性能訓(xùn)練的結(jié)果。

實(shí)驗(yàn)結(jié)果

該研究在多種模擬環(huán)境中實(shí)現(xiàn)了顯著的訓(xùn)練加速：在單個 NVIDIA A100 GPU 上，Ant 和 Humanoid 環(huán)境可以分別在 20 秒和 4 分鐘內(nèi)實(shí)現(xiàn)高性能運(yùn)動，而 ANYmal 的訓(xùn)練也只需要不到 2 分鐘的時間，使用 AMP 的 Humanoid 角色動畫訓(xùn)練需要 6 分鐘，使用 Shadow Hand 旋轉(zhuǎn)立方體需 35 分鐘。

研究者首先進(jìn)行了「Ant 實(shí)驗(yàn)」，在這個環(huán)境中，智能體被訓(xùn)練在平坦的地面上運(yùn)行。實(shí)驗(yàn)發(fā)現(xiàn)，隨著智能體數(shù)量的增加，訓(xùn)練時間如預(yù)期中一樣在減少。也就是說，環(huán)境數(shù)量由 256 個改變?yōu)?8192 個，增加了 5 個數(shù)量級，達(dá)到 7000 個獎勵數(shù)量級的訓(xùn)練時間從 1000 秒 (約 16.6 分鐘) 減少到了 100 秒(約 1.6 分鐘)。但 Ant 在單個 GPU 上僅用 20 秒就可以達(dá)到 3000 次運(yùn)動效果。

「Ant」是模擬最簡單的環(huán)境之一，每秒并行環(huán)境步驟的數(shù)量可能高達(dá) 700K。當(dāng)環(huán)境數(shù)量從 8192 增加到 16384 時，由于水平長度減少，研究者沒能再觀測到增益。

研究者還為 ANYmal 開發(fā)了一個粗糙地面的運(yùn)動任務(wù)，并通過將訓(xùn)練好的策略轉(zhuǎn)移到實(shí)體機(jī)器人上來驗(yàn)證了該方法的有效性。機(jī)器人學(xué)習(xí)在不平坦的地面、斜坡、樓梯和障礙物上行走，除了對平坦地形環(huán)境的觀察外，它還需要接收機(jī)器人周圍地形高度的測量數(shù)據(jù)。

該研究使用不對稱的 actor-critic 和域隨機(jī)化重現(xiàn) OpenAI Shadow Hand 立方體的訓(xùn)練設(shè)置。實(shí)驗(yàn)結(jié)果表明，該研究在 A100 上，前饋連續(xù) 20 次成功和 LSTM 網(wǎng)絡(luò)連續(xù) 37 次成功，平均分別在大約 1 小時和 6 小時內(nèi)獲得與當(dāng)時 OpenAI 實(shí)驗(yàn)結(jié)果相似的性能，成功容差為 0.4 rad。相比之下，使用傳統(tǒng) RL 訓(xùn)練設(shè)置，在 CPU 集群（384 個 CPU，每個 CPU 具有 16 個內(nèi)核）和 8 個 NVIDIA V100 GPU、MuJoCo 的組合上，OpenAI 的工作分別需要 30 小時和 17 小時。

值得一提的是， OpenAI 僅顯示 1 個隨機(jī)種子的結(jié)果，相比之下，該研究中最好的種子在短短 2.5 小時內(nèi)使用 LSTM 實(shí)現(xiàn)了 37 次連續(xù)成功。

該研究還在 ANYmal 和 TriFinger 上演示了模擬到現(xiàn)實(shí)的傳輸結(jié)果，這進(jìn)一步展示了該研究的模擬器執(zhí)行高保真接觸豐富操作任務(wù)的能力。

與吳恩達(dá)共話ML未來發(fā)展，2021亞馬遜云科技中國峰會可「玩」可「學(xué)」

2021亞馬遜云科技中國峰會「第二站」將于9月9日-9月14日全程在線上舉辦。對于AI開發(fā)者來說，9月14日舉辦的「人工智能和機(jī)器學(xué)習(xí)峰會」最值得關(guān)注。

當(dāng)天上午，亞馬遜云科技人工智能與機(jī)器學(xué)習(xí)副總裁Swami Sivasubramanian 博士與 AI 領(lǐng)域著名學(xué)者、Landing AI 創(chuàng)始人吳恩達(dá)（Andrew Ng ）博士展開一場「爐邊談話」。

不僅如此，「人工智能和機(jī)器學(xué)習(xí)峰會」還設(shè)置了四大分論壇，分別為「機(jī)器學(xué)習(xí)科學(xué)」、「機(jī)器學(xué)習(xí)的影響」、「無需依賴專業(yè)知識的機(jī)器學(xué)習(xí)實(shí)踐」和「機(jī)器學(xué)習(xí)如何落地」，從技術(shù)原理、實(shí)際場景中的應(yīng)用落地以及對行業(yè)領(lǐng)域的影響等多個方面詳細(xì)闡述了機(jī)器學(xué)習(xí)的發(fā)展。

點(diǎn)擊閱讀原文，立即報名。