環境生成模型（generative models of environments），也可以叫世界模型（world model），在「通用智能體規劃」和「推理環境」中的關鍵組成部分，相比傳統強化學習采樣效率更高。

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

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

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

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

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

DIAMOND在Atari 100k基準測試中達到了1.46的平均人類標準化分數（mean human

normalized score），也是完全在世界模型內訓練智能體的最佳成績。

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

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

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

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

強化學習和世界模型

我們可以把環境看作是一個復雜的系統，智能體在這個系統中通過執行動作來探索并接收反饋（獎勵）。

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

為此，研究人員構建了一個世界模型來模擬環境的行為，讓智能體在模擬環境中進行訓練，這樣可以更高效地利用數據，提高學習速度。

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

基于評分的擴散模型

擴散模型是一類受非平衡熱力學啟發的生成模型，通過逆轉加噪過程來生成樣本。

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

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

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

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

用于世界建模的擴散模型

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

DIAMOND

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

模型的核心是預測環境的下一個狀態，為了訓練該網絡，需要提供一系列的數據，包括過去的觀察結果和動作，網絡的目標是從當前的狀態和動作中預測出下一個狀態。

在訓練過程中，會逐漸向數據中加入噪聲，模擬環境的不確定性；然后，網絡需要學會從這些帶有噪聲的數據中恢復出原始的、清晰的下一個狀態，整個過程就像是在一堆雜亂無章的信息中找到規律，預測出接下來可能發生的事情。

為了幫助網絡更好地學習和預測，DIAMOND使用了一種叫做U-Net的神經網絡結構。這種結構特別適合處理圖像數據，因為它可以捕捉到圖像中的復雜模式。我們還使用了一種特殊的技術，叫做自適應組歸一化，這有助于網絡在處理不同噪聲水平的數據時保持穩定。

最后使用歐拉方法來生成預測結果，不需要復雜的計算，在大多數情況下都可以提供足夠準確的預測。

在想象中強化學習

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

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

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

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

Atari 100k基準結果

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

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

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

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