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

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

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

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

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

項(xiàng)目鏈接：https://diamond-wm.github.io

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

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

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

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

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

不過該方法在模擬世界模型時(shí)，在部分場(chǎng)景下仍然會(huì)失效。

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

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

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

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

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

基于評(píng)分的擴(kuò)散模型

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

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

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

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

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

用于世界建模的擴(kuò)散模型

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

DIAMOND

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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