AXIOM：對象模型分鐘級策略學習開源，10000步內掌握多種游戲

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

AXIOM：借助擴展對象中心模型實現分鐘級游戲策略學習

https://arxiv.org/pdf/2505.24784

https://github.com/VersesTech/axiom

摘要

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

1 引言

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

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

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

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

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

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

2 方法

2.1 模型的增長與剪枝

快速結構學習。秉承快速結構學習 [23] 的理念，AXIOM 使用一種在線增長啟發式方法動態擴展其全部四個混合模塊（sMM、iMM、tMM、rMM）：依次處理每一個新數據點，判斷它是由現有成分最好地解釋，還是應當創建一個新的成分，并隨后更新所選成分的參數。

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

2.2 規劃

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

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

3 結果

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

基準測試環境

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

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

基線方法

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

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

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

獎勵表現

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

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

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

綜合來看，這表明 AXIOM 的以物體為中心的世界模型，結合其快速的在線結構學習和推理算法，可以顯著減少實現高性能所需的交互次數。

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

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

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

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

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

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

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

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

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

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

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

這種自動簡化揭示了實現最優性能所需的最簡動力學集合，使 AXIOM 的決策過程變得透明且穩健。

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

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

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

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

有關更多細節，請參見附錄 E.3。

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

A 完整模型細節

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

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

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

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

這種方法受到最近提出的快速結構學習方法[23] 的啟發，后者旨在從零開始自動學習數據集的層次化生成模型。

符號說明前言

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

A.1 生成模型

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

A.2 Slot Mixture Model (sMM)

A.3 運動與存在潛變量

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

最終效果是：那些被推斷為正在運動且存在的 slot 保持完整的精度，而其他任何組合狀態都會將該 slot 的輸入協方差壓縮至 0，從而在參數學習中去除其充分統計量的影響。

A.4 相互作用變量

A.7 轉移混合模型

A.8 循環混合模型

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

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

多元正態分布成分的參數配備了NIW 先驗（正態-逆-威沙特先驗），而離散的分類似然（Categorical likelihood）的參數則配備了狄利克雷先驗。與 AXIOM 的所有其他模塊一樣，我們為的混合權重配備了一個截斷的 stick-breaking 先驗，其中最后一個第 M個偽計數參數用于調節添加新 rMM 成分的傾向性。

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

固定距離變體（Fixed distance variant）

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  
  

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

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

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

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

  
  

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

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

  
  

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

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

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

  規劃展開次數與采樣數量（Planning rollouts and samples）

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

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

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

  
  

  E.3 擾動實驗（Perturbations）

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

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

  評估模型魯棒性的方法

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

  我們的結果還表明：

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

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

  
  

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

  具體來說：

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

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

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

  
  

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

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

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

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

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

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

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

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

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

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

  • 顏色擾動

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

  • 形狀擾動

    （Shape perturbation）：將基本圖形從正方形變為圓形和三角形（見圖10c）。

    
    

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

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

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

后續在這些開創性工作基礎上的改進包括：

• Latent Slot Diffusion

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

• SlotSSM [48]

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

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

• FOCUS

  ：驗證了此類方法在機器人操作任務中低數據量下的泛化優勢 [49]；

• OC-STORM 和 SSWM

  ：利用以對象為中心的信息來預測環境的動力學和獎勵 [14, 50]；

• SPARTAN

  ：提出使用大型 Transformer 架構來識別稀疏的局部因果模型，以準確預測未來物體狀態 [13]。

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

貝葉斯推理（Bayesian Inference）

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

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

我們在時間混合建模方面的做法，與近期關于結構學習高斯混合模型的研究有概念上的相似之處，這些研究在強化學習背景下自適應地確定感知和轉移建模所需的成分數量 [52]。

AXIOM 的模型與最初的快速結構學習方法 [23] 的一個重要區別在于：

• AXIOM 使用了更結構化的先驗（表現為 sMM 的以對象為中心的因子化和 tMM 的分段線性結構）；

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

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

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

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