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

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

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

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

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

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

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

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

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

蒙特卡洛樹擴(kuò)散（MCTD）

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

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

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

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

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

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

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

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

MCTD 的效果得到了實(shí)驗(yàn)的驗(yàn)證

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

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

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

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

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

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

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

 該團(tuán)隊(duì)表示：「未來還將探索自適應(yīng)計(jì)算分配、基于學(xué)習(xí)的元?jiǎng)幼鬟x擇和獎(jiǎng)勵(lì)塑造，以進(jìn)一步提高性能，為更具可擴(kuò)展性和靈活性的 System 2 規(guī)劃鋪平道路?！?/span>