單GPU運行數千環(huán)境、800萬步模擬只需3秒,斯坦福開發(fā)超強游戲引擎
現階段,AI 智能體仿佛無所不能,玩游戲、模仿人類完成各種任務,而這些智能體基本是在復雜環(huán)境中訓練而成的。不僅如此,隨著學習任務變得越來越復雜,模擬環(huán)境的復雜性也隨之增加,從而增加了模擬環(huán)境的成本。
即使擁有超級計算規(guī)模資源的公司和機構,訓練好一個可用的智能體也可能需要數天的時間才能完成。
這阻礙了該領域的進展,降低了訓練先進 AI 智能體的實用性。為了解決環(huán)境模擬的高成本問題,最近的研究努力從根本上重新設計模擬器,以在訓練智能體時實現更高的效率。這些工作共享批量模擬的思想,即在單個模擬器引擎內同時執(zhí)行許多獨立的環(huán)境(訓練實例)。
本文,來自斯坦福大學等機構的研究者,他們提出了一個名為 Madrona 的強化學習游戲引擎,可以在單個 GPU 上并行運行數千個環(huán)境,將智能體的訓練時間從幾小時縮減到幾分鐘。
- 論文地址:https://madrona-engine.github.io/shacklett_siggraph23.pdf
- 論文主頁:https://madrona-engine.github.io/
具體而言,Madrona 是一款研究型游戲引擎,專為創(chuàng)建學習環(huán)境而設計,可以在單個 GPU 上同時運行數千個環(huán)境實例,并且以極高的吞吐量(每秒數百萬個聚合步驟)執(zhí)行。Madrona 的目標是讓研究人員更輕松地為各種任務創(chuàng)建新的高性能環(huán)境,從而使 AI 智能體訓練的速度提高幾個數量級。
Madrona 具有以下特點:
- GPU 批量模擬:單個 GPU 上可運行數千個環(huán)境;
- 實體組件系統 (ECS) 架構;
- 可與 PyTorch 輕松互操作。
Madrona 環(huán)境示例:
上面我們已經提到,該研究利用了 ECS 設計原則,其具體過程如下:
研究者借助 Madrona 框架,實現了多個學習環(huán)境,結果表明,相較于開源 CPU 基線,GPU 的速度提升達到了兩到三個數量級,相較于在 32 線程 CPU 上運行的強基線,速度提升為 5-33 倍。此外,該研究還在該框架中實現了 OpenAI 的「hide and seek 3D」環(huán)境,每個模擬步驟都執(zhí)行剛體物理學和光線追蹤,在單個 GPU 上實現了每秒超過 190 萬個 step 速度。
作者之一、斯坦福大學計算機科學副教授 Kayvon Fatahalian 表示,在一款讓多個智能體玩烹飪游戲 Overcooked 上,借助 Madrona 游戲引擎,模擬 800 萬個環(huán)境步驟的時間從一小時縮短到三秒。
目前,Madrona 需要使用 C++ 來編寫游戲邏輯。Madrona 僅提供了可視化渲染支持,雖然它可以同時模擬數千個環(huán)境,但可視化器一次只能查看一個環(huán)境。
基于 Madrona 搭建的環(huán)境模擬器有哪些?
Madrona 本身不是一個 RL 環(huán)境模擬器,而是一個游戲引擎或框架。開發(fā)者借助它可以更容易地實現自己的新的環(huán)境模擬器,從而通過在 GPU 上運行批次模擬并將模擬輸出與學習代碼緊密結合來實現高性能。
下面是基于 Madrona 搭建的一些環(huán)境模擬器。
Madrona Escape Room
Madrona Escape Room 是一個簡單的 3D 環(huán)境,使用了 Madrona 的 ECS API 以及物理和渲染功能。在這個簡單任務中,智能體必須學習按下紅色按鈕并推動其他顏色的箱子以通過一系列房間。
Overcooked AI
Overcooked AI 環(huán)境是一個基于協作電子游戲的多智能體學習環(huán)境(多人協作烹飪游戲),這里對它進行了高通量 Madrona 重寫。
圖源:https://store.epicgames.com/zh-CN/p/overcooked
Hide and Seek
2019 年 9 月,OpenAI 智能體上演了捉迷藏大戰(zhàn),自創(chuàng)套路與反套路。這里使用 Madrona 對「Hide and Seek」環(huán)境進行了復現。
Hanabi
Hanabi 是一個基于 Madrona 游戲引擎的 Hanabi 紙牌游戲的實現,也是一個協作式 Dec-POMDP。該環(huán)境基于 DeepMind 的 Hanabi 環(huán)境,并支持部分 MAPPO 實現。
Cartpole
Cartpole 是一個典型的 RL 訓練環(huán)境,它與構建在 Madrona 游戲引擎之上的 gym 實現有相同的動力學。
GitHub 地址:https://github.com/shacklettbp/madrona
Overcooked 烹飪游戲:一分鐘內訓練最佳智能體
Overcooked in Thousands of Kitchens: Training Top Performing Agents in Under a Minute
論文作者之一、斯坦福大學本科生 Bidipta Sarkar 撰寫博客詳細介紹了訓練智能體玩 Overcooked 烹飪游戲的過程。Overcooked 是一款流行的烹飪游戲, 也可以作為協作多智能體研究的基準。
在 Sarkar 的 RL 研究中,模擬虛擬環(huán)境的高成本對他來說始終是訓練智能體的一大障礙。
就 Overcooked 烹飪游戲而言,大約需要 800 萬步的游戲經驗,才能訓練一對在 Overcooked 狹窄房間布局(下圖)中收斂到穩(wěn)定均衡策略的智能體。Overcooked 的開源實現使用 Python 編寫,在 8 核 AMD CPU 上每秒運行 2000 步,因此生成必要的智能體經驗需要花費 1 個小時以上。
相比之下,在英偉達 A40 GPU 上執(zhí)行訓練所需的所有其他操作(包括所有 800 萬個模擬步驟的策略推理、策略訓練的反向傳播)僅需不到 1 分鐘的時間。很顯然,訓練 Overcooked 智能體受限于 Overcooked 環(huán)境模擬器的速度。
考慮到 Overcooked 是一個簡單的環(huán)境,讓模擬速度難住似乎很愚蠢。因此 Sarkar 試著看看 Overcooked 環(huán)境模擬的速度是否可以提升,這就需要用到 Madrona 游戲引擎。
利用 Madrona 游戲引擎,Sarkar 得到了一個原始 Overcooked Python 實現的即插即用的 GPU 加速版替代。當并行模擬 1000 個 Overcooked 環(huán)境時,GPU 加速后的實現在 A40 GPU 上每秒可以生成 350 萬步經驗。
作為結果,模擬 800 萬個環(huán)境步驟的時間從 1 小時縮短至了 3 秒,從而可以使用 A40 GPU 在短短 1 分鐘內訓練一個策略。
該模擬器的速度為在 Overcooked 中執(zhí)行廣泛的超參數掃描打開了新的可能性,尤其是在以往訓練單個策略所需的時間內有了訓練多個策略的可能。
最后,Sarkar 意識到與創(chuàng)建 GPU 加速環(huán)境的現有替代方案(如 PyTorch、Taichi Lang、Direct CUDA C++)相比,將 Overcooked 移植到 Madrona 的過程更加地順利。
博客詳情:https://bsarkar321.github.io/blog/overcooked_madrona/index.html