NeurIPS 2023 Spotlight | 面向通用決策場景的MCTS基準框架：LightZero

這一部分，本文首先會概述講解LightZero，然后深入討論設計通用MCTS算法將會面臨的環境挑戰，并提出模塊化的訓練流程設計方案和高度可擴展的MCTS工具鏈。

圖1：LightZero 概覽。左側部分描繪了蒙特卡羅樹搜索（MCTS）的發展過程，而右側則展示了各類強化學習（RL）環境。LightZero 融合并拓展了 MCTS/MuZero 子領域的最新研究成果，并將這些成果有效地應用于多種不同的環境中。

如圖1所示，LightZero是首個匯集了MCTS/MuZero領域幾乎全部最新算法進展的Benchmark測試平臺。更準確地說，LightZero匯集了源自AlphaZero/MuZero的9項關鍵算法和超過20種決策環境，為其設計了標準化的訓練和部署接口，以適應多種決策環境的需要。

與原有其他各類定制化算法框架相比，LightZero提供了一種統一的視角和用戶接口。這種統一性使得對比和分析各類技術組合成為可能，為MCTS相關算法和應用的研究提供一個強大且通用的基準平臺。

圖2：MCTS+RL方法和model-free RL方法（例如 PPO）在六個環境挑戰及數據效率維度上的定性對比結果。LightZero中將通用決策求解器的關鍵能力進行了分類，包括：多模態觀察空間、復雜動作空間、環境固有隨機性、對先驗知識的依賴程度、仿真成本、探索難度和數據效率。圖中的每條曲線表示一個算法在這六個類別上的評分。評分為1表示該算法在此維度上表現欠佳，只適用于有限場景，而更高的評分則意味著更廣泛的應用范圍和更優的性能。特別地，model-free RL方法不依賴仿真，對先驗知識的依賴度較小，因此在這些領域得分較高。請注意，LightZero在此上下文中指的是在 LightZero 框架內，將各種技術和超參數設置進行最優組合的特殊算法變體。關于定性評分規則的詳細信息，請參見原論文附錄 D。

LightZero對各類RL環境進行了統一的分析，總結出設計通用高效MCTS算法的六個關鍵挑戰（圖2）。具體來說，這六種環境挑戰分別是：

1）多模態觀察空間（Multi-modal observation spaces），需要智能體能夠提取并有效融合不同形式的信息如低維向量、視覺圖像和復雜結構化觀察空間等。

2）復雜動作空間（Complex action space），需要智能體能在生成多樣化決策信號上表現出色，包括離散動作選擇、連續控制和混合結構的動作空間等。

3）依賴先驗知識（Reliance on prior knowledge），這是AlphaZero等方法的一個主要缺點。這些方法本質上需要能夠訪問完美的模擬器和環境的具體規則。相反，MuZero及派生方法通過學習環境模型來替代模擬器和相關先驗，解決了這一限制。

4）環境固有隨機性（Inherent stochasticity）在基于樹搜索的規劃方法中提出了根本性挑戰。環境動力學的不確定性和部分可觀測的狀態空間都可能導致規劃軌跡的錯位，產生大量無用或沖突的搜索結果。

5）模擬成本（Simulation cost）是MCTS-style方法的時間消耗的主要來源。同時，如果算法在模擬過程中未能訪問所有必要的動作，算法的性能會大幅下降。

6）探索困難（Hard exploration）是一個經常被忽視的關鍵挑戰。雖然搜索樹可以通過減少探索范圍來提高效率，但MCTS-style的方法在有大量非終止軌跡（即沒有有效的獎勵信號的軌跡）的環境中容易遇到困難，如迷宮等。

圖3：LightZero訓練流程中的四個核心子模塊。上下文交換器（Context Exchanger）負責在各個子模塊之間傳輸配置、模型和軌跡。

MCTS系列方法雖然表現優秀，但其實現復雜性也相當高，這大大限制了其應用場景。與DQN/PPO等model-free RL算法不同，MCTS方法在每次智能體-環境交互中都需要使用搜索樹進行多步模擬。此外，為了提升訓練數據質量，MuZero Unplugged引入了數據重分析機制，通過使用新模型重新計算舊數據，以獲取改進的訓練目標。然而，這兩種技術都需要多次調用模擬器或神經網絡，增加了分布式訓練和通信拓撲等方面的復雜性。因此，對算法進行集成以簡化整個框架是必要的。

LightZero中設計的訓練流程（如圖3所示）包括四個核心子模塊。

• 首先，LightZero同時支持online RL和offline RL訓練。兩者的主要區別在于是采用在線交互數據收集器還是直接利用離線數據集。

• 其次，LightZero對相關訓練流程進行了重構，將其分解為四個核心子模塊，遵循高內聚、低耦合（high cohesion and low coupling）的設計原則。

• 數據收集器（Data collector）負責利用策略網絡和搜索樹進行有效的動作選擇，同時包括各種探索策略，以及數據預處理和打包操作。
  
  
• 數據整理器（Data arranger）在MCTS中扮演著關鍵角色，它能有效地存儲和準備用于訓練的有價值數據。此模塊涉及到數據重分析技術，用以校正off-policy或離線數據。優先級采樣的改良版確保了訓練的小批量數據具有足夠的多樣性和較高的學習潛能。吞吐量限制器負責控制添加和采樣數據的比例，以在固定的通信帶寬內實現最優的數據利用效果。
  
  
• 智能體學習器（Agent learner）負責多個網絡的訓練，可以利用自監督表征學習，基于模型的rollout，基于值分布的強化學習和價值函數歸一化等優化技術進行增強。這些技術有助于策略的提升，進而提高智能體的整體性能。
  
  
• 智能體評估器（Agent evaluator）提供各類針對強化學習的評價指標以監控訓練過程并評估策略行為。同時，這個模塊集成了一些模型推理時常用的技巧，如beam search，以提升測試性能。

LightZero基于上述設計提供一系列標準接口和工具，讓研究者和工程師能更有效地開發算法或進行系統優化。其中，MCTS的探索策略和世界模型的對齊問題至關重要，將在后續章節詳細探討。對于多向量環境和搜索樹的并行方案也是值得深究的話題，一個簡易的對比示意圖如下所示，具體的數據流和開銷分析請參考原論文附錄 E。

圖4：各種不同決策算法在數據收集時的并行化設計方案對比圖。MCTS 系列算法需要更加深度優化的并行策略。

為了測試不同算法之間的性能差異和LightZero中的算法變體作為通用決策求解器的能力，實驗部分在各種強化學習環境中進行了廣泛的對比。一共評估了下面這些算法變體，包括AlphaZero、MuZero、EfficientZero、Sampled MuZero、Stochastic MuZero、Gumbel MuZero和LightZero中綜合改進的算法變體。下面的部分展示了其中一些Benchmark結果作為例子。有關環境和算法的詳細設置，更詳盡全面的 Benchmark結果和相關分析，請參閱論文的附錄 B、C、G。

圖5：圖示為在六大代表性Atari 環境下（圖像輸入），LightZero集成的各種算法的性能對比。橫坐標代表環境步驟（Env Steps），縱坐標表示在20個 episode 中的平均收益（Return）。在此上下文中，MuZero w/ SSL代表加入了自監督損失的原始 MuZero 算法。EfficientZero則指通過自監督損失及value_prefix增強的MuZero算法。Sampled EfficientZero則在EfficientZero的基礎上引入了與采樣相關的改進技術。

圖6：圖示為AlphaZero與MuZero在Connect4和Gomoku中的性能比較。AlphaZero相比于MuZero表現出顯著的樣本效率優勢，說明在環境模擬器可直接使用時，AlphaZero將有明顯優勢。然而，即便在無模擬器的環境下，MuZero仍可產生相當的結果，這充分體現了其廣泛的適應性。

圖7：上部：圖示為Sampled EfficientZero在連續動作空間環境中，采用不同策略建模技術的性能比較。下部：圖示為該算法在MuJoCo連續動作空間環境中，應用各種策略建模方法的性能比較。隨著動作空間尺寸的增加，前者（連續空間離散化版本）表現逐漸下降，而后者表現則相對穩定。

圖8：圖示為在不同模擬成本下，Gumbel MuZero和MuZero的性能比較。在模擬次數有限時，Gumbel MuZero相對于MuZero 表現出顯著的性能優勢，突顯了其在設計低時延MCTS智能體上的潛力。針對Gomoku（棋盤尺寸=6），評估了sim={20, 10}；針對LunarLander-v2，評估了sim={20, 10, 5}；針對Atari Games，評估了sim={50, 16, 2}。

圖9：圖示為在不同隨機性等級（num_chances=2 和 5）的2048環境中，Stochastic MuZero和MuZero的性能比較。在環境的狀態轉移存在顯著隨機性的環境中，Stochastic MuZero略超MuZero。然而，隨著隨機性等級的升高，Stochastic MuZero的性能也開始受限。

圖10：左圖：圖示為在MiniGrid-KeyCorridorS3R3-v0環境中，不同探索策略的性能比較（收集階段的回報）。利用好奇心機制探索狀態空間的IntrinsicExploration策略表現出較高的樣本效率。右圖：圖示為同樣環境中，LightZero實施的算法的性能比較。在特性為高維向量觀察及稀疏獎勵的環境中，自監督學習損失有助于模型對齊，但預測value_prefix 可能帶來挑戰，對學習造成阻礙。

圖11：圖示為在多智能體對抗和協作環境GoBigger的T2P2和T2P3場景中，MuZero和EfficientZero（均在獨立學習模式下進行訓練）的性能比較。兩種算法在與內置機器人對戰時均能穩定收斂，其樣本效率表現出顯著優勢，相對于其他非 MCTS方法提升了約六倍。

圖12：自監督一致性損失在不同環境中的影響，涉及到各種類型的觀測。從左到右，性能比較涉及到標準圖像輸入、緊湊向量輸入和棋盤圖像輸入，考慮了有和沒有一致性損失的情況。實驗顯示，一致性損失對于標準圖像輸入來說是至關重要的。

10: 如圖12所示，對于標準圖像輸入，一致性損失顯得至關重要。如果無視一致性損失，模型性能會顯著下滑，這揭示了學習針對高維輸入的動態模型的挑戰性。然而，對于LunarLander這樣的向量輸入環境，一致性損失只提供了輕微的優勢，這暗示在緊湊的向量觀測上實施dynamic model學習相對較為簡單。在特定的二維輸入環境，如井字游戲（TicTacToe），一致性損失依然較大，凸顯出在潛在狀態輸出間實現一致性的難度。此外，如果在一致性損失中使用不適當的超參數，可能會導致訓練無法收斂。綜合來看，LightZero中的實驗表明，一致性損失的有效性高度依賴于特定的觀察屬性。對于棋盤游戲，未來的研究方向應考慮探索適當的損失函數，以確保訓練過程中的對齊性。

另外，為了進一步拓展決策智能算法和模型的應用范圍，LightZero項目也已經在Hugging Face Model Zoo 中設立了專門的模型空間：Hugging Face Model Zoo for LightZero。這個專門為LightZero打造的模型庫集合，旨在為研究者和開發者提供一個共享、交流研究成果，并進行快速接入和實驗的模型平臺，其中包含了最前沿MCTS系列算法模型，并針對特定應用場景進行優化。致力于讓研究者能在一個標準化的環境下快速復現、測試最新的基于MCTS的決策智能算法，并進行深度的比較和評估。

Hugging Face Model Zoo for LightZero空間：

https://huggingface.co/OpenDILabCommunity