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