世界模型也擴散！訓練出的智能體竟然不錯

世界模型提供了一種以安全且樣本高效的方式訓練強化學習智能體的方法。近期，世界模型主要對離散潛在變數序列進行操作來模擬環境動態。

然而，這種壓縮為緊湊離散表徵的方法可能會忽略對強化學習很重要的視覺細節。另一方面，擴散模型已成為影像生成的主要方法，對離散潛在模型提出了挑戰。

這種典範轉移的推動，來自日內瓦大學、愛丁堡大學、微軟研究院的研究者聯合提出一種在擴散世界模型中訓練的強化學習智能體——DIAMOND （DIffusion As a Model Of eNvironment Dreams）。

• #論文網址：https://arxiv .org/abs/2405.12399
• #專案網址：https://github.com/eloialonso/diamond
• #論文標題：Diffusion for World Modeling: Visual Details Matter in Atari

在Atari 100k基準測試中， DIAMOND+獲得了1.46的平均人類歸一化得分（HNS）。這意味著在世界模型中訓練的智能體的SOTA水平完全可以在世界模型中訓練的智能體的SOTA水平進行完全訓練。該研究提供了穩定性分析來說明DIAMOND的設計選擇對於確保擴散世界模型的長期高效穩定是必要的。

除了在影像空間中操作的好處是使擴散世界模型能夠成為環境的直接代表，從而提供對世界模型和智能體行為更深入的了解。特別地，該研究發現某些遊戲中表現的提高源於對關鍵視覺細節的更好建模。

方法介紹

接下來，本文介紹了 DIAMOND，這是一種在擴散世界模型中訓練的強化學習智能體。具體來說，研究者基於 2.2 節引入的漂移和擴散係數 f 和 g，這兩個係數對應於特定的擴散範式選擇。此外，該研究還選擇了基於 Karras 等人提出的 EDM 公式。

首先定義一個擾動核，，其中， 是一個與擴散時間相關的實值函數，稱為雜訊時間表。這對應於將漂移和擴散係數設為 和。

接著使用Karras 等人（2022）引入的網路預處理，同時參數化公式（5）中的，作為雜訊觀測值和神經網路 預測價值的加權與： 

得到公式（6）

其中為了簡潔定義，包含所有條件變數。 

預處理器的選擇。選擇預處理器和，以保持網路輸入和輸出在任何雜訊等級 下的單位變異數。  是雜訊等級的經驗轉換，  由  與資料分佈的標準差  給出，公式為

 結合公式5 和6，得到訓練目標：

#該研究使用標準的U-Net 2D 來建構向量場，並保留一個包含過去L 個觀測和動作的緩衝區，以此來條件化模型。接下來他們將這些過去的觀測以通道方式與下一個帶雜訊觀測拼接，並透過自適應組歸一化層將動作輸入到 U-Net 的殘差區塊中。如同在第 2.3 節和附錄 A 中討論的，有許多可能的採樣方法可以從訓練好的擴散模型中產生下一個觀測。雖然該研究發布的程式碼庫支援多種採樣方案，但該研究發現歐拉方法在不需要額外的NFE（函數評估次數）以及避免了高階採樣器或隨機採樣的不必要複雜性的情況下是有效的。

實驗

為了全面評估DIAMOND，該研究使用了公認的Atari 100k 基準測試，該基準測試包括26 個遊戲，用於測試智能體的廣泛能力。對於每個遊戲，智能體只允許在環境中進行 100k 次操作，這大約相當於人類 2 小時的遊戲時間，以便在評估前學習玩遊戲。作為參考，沒有限制的 Atari 智能體通常訓練 5000 萬步，相當於經驗的 500 倍增加。研究者從頭開始在每個遊戲上用 5 個隨機種子訓練 DIAMOND。每次運行大約使用 12GB 的 VRAM，在單一 Nvidia RTX 4090 上大約需要 2.9 天（總計 1.03 個 GPU 年）。

表1 比較了在世界模型中訓練智能體的不同分數：

圖2 中提供了平均值和IQM（ Interquartile Mean ）置信區間：

結果表明， DIAMOND 在基準測試中表現強勁，超過人類玩家在11 個遊戲中的表現，並達到了1.46 的HNS 得分，這是完全在世界模型中訓練的智能體的新紀錄。研究還發現，DIAMOND 在需要捕捉細節的環境中表現特別出色，例如 Asterix、Breakout 和 Road Runner。

為了研究擴散變數的穩定性，研究分析了自迴歸產生的想像軌跡（imagined trajectory），如下圖3 所示：

研究發現有些情況需要迭代求解器將取樣過程驅動到特定模式，如圖4 所示的拳擊遊戲：

如圖5 所示，與IRIS 想像的軌跡相比，DIAMOND 想像的軌跡通常具有更高的視覺質量，並且更符合真實環境。

有興趣的讀者可以閱讀論文原文，了解更多研究內容。

以上是世界模型也擴散！訓練出的智能體竟然不錯的詳細內容。更多資訊請關注PHP中文網其他相關文章！