DeepSeek R1與OpenAI o系列模型的爆發(fā)，掀起了人工智能領域的新一輪浪潮。

甚至讓我們看到了通往AGI的曙光。

這些性能強大模型的背后都離不開一個關鍵技術——強化學習（Reinforcement Learning）。

谷歌的DeepMind團隊在這方面一直很強，轟動全球的「人機大戰(zhàn)」（李世石對戰(zhàn)AlphaGo）就出自他們團隊之手。

近日，DeepMind研究團隊又整了個大的！

他們使用改進的強化學習技術，在類似Minecraft（我的世界）游戲中，使智能體的游戲水平甚至超過了人類專家！

論文地址：https://arxiv.org/pdf/2502.01591

強化學習通過讓AI不斷嘗試和犯錯來學習，就像人在學習新技能一樣。

AI在環(huán)境中行動，根據(jù)結果的好壞來調(diào)整自己的策略。

那些在環(huán)境中一邊嘗試，一邊學習，收集數(shù)據(jù)（觀察和獎勵），然后更新自己策略的強化學習算法被稱作在線強化學習。

在線強化學習算法通常不預先「建模」，AI看到什么就做什么，這種方法稱為無模型強化學習（MFRL）。

但這種方法需要收集大量環(huán)境數(shù)據(jù)。

因此，有人提出了基于模型的強化學習（MBRL），很顯然它試圖減少訓練所需的數(shù)據(jù)量。

MBRL會先學習一個「世界模型」（WM），就像在腦海中模擬出一個世界，然后在這個模擬的世界里進行「想象」和「規(guī)劃」。好比AI自己先在腦子里模擬一下，然后再行動。

為了評估強化學習算法的效率，大家通常使用Atari-100k基準測試，查看算法在Atari游戲中使用不超過10萬幀訓練數(shù)據(jù)時的表現(xiàn)。

但是，Atari游戲的確定性比較高，AI很容易記住一些操作，而不是真正學會泛化。而且，Atari游戲往往只側重于一兩種技能，不能很好地考察AI的綜合能力。

為了訓練出能力更全面的AI，谷歌DeepMind的研究團隊選擇了Crafter這個環(huán)境。

Crafter是一個2D版的《我的世界》，具體來說，他們用的是Craftax-classic環(huán)境，它是Crafter的快速復刻版。

Craftax-classic環(huán)境有幾個很好的特點：

1. 每次游戲的環(huán)境都是隨機生成的，AI需要應對不同的挑戰(zhàn)。
2. AI只能看到局部視野，就好像只能看到屏幕的一部分，而不是整個地圖。
3. 這是一個以成就層級來設定獎勵信號的體系，需要進行深入且廣泛的探索才能達成。

DeepMind研究團隊的這篇論文主要研究了如何在Craftax-classic環(huán)境中改進基于Transformer世界模型（TWM）的強化學習方法。

研究人員主要從三個方面入手：如何使用TWM、如何將圖像轉換成TWM的輸入以及如何訓練TWM。

結果讓人驚喜！

研究團隊的方法讓智能體在僅用100萬步環(huán)境交互的情況下，就取得了Craftax-classic 67.42%的獎勵和 27.91%的得分，這比之前的最佳研究成果（SOTA）——53.20%的獎勵和19.4%的得分——都有了顯著提升。

智能體的表現(xiàn)甚至超越了人類專家！相當炸裂。

下圖1展示了論文中提出的MBRL方法在Craftax-classic環(huán)境中的性能，其中左圖展示了不同算法隨著環(huán)境交互步驟的增加所獲得的獎勵。

中間圖展示了Craftax-classic的畫面，即一個63x63像素的圖像，包含智能體周圍的地圖以及智能體的生命值和物品欄。右圖顯示了NNT（最近鄰標記器）提取的 64 個不同的圖像塊，這些圖像塊用于TWM（Transformer 世界模型）的輸入。

相關工作

基于模型的強化學習（MBRL）一般分為背景規(guī)劃和決策時規(guī)劃兩種。

背景規(guī)劃在模型外部（即決策發(fā)生之前）使用世界模型（WM）生成想象軌跡，進而訓練策略。而決策時規(guī)劃在做決策的時刻利用WM進行前瞻搜索。

由于決策時規(guī)劃速度較慢，特別是使用大型世界模型如Transformer模型時。所以，論文側重于背景規(guī)劃。

另一個需要關心的問題是世界模型（WM）。世界模型分為生成式世界模型和非生成式世界模型。

生成式世界模型可以生成（或想象）未來的觀察結果，并以此輔助策略學習，從而提高強化學習的效率，而非生成式世界模型則僅使用自預測損失進行訓練。

相對之下，生成式世界模型更適合背景規(guī)劃，因為它能方便地將真實和想象數(shù)據(jù)結合進行策略學習。

訓練方法

研究團隊首先建立了一個基線模型。

這個基線模型在環(huán)境中訓練了100萬步后，達到了46.91%的獎勵和15.60%的分數(shù)。

研究團隊在此基礎上進行了兩方面的改進：一是增大模型規(guī)模，二是在策略中加入RNN（具體來說是GRU）來提供記憶能力。

有趣的是，他們發(fā)現(xiàn)如果只是單純增大模型規(guī)模，反而會降低性能。但當更大的模型與精心設計的RNN相結合時，性能會有所提升。

對于RNN，研究團隊認為保持隱藏狀態(tài)的維度足夠小至關重要，這樣才能讓記憶只關注當前圖像無法提取的、真正重要的歷史信息。

通過這些結構上的改動，模型的獎勵指標提升到了55.49%，得分則達到了16.77%。性能上超過了更為復雜、運行也更慢的DreamerV3（其獎勵為53.20%，得分為 14.5%）。

在性能提升的同時，成本也降低了。模型在一張A100 GPU上對環(huán)境進行100萬步訓練僅需約15分鐘。

使用Dyna方法進行預熱

接下來就是論文提出的核心改進之一，即如何有效地利用真實環(huán)境數(shù)據(jù)和世界模型（WM）生成的虛擬數(shù)據(jù)來訓練智能體。

與通常只使用世界模型生成的軌跡進行策略訓練不同，研究者受到Dyna算法的啟發(fā)，改進之后以適應深度強化學習。

Dyna方法將真實環(huán)境中的軌跡數(shù)據(jù)和TWM生成的想象軌跡數(shù)據(jù)混合起來訓練智能體。將世界模型視為一種生成式數(shù)據(jù)增強

智能體首先與環(huán)境交互，收集真實的軌跡數(shù)據(jù)，并立即用這些數(shù)據(jù)來更新策略。之后，智能體使用世界模型在想象中生成軌跡，并用這些想象數(shù)據(jù)來進一步更新策略。

這種混合使用真實數(shù)據(jù)和虛擬數(shù)據(jù)的機制，可以被看作是一種生成式數(shù)據(jù)增強的方式。

論文強調(diào)，世界模型的準確性對于策略學習至關重要。

為了確保世界模型足夠準確，避免其不準確的預測「污染」訓練數(shù)據(jù)，研究者提出在開始使用想象軌跡訓練策略之前，先讓智能體與環(huán)境交互一段時間。

這個過程被稱為「預熱（warmup）」。具體來說，就是只有在智能體與環(huán)境交互達到一定步數(shù)之后，才開始使用世界模型生成的軌跡來訓練。

實驗表明，移除預熱步驟會導致獎勵大幅下降，從67.42%降至33.54%。此外，僅僅使用想象數(shù)據(jù)訓練策略也會導致性能下降到55.02%。

圖像塊最近鄰分詞器

不同于常規(guī)的使用VQ-VAE方法來在圖像和tokens之間進行轉換。

在論文中，研究團隊提出了一種新的圖像標記化（tokenization）方法。用于將圖像轉換為Transformer世界模型（TWM）可以處理的token輸入。

研究團隊利用Craftax-classic環(huán)境的一個特點，即每個觀察圖像都是由9x9個7x7大小的圖像塊組成。

因此，他們首先將圖像分解為這些不重疊的圖像塊(patches)。然后，獨立地將每個圖像塊編碼為token。

在將圖像塊分解的基礎上，論文使用一個更簡單的最近鄰標記器（Nearest-Neighbor Tokenizer，NNT）來代替?zhèn)鹘y(tǒng)的VQ-VAE3。

NNT的編碼過程類似于最近鄰分類器，它將每個圖像塊與一個代碼本中的條目進行比較。

如果圖像塊與代碼本中最近的條目之間的距離小于一個閾值，則將該條目的索引作為token；否則，將該圖像塊作為一個新的代碼添加到代碼本中。

與VQ-VAE 不同，NNT的代碼本一旦添加條目，就不再更新。這意味著代碼本是靜態(tài)的，但可以不斷增長。

解碼時，NNT只是簡單地返回代碼本中與token索引相對應的代碼 (圖像塊)。

這種靜態(tài)但不斷增長的代碼本使得TWM的目標分布更加穩(wěn)定，大大簡化了TWM的在線學習過程。

實驗結果顯示，在圖像塊分解的基礎上，用NNT替換VQ-VAE可以顯著提高智能體的獎勵，從58.92%提升到64.96%。

不過，NNT對圖像塊的大小比較敏感，如果圖像塊的大小不合適，可能會影響性能。而且如果圖像塊內(nèi)部的視覺變化很大，NNT可能會生成一個非常大的代碼本。

塊狀教師強制

在通常的做法中，TWM通常采用教師強制來訓練，論文的研究團隊則提出了一種更有效的替代方案，稱為塊狀教師強制（block teacher forcing, BTF）。

這個方案同時修改了TWM的監(jiān)督方式和注意力機制：當給定前面的全部token后，BTF會并行預測下一時間步中的所有潛在token，從而不再依賴當前時間步已生成的token。

下圖2清晰地展示了BTF如何通過改變注意力模式和監(jiān)督方式來改進TWM的訓練。

傳統(tǒng)的教師強制自回歸地預測每個token，而BTF則并行預測同一時間步的所有token，從而提高了訓練速度和模型準確性。

實驗表明，與完全自回歸（AR）的方法相比，BTF能得到更準確的TWM。

在本實驗中，BTF將獎勵從64.96%提升到了67.42%，從而獲得了表現(xiàn)最優(yōu)的基于模型的強化學習（MBRL）智能體。

一舉超越了人類專家的表現(xiàn)！（見表1）

實驗結果

性能階梯

在論文中，智能體按照研究者所采用的改進措施進行排序，形成了一個「MBRL 階梯」，具體如下：

M1: Baseline：基準MBRL智能體，使用了VQ-VAE進行tokenization，其獎勵為31.93%，優(yōu)于IRIS的25.0%。

M2: M1+Dyna：在M1的基礎上，使用Dyna方法，即混合使用真實環(huán)境數(shù)據(jù)和TWM生成的想象數(shù)據(jù)來訓練策略，獎勵提升至43.36%。

M3: M2+patches：在M2的基礎上，將VQ-VAE的tokenization過程分解到各個圖像塊 (patches) 上，獎勵進一步提升至 58.92%。

M4: M3 + NNT：在M3的基礎上，用最近鄰標記器 (NNT) 替換VQ-VAE，獎勵提升至 64.96%。

M5: M4 + BTF：在M4的基礎上，引入塊教師強制 (BTF)，最終的獎勵達到67.42%(±0.55)，成為論文中最佳的 MBRL智能體。

下圖3清晰地展示了每一步改進帶來的性能提升。

與現(xiàn)有方法比較

研究團隊這次性能最優(yōu)的模型M5創(chuàng)造了新的SOTA成績，獎勵達到67.42%，得分達到27.91%。

這是首次超過人類專家的平均獎勵水平（該人類水平基于5名專家玩家玩了100局所測得）。

需要指出的是，盡管模型在獎勵上已超越了人類專家，但得分仍明顯低于人類專家水平。

消融實驗

實驗表明，當NNT使用7×7大小的圖像塊時效果最佳，使用較小 (5×5) 或較大 (9×9) 的圖像塊時，性能會有所下降，但仍然具有競爭力。

如果不使用量化，而是讓TWM重建連續(xù)的7×7圖像塊，性能會大幅下降。

研究者發(fā)現(xiàn)，移除「MBRL 階梯」中的任何一個步驟，都會導致模型性能下降，這表明論文提出的每個改進都至關重要。

下圖5可視化地展示消融研究的結果，驗證了論文提出的各個改進措施的重要性。

模型如果過早地開始在想象數(shù)據(jù)上訓練，性能會因TWM的不準確而崩潰。只有在智能體與環(huán)境交互足夠長時間，并獲得足夠數(shù)據(jù)來訓練可靠的WM后，使用想象數(shù)據(jù)進行訓練才是有效的。

去除MFRL智能體中的RNN或使用較小的模型都會導致模型性能下降。

比較TWM的生成序列

研究者比較了三種不同的世界模型 (TWM) 的生成軌跡質量，這些模型分別是M1（基線模型）、M3（加入了Dyna和圖像塊分解的模型）以及M5（最佳模型，包含了所有改進）。

為了進行評估，研究者首先構建了一個包含160條軌跡的評估數(shù)據(jù)集，每條軌跡長度為20。然后，他們使用每個TWM模型，從相同的起始狀態(tài)和動作序列出發(fā)，生成對應的想象軌跡。

評估的關鍵指標是通過訓練一個CNN符號提取器，來預測真實軌跡和TWM生成軌跡中的符號，并計算預測的符號準確率。

這種方法能夠深入了解模型在多大程度上捕捉到了游戲的核心動態(tài)。

定量評估

通過定量評估，研究團隊發(fā)現(xiàn)符號準確率隨著TWM生成軌跡步數(shù)的增加而下降，這種下降是由于誤差的累積導致的。

M5模型由于其采用了最近鄰標記器（NNT），保持了所有時間步中最高的符號準確率，表明其能夠更好地捕捉游戲動態(tài)，并且NNT使用的靜態(tài)代碼本簡化了TWM的學習過程。

定性評估與分析

除了定量評估外，研究團隊還對TWM生成的軌跡進行了定性分析。

通過視覺檢查，他們觀察到了三種現(xiàn)象：地圖不一致性、符合游戲規(guī)則的幻覺以及不符合游戲規(guī)則的幻覺。

M1模型在地圖和游戲動態(tài)方面都存在明顯的錯誤，而M3和M5模型能夠生成一些符合游戲規(guī)則的幻覺，例如出現(xiàn)怪物和生命值變化。

M3模型仍然會產(chǎn)生一些不符合游戲規(guī)則的幻覺，例如怪物突然消失或生成的動物外觀錯誤，而M5模型則很少出現(xiàn)這種不合理的幻覺。

定性分析表明，盡管所有模型都存在一定的誤差，但M5模型在保持游戲動態(tài)一致性方面明顯優(yōu)于其他模型，體現(xiàn)了其學習到的世界模型質量的提升。

下圖6表明，NNT和BTF等改進措施對于提高TWM學習效果的重要性，最終促進了MBRL智能體性能的提升。

Craftax完整版本測試結果

研究團隊還比較了多種智能體在Craftax的完整版本（Craftax Full）上的性能。相比 Craftax-classic，這個完整版在關卡數(shù)量和成就設置上都有顯著提升，難度更高。

此前的最佳智能體只能達到2.3%的獎勵，而DeepMind團隊的MFRL智能體取得了4.63%的獎勵，MBRL智能體則更是將獎勵提高到5.44%，再次刷新了SOTA紀錄。

這些結果表明，DeepMind團隊所采用的訓練方法能夠推廣到更具挑戰(zhàn)性的環(huán)境。

結論與下一步工作

在本論文，DeepMind研究團隊提出了三項針對基于Transformer世界模型 (TWM) 的視覺MBRL智能體的改進措施。

這些改進應用于背景規(guī)劃：帶有預熱的 Dyna 方法、圖像塊最近鄰標記化 (NNT)以及塊教師強制 (BTF)。

這些改進措施結合起來，使得MBRL智能體在Craftax-classic基準測試中取得了顯著更高的獎勵和分數(shù)，首次超越了人類專家的獎勵水平。

論文提出的技術也成功地推廣到了更具挑戰(zhàn)性的Craftax(full) 環(huán)境中，取得了新的SOTA結果。

下一步工作

DeepMind研究團隊未來將研究如何將這些技術推廣到Craftax之外的其他環(huán)境，以驗證其通用性。探索使用優(yōu)先經(jīng)驗回放來加速TWM的訓練，以提高數(shù)據(jù)利用效率。

團隊還考慮將大型預訓練模型 (如SAM和Dino-V2) 的能力與當前的標記器結合起來，從而獲得更穩(wěn)定的代碼本，并減少對圖像塊大小和表觀變化的敏感性。

為探究無法生成未來像素的非重構型世界模型，團隊還計劃改造策略網(wǎng)絡，使其能夠直接接收TWM生成的潛變量 token，而不是像素。