?作者?|?張倩

來源?|?機器之心

在強化學習研究中，一個實驗就要跑數天或數周，有沒有更快的方法？近日，來自 SalesForce 的研究者提出了一種名為 WarpDrive（曲率引擎）的開源框架，它可以在一個 V100 GPU 上并行運行、訓練數千個強化學習環境和上千個智能體。實驗結果表明，與 CPU+GPU 的 RL 實現相比，WarpDrive 靠一個 GPU 實現的 RL 要快幾個數量級。

深度強化學習是一種訓練 AI 智能體的強大方法。然而，如何將強化學習應用于多個智能體或高維狀態、觀察結果、行動空間等復雜場景仍然是一個具有挑戰性的問題。多智能體系統是強化學習研究和應用的前沿，尤其是那些存在多個智能體交互的系統。它們對于解決經濟、自動駕駛、對話智能體、機器人學以及其他很多領域的工程和科學問題具有重要意義。

在這一問題中，訓練智能體的時間是一大挑戰。如果訓練的智能體比較多，一個強化學習實驗往往要花費數天或數周才能完成。主要原因在于：在在線訓練中，強化學習需要反復運行模擬、訓練智能體模型。此外，學習所需的重復次數會隨著問題的復雜度呈指數級增長。這在無模型設置中是最突出的，在這種設置中，強化學習智能體在訓練時通常對當前的模擬或任務沒有任何初始知識。由于深度強化學習實現通常將基于 CPU 的模擬和 GPU 神經網絡模型結合在一起，現實世界中的訓練時間可能會非常長。

為了提高強化學習系統的吞吐量，來自 Salesforce 的研究者提出了一種名為 WarpDrive 的開源框架，來構建極其快速的（多智能體）深度強化學習（MADRL）pipeline。

論文鏈接：

https://arxiv.org/pdf/2108.13976.pdf

項目地址：

https://github.com/salesforce/warp-drive

WarpDrive 在一個 GPU 上端到端地運行整個 MADRL 工作流，使用單個數據存儲進行模擬 roll-out、推理和訓練。這將最小化通信和復制的成本，顯著提升模擬采樣和學習速率。WarpDrive 還利用 GPU 的并行能力，并行地運行模擬和每次模擬中的智能體。綜合來看，這些設計選擇使得運行數千個并發模擬成為可能，每個模擬包含數千個智能體，在非常大的批次上進行訓練。

論文中的基準測試表明，與混合 CPU 模擬和 GPU 模型的一般實現相比，WarpDrive 實現的 RL 要快幾個數量級。比如，在有 2000 個模擬、1000 個智能體的 Tag 模擬中，WarpDrive 的吞吐量實現了至少 100 倍的性能提升。

此外，研究者還將 WarpDrive 構建得非常靈活、易用、易擴展，并在現有庫的基礎上進行構建和補充，如 PyCUDA 和 PyTorch。在這個初始版本中，WarpDrive 使用在 CUDA C 中實現的 step 和 reset 方法的模擬，使用 OpenAI gym-API 風格。這使得用戶在實現模擬的時候能夠使用 CUDA 編程的全部特征集。WarpDrive 在 PyCUDA 的基礎上構建，可以方便地訪問用 CUDA C 編寫的 GPU 數據和模擬方法。它還提供了一個輕量級的 Python 包裝器來自動構建 gym 風格的環境對象，實現模擬 API。因此，WarpDrive 可以讓用戶輕松地創建和擴展自定義 RL pipeline，以最大化 GPU 的用途。

作者表示，WarpDrive 這一名字的靈感來自科幻小說中的超光速宇宙飛船推進系統——曲率引擎，后者可以通過對時空本身的改造來驅動飛船。

在這篇論文中，研究者概述了 WarpDrive 的設計，并展示了它在基準模擬中的吞吐量和可擴展性。最后，他們總結了未來的發展方向。

用 WarpDrive 加速強化學習

WarpDrive 提供了一個框架和一些高質量的工具，幫助用戶快速構建快捷而靈活的多智能體強化學習系統。圖 3 顯示了 WarpDrive 的一個基本用例。

接下來，研究者自下而上地介紹了 WarpDrive 的設計和組件，以及它們在加速強化學習方面的作用。

CUDA 計算

WarpDrive 專注于 CUDA（Compute Unified Device Architecture），這是一個流行的平臺和編程模型，允許用戶在（CUDAenabled）GPU 硬件上運行程序。CUDA 程序也被稱為計算 kernel。CUDA API 可以直接訪問 GPU 的虛擬指令集和并行計算元素。

GPU 的一個關鍵特性是它們可以并行運行許多計算線程。線程組成線程塊，多個線程塊組成一個網格結構。一個 CUDA kernel 可以訪問并定義這些線程的并行計算，見下圖 2。

DataManager 和 FunctionManager

圖 2 顯示了 WarpDrive 框架的一些細節。所有在 GPU 上運行的 CUDA 程序都需要通過 CPU 觸發。按照 CUDA 的規則，CPU 被稱為主機（host），GPU 被稱為設備。任何 CUDA 程序的執行都包括三個主要步驟：

主機到設備的傳輸（push）：將輸入數據從主機內存復制到設備內存，例如在訓練開始時。

加載 CUDA 函數，在 GPU 上執行并緩存數據以提高性能。

設備到主機的傳輸（pull）：將數據從設備內存復制回主機，例如訓練結束的時候。

按照這個模式，WarpDrive 實現了兩個重要的 Python 類（駐留在 CPU 上）——DataManager 和 FunctionManager，以方便主機和設備之間所有與 RL 相關的通信和交互。DataManager 提供了一個簡易 API 來處理與 RL 相關的主機和設備之間的所有數據傳輸（pull 和 push）。FunctionManager 允許用戶從 CPU 調用 CUDA kernel，并在 GPU 上執行它們。

WarpDrive RL 工作流

使用 WarpDrive，在 GPU 上收集 rollout 和訓練模型的一個典型 RL 工作流如下所示：

1. 初始化并重置環境對象后，將主機上的所有數據復制到設備上。DataManager 提供 API 方法來執行此 push 操作。從主機復制到設備的數據可能包括環境配置參數、在第一次重置結束時創建的數據數組，以及觀察、動作、獎勵和「完成」標志的占位符。DataManager 還幫助維護變量的副本，這些變量需要在每次重置時重新初始化。在這點之后，就不再從主機向設備推送數據了。

2. FunctionManager 提供 API 方法來初始化和調用 CUDA C kernel 函數，這些函數用于從主機節點執行環境步驟。這些函數只在設備上執行，所有的數據數組都被就地修改。出于可視化或分析的目的，主機可能會不時地拉取（pulled）數據，但所有數據本質上只能在訓練期間駐留在 GPU 上。

3.WarpDrive 還包含一個 Sampler 類，用于對動作進行采樣，以便逐步遍歷環境。使用策略模型計算出的概率對動作進行抽樣。WarpDrive 的采樣器在每個智能體線程上并行運行，運行速度比等效的 PyTorch 實現快 2 倍。詳見實驗部分。

4. 在 GPU 中，研究者并行執行多個環境副本。每個環境運行在一個單獨的線程塊上。由于一個典型的 GPU 有數千個塊，在一個 GPU 上并行執行數千個環境也是可行的。

5. 任何環境都可能達到終端狀態并被「完成」。WarpDrive 提供了一個 EnvironmentReset 類，用于自動識別和重置那些已完成的環境。此時，那些環境也被重置并給出（新的）初始數據。例如，它們可能使用在初始重置時復制的初始數據數組。

6. 在每個環境中，每個智能體還可以在智能體獨占（agent-exclusive）的線程上執行它自己的邏輯。圖 2 顯示了一個示例，其中智能體 i、j 和 k 分別在單個線程 i、j 和 k 上并行操作。這在多智能體 RL 中非常有用，因為可以完全將智能體操作并行化，因此模擬時間復雜度保持不變，即使智能體數量增加（達到可用線程的數量）。

7. 一旦從幾個環境和智能體中收集到的數據被送到訓練數據 batch 中，也可以用 WarpDrive 執行端到端訓練。這次推出的 WarpDrive 包括一個示例訓練腳本和 Trainer 類，它目前實現了 Advantage Actor Critic（A2C）、近端策略優化（PPO）RL 算法和一個全連接神經網絡策略模型。Trainer 基于 PyTorch 構建，并從 CPU 主機調用所有 CUDA kernel 和 PyTorch 操作。然而，PyTorch 可以直接訪問和解釋以 torch 形式存儲在設備上的 WarpDrive 數據 batch（狀態、動作、獎勵和完成標志）。這使得用戶可以計算損失并修改模型參數，同時消除數據復制。考慮到 WarpDrive 的模塊化特性，這一框架很容易擴展。

這些設計選擇使得 WarpDrive 非常高效。

實驗結果

研究者使用 Tag 游戲演示了 WarpDrive 的工作流，并對其性能進行了測試。他們比較了以下兩種情況的性能：

• 使用 CPU-simulations + GPU-agent 的模型；

• 在單個 Nvidia V100 GPU 上運行 WarpDrive。

測試結果表明，與使用 CPU-simulation + GPU-agent 的模型相比，WarpDrive 的吞吐量要高幾個數量級。

#投 稿?通 道#

?讓你的文字被更多人看到?

如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。

總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。?

PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學術熱點剖析、科研心得或競賽經驗講解等。我們的目的只有一個，讓知識真正流動起來。

📝?稿件基本要求：

? 文章確系個人原創作品，未曾在公開渠道發表，如為其他平臺已發表或待發表的文章，請明確標注?

? 稿件建議以?markdown?格式撰寫，文中配圖以附件形式發送，要求圖片清晰，無版權問題

? PaperWeekly 尊重原作者署名權，并將為每篇被采納的原創首發稿件，提供業內具有競爭力稿酬，具體依據文章閱讀量和文章質量階梯制結算

📬?投稿通道：

? 投稿郵箱：hr@paperweekly.site?

? 來稿請備注即時聯系方式（微信），以便我們在稿件選用的第一時間聯系作者

? 您也可以直接添加小編微信（pwbot02）快速投稿，備注：姓名-投稿

△長按添加PaperWeekly小編

🔍

現在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關注」訂閱我們的專欄吧

·

總結

以上是生活随笔為你收集整理的一块V100运行上千个智能体、数千个环境，这个曲率引擎框架实现RL百倍提速的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。