一塊V100運行上千個智能體、數(shù)千個環(huán)境，這個「曲率引擎」框架實現(xiàn)RL百倍提速

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

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

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

為了提高強化學(xué)習(xí)系統(tǒng)的吞吐量，來自 Salesforce 的研究者提出了一種名為 WarpDrive 的開源框架，來構(gòu)建極其快速的（多智能體）深度強化學(xué)習(xí)（MADRL）pipeline。

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

項目地址：https://github.com/salesforce/warp-drive

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

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

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

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

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

用 WarpDrive 加速強化學(xué)習(xí)

WarpDrive 提供了一個框架和一些高質(zhì)量的工具，幫助用戶快速構(gòu)建快捷而靈活的多智能體強化學(xué)習(xí)系統(tǒng)。圖 3 顯示了 WarpDrive 的一個基本用例。

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

CUDA 計算

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

GPU 的一個關(guān)鍵特性是它們可以并行運行許多計算線程。線程組成線程塊，多個線程塊組成一個網(wǎng)格結(jié)構(gòu)。一個 CUDA kernel 可以訪問并定義這些線程的并行計算，見下圖 2。

DataManager 和 FunctionManager

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

主機到設(shè)備的傳輸（push）：將輸入數(shù)據(jù)從主機內(nèi)存復(fù)制到設(shè)備內(nèi)存，例如在訓(xùn)練開始時。

加載 CUDA 函數(shù)，在 GPU 上執(zhí)行并緩存數(shù)據(jù)以提高性能。

設(shè)備到主機的傳輸（pull）：將數(shù)據(jù)從設(shè)備內(nèi)存復(fù)制回主機，例如訓(xùn)練結(jié)束的時候。

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

WarpDrive RL 工作流

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

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

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

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

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

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

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

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

這些設(shè)計選擇使得 WarpDrive 非常高效。

實驗結(jié)果

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

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

在單個 Nvidia V100 GPU 上運行 WarpDrive。

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