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

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

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

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

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

項(xiàng)目地址：https://github.com/salesforce/warp-drive

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

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

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

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

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

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

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

接下來，研究者自下而上地介紹了 WarpDrive 的設(shè)計(jì)和組件，以及它們?cè)诩铀購?qiáng)化學(xué)習(xí)方面的作用。

CUDA 計(jì)算

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

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

DataManager 和 FunctionManager

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

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

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

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

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

WarpDrive RL 工作流

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

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

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

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

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

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

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

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

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

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

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

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

在單個(gè) Nvidia V100 GPU 上運(yùn)行 WarpDrive。

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