AXIOM: Learning to Play Games in Minutes withExpanding Object-Centric Models

AXIOM：借助擴展對象中心模型實現(xiàn)分鐘級游戲策略學(xué)習(xí)

https://arxiv.org/pdf/2505.24784

https://github.com/VersesTech/axiom

摘要

當前的深度強化學(xué)習(xí)（DRL）在多個領(lǐng)域?qū)崿F(xiàn)了最先進的性能，但在數(shù)據(jù)效率方面仍不如人類學(xué)習(xí)，因為人類能夠利用關(guān)于物體及其相互作用的核心先驗知識。主動推理（active inference）提供了一個有原則的框架，將感知信息與先驗知識結(jié)合，用于學(xué)習(xí)世界模型，并量化自身信念和預(yù)測的不確定性。然而，主動推理模型通常為特定任務(wù)手工設(shè)計，缺乏DRL所具備的跨領(lǐng)域靈活性。為此，我們提出了一種新架構(gòu)AXIOM，它整合了關(guān)于以物體為中心的動力學(xué)和交互的最簡但具有表達能力的核心先驗，以加速在低數(shù)據(jù)場景下的學(xué)習(xí)。AXIOM結(jié)合了貝葉斯方法的數(shù)據(jù)效率和可解釋性，以及DRL的跨任務(wù)泛化能力。該模型將場景表示為物體的組合，其動力學(xué)建模為分段線性軌跡，捕捉稀疏的物體間交互。生成模型結(jié)構(gòu)通過從單個事件中增長并學(xué)習(xí)混合模型，并通過貝葉斯模型簡化定期優(yōu)化，從而實現(xiàn)泛化。AXIOM在僅10,000次交互步驟內(nèi)掌握多種游戲，參數(shù)數(shù)量遠少于DRL方法，且無需基于梯度優(yōu)化的計算開銷。

1 引言

強化學(xué)習(xí)（RL）作為一種靈活的框架，在復(fù)雜任務(wù)中取得了顯著成功。然而，現(xiàn)有方法存在多個缺陷：需要大量訓(xùn)練數(shù)據(jù)、依賴大容量經(jīng)驗回放緩沖區(qū)，并專注于最大化累積獎勵而缺乏結(jié)構(gòu)化探索。這與人類學(xué)習(xí)形成對比——人類依靠核心先驗快速泛化到新任務(wù)。核心先驗代表了塑造感知與學(xué)習(xí)的基本組織原則（超先驗），是構(gòu)建更復(fù)雜知識結(jié)構(gòu)的基礎(chǔ)。例如，這些先驗使人類能直觀理解“物體在無外力時沿平滑軌跡運動”，并有助于因果推理、掌握行為與結(jié)果之間的關(guān)系。將視覺場景分解為物體表示已被證明在樣本效率、泛化性和魯棒性方面具有潛力。這些問題天然適合貝葉斯智能體架構(gòu)，如主動推理（active inference），它為整合先驗知識提供了理論基礎(chǔ)，支持持續(xù)適應(yīng)而不遺忘舊知識。已有研究指出這種方法與人類認知過程高度一致，其中信念隨新證據(jù)不斷更新。盡管有這些理論優(yōu)勢，主動推理的應(yīng)用多局限于小規(guī)模任務(wù)，依賴精心設(shè)計的先驗，難以達到DRL在多樣領(lǐng)域的廣泛適用性。

為彌合這一差距，我們提出了一種新的主動推理架構(gòu)，融合了最小但具表達性的關(guān)于物體及交互的核心先驗。具體來說，我們提出了AXIOM（基于對象中心模型的主動擴展推理），其包含三個關(guān)鍵組件：(1) 高斯混合模型，將視覺輸入解析為以物體為中心的表示，并自動擴展以容納新物體；(2) 轉(zhuǎn)移混合模型，發(fā)現(xiàn)運動原型（如下落、滑動、彈跳）；(3) 多物體潛在特征上的稀疏關(guān)系混合模型，學(xué)習(xí)由物體狀態(tài)、動作、獎勵和動態(tài)模式共同驅(qū)動的因果相關(guān)交互。AXIOM的學(xué)習(xí)算法具有三重效率：第一，它采用變分貝葉斯更新逐幀順序?qū)W習(xí)，無需經(jīng)驗回放緩沖或梯度計算，支持對數(shù)據(jù)分布變化的在線適應(yīng)；第二，其混合結(jié)構(gòu)可通過添加新成分解釋新數(shù)據(jù)、合并冗余成分降低模型復(fù)雜度，實現(xiàn)快速結(jié)構(gòu)學(xué)習(xí)；第三，通過維護參數(shù)后驗分布，AXIOM可以在策略選擇中引入信息尋求目標，從而進行考慮不確定性的探索。

為驗證模型效果，我們引入Gameworld 10k基準測試，一套專為評估智能體在10,000次交互內(nèi)玩不同像素游戲效率的新環(huán)境。許多現(xiàn)有RL基準（如Arcade Learning Environment或MuJoCo）強調(diào)長視野信用分配、復(fù)雜物理或視覺復(fù)雜性，往往掩蓋了快速學(xué)習(xí)與泛化的真正挑戰(zhàn)。因此，Gameworld 10k中的每個游戲均遵循類似模式：視覺場景中包含多個物體、一個可控得分玩家物體，以及遵循連續(xù)軌跡且交互稀疏的其他物體。我們設(shè)計了10款視覺元素簡化的游戲（使用不同形狀大小的單色精靈圖），以聚焦于動態(tài)建模與控制機制的研究，而非過度復(fù)雜的物體分割模型。Gameworld環(huán)境還允許精確控制游戲特征與動力學(xué)，便于測試系統(tǒng)如何適應(yīng)對游戲因果或視覺結(jié)構(gòu)的稀疏干預(yù)（如物體形狀與顏色）。在該基準上，我們的智能體在低數(shù)據(jù)環(huán)境下（10,000步）表現(xiàn)優(yōu)于主流強化學(xué)習(xí)模型，且不依賴任何基于梯度的優(yōu)化方法。雖然我們尚未將AXIOM部署到RL文獻中常見的復(fù)雜控制任務(wù)規(guī)模，但我們的成果標志著朝向構(gòu)建具備緊湊、可解釋世界模型并在多領(lǐng)域?qū)崿F(xiàn)快速決策智能體的重要進展。我們的主要貢獻如下：

• 提出AXIOM：一種全新的以物體為中心的主動推理智能體，具備在線學(xué)習(xí)、可解釋、樣本高效、適應(yīng)性強且計算成本低的特點。

• 為展示AXIOM的有效性，我們引入了一個可修改的新基準套件，專門針對具有物體與稀疏交互的環(huán)境中樣本高效學(xué)習(xí)問題。

• 我們展示了這種無需梯度的方法在樣本效率與絕對性能上均可超越最先進的深度學(xué)習(xí)方法，且在線學(xué)習(xí)方案對環(huán)境擾動具有魯棒性。

2 方法

2.1 模型的增長與剪枝

快速結(jié)構(gòu)學(xué)習(xí)。秉承快速結(jié)構(gòu)學(xué)習(xí) [23] 的理念，AXIOM 使用一種在線增長啟發(fā)式方法動態(tài)擴展其全部四個混合模塊（sMM、iMM、tMM、rMM）：依次處理每一個新數(shù)據(jù)點，判斷它是由現(xiàn)有成分最好地解釋，還是應(yīng)當創(chuàng)建一個新的成分，并隨后更新所選成分的參數(shù)。

貝葉斯模型縮減（BMR）。每 ΔTBMR=500 幀，我們最多采樣 2000 個已使用的 rMM 成分，使用祖先采樣從模型生成的數(shù)據(jù)中計算它們之間的互期望對數(shù)似然，并貪婪地測試合并候選。若某一合并操作能減少在剩余變量采樣數(shù)據(jù)條件下，關(guān)于獎勵和下一 tMM 開關(guān)狀態(tài)的多項分布的期望自由能，則接受該合并；否則回滾。BMR 使得 AXIOM 能夠從單一事件中泛化動力學(xué)規(guī)律，例如通過合并多個單事件聚類，學(xué)習(xí)到當球撞擊屏幕底部時會獲得負獎勵（見第 3 節(jié)，圖 4a）。

2.2 規(guī)劃

AXIOM 使用主動推理（active inference）進行規(guī)劃 [33]；它在不同的策略（動作序列）條件下展開對未來軌跡的預(yù)測，然后使用期望自由能（expected free energy）對這些策略進行推理，所選擇的策略 是使期望自由能最小化的那個策略：

每時間步的期望效用在規(guī)劃時通過已學(xué)習(xí)的模型和 slot 潛變量進行評估，并在整個規(guī)劃視野中隨時間步累加。期望信息增益（公式(10)右邊第二項）是根據(jù) rMM 的后驗 Dirichlet 計數(shù)計算得出的，用于衡量采取當前考慮中的策略時，能在多大程度上獲得關(guān)于 rMM 開關(guān)狀態(tài)的信息。有關(guān)規(guī)劃的更多細節(jié)見附錄 A.11。

3 結(jié)果

為了評估 AXIOM，我們在 Gameworld 環(huán)境中將其與兩種最先進的基于樣本高效、像素輸入的深度強化學(xué)習(xí)基線方法進行比較：BBF 和 DreamerV3。

基準測試環(huán)境

Gameworld 環(huán)境的設(shè)計目標是讓人類學(xué)習(xí)者在幾分鐘內(nèi)就能解決，從而確保學(xué)習(xí)過程不依賴于脆弱的探索機制或復(fù)雜的信用分配。該套件包括 10 個多樣化的游戲，借助大型語言模型生成，靈感來自 ALE 和經(jīng)典電子游戲，同時保持了輕量級和結(jié)構(gòu)化的設(shè)計。

Gameworld 環(huán)境可在 https://github.com/VersesTech/gameworld 獲取。圖 2 展示了其中包含的游戲的多樣性與視覺簡潔性。為評估魯棒性，Gameworld 10k 支持受控干預(yù)，如物體顏色或形狀的變化，以測試智能體在表面領(lǐng)域變化下的泛化能力。

基線方法

BBF [34] 建立在 SR-SPR [35] 的基礎(chǔ)上，代表了目前最樣本高效的無模型方法之一。我們針對 Gameworld 10k 套件對其預(yù)處理進行了調(diào)整，將幀跳過替換為對連續(xù)兩幀的最大池化；其余所有已發(fā)表的超參數(shù)保持不變。

其次，DreamerV3 [36] 是一種基于世界模型的智能體，在僅使用像素輸入的游戲和控制任務(wù)中表現(xiàn)優(yōu)異；我們使用已發(fā)布的設(shè)置，但將訓(xùn)練比設(shè)為 1024，批量大小為 16（有效訓(xùn)練比為 64:1）。

我們選擇這兩個基線是因為它們代表了從原始像素中進行樣本高效學(xué)習(xí)的最先進水平。請注意，對于 BBF 和 DreamerV3，我們將圖像幀分別縮放至 84×84 和 96×96 像素（遵循已發(fā)布實現(xiàn)），而 AXIOM 則直接操作 Gameworld 的完整分辨率 210×160 幀。

獎勵表現(xiàn)

圖 3 顯示了在 Gameworld 10k 套件上，從第 0 步到第 10000 步的每步獎勵的 1000 步移動平均值（均值 ± 1 標準差，共 10 個種子）。表 1 顯示了 AXIOM、BBF 和 DreamerV3 在 10k 交互步結(jié)束時所達到的累積獎勵。

AXIOM 在每一個 Gameworld 環(huán)境中都達到了高于或與 BBF 和 DreamerV3 相當?shù)钠骄鄯e獎勵。值得注意的是，AXIOM 不僅在多個游戲中達到了更高的峰值得分，而且收斂速度明顯更快，通常在前 5000 步就已獲得大部分最終獎勵，而 BBF 和 DreamerV3 幾乎需要完整的 10000 步才能達到類似效果。

對于那些在 10k 步時 BBF 和 Dreamer 表現(xiàn)接近隨機的游戲，我們確認其性能最終確實有所提升，排除了這些游戲本身對這些架構(gòu)來說本質(zhì)上過于困難的可能性（詳見附錄 E.1）。

綜合來看，這表明 AXIOM 的以物體為中心的世界模型，結(jié)合其快速的在線結(jié)構(gòu)學(xué)習(xí)和推理算法，可以顯著減少實現(xiàn)高性能所需的交互次數(shù)。

固定交互距離可以帶來更高的累積獎勵，因為智能體無需花費動作去學(xué)習(xí)該距離，但這樣做要求為每個游戲單獨調(diào)優(yōu)這一交互距離。這說明了如何將關(guān)于特定領(lǐng)域的額外知識融入像 AXIOM 這樣的貝葉斯模型中，以進一步提高樣本效率。

引入公式 (10) 中的信息增益項可以使智能體在某些游戲中更快地獲得獎勵（例如 Bounce），但在其他游戲（例如 Gold）中反而會導(dǎo)致平均獎勵增長更慢，因為它鼓勵訪問信息豐富但獎勵為負的狀態(tài)。

BMR 對需要空間泛化的游戲（如 Gold 和 Hunt）至關(guān)重要，但在 Cross 游戲中卻會損害性能，因為早期合并聚類會削弱信息增益項并抑制探索。有關(guān)詳細討論請參見附錄 E.2。

計算成本 表 2 比較了在單塊 A100 GPU 上測得的模型大小以及每步訓(xùn)練時間（模型更新與規(guī)劃）。

盡管由于使用了大量基于模型的 rollout，AXIOM 存在一定的規(guī)劃開銷，但其模型更新遠比 BBF 更高效，在每樣本的墻上時間（wall-clock time）方面具有更優(yōu)的權(quán)衡。

AXIOM 的以物體為中心的擴展模型會根據(jù)環(huán)境復(fù)雜度收斂到一個足夠的復(fù)雜度，而 BBF 和 DreamerV3 的模型大小則是固定的（且大得多）。

可解釋性 與傳統(tǒng)的深度強化學(xué)習(xí)方法不同，AXIOM 擁有一個結(jié)構(gòu)化的、以物體為中心的模型，其潛在變量和參數(shù)可以用人類可讀的方式直接解釋（例如：形狀、顏色、位置）。AXIOM 的轉(zhuǎn)移混合模型還將復(fù)雜的軌跡分解為更簡單的線性子序列。

圖 4a 展示了在 Impact 游戲中，由 AXIOM 想象出的軌跡以及基于獎勵條件的 rMM 聚類結(jié)果。圖 4a 中間面板展示了潛在空間中的想象軌跡，可以直接從對應(yīng)物體的顏色和位置來解讀。

由于循環(huán)混合模型（rMM）將開關(guān)狀態(tài)建立在各種與游戲和物體相關(guān)的特征之上，我們可以將這些開關(guān)變量設(shè)定為依賴于不同的游戲特征，并通過可視化展示 rMM 所學(xué)到的關(guān)聯(lián)關(guān)系（例如，獎勵與空間位置之間的聯(lián)系）。

圖 4a 最右側(cè)面板展示了與獎勵（綠色）和懲罰（紅色）相關(guān)的 rMM 聚類在空間中的分布。這些聚類的分布解釋了 AXIOM 對“在空間中的哪些位置會遇到獎勵”的信念，例如當玩家未接到球時預(yù)期會在屏幕底部出現(xiàn)懲罰（圖 4a 最右下面板的紅色聚類）。

圖 4b 顯示了在訓(xùn)練過程中活躍的 rMM 成分數(shù)量急劇下降。通過主動合并聚類以最小化簡化模型所對應(yīng)的期望自由能，貝葉斯模型縮減（BMR）在保持或提升性能的同時提高了計算效率（見表 1）。這種合并后的成分使得模型能夠超越訓(xùn)練數(shù)據(jù)進行插值，從而增強泛化能力。

這種自動簡化揭示了實現(xiàn)最優(yōu)性能所需的最簡動力學(xué)集合，使 AXIOM 的決策過程變得透明且穩(wěn)健。

圖 4c 表明，隨著訓(xùn)練的推進，每步的信息增益逐漸減少，而期望效用則逐步上升，反映出世界模型變得可靠后，智能體從探索階段轉(zhuǎn)向了利用階段。

干擾魯棒性 最后，我們在系統(tǒng)性地對游戲機制施加干擾的情況下測試了 AXIOM 的表現(xiàn)。在這里，我們在第 5000 步時對每個物體的顏色或形狀進行了擾動。

圖 4d 顯示，AXIOM 對形狀擾動具有很強的魯棒性，因為它仍能通過 iMM 正確推斷出物體類型。而在面對顏色擾動時，AXIOM 會新增身份類型并重新學(xué)習(xí)其動力學(xué)特性，導(dǎo)致性能略有下降，隨后逐漸恢復(fù)。

由于 AXIOM 的世界模型具有良好的可解釋性結(jié)構(gòu)，我們可以預(yù)先為其注入關(guān)于可能顏色擾動的知識，在 iMM 推理步驟中僅使用形狀信息，并在之后根據(jù)形狀重新映射被擾動的 slot，從而恢復(fù)性能。

有關(guān)更多細節(jié)，請參見附錄 E.3。

局限性與未來工作。我們的工作受到這樣一個事實的限制：核心先驗本身是由人工設(shè)計的，而不是從數(shù)據(jù)中自主發(fā)現(xiàn)的。未來的工作將聚焦于開發(fā)能夠從數(shù)據(jù)中自動推斷這類核心先驗的方法，這將使我們的方法能夠應(yīng)用于更復(fù)雜的領(lǐng)域，如 Atari 或 Minecraft [36]，在這些領(lǐng)域中，潛在的生成過程不那么透明，但仍受類似的因果原理所支配。我們相信，這一方向代表了構(gòu)建自適應(yīng)智能體的關(guān)鍵一步，這樣的智能體能夠在無需顯式設(shè)計領(lǐng)域知識的情況下，快速構(gòu)建新環(huán)境的結(jié)構(gòu)化模型。

A 完整模型細節(jié)

AXIOM 的世界模型是一個隱馬爾可夫模型（hidden Markov model），其具有以物體為中心的潛在狀態(tài)空間。該模型本身包含兩個主要組成部分：

1. 以物體為中心、類似 slot-attention 的似然模型
2. 遞歸切換狀態(tài)空間模型（recurrent switching state space model）。

這一遞歸切換狀態(tài)空間模型被應(yīng)用于由似然模型識別出的每一個物體或“slot”，并使用分段線性軌跡來建模每個物體的動力學(xué)行為。

不同于大多數(shù)其他潛在狀態(tài)空間模型（包括其他以物體為中心的模型），AXIOM 的一個顯著特點是其可適應(yīng)的復(fù)雜度——它通過在線增長和剪枝機制（見算法1 和 算法2）來迭代擴展與縮減模型，使其結(jié)構(gòu)能夠匹配其所交互世界的復(fù)雜度。這包括自動推斷場景中物體的數(shù)量，以及描述所有物體運動所需的動態(tài)模式數(shù)量。

這種方法受到最近提出的快速結(jié)構(gòu)學(xué)習(xí)方法[23] 的啟發(fā)，后者旨在從零開始自動學(xué)習(xí)數(shù)據(jù)集的層次化生成模型。

符號說明前言

大寫加粗符號表示矩陣或向量值隨機變量的集合，小寫加粗符號表示多元變量。

A.1 生成模型

該模型將感知與動力學(xué)因式分解為兩個獨立的生成模塊：

A.2 Slot Mixture Model (sMM)

A.3 運動與存在潛變量

這個二值門控變量隨后調(diào)節(jié)與身份模型（iMM）、轉(zhuǎn)移混合模型（tMM）和遞歸混合模型（rMM）相關(guān)的各種似然函數(shù)的輸入精度，從而有效地對這些模型在未跟蹤或不存在的 slot 上的學(xué)習(xí)過程進行“掩碼”處理。

最終效果是：那些被推斷為正在運動且存在的 slot 保持完整的精度，而其他任何組合狀態(tài)都會將該 slot 的輸入?yún)f(xié)方差壓縮至 0，從而在參數(shù)學(xué)習(xí)中去除其充分統(tǒng)計量的影響。

A.4 相互作用變量

A.7 轉(zhuǎn)移混合模型

A.8 循環(huán)混合模型

循環(huán)混合模型（Recurrent Mixture Model，簡稱 rMM）用于直接從當前槽級別的特征中推斷轉(zhuǎn)移模型的開關(guān)狀態(tài)。這種將開關(guān)狀態(tài)依賴于連續(xù)特征的建模方式，與循環(huán)切換線性動態(tài)系統(tǒng)（recurrent Switching Linear Dynamical System，rSLDS）[19] 中所使用的結(jié)構(gòu)相同。然而，與 rSLDS 不同的是：rSLDS 通常使用一種判別式映射（例如 softmax 或 stick-breaking 參數(shù)化形式）來從連續(xù)狀態(tài)中推斷開關(guān)狀態(tài)，而 rMM 則通過一個在混合連續(xù)-離散槽狀態(tài)上的混合模型[30] 來生成式地恢復(fù)這種依賴關(guān)系。

通過這種方式，“選擇”用于對 tMM 進行條件建模的開關(guān)狀態(tài)，實際上是通過對具有特定條件關(guān)系（在此背景下，是一種聯(lián)合混合似然關(guān)系）的其他隱變量和觀測變量進行推理而自然產(chǎn)生的。具體來說，rMM 將連續(xù)變量和離散變量的分布建模為由另一個每槽隱含分配變量 驅(qū)動的混合模型。rMM 定義了一個關(guān)于連續(xù)和離散槽特定信息元組的混合似然。

多元正態(tài)分布成分的參數(shù)配備了NIW 先驗（正態(tài)-逆-威沙特先驗），而離散的分類似然（Categorical likelihood）的參數(shù)則配備了狄利克雷先驗。與 AXIOM 的所有其他模塊一樣，我們?yōu)榈幕旌蠙?quán)重配備了一個截斷的 stick-breaking 先驗，其中最后一個第 M個偽計數(shù)參數(shù)用于調(diào)節(jié)添加新 rMM 成分的傾向性。

另外，我們還使用了門控變量 來過濾用于動態(tài)學(xué)習(xí)的槽：通過放大那些未被推斷為正在移動且存在的槽所對應(yīng)的協(xié)方差，從而降低其在模型中的影響。

固定距離變體（Fixed distance variant）

Aviate（飛行）
在這個環(huán)境中，玩家控制一只小鳥，需要穿越一系列垂直的管道。小鳥在重力作用下下落，玩家可以通過執(zhí)行“拍打”動作讓小鳥跳躍。目標是引導(dǎo)小鳥穿過管道之間的狹窄水平空隙，避免與管道結(jié)構(gòu)的任何部分或屏幕的頂部/底部邊緣發(fā)生碰撞。一旦與管道碰撞，或飛出屏幕頂部或底部，將獲得負獎勵，并結(jié)束游戲。

Bounce（反彈）
這個環(huán)境模擬了一個簡化版的經(jīng)典游戲 Pong，玩家控制一個球拍，與由 AI 控制的對手對打一個球。玩家有三個離散動作：向上移動球拍、向下移動球拍，或保持不動，這會影響球接觸時的垂直軌跡。目標是將球擊過對手的球拍以得分（+1 獎勵），同時防止對手擊中自己（-1 獎勵）。游戲是回合制的，一方得分后游戲重置。

Cross（穿越）
受經(jīng)典 Atari 游戲 Freeway 啟發(fā)，該環(huán)境要求玩家扮演一個黃色方塊，穿越多車道的道路而不被汽車撞到。玩家有三個離散動作：向上移動、向下移動或保持原地不動，用于在八個不同的車道之間進行垂直移動。各種顏色和速度的汽車持續(xù)從左向右行駛并在屏幕邊緣循環(huán)出現(xiàn)。目標是到達屏幕頂部以獲得正獎勵；若與汽車碰撞，則玩家會被重置回屏幕底部并受到負獎勵。

Driver（駕駛）
這是一個基于車道的駕駛模擬環(huán)境，玩家從俯視視角控制一輛汽車，在多車道道路上行駛。玩家可以選擇三個離散動作：停留在當前位置、向左移動或向右移動，從而實現(xiàn)變道。目標是盡可能行駛更遠的距離，避免與其他出現(xiàn)在車道上并以不同速度向下行駛的車輛發(fā)生碰撞。與另一輛車碰撞會帶來負獎勵并結(jié)束游戲。

Explode（爆炸）
這個游戲靈感來自街機經(jīng)典游戲 Kaboom!，玩家控制屏幕底部的一個水平水桶，任務(wù)是接住從移動轟炸機上掉落的炸彈。玩家可以執(zhí)行三個離散動作：保持靜止、向左移動或向右移動，以便精確地在水平方向上攔截下落的投射物。轟炸機會持續(xù)在屏幕頂部來回移動，并周期性釋放炸彈，這些炸彈在下落過程中會加速。成功用水桶接到炸彈可獲得正獎勵，而讓炸彈掉出屏幕則會導(dǎo)致負獎勵。

Fruits（水果）在這個游戲中，玩家扮演一個角色，必須在躲避危險巖石的同時收集下落的水果。玩家可以執(zhí)行三個離散動作之一：向左移動、向右移動或保持不動，控制屏幕底部的水平移動。各種顏色的水果從頂部掉落，被玩家的“隱形籃子”接住時會獲得正獎勵。相反，以深灰色矩形表示的巖石如果被接住，則會結(jié)束游戲并帶來負獎勵。

Gold（黃金）在這個游戲中，玩家控制一個由黃色方塊表示的角色，以俯視視角在一個草地上移動，目標是收集金幣并避開狗。玩家可選擇五個離散動作之一：靜止不動、向上、向右、向下或向左移動，從而實現(xiàn)靈活的屏幕導(dǎo)航。金幣是靜態(tài)的可收集物品，接觸時給予正獎勵；而狗則在屏幕上動態(tài)移動，作為障礙物存在，若發(fā)生碰撞則結(jié)束游戲并帶來負獎勵。

Hunt（狩獵）這個游戲采用類似網(wǎng)格的多車道環(huán)境，玩家以俯視視角控制角色進行導(dǎo)航。玩家有四個離散動作可供選擇：向左、向右、向上或向下移動，實現(xiàn)在游戲區(qū)域內(nèi)的二維自由移動。屏幕上不斷出現(xiàn)沿車道水平移動的物品和障礙物。玩家的目標是收集有益物品以獲得正獎勵，同時巧妙規(guī)避有害障礙物，避免受到負獎勵的影響，鼓勵策略性路徑規(guī)劃。

Impact（沖擊）該環(huán)境模擬經(jīng)典街機游戲《打磚塊》（Breakout），玩家控制屏幕底部的一個水平球拍，用來反彈球以擊碎磚墻。玩家有三個離散動作：將球拍向左移動、向右移動或保持靜止。目標是通過用球擊中磚塊來將其全部清除，每擊碎一塊磚獲得一次正獎勵。如果球漏過球拍，玩家將獲得負獎勵且游戲重置。當所有磚塊都被清除后游戲結(jié)束。

Jump（跳躍）這是一個橫向卷軸式的無盡奔跑游戲，玩家控制一個角色持續(xù)向前奔跑，并面對各種障礙物。玩家有兩個離散動作：不采取任何操作，或發(fā)起跳躍以避開不同類型的障礙物。與障礙物碰撞將導(dǎo)致負獎勵，并立即重置游戲。

E 補充結(jié)果與消融實驗E.1 在 100K 步長下的基線性能

將實際運行時間延長至 100,000 次交互步數(shù)后，模型基礎(chǔ)方法（model-based）與無模型方法（model-free）之間的對比更加明顯。

• 在Hunt游戲中，DreamerV3 在整個訓(xùn)練過程中幾乎沒有任何顯著進展，其表現(xiàn)始終接近于隨機策略的基準，而BBF則持續(xù)提升，并最終達到了與我們的以對象為中心的智能體在僅 10,000 步內(nèi)所達到的平均回合回報相當?shù)乃健?/p>

• 在Gold游戲中，兩個基線方法在 100,000 步內(nèi)確實實現(xiàn)了學(xué)習(xí)，但它們的最終性能仍低于我們的智能體在短短 10,000 步內(nèi)所達到的表現(xiàn)水平（見圖6）。

  
  

  E.2 消融實驗無信息增益（No information gain）

  當禁用信息增益時，我們得到了圖7中的紫色曲線。總體來看，從表面上看，信息增益對大多數(shù)游戲的影響似乎不大。然而，這其實是可以預(yù)期的：如圖4c所示，在Explode游戲中，信息增益僅在前幾百步中驅(qū)動性能提升，之后便由期望效用主導(dǎo)。

  從累積獎勵的角度來看，信息增益實際上在許多因玩家與物體互動而產(chǎn)生負獎勵的游戲中損害了表現(xiàn)。這是因為這些交互事件在初期會被預(yù)測為信息量豐富，從而鼓勵智能體多次嘗試經(jīng)歷它們。這一點在Cross游戲中尤為明顯：在經(jīng)歷了最初幾次碰撞后，沒有信息增益的消融模型立即決定完全不再嘗試穿越道路。

  圖8可視化了rMM生成的聚類結(jié)果，清楚地展示了在Cross游戲中，沒有信息增益會扼殺探索行為。因此，我們認為信息增益在困難探索任務(wù)中將發(fā)揮更重要的作用，這是一個值得未來研究的方向。

  
  

  無貝葉斯模型簡化（No Bayesian Model Reduction, BMR）

  圖7中的橙色曲線顯示了禁用貝葉斯模型簡化（BMR）的影響。BMR 對 Gold 和 Hunt 這兩個允許玩家在二維區(qū)域內(nèi)自由移動的游戲具有關(guān)鍵影響。在這些情況下，BMR 能夠通過合并相似聚類來實現(xiàn)對動態(tài)和物體交互的空間泛化能力。

  
  

  唯一的例外再次是 Cross 游戲：在這里，禁用 BMR 反而產(chǎn)生了表現(xiàn)最好的智能體。這一現(xiàn)象再次與信息增益的相互作用有關(guān)。由于 BMR 會合并相似的聚類，未發(fā)生碰撞的向上移動會被歸入一個經(jīng)常訪問的單一聚類。這使得該聚類從信息增益的角度來看變得不那么“有信息量”，導(dǎo)致智能體更傾向于先去與不同的車輛發(fā)生碰撞。

  然而，當禁用 BMR 時，每個空間位置都會擁有自己的聚類，智能體會被吸引去訪問那些較少觀察到的位置，例如屏幕頂部。如果我們繪制圖8c中最終的 rMM 聚類結(jié)果，也可以定性地看到這種差異。

  這引出了一個問題：在學(xué)習(xí)過程中何時進行 BMR 最合適？顯然，BMR 對于將已觀察到的事件泛化到新情境至關(guān)重要，但如果在學(xué)習(xí)早期階段就啟用它，反而可能阻礙學(xué)習(xí)過程。這種相互作用機制仍需進一步研究。

  規(guī)劃展開次數(shù)與采樣數(shù)量（Planning rollouts and samples）

  在規(guī)劃階段，我們在每個時間步都進行 rollout（軌跡展開）采樣，因此存在一個明顯的權(quán)衡關(guān)系：即策略數(shù)量與每次 rollout 的采樣數(shù)量之間的平衡——既要考慮計算時間成本（見圖5），又要保證所找到計劃的質(zhì)量。

  我們進行了網(wǎng)格搜索，調(diào)整 rollout 數(shù)量 [64, 128, 256, 512] 和每條 rollout 的采樣次數(shù) [1, 3, 5]，并評估了三個隨機種子下的表現(xiàn)。結(jié)果顯示在圖9中，表明性能差異并不顯著，但平均而言，更多的 rollout 和每條 rollout 中抽取多個樣本略優(yōu)于單一樣本的情況。

  因此，在我們的主要評估中，我們使用了512 條策略，每條策略采樣3 個樣本。但圖5和圖9的結(jié)果表明，當計算時間受限時，將策略數(shù)量減少到 128 或 64 是一種可行的方式，可以在不犧牲性能的前提下提高效率。

  
  

  E.3 擾動實驗（Perturbations）

  Gameworld 10k 基準測試的一個優(yōu)勢在于它能夠在不同環(huán)境中施加統(tǒng)一的擾動，從而讓我們可以量化不同模型對視覺特征變化的魯棒性。

  在當前實驗中，我們引入了兩種類型的擾動：

  評估模型魯棒性的方法

  我們在訓(xùn)練中途（即第5,000步）應(yīng)用每種擾動，并繪制出 Axiom、Dreamer 和 BBF 在各游戲中平均獎勵的變化曲線（見圖11）。

  我們的結(jié)果還表明：

  重映射槽身份擾動（Remapped Slot Identity Perturbations）

  在圖11中以紫色線表示的這項擾動實驗，旨在展示 AXIOM 世界模型的“白盒”性質(zhì)及其可解釋性。

  
  

  在這個實驗中，我們首先像之前一樣執(zhí)行標準的“顏色擾動”，然后我們將關(guān)于“物體顏色不可靠”的知識編碼進 AXIOM 的世界模型中。

  具體來說：

  在實踐中，這意味著顏色發(fā)生變化的對象不會被快速分配新的身份，從而使得相同的基于身份的動力學(xué)（即 rMM 中的聚類）仍然可以用于預(yù)測和解釋這些對象的行為，即使它們的顏色發(fā)生了變化。

  這也解釋了為什么在某些游戲中，使用這種“顏色重映射”技巧后，擾動幾乎沒有任何負面影響，尤其是在那些可以通過形狀輕松推斷對象身份的游戲（如 Explode）中。

  圖12 展示了有和沒有“重映射技巧”的 iMM 身份槽；圖11d 則展示了所有游戲的性能影響。

  
  

  對于某些游戲中具有相同形狀的對象（例如 Hunt 中的獎勵物和障礙物，或 Fruits 中的水果和巖石），這種重映射技巧無效，因為僅憑形狀信息不足以推斷對象類型并據(jù)此對動力學(xué)建模。

  在這種情況下，可以考慮使用更多特征來推斷對象身份，比如位置或動力學(xué)特征。然而，擴展我們的模型以整合這些信息以進一步提升魯棒性，將是我們未來的工作方向。

• • 因為 AXIOM 學(xué)到的潛在對象特征可以直接解釋為幀中物體的顏色，

  • 我們可以在提取對象身份的推理步驟中（即 iMM 的推理步驟），去除與顏色相關(guān)的潛在維度的影響，

  • 并僅使用形狀信息來進行對象類型推斷。

  • BBF 和 Dreamer 對形狀變化也具有魯棒性

  • 對于顏色擾動，Dreamer（與 Axiom 類似）有時會出現(xiàn)短暫的性能下降（例如在 Explode 游戲中），但隨后也會恢復(fù)。

  • 相比之下，BBF 對這兩種擾動都幾乎沒有受到影響。我們推測這種魯棒性源于擾動是在 BBF 尚未收斂時就施加的，因此視覺特征的變化對其學(xué)習(xí)過程影響較小。

  • 在形狀擾動下，Axiom 表現(xiàn)出較強的魯棒性。我們認為這歸功于其身份模型（iMM），該模型即使在物體外觀發(fā)生改變的情況下，仍能成功地將新形狀映射到已有的對象身份上。

  • 在顏色擾動下，Axiom 的表現(xiàn)通常會下降——這是因為身份模型最初會將擾動后的精靈視為新的對象——但隨后它迅速恢復(fù)，因為它會將這些新身份重新分配給之前學(xué)習(xí)的動力學(xué)模式。

  • 顏色擾動

    （Color perturbation）：改變所有精靈（sprite）和背景的顏色（見圖10b）；

  • 形狀擾動

    （Shape perturbation）：將基本圖形從正方形變?yōu)閳A形和三角形（見圖10c）。

    
    

深度強化學(xué)習(xí)的早期突破，例如使用深度 Q 網(wǎng)絡(luò)玩 Atari 游戲 [39]，是無模型的（model-free），并且需要在數(shù)百萬幀圖像上進行訓(xùn)練才能達到人類水平的表現(xiàn)。為此，近年來的研究開始轉(zhuǎn)向基于模型的強化學(xué)習(xí)（model-based reinforcement learning），通過學(xué)習(xí)環(huán)境的“世界模型”來減少所需的環(huán)境交互次數(shù)，從而實現(xiàn)更好的泛化能力 [40, 41]。

一個著名的例子是 Dreamer 系列模型，它依賴于連續(xù)和離散狀態(tài)空間的混合結(jié)構(gòu)來建模環(huán)境動力學(xué) [36, 42, 43]。這類世界模型模擬了人類認知的某些方面，例如對物理現(xiàn)象的直覺理解和物體追蹤 [5, 7]。為此，人們也在這一類架構(gòu)中引入先驗知識，使得特定類型的世界結(jié)構(gòu)可以被更快、更好地學(xué)習(xí)。例如，在物體層面上建模交互已被證明可以在多個任務(wù)中提升樣本效率、泛化能力和魯棒性 [9–12]。

近年來，隨著 IODINE [44] 和 Slot Attention [45] 等模型的提出，物體分割領(lǐng)域也取得了顯著進展。這些模型利用自注意力機制的優(yōu)勢，強制槽位隱變量之間競爭解釋圖像像素的能力。Slot Attention 中使用的自注意力形式與用于擬合高斯混合模型的 E 步和 M 步密切相關(guān) [46, 47]，這也啟發(fā)了我們：AXIOM 通過槽混合模型（sMM）的推理與學(xué)習(xí)從圖像中分割出物體。

后續(xù)在這些開創(chuàng)性工作基礎(chǔ)上的改進包括：

• Latent Slot Diffusion

  ：使用擴散模型進一步提升了原始方法；

• SlotSSM [48]

  ：不僅將物體分解作為圖像分割的歸納偏置，還用于視頻預(yù)測。

最近一些提出以對象為中心、基于模型的方法包括：

• FOCUS

  ：驗證了此類方法在機器人操作任務(wù)中低數(shù)據(jù)量下的泛化優(yōu)勢 [49]；

• OC-STORM 和 SSWM

  ：利用以對象為中心的信息來預(yù)測環(huán)境的動力學(xué)和獎勵 [14, 50]；

• SPARTAN

  ：提出使用大型 Transformer 架構(gòu)來識別稀疏的局部因果模型，以準確預(yù)測未來物體狀態(tài) [13]。

與 OC-STORM 不同的是，后者使用預(yù)訓(xùn)練視覺基礎(chǔ)模型和分割掩碼提取對象特征，而 AXIOM 則無需對象級別的監(jiān)督即可在線識別并分割物體（盡管目前僅測試了單色多邊形等簡單對象）。此外，AXIOM 還能在線增長和修剪其以對象為中心的狀態(tài)空間，并像 OC-STORM 一樣使用其世界模型生成的軌跡進行規(guī)劃。

貝葉斯推理（Bayesian Inference）

我們的模型中的推理、學(xué)習(xí)和規(guī)劃源自主動推理框架（active inference framework），該框架允許我們將貝葉斯原理與強化學(xué)習(xí)相結(jié)合，通過最小化預(yù)期自由能（expected free energy）來平衡獎勵最大化與信息增益 [15, 16]。

為了學(xué)習(xí)環(huán)境的結(jié)構(gòu)，我們借鑒了快速結(jié)構(gòu)學(xué)習(xí)方法[24]：首先向模型中添加混合成分 [51]，然后使用貝葉斯模型簡化（Bayesian model reduction）對其進行剪枝 [21, 22, 24]。

我們在時間混合建模方面的做法，與近期關(guān)于結(jié)構(gòu)學(xué)習(xí)高斯混合模型的研究有概念上的相似之處，這些研究在強化學(xué)習(xí)背景下自適應(yīng)地確定感知和轉(zhuǎn)移建模所需的成分數(shù)量 [52]。

AXIOM 的模型與最初的快速結(jié)構(gòu)學(xué)習(xí)方法 [23] 的一個重要區(qū)別在于：

• AXIOM 使用了更結(jié)構(gòu)化的先驗（表現(xiàn)為 sMM 的以對象為中心的因子化和 tMM 的分段線性結(jié)構(gòu)）；

• 并且使用了連續(xù)混合模型似然函數(shù)，而非純粹離散的似然。

我們所使用的轉(zhuǎn)移混合模型（tMM）是一種截斷的無限切換線性動態(tài)系統(tǒng)（SLDS）[29, 53, 54]。特別是，我們依賴于一種稱為循環(huán) SLDS（recurrent SLDS）的最新形式 [19]，它引入了開關(guān)狀態(tài)對連續(xù)狀態(tài)的依賴關(guān)系，從而解決了標準 SLDS 的兩個關(guān)鍵限制：狀態(tài)無關(guān)的轉(zhuǎn)移和上下文盲的動力學(xué)建模。

我們的創(chuàng)新之處在于如何處理 rSLDS 中的循環(huán)連接：我們采用了一種生成式模型（generative model）而非判別式模型來建模開關(guān)狀態(tài)。這使得我們可以更加靈活地根據(jù)各種信息源（包括連續(xù)和離散信息）對開關(guān)狀態(tài)進行條件建模，并且開關(guān)狀態(tài)的依賴關(guān)系在連續(xù)特征上是二次的（quadratic），從而克服了原始 rSLDS 中使用 softmax 似然時對線性可分性的內(nèi)在假設(shè) [19, 55]。

原文鏈接：https://arxiv.org/pdf/2505.24784