擴散模型（Diffusion Model）通過利用大規(guī)模離線數(shù)據(jù)對軌跡分布進行建模，能夠生成復(fù)雜的軌跡。與傳統(tǒng)的自回歸規(guī)劃方法不同，基于擴散的規(guī)劃器通過一系列去噪步驟可以整體生成完整軌跡，無需依賴前向動力學模型，有效解決了前向模型的關(guān)鍵局限性，特別適用于具有長周期或稀疏獎勵的規(guī)劃任務(wù)。

盡管擴散模型具有這些優(yōu)勢，但如何通過利用額外的測試時間計算（TTC）來有效提高規(guī)劃精度仍然是一個懸而未決的問題。一種潛在的方法是增加去噪步驟的數(shù)量，或者增加采樣次數(shù)。然而，已知增加去噪步驟帶來的性能提升會迅速趨于平緩，而通過多個樣本進行獨立隨機搜索的效率非常低，因為它們無法利用其他樣本的信息。

另一方面，蒙特卡洛樹搜索（MCTS）則具有強大的 TTC 可擴展性。通過利用迭代模擬，MCTS 可以根據(jù)探索性反饋改進決策并進行調(diào)整，使其可以隨著計算量的增加而有效地提升規(guī)劃準確度。

這種能力使 MCTS 成為了許多 System 2 推理任務(wù)的基石，例如求解數(shù)學問題和生成程序。

然而，與基于擴散的規(guī)劃器不同，傳統(tǒng)的 MCTS 依賴于前向模型來執(zhí)行樹 rollout。這也就繼承了其局限性，包括失去全局一致性。

除了局限于離散動作空間之外，生成的搜索樹在深度和寬度上都可能變得過大。這會導致計算需求變得非常大，特別是當任務(wù)場景涉及到長遠規(guī)劃和大動作空間時。

那么，關(guān)鍵的問題來了：為了克服擴散模型和 MCTS 各自的缺陷，同時提升基于擴散的規(guī)劃的 TTC 可擴展性，可以將擴散模型與 MCTS 組合起來嗎？又該怎么去組合它們？

近日，Yoshua Bengio 和 Sungjin Ahn 領(lǐng)導的一個團隊為上述問題提供了一個答案，提出了蒙特卡洛樹擴散（MCTD）。這是一種將擴散模型的生成力量與 MCTS 的自適應(yīng)搜索功能相結(jié)合的全新框架。該團隊表示：「我們的方法將去噪（denoising）重新概念化為一個樹結(jié)構(gòu)過程，允許對部分去噪的規(guī)劃進行迭代評估、修剪和微調(diào)。」

• 論文標題：Monte Carlo Tree Diffusion for System 2 Planning
• 論文地址：https://arxiv.org/pdf/2502.07202v1

蒙特卡洛樹擴散（MCTD）

簡單來說，MCTD = 擴散模型 + MCTS。該框架整合了基于擴散的軌跡生成以及 MCTS 的迭代搜索能力，可實現(xiàn)更加高效和可擴展的規(guī)劃。

具體方法上，MCTD 有三項創(chuàng)新。

• 第一，MCTD 將去噪（denoising）過程重構(gòu)成了一種基于樹的 rollout 過程，于是便能在維持軌跡連貫性的同時實現(xiàn)半自回歸的因果規(guī)劃。
• 第二，其引入了引導層級作為元動作（meta-action），從而可實現(xiàn)「探索」與「利用」的動態(tài)平衡，進而確保在擴散框架內(nèi)實現(xiàn)自適應(yīng)和可擴展的軌跡優(yōu)化。
• 第三，其采用的模擬機制是快速跳躍去噪（fast jumpy denoising）。從名字也能看出來，該機制的效率肯定很高 —— 不使用成本高昂的前向模型 rollout 即可有效估計軌跡質(zhì)量。

基于這些創(chuàng)新，便可以在擴散過程中實現(xiàn) MCTS 的四大步驟，即選擇、擴展、模擬和反向傳播，從而有效地將結(jié)構(gòu)化搜索與生成式建模組合到了一起。

蒙特卡羅樹擴散（MCTD）的兩種視角。

上圖的（a）為 MCTS 視角：展示了 MCTD 一輪的四個步驟 —— 選擇、擴展、模擬和反向傳播 —— 在一個部分去噪樹上的過程。每個節(jié)點對應(yīng)一個部分去噪的子軌跡，邊標記為二元引導級別（0 = 無引導，1 = 有引導）。在新節(jié)點擴展后，執(zhí)行「跳躍」去噪以快速估計其值，然后沿著樹中的路徑反向傳播。

上圖的（b）為擴散視角：同一過程被視為在去噪深度（縱軸）和規(guī)劃范圍（橫軸）上的部分去噪。每個彩色塊表示在特定噪聲水平下的部分去噪規(guī)劃，顏色越深表示噪聲越高。不同的擴展（0 或 1）在規(guī)劃方向上創(chuàng)建分支，代表替代的軌跡優(yōu)化。值得注意的是，整行同時去噪，但去噪水平不同。

MCTD 框架將這兩種視角統(tǒng)一了起來。整體的算法過程如下所示：

MCTD 的效果得到了實驗的驗證

該團隊也通過實驗驗證了蒙特卡洛樹擴散的效果。他們采用的評估任務(wù)套件是 Offline Goal-conditioned RL Benchmark（OGBench），其中涉及包括迷宮導航在內(nèi)的多種任務(wù)以及多種機器人形態(tài)和機器臂操作。

下表 1 展示了質(zhì)點和機器蟻在中、大、巨型迷宮中的成功率，可以看到 MCTD 的表現(xiàn)遠超其它方法。

下圖展示了三個規(guī)劃器的規(guī)劃結(jié)果以及實際的  rollout。

三種規(guī)劃器 ——Diffuser、Diffusion Forcing 和 MCTD 生成的規(guī)劃與實際展開的對比。雖然 Diffuser 和 Diffusion Forcing 未能生成成功的軌跡規(guī)劃，但 MCTD 通過自適應(yīng)優(yōu)化其規(guī)劃取得了成功。

在點陣迷宮中等任務(wù)中，使用二元引導集 {無引導，引導} 的 MCTD 樹搜索過程可視化展示。每個節(jié)點對應(yīng)一個部分去噪的軌跡，其中左圖顯示帶噪聲的部分規(guī)劃，右圖顯示快速去噪后的規(guī)劃。搜索通過選擇無引導或引導來擴展子節(jié)點，評估每個新生成的規(guī)劃，并最終收斂到高亮的葉節(jié)點作為解決方案。

下面兩個表格則展示了不同方法在機器臂方塊操作任務(wù)以及視覺點迷宮任務(wù)上的結(jié)果。

總體而言，MCTD 在長期任務(wù)上的表現(xiàn)優(yōu)于現(xiàn)有方法，可實現(xiàn)卓越的可擴展性，并得到高質(zhì)量的解決方案。

 該團隊表示：「未來還將探索自適應(yīng)計算分配、基于學習的元動作選擇和獎勵塑造，以進一步提高性能，為更具可擴展性和靈活性的 System 2 規(guī)劃鋪平道路?！?/span>