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

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

這些性能強(qiáng)大模型的背后都離不開一個(gè)關(guān)鍵技術(shù)——強(qiáng)化學(xué)習(xí)（Reinforcement Learning）。

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

近日，DeepMind研究團(tuán)隊(duì)又整了個(gè)大的！

他們使用改進(jìn)的強(qiáng)化學(xué)習(xí)技術(shù)，在類似Minecraft（我的世界）游戲中，使智能體的游戲水平甚至超過了人類專家！

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

強(qiáng)化學(xué)習(xí)通過讓AI不斷嘗試和犯錯(cuò)來學(xué)習(xí)，就像人在學(xué)習(xí)新技能一樣。

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

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

在線強(qiáng)化學(xué)習(xí)算法通常不預(yù)先「建?！?，AI看到什么就做什么，這種方法稱為無模型強(qiáng)化學(xué)習(xí)（MFRL）。

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

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

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

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

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

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

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

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

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

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

研究人員主要從三個(gè)方面入手：如何使用TWM、如何將圖像轉(zhuǎn)換成TWM的輸入以及如何訓(xùn)練TWM。

結(jié)果讓人驚喜！

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

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

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

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

相關(guān)工作

基于模型的強(qiáng)化學(xué)習(xí)（MBRL）一般分為背景規(guī)劃和決策時(shí)規(guī)劃兩種。

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

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

另一個(gè)需要關(guān)心的問題是世界模型（WM）。世界模型分為生成式世界模型和非生成式世界模型。

生成式世界模型可以生成（或想象）未來的觀察結(jié)果，并以此輔助策略學(xué)習(xí)，從而提高強(qiáng)化學(xué)習(xí)的效率，而非生成式世界模型則僅使用自預(yù)測損失進(jìn)行訓(xùn)練。

相對之下，生成式世界模型更適合背景規(guī)劃，因?yàn)樗芊奖愕貙⒄鎸?shí)和想象數(shù)據(jù)結(jié)合進(jìn)行策略學(xué)習(xí)。

訓(xùn)練方法

研究團(tuán)隊(duì)首先建立了一個(gè)基線模型。

這個(gè)基線模型在環(huán)境中訓(xùn)練了100萬步后，達(dá)到了46.91%的獎(jiǎng)勵(lì)和15.60%的分?jǐn)?shù)。

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

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

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

通過這些結(jié)構(gòu)上的改動，模型的獎(jiǎng)勵(lì)指標(biāo)提升到了55.49%，得分則達(dá)到了16.77%。性能上超過了更為復(fù)雜、運(yùn)行也更慢的DreamerV3（其獎(jiǎng)勵(lì)為53.20%，得分為 14.5%）。

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

使用Dyna方法進(jìn)行預(yù)熱

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

與通常只使用世界模型生成的軌跡進(jìn)行策略訓(xùn)練不同，研究者受到Dyna算法的啟發(fā)，改進(jìn)之后以適應(yīng)深度強(qiáng)化學(xué)習(xí)。

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

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

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

論文強(qiáng)調(diào)，世界模型的準(zhǔn)確性對于策略學(xué)習(xí)至關(guān)重要。

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

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

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

圖像塊最近鄰分詞器

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

在論文中，研究團(tuán)隊(duì)提出了一種新的圖像標(biāo)記化（tokenization）方法。用于將圖像轉(zhuǎn)換為Transformer世界模型（TWM）可以處理的token輸入。

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

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

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

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

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

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

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

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

實(shí)驗(yàn)結(jié)果顯示，在圖像塊分解的基礎(chǔ)上，用NNT替換VQ-VAE可以顯著提高智能體的獎(jiǎng)勵(lì)，從58.92%提升到64.96%。

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

塊狀教師強(qiáng)制

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

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

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

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

實(shí)驗(yàn)表明，與完全自回歸（AR）的方法相比，BTF能得到更準(zhǔn)確的TWM。

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

一舉超越了人類專家的表現(xiàn)?。ㄒ姳?）

實(shí)驗(yàn)結(jié)果

性能階梯

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

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

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

M3: M2+patches：在M2的基礎(chǔ)上，將VQ-VAE的tokenization過程分解到各個(gè)圖像塊 (patches) 上，獎(jiǎng)勵(lì)進(jìn)一步提升至 58.92%。

M4: M3 + NNT：在M3的基礎(chǔ)上，用最近鄰標(biāo)記器 (NNT) 替換VQ-VAE，獎(jiǎng)勵(lì)提升至 64.96%。

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

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

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

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

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

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

消融實(shí)驗(yàn)

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

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

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

下圖5可視化地展示消融研究的結(jié)果，驗(yàn)證了論文提出的各個(gè)改進(jìn)措施的重要性。

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

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

比較TWM的生成序列

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

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

評估的關(guān)鍵指標(biāo)是通過訓(xùn)練一個(gè)CNN符號提取器，來預(yù)測真實(shí)軌跡和TWM生成軌跡中的符號，并計(jì)算預(yù)測的符號準(zhǔn)確率。

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

定量評估

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

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

定性評估與分析

除了定量評估外，研究團(tuán)隊(duì)還對TWM生成的軌跡進(jìn)行了定性分析。

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

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

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

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

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

Craftax完整版本測試結(jié)果

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

此前的最佳智能體只能達(dá)到2.3%的獎(jiǎng)勵(lì)，而DeepMind團(tuán)隊(duì)的MFRL智能體取得了4.63%的獎(jiǎng)勵(lì)，MBRL智能體則更是將獎(jiǎng)勵(lì)提高到5.44%，再次刷新了SOTA紀(jì)錄。

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

結(jié)論與下一步工作

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

這些改進(jìn)應(yīng)用于背景規(guī)劃：帶有預(yù)熱的 Dyna 方法、圖像塊最近鄰標(biāo)記化 (NNT)以及塊教師強(qiáng)制 (BTF)。

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

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

下一步工作

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

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

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