智慧體零樣本解決未見過人類設計環境！全靠這個開放式物理RL環境空間
机器之心發表於2024-11-24
當物理推理能力進化後，通用強化學習智慧體能在2D物理環境中執行多樣化任務了。
在機器學習領域，開發一個在未見過領域表現出色的通用智慧體一直是長期目標之一。一種觀點認為，在大量離線文字和影片資料上訓練的大型 transformer 最終可以實現這一目標。不過，在離線強化學習（RL）設定中應用這些技術往往會將智慧體能力限制在資料集內。另一種方法是使用線上 RL，其中智慧體透過環境互動自己收集資料。然而，除了一些明顯的特例外，大多數 RL 環境都是一些狹窄且同質化的場景，限制了訓練所得智慧體的泛化能力。近日，牛津大學的研究者提出了 Kinetix 框架，它可以表徵 2D 物理環境中廣闊的開放式空間，並用來訓練通用智慧體。- 論文地址：https://arxiv.org/pdf/2410.23208
- 專案主頁：https://kinetix-env.github.io/
- 論文標題：Kinetix: Investigating the Training of General Agents through Open-Ended Physics-Based Control Tasks
Kinetix 涵蓋的範圍足夠廣，可以表徵機器人任務（如抓取和移動）、經典的 RL 環境（如 Cartpole、Acrobot 和 Lunar）、電子遊戲（Pinball）和其他很多工，如下圖 1 所示。此外，為了後端執行 Kinetix，研究者開發了一種硬體加速物理引擎 Jax2D，它能夠高效地模擬訓練智慧體所需的數十億次環境互動。他們表示，透過從可表徵的 2D 物理問題空間中隨機取樣 Kinetix 環境，可以幾乎無限地生成有意義的多樣化訓練任務。研究者發現，在這些環境中訓練的 RL 智慧體表現出了對一般機械特性的理解，並能夠零樣本地解決未見過的手工環境。他們進一步分析了在特定困難環境中微調該通用智慧體能帶來哪些好處，結果發現與白板智慧體相比，這樣做能夠大大減少學習特定任務所需的樣本數量。同時，微調還帶來了一些新能力，包括解決專門訓練過的智慧體無法取得進展的任務。Kinetix 是一個大型開放式 RL 環境，完全在 JAX 中實現。為了支援 Kinetix，研究團隊開發了基於脈衝的 2D 剛體物理引擎 ——Jax2D，完全用 JAX 編寫，構成了 Kinetix 基準測試的基礎。研究團隊透過僅模擬幾個基本元件來將 Jax2D 設計得儘可能具有表達能力。為此，Jax2D 場景僅包含 4 個獨特的實體：圓形、（凸）多邊形、關節和推進器。從這些簡單的構建塊中，可以表徵出多種多樣的不同物理任務。Jax2D 與 Brax 等其他基於 JAX 的物理模擬器的主要區別在於 Jax2D 場景幾乎完全是動態指定的，這意味著每次模擬都會執行相同的底層計算圖，使得能夠透過 JAX vmap 操作並行處理不同任務，這是在多工 RL 環境中利用硬體加速功能的關鍵組成部分。相比之下，Brax 幾乎完全是靜態指定的。動作空間 Kinetix 支援多離散和連續動作空間。在多離散動作空間中，每個電機和推進器可以不活動，也可以在每個時間步以最大功率啟用，電機可以向前或向後執行。 使用符號觀察，其中每個實體（形狀、關節或推進器）由一系列物理屬性值（包括位置、旋轉和速度）定義。然後將觀察定義為這些實體的集合，允許使用排列不變的網路架構，例如 transformer。這種觀察空間使環境完全可觀察，從而無需具有記憶的策略。還提供基於畫素的觀察和符號觀察的選項，它可以簡單地連線和展平實體資訊。為了實現通用智慧體的目標，該研究選擇了一個簡單但具有高度表達力的獎勵函式，該函式在所有環境中保持固定。每個場景必須包含一個綠色形狀和一個藍色形狀 - 目標只是使這兩個形狀發生碰撞，此時該情節以 + 1 獎勵結束。場景還可以包含紅色形狀，如果它們與綠色形狀碰撞，將會以 -1 獎勵終止該情節。如圖 1 所示，這些簡單且可解釋的規則允許表示大量語義上不同的環境。Kinetix 的表現力、多樣性和速度使其成為研究開放性的理想環境，包括通用智慧體、UED 和終身學習。為了使其對智慧體訓練和評估發揮最大作用，該研究提供了一個啟發式環境生成器、一組手工設計的級別以及描述環境複雜性的環境分類法。 環境生成器 Kinetix 的優勢在於它可以表示環境的多樣性。然而，這個環境集包含許多退化的情況，如果簡單地取樣，它們可能會主導分佈。因此，該研究提供了一個隨機級別生成器，旨在最大程度地提高表達能力，同時最大限度地減少簡併級別的數量。確保每個關卡都具有完全相同的綠色和藍色形狀，以及至少一個可控方面（電機或推進器）。研究者在程式生成的 Kinetix 關卡上進行訓練，後者從靜態定義分佈中抽取。他們將來自該分佈的取樣關卡上的訓練稱為 DR。主要評估指標是在手動 holdout 關卡的解決率。智慧體不會在這些關卡上訓練，但它們確實存在於該訓練分佈的支援範圍內。由於所有關卡都遵循相同的底層結構並完全可觀察，因此理論上可以學習一種在分佈內所有關卡上表現最佳的策略。為了選擇要訓練的關卡，研究者使用了 SOTA UED 演算法 SFL，它定期在隨機生成的關卡上執行大量 rollout，然後選擇具有高學習能力的子集，並在固定時間內對它進行訓練，最後再次選擇新的關卡。同時，研究者使用 PLR 和 ACCEL 進行了初步實驗，但發現這些方法相較於 DR 沒有任何改進。下圖 2 是訓練所用的基於 transforme r 的架構。可以看到，一個場景被分解為它的組成實體，然後透過網路傳遞。該網路由 L 層的自注意力和訊息傳遞組成，K 個完全連線層緊隨其後。其中為了以置換不變的方式處理觀察結果，研究者將每個實體表徵為向量 v，其中包含物理屬性，比如摩擦、質量和旋轉。在下圖 3 中，研究者分別在 S、M 和 L 大小的環境中訓練 SFL。在每種情況下，訓練環境（隨機）具有相應的大小，而研究者使用相應的 holdout 集來評估智慧體的泛化能力。可以看到，在每種情況下，智慧體的效能都會在訓練過程中提高，這表明它正在學習一種可以應用於未見過環境的通用策略。接下來，研究者透過探索學得的通用智慧體在受限目標遵循設定中的行為，仔細探究了它的零樣本效能。具體來講，他們建立的關卡在其中心具有單一形態（一組與電機連線幷包含綠色形狀的形狀），目標（藍色形狀）固定在關卡頂部，並且位置 x 是隨機的。研究者測量了目標位置 x 與可控形態位置 x 之間的關聯，如下圖 4 所示。其中最佳智慧體的行為表現為高相關性，因此會在對角線上表現出高發生率。他們還評估了在隨機 M 關卡上訓練 50 億時間步的隨機智慧體和通用智慧體。正如預期的那樣，隨機智慧體在可控形態和目標位置之間沒有表現出相關性，而經過訓練的智慧體表現出了正相關性，表明它可以將操縱形態到目標位置。本節中，研究者探究了在使用給定有限樣本數量來微調 holdout 任務時，通用智慧體的效能。在下圖 5 中，他們為 L holdout 集中的每個關卡訓練了單獨的專用智慧體，並將它們與微調通用智慧體進行比較。研究者繪製了四個選定環境的學習曲線，以及整個 holdout 集的總體效能曲線。在其中三個關卡上，微調智慧體的表現遠遠優於從頭開始訓練，尤其是對於 Mujoco-Hopper-Hard 和 Mujoco-Walker-Hard，微調智慧體能夠完全勝任這些關卡，而白板智慧體無法始終如一地做到這一點。