在機器學(xué)習(xí)領(lǐng)域，開發(fā)一個在未見過領(lǐng)域表現(xiàn)出色的通用智能體一直是長期目標(biāo)之一。一種觀點認為，在大量離線文本和視頻數(shù)據(jù)上訓(xùn)練的大型 transformer 最終可以實現(xiàn)這一目標(biāo)。

不過，在離線強化學(xué)習(xí)（RL）設(shè)置中應(yīng)用這些技術(shù)往往會將智能體能力限制在數(shù)據(jù)集內(nèi)。另一種方法是使用在線 RL，其中智能體通過環(huán)境交互自己收集數(shù)據(jù)。

然而，除了一些明顯的特例外，大多數(shù) RL 環(huán)境都是一些狹窄且同質(zhì)化的場景，限制了訓(xùn)練所得智能體的泛化能力。

近日，牛津大學(xué)的研究者提出了 Kinetix 框架，它可以表征 2D 物理環(huán)境中廣闊的開放式空間，并用來訓(xùn)練通用智能體。

• 論文地址：https://arxiv.org/pdf/2410.23208
• 項目主頁：https://kinetix-env.github.io/
• 論文標(biāo)題：Kinetix: Investigating the Training of General Agents through Open-Ended Physics-Based Control Tasks

Kinetix 涵蓋的范圍足夠廣，可以表征機器人任務(wù)（如抓取和移動）、經(jīng)典的 RL 環(huán)境（如 Cartpole、Acrobot 和 Lunar）、電子游戲（Pinball）和其他很多任務(wù)，如下圖 1 所示。

此外，為了后端運行 Kinetix，研究者開發(fā)了一種硬件加速物理引擎 Jax2D，它能夠高效地模擬訓(xùn)練智能體所需的數(shù)十億次環(huán)境交互。他們表示，通過從可表征的 2D 物理問題空間中隨機采樣 Kinetix 環(huán)境，可以幾乎無限地生成有意義的多樣化訓(xùn)練任務(wù)。

研究者發(fā)現(xiàn)，在這些環(huán)境中訓(xùn)練的 RL 智能體表現(xiàn)出了對一般機械特性的理解，并能夠零樣本地解決未見過的手工環(huán)境。

他們進一步分析了在特定困難環(huán)境中微調(diào)該通用智能體能帶來哪些好處，結(jié)果發(fā)現(xiàn)與白板智能體相比，這樣做能夠大大減少學(xué)習(xí)特定任務(wù)所需的樣本數(shù)量。

同時，微調(diào)還帶來了一些新能力，包括解決專門訓(xùn)練過的智能體無法取得進展的任務(wù)。

Kinetix 詳解

Kinetix 是一個大型開放式 RL 環(huán)境，完全在 JAX 中實現(xiàn)。

Jax2D

為了支持 Kinetix，研究團隊開發(fā)了基于脈沖的 2D 剛體物理引擎 ——Jax2D，完全用 JAX 編寫，構(gòu)成了 Kinetix 基準(zhǔn)測試的基礎(chǔ)。研究團隊通過僅模擬幾個基本組件來將 Jax2D 設(shè)計得盡可能具有表達能力。

為此，Jax2D 場景僅包含 4 個獨特的實體：圓形、（凸）多邊形、關(guān)節(jié)和推進器。從這些簡單的構(gòu)建塊中，可以表征出多種多樣的不同物理任務(wù)。

Jax2D 與 Brax 等其他基于 JAX 的物理模擬器的主要區(qū)別在于 Jax2D 場景幾乎完全是動態(tài)指定的，這意味著每次模擬都會運行相同的底層計算圖，使得能夠通過 JAX vmap 操作并行處理不同任務(wù)，這是在多任務(wù) RL 環(huán)境中利用硬件加速功能的關(guān)鍵組成部分。相比之下，Brax 幾乎完全是靜態(tài)指定的。

Kinetix：RL 環(huán)境規(guī)范

動作空間 Kinetix 支持多離散和連續(xù)動作空間。在多離散動作空間中，每個電機和推進器可以不活動，也可以在每個時間步以最大功率激活，電機可以向前或向后運行。 

• 觀察空間 

使用符號觀察，其中每個實體（形狀、關(guān)節(jié)或推進器）由一系列物理屬性值（包括位置、旋轉(zhuǎn)和速度）定義。然后將觀察定義為這些實體的集合，允許使用排列不變的網(wǎng)絡(luò)架構(gòu)，例如 transformer。這種觀察空間使環(huán)境完全可觀察，從而無需具有記憶的策略。還提供基于像素的觀察和符號觀察的選項，它可以簡單地連接和展平實體信息。

• 獎勵 

為了實現(xiàn)通用智能體的目標(biāo)，該研究選擇了一個簡單但具有高度表達力的獎勵函數(shù)，該函數(shù)在所有環(huán)境中保持固定。每個場景必須包含一個綠色形狀和一個藍色形狀 - 目標(biāo)只是使這兩個形狀發(fā)生碰撞，此時該情節(jié)以 + 1 獎勵結(jié)束。場景還可以包含紅色形狀，如果它們與綠色形狀碰撞，將會以 -1 獎勵終止該情節(jié)。如圖 1 所示，這些簡單且可解釋的規(guī)則允許表示大量語義上不同的環(huán)境。

Kinetix 的表現(xiàn)力、多樣性和速度使其成為研究開放性的理想環(huán)境，包括通用智能體、UED 和終身學(xué)習(xí)。為了使其對智能體訓(xùn)練和評估發(fā)揮最大作用，該研究提供了一個啟發(fā)式環(huán)境生成器、一組手工設(shè)計的級別以及描述環(huán)境復(fù)雜性的環(huán)境分類法。 

環(huán)境生成器 Kinetix 的優(yōu)勢在于它可以表示環(huán)境的多樣性。然而，這個環(huán)境集包含許多退化的情況，如果簡單地采樣，它們可能會主導(dǎo)分布。因此，該研究提供了一個隨機級別生成器，旨在最大程度地提高表達能力，同時最大限度地減少簡并級別的數(shù)量。確保每個關(guān)卡都具有完全相同的綠色和藍色形狀，以及至少一個可控方面（電機或推進器）。

實驗結(jié)果

研究者在程序生成的 Kinetix 關(guān)卡上進行訓(xùn)練，后者從靜態(tài)定義分布中抽取。他們將來自該分布的采樣關(guān)卡上的訓(xùn)練稱為 DR。主要評估指標(biāo)是在手動 holdout 關(guān)卡的解決率。智能體不會在這些關(guān)卡上訓(xùn)練，但它們確實存在于該訓(xùn)練分布的支持范圍內(nèi)。由于所有關(guān)卡都遵循相同的底層結(jié)構(gòu)并完全可觀察，因此理論上可以學(xué)習(xí)一種在分布內(nèi)所有關(guān)卡上表現(xiàn)最佳的策略。

為了選擇要訓(xùn)練的關(guān)卡，研究者使用了 SOTA UED 算法 SFL，它定期在隨機生成的關(guān)卡上執(zhí)行大量 rollout，然后選擇具有高學(xué)習(xí)能力的子集，并在固定時間內(nèi)對它進行訓(xùn)練，最后再次選擇新的關(guān)卡。同時，研究者使用 PLR 和 ACCEL 進行了初步實驗，但發(fā)現(xiàn)這些方法相較于 DR 沒有任何改進。

架構(gòu)

下圖 2 是訓(xùn)練所用的基于 transforme r 的架構(gòu)。可以看到，一個場景被分解為它的組成實體，然后通過網(wǎng)絡(luò)傳遞。該網(wǎng)絡(luò)由 L 層的自注意力和消息傳遞組成，K 個完全連接層緊隨其后。

其中為了以置換不變的方式處理觀察結(jié)果，研究者將每個實體表征為向量 v，其中包含物理屬性，比如摩擦、質(zhì)量和旋轉(zhuǎn)。

零樣本結(jié)果

在下圖 3 中，研究者分別在 S、M 和 L 大小的環(huán)境中訓(xùn)練 SFL。在每種情況下，訓(xùn)練環(huán)境（隨機）具有相應(yīng)的大小，而研究者使用相應(yīng)的 holdout 集來評估智能體的泛化能力。

可以看到，在每種情況下，智能體的性能都會在訓(xùn)練過程中提高，這表明它正在學(xué)習(xí)一種可以應(yīng)用于未見過環(huán)境的通用策略。

接下來，研究者通過探索學(xué)得的通用智能體在受限目標(biāo)遵循設(shè)置中的行為，仔細探究了它的零樣本性能。具體來講，他們創(chuàng)建的關(guān)卡在其中心具有單一形態(tài)（一組與電機連接并包含綠色形狀的形狀），目標(biāo)（藍色形狀）固定在關(guān)卡頂部，并且位置 x 是隨機的。

研究者測量了目標(biāo)位置 x 與可控形態(tài)位置 x 之間的關(guān)聯(lián)，如下圖 4 所示。其中最佳智能體的行為表現(xiàn)為高相關(guān)性，因此會在對角線上表現(xiàn)出高發(fā)生率。他們還評估了在隨機 M 關(guān)卡上訓(xùn)練 50 億時間步的隨機智能體和通用智能體。

正如預(yù)期的那樣，隨機智能體在可控形態(tài)和目標(biāo)位置之間沒有表現(xiàn)出相關(guān)性，而經(jīng)過訓(xùn)練的智能體表現(xiàn)出了正相關(guān)性，表明它可以將操縱形態(tài)到目標(biāo)位置。

微調(diào)結(jié)果

本節(jié)中，研究者探究了在使用給定有限樣本數(shù)量來微調(diào) holdout 任務(wù)時，通用智能體的性能。在下圖 5 中，他們?yōu)?L holdout 集中的每個關(guān)卡訓(xùn)練了單獨的專用智能體，并將它們與微調(diào)通用智能體進行比較。

研究者繪制了四個選定環(huán)境的學(xué)習(xí)曲線，以及整個 holdout 集的總體性能曲線。在其中三個關(guān)卡上，微調(diào)智能體的表現(xiàn)遠遠優(yōu)于從頭開始訓(xùn)練，尤其是對于 Mujoco-Hopper-Hard 和 Mujoco-Walker-Hard，微調(diào)智能體能夠完全勝任這些關(guān)卡，而白板智能體無法始終如一地做到這一點。

更多技術(shù)細節(jié)和實驗結(jié)果請參考原論文。