環(huán)境生成模型（generative models of environments），也可以叫世界模型（world model），在「通用智能體規(guī)劃」和「推理環(huán)境」中的關(guān)鍵組成部分，相比傳統(tǒng)強化學(xué)習(xí)采樣效率更高。

但世界模型主要操作一系列離散潛在變量（discrete latent variables）以模擬環(huán)境動態(tài)，但這種壓縮緊湊的離散表征有可能會忽略那些在強化學(xué)習(xí)中很重要的視覺細節(jié)。

日內(nèi)瓦大學(xué)、愛丁堡大學(xué)的研究人員提出了一個在擴散世界模型中訓(xùn)練的強化學(xué)習(xí)智能體DIAMOND（DIffusion As a Model Of eNvironment Dreams），文中分析了使擴散模型適應(yīng)于世界建模（world modeling）所需的設(shè)計要素，并展示了如何通過改善視覺細節(jié)來提高智能體的性能。

論文鏈接：https://arxiv.org/pdf/2405.12399

代碼鏈接：https://github.com/eloialonso/diamond

項目鏈接：https://diamond-wm.github.io

DIAMOND在Atari 100k基準(zhǔn)測試中達到了1.46的平均人類標(biāo)準(zhǔn)化分?jǐn)?shù)（mean human

normalized score），也是完全在世界模型內(nèi)訓(xùn)練智能體的最佳成績。

此外，在圖像空間中操作還有一個好處是，擴散世界模型能夠成為環(huán)境的即插即用替代品，更方便地深入分析世界模型和智能體行為。

在項目主頁，研究人員還展示了智能體玩CS: GO的畫面，先收集了87小時人類玩家的視頻；然后用兩階段管道（two-stage pipeline:）以低分辨率執(zhí)行動態(tài)預(yù)測，降低訓(xùn)練成本；將擴散模型從Atari的4.4M參數(shù)擴展（scaling）到 CS: GO 的381M；最后對上采樣器使用隨機采樣（stochastic sampling）來提高視覺生成質(zhì)量。

模型在RTX 4090上訓(xùn)練了12天，并且可以在RTX 3090上以約10 FPS的速度運行。

不過該方法在模擬世界模型時，在部分場景下仍然會失效。

強化學(xué)習(xí)和世界模型

我們可以把環(huán)境看作是一個復(fù)雜的系統(tǒng)，智能體在這個系統(tǒng)中通過執(zhí)行動作來探索并接收反饋（獎勵）。

智能體不能直接知道環(huán)境的具體狀態(tài)，只能通過圖像觀測來理解環(huán)境，最終的目標(biāo)是教會智能體一個策略，使其能夠根據(jù)所看到的圖像來決定最佳的行動方式，以獲得最大的長期獎勵。

為此，研究人員構(gòu)建了一個世界模型來模擬環(huán)境的行為，讓智能體在模擬環(huán)境中進行訓(xùn)練，這樣可以更高效地利用數(shù)據(jù)，提高學(xué)習(xí)速度。

整個訓(xùn)練過程包括收集真實世界中的數(shù)據(jù)，用這些數(shù)據(jù)來訓(xùn)練世界模型，然后讓智能體在世界模型中進行訓(xùn)練，類似于在一個虛擬的環(huán)境中進行練習(xí)一樣，也可以稱之為「想象中的訓(xùn)練」（imagination）。

基于評分的擴散模型

擴散模型是一類受非平衡熱力學(xué)啟發(fā)的生成模型，通過逆轉(zhuǎn)加噪過程來生成樣本。

假設(shè)有一個由連續(xù)時間變量τ索引的擴散過程，其中τ的取值范圍是0到T，然后有一系列的分布，以及邊界條件：在τ=0時，分布是數(shù)據(jù)的真實分布，而在τ=T時，分布是一個易于處理的無結(jié)構(gòu)先驗分布，比如高斯分布。

為了逆轉(zhuǎn)正向的加噪過程，需要定義漂移系數(shù)和擴散系數(shù)的函數(shù)，以及估計與過程相關(guān)的未知得分函數(shù)；在實踐中，可以使用一個單一的時間依賴得分模型來估計這些得分函數(shù)。

不過在任意時間點估計得分函數(shù)并不簡單，現(xiàn)有的方法使用得分匹配作為目標(biāo)，可以在不知道潛在得分函數(shù)的情況下，從數(shù)據(jù)樣本中訓(xùn)練得分模型。

為了獲得邊際分布的樣本，需要模擬從時間0到時間τ的正向過程，然后通過一個高斯擾動核到清潔數(shù)據(jù)樣本，在一步之內(nèi)解析地到達正向過程的任何時間τ；由于核是可微的，得分匹配簡化為一個去噪得分匹配目標(biāo)（denoising score matching），這時目標(biāo)變成了一個簡單的L2重建損失，其中包含了一個時間依賴的重參數(shù)化項。

用于世界建模的擴散模型

世界模型需要一個條件生成模型來模擬環(huán)境的動態(tài)，即給定過去的狀態(tài)和動作，預(yù)測下一個狀態(tài)的概率分布，可以看作是部分可觀察馬爾可夫決策過程（POMDP），通過在歷史數(shù)據(jù)上訓(xùn)練一個條件生成模型，來預(yù)測環(huán)境的下一個狀態(tài)，雖然理論上可以采用任意常微分方程（ODE）或隨機微分方程（SDE）求解器，但在生成新的觀察結(jié)果時，需要在采樣質(zhì)量和計算成本之間做出權(quán)衡。

DIAMOND

DIAMOND模型有兩個重要的參數(shù)，一個是漂移系數(shù)，決定了系統(tǒng)隨時間變化的趨勢；另一個是擴散系數(shù)，決定了噪聲的強度，兩個系數(shù)共同調(diào)節(jié)可以使模型更好地模擬真實世界的變化。

模型的核心是預(yù)測環(huán)境的下一個狀態(tài)，為了訓(xùn)練該網(wǎng)絡(luò)，需要提供一系列的數(shù)據(jù)，包括過去的觀察結(jié)果和動作，網(wǎng)絡(luò)的目標(biāo)是從當(dāng)前的狀態(tài)和動作中預(yù)測出下一個狀態(tài)。

在訓(xùn)練過程中，會逐漸向數(shù)據(jù)中加入噪聲，模擬環(huán)境的不確定性；然后，網(wǎng)絡(luò)需要學(xué)會從這些帶有噪聲的數(shù)據(jù)中恢復(fù)出原始的、清晰的下一個狀態(tài)，整個過程就像是在一堆雜亂無章的信息中找到規(guī)律，預(yù)測出接下來可能發(fā)生的事情。

為了幫助網(wǎng)絡(luò)更好地學(xué)習(xí)和預(yù)測，DIAMOND使用了一種叫做U-Net的神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)。這種結(jié)構(gòu)特別適合處理圖像數(shù)據(jù)，因為它可以捕捉到圖像中的復(fù)雜模式。我們還使用了一種特殊的技術(shù)，叫做自適應(yīng)組歸一化，這有助于網(wǎng)絡(luò)在處理不同噪聲水平的數(shù)據(jù)時保持穩(wěn)定。

最后使用歐拉方法來生成預(yù)測結(jié)果，不需要復(fù)雜的計算，在大多數(shù)情況下都可以提供足夠準(zhǔn)確的預(yù)測。

在想象中強化學(xué)習(xí)

比如說，我們正在訓(xùn)練一個智能體如何在一個虛擬世界中行動：智能體需要「獎勵模型」告訴它做得好不好，需要「終止模型」告訴他什么時候游戲結(jié)束。

智能體有兩個部分：一個部分告訴它該怎么做（actor），用REINFORCE方法來訓(xùn)練；另一個部分告訴它做得怎么樣（critic ），用λ-回報的貝爾曼誤差的方法來訓(xùn)練。

讓智能體在一個完全由計算機生成的世界中進行訓(xùn)練，這樣就可以在不真實接觸環(huán)境的情況下學(xué)習(xí)和成長。

只需要在真實環(huán)境中收集一些數(shù)據(jù)；每次收集完數(shù)據(jù)后，都會更新智能體的虛擬世界，然后讓模型在這個更新后的世界中繼續(xù)訓(xùn)練；整個過程不斷重復(fù)，直到智能體學(xué)會如何在虛擬世界中更好地行動。

Atari 100k基準(zhǔn)結(jié)果

Atari 100k包括了26個不同的電子游戲，每個游戲都要求模型具有不同的能力。

在測試中，智能體在開始真正玩游戲之前，只能在游戲中嘗試100,000次動作，大概相當(dāng)于人類玩2個小時的游戲時間，而其他無限嘗試的游戲智能體通常會嘗試5億次動作，多了500倍。

為了更容易與人類玩家的表現(xiàn)進行比較，使用人類歸一化得分（HNS）指標(biāo)，結(jié)果顯示，DIAMOND的表現(xiàn)非常出色，在11個游戲中超過了人類玩家的表現(xiàn)，基本實現(xiàn)了超越人類的水平，平均得分為1.46，在所有世界模型訓(xùn)練的智能體中是最高的。

DIAMOND在某些游戲中的表現(xiàn)尤其好，要求智能體能夠捕捉到細節(jié)，比如《阿斯特里克斯》、《打磚塊》和《公路賽跑者》。