現階段,AI 智慧體仿佛無所不能,玩遊戲、模仿人類完成各種任務,而這些智慧體基本是在複雜環境中訓練而成的。不僅如此,隨著學習任務變得越來越複雜,模擬環境的複雜性也隨之增加,進而增加了模擬環境的成本。
即使擁有超級運算規模資源的公司和機構,訓練好一個可用的智慧體也可能需要數天的時間才能完成。
這阻礙了該領域的進展,降低了訓練先進 AI 智慧體的實用性。為了解決環境模擬的高成本問題,最近的研究努力從根本上重新設計模擬器,以便在訓練智慧體時實現更高的效率。這些工作共用批次模擬的思想,即在單個模擬器引擎內同時執行許多獨立的環境(訓練實例)。
本文,來自史丹佛大學等機構的研究者,他們提出了一個名為 Madrona 的強化學習遊戲引擎,可以在單個 GPU 上並行運行數千個環境,將智慧體的訓練時間從幾小時縮減到幾分鐘。
- 論文連結:https://madrona-engine.github.io/shacklett_siggraph23.pdf
- 論文主頁:https://madrona-engine.github.io/
具體而言,Madrona 是一款研究型遊戲引擎,專為創建學習環境而設計,可以在單個 GPU 上同時運行數千個環境實例,並且以極高的輸送量(每秒數百萬個聚合步驟)執行。Madrona 的目標是讓研究人員更輕鬆地為各種任務創建新的高性能環境,進而使 AI 智慧體訓練的速度提高幾個數量級。
Madrona 具有以下特點:
- GPU 批次模擬:單個 GPU 上可運行數千個環境;
- 實體元件系統 (ECS) 架構;
- 可與 PyTorch 輕鬆互動操作。
Madrona 環境示範:
上面我們已經提到,該研究利用了 ECS 設計原則,其具體過程如下:
研究者借助 Madrona 框架,實現了多個學習環境,結果表明,相較於開源 CPU 基線,GPU 的速度提升達到了兩到三個數量級,相較於在 32 執行緒 CPU 上運行的強基線,速度提升為 5-33 倍。此外,該研究還在該框架中實現了 OpenAI 的「hide and seek 3D」環境,每個模擬步驟都執行剛體物理學和光線追蹤,在單個 GPU 上實現了每秒超過 190 萬個 step 速度。
作者之一、史丹佛大學電腦科學副教授 Kayvon Fatahalian 表示,在一款讓多個智慧體玩烹飪遊戲 Overcooked 上,借助 Madrona 遊戲引擎,模擬 800 萬個環境步驟的時間從一小時縮短到三秒。
A few initial examples…
— Kayvon Fatahalian (@kayvonf) August 7, 2023
We’ve reimplemented @OpenAI’s hide and seek env to run at 2 million steps/sec on a single RTX 4090, Hanabi at 40M. Overcooked AI agents now train in under a minute.@bidiptas13 describes his Overcooked port here:https://t.co/P5z0IyHz8w
2/
目前,Madrona 需要使用 C++ 來編寫遊戲邏輯。Madrona 僅提供了視覺化算繪支援,雖然它可以同時模擬數千個環境,但視覺化器一次只能查看一個環境。
基於 Madrona 搭建的環境模擬器有哪些?
Madrona 本身不是一個 RL 環境模擬器,而是一個遊戲引擎或框架。開發者借助它可以更容易地實現自己的新的環境模擬器,進而透過在 GPU 上運行批次模擬並將模擬輸出與學習程式碼緊密結合來實現高性能。
下面是基於 Madrona 搭建的一些環境模擬器。
Madrona Escape Room
Madrona Escape Room 是一個簡單的 3D 環境,使用了 Madrona 的 ECS API 以及物理和算繪功能。在這個簡單任務中,智慧體必須學習按下紅色按鈕並推動其他顏色的箱子以通過一系列房間。
Overcooked AI
Overcooked AI 環境是一個基於協作電子遊戲的多智慧體學習環境(多人協作烹飪遊戲),這裡對它進行了高通量 Madrona 重寫。
Hide and Seek
2019 年 9 月,OpenAI 智慧體上演了捉迷藏攻防大戰,自創一套方式。這裡使用 Madrona 對「Hide and Seek」環境進行了再現。
Hanabi
Hanabi 是一個基於 Madrona 遊戲引擎的 Hanabi 紙牌遊戲的實現,也是一個協作式 Dec-POMDP。該環境基於 DeepMind 的 Hanabi 環境,並支援部分 MAPPO 實現。
Cartpole
Cartpole 是一個典型的 RL 訓練環境,它與構建在 Madrona 遊戲引擎之上的 gym 實現有相同的動力學。
- GitHub 地址:https://github.com/shacklettbp/madrona
Overcooked 烹飪遊戲:一分鐘內訓練最佳智慧體
論文作者之一、史丹佛大學本科生 Bidipta Sarkar 撰寫部落格詳細介紹了訓練智慧體玩 Overcooked 烹飪遊戲的過程。Overcooked 是一款流行的烹飪遊戲, 也可以作為協作多智慧體研究的基準。
在 Sarkar 的 RL 研究中,模擬虛擬環境的高成本對他來說始終是訓練智慧體的一大障礙。
就 Overcooked 烹飪遊戲而言,大約需要 800 萬步的遊戲經驗,才能訓練一對在 Overcooked 狹窄房間佈局(下圖)中收斂到穩定均衡策略的智慧體。Overcooked 的開源實現使用 Python 編寫,在 8 核 AMD CPU 上每秒運行 2000 步,因此生成必要的智慧體經驗需要花費 1 個小時以上。
相比之下,在NVIDIA A40 GPU 上執行訓練所需的所有其他操作(包括所有 800 萬個模擬步驟的策略推理、策略訓練的反向傳播)僅需不到 1 分鐘的時間。很顯然,訓練 Overcooked 智慧體受限於 Overcooked 環境模擬器的速度。
考慮到 Overcooked 是一個簡單的環境,讓模擬速度難住似乎很愚蠢。因此 Sarkar 試著看看 Overcooked 環境模擬的速度是否可以提升,這就需要用到 Madrona 遊戲引擎。
利用 Madrona 遊戲引擎,Sarkar 得到了一個原始 Overcooked Python 實現的隨插即用的 GPU 加速版替代。當並行模擬 1000 個 Overcooked 環境時,GPU 加速後在 A40 GPU 上每秒可以生成 350 萬步經驗。
作為結果,模擬 800 萬個環境步驟的時間從 1 小時縮短至了 3 秒,進而可以使用 A40 GPU 在短短 1 分鐘內訓練一個策略。
該模擬器的速度為在 Overcooked 中執行廣泛的超參數掃描打開了新的可能性,尤其是在以往訓練單個策略所需的時間內有了訓練多個策略的可能。
最後,Sarkar 意識到與創建 GPU 加速環境的現有替代方案(如 PyTorch、Taichi Lang、Direct CUDA C++)相比,將 Overcooked 移植到 Madrona 的過程更加地順利。
- 部落格連結:Overcooked in Thousands of Kitchens: Training Top Performing Agents in Under a Minute
- 資料來源:Madrona Engine