Monte Carlo Tree Diffusion for System 2 Planning
Authors: Jaesik Yoon
Hyeonseo Cho
Doojin Baek
Yoshua Bengio
Sungjin Ahn
Venue: ICML 2025 spotlightposter (2025)
Keywords: Diffusion
MCTS
Long-term Planning
Offline RL
Goal-conditioned RL
Inference-Time Scaling
Primary area: reinforcement_learning->planning
TL;DR: Monte Carlo Tree Diffusion (MCTD) merges MCTS with diffusion-based partial denoising, enabling superior long-horizon planning by adaptively exploring and refining trajectories.

Abstract: Diffusion models have recently emerged as a powerful tool for planning. However, unlike Monte Carlo Tree Search (MCTS)—whose performance naturally improves with inference-time computation scaling—standard diffusion‐based planners offer only limited avenues for the scalability. In this paper, we introduce Monte Carlo Tree Diffusion (MCTD), a novel framework that integrates the generative strength of diffusion models with the adaptive search capabilities of MCTS. Our method reconceptualizes denoising as a tree‐structured process, allowing partially denoised plans to be iteratively evaluated, pruned, and refined. By selectively expanding promising trajectories while retaining the flexibility to revisit and improve suboptimal branches, MCTD achieves the benefits of MCTS such as controlling exploration-exploitation trade-offs within the diffusion framework. Empirical results on challenging long‐horizon tasks show that MCTD outperforms diffusion baselines, yielding higher‐quality solutions as inference-time computation increases.

Monte Carlo Tree Diffusion for System Planning
Jaesik Yoon Hyeonseo Cho
Doojin Back'
Yoshua Bengio Sungjin Ahn
Abstract
error accumulation (Hafner et al., 2022; Hamed al., 2024), making it particularly well-suited for planning tasks with long horizons or sparse rewards.
Diffusion models have recently emerged as a powerful tool for planning. However, unlike Monte Carlo Tree Search whose per- formance naturally improves with inference-time scaling-standann diffusion-based planners of- fer only limited avenues for scalability. In this paper, we introduce Monte Carlo Tree Diffu- sion (MCTD), a novel framework that integrates the generative strength of diffusion models with the adaptive search capabilities of MCTS. Our method reconceptualizes the denoising as tree- structured process, allowing partially denoised plans to be iteratively evaluated, pruned, and re- fined. By selectively expanding promising tra- jectories while retaining the flexibility re- visit and improve suboptimal branches, MCTD achieves the benefits of MCTS such as control- ling exploration-explootation trade-offs within the diffusion framework. Empirical results on chal- lenging long-horizon tasks show that MCTD outperforms diffusion baselines, yielding higher- quality solutions as inference-time computation increases.
Despite their strengths, it remains uncertain how diffusion- based planners can effectively enhance planning accuracy through the computation scaling on the inference time-a property referred to inference-time scalability, One poten- tial approach is to increase the number of denoising steps or, alternatively, draw additional samples (Zhou et al., 2024). However, it is known that performance gains from increas- ing denoising steps plateau quickly (Karras et al., 2022; Song et al., 2021a;b) and independent random searches with multiple samples are highly inefficient as they fail to leverage information from other samples. Moreover, how to effectively manage the applorationeeppootatiie tradeoff within this framework also remains unclear.
In contrast, Monte Carlo Tree Search (MCTS) (Coulom, 2007), a widely adopted planning method, demonstrates robust inference-time scalability. By leveraging iterative simulations, MCTS refines decisions and adapts based on ex- ploratory feedback, making highly effective in improving planning accuracy as more computation is allocated. This capability has established MCTS as cornerstone in many System 2 reasoning tasks, such as mathematical problem- solving (Guan et al., 2025; Zhang et al., 2024a) and program synthesis (Brandfonbrener et al., 2024). However, unlike diffusion- -based planners, traditional MCTS relies on a for- ward model for tree rollouts, inheriting its limitations includ- ing losing global consistency. In addition to being restricted to discrete action spaces, the resulting search tree can grow excessively large in both depth and width. This leads to significant computational demands, particularly in scenarios involving long horizons and large action spaces.
1. Introduction
Diffusion models have recently emerged a powerful ap- proach to planning, enabling the generation of complex trajectories by modeling trajectory distributions using large- scale offline data (Janner et al., 2022; Ajay et al., 2023; Zhou et al., 2024: Chen et al., 2024a;e:b). Unlike traditional autoregressive planning methods, diffusion-basee planners, such as Diffuser (Janner et al., 2022), generate entire tra- jectories holistically through series of denoising steps, eliminating the need for a forward dynamics model. This approach effectively addresses key limitations of forward models, such poor long-term dependency modeling and
This raises a crucial question: how can we combine the strengths of Diffuser and MCTS overcome their lim- itations and enhance the inference-time scalability of diffusion-based planning? To address this, we propose Monte Carlo Tree Diffusion (MCTD), framework that integrates diffusion- based trajectory generation with the it- erative search capabilities of MCTS for more efficient and scalable planning.
'KAIST ?SAP 'Mila "New York University. Correspon- dence Jaesik Yoon ejaesik.yoone kc. Sungjin Ahn <sungjin ahn kr.>.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
MCTD builds on three key innovations. First, it restructures denoising into a tree-based rollout process, enabling semi-

---

Monte Carlo Tree Diffusion
autoregressive causal planning while maintaining trajectory coherence. Second, it introduces guidance levels as meta- actions to dynamically balance exploration and exploitation, ensuring adaptive and scalable trajectory refinmmen within the diffusion framework. Third, it employs fast jumpy de- noising as a simulation mechanism, efficiently estimating trajectory quality without costly forward model rollouts. These innovations enable the four steps of MCTS (Selec- tion, Expansion, Simulation, and Backpropagation) within diffusion planning, effectively bridging structured search with generative modeling. Experimental results show that MCTD outperforms existing approaches in long-horizon tasks, achieving superior scalability and solution quality.
is high (e.g., future plans), without requiring complete transition from full noise to no noise across the entire tra- jectory. Such token-level control is particularly beneficial in domains that demand causal consistency, such as long- horizon planning problems.
2.2. Monte Carlo Tree Search
Monte Carlo Tree Search (MCTS) is a planning algo- rithm that combines tree search with stochastic sim- ulations to effectively balance exploration and ex- ploitation (Coulom, 2007). Typically, MCTS proceeds in four stages: selection, expansion, simulation, and backpropagation. During selection, the algorithm starts at root node corresponding to the current state (or partial plan) and descends the tree according to policy such as Upper Confidence Bounds for Trees (UCT) (Kocsis & Szepesvári, 2006), aiming to choose the most promising child at each step or explore unexperienced states for balanced. Once a leaf or expandable node reached, expansion adds new child nodes representing previously unexplored actions or plans.
The main contributions of this paper are as follows: First, to the best of our knowledge, this the first work propose an MCCTS-integraan diffusion planning framework that explic- itly incorporates the four steps of MCTS, providing an ef- fective inference-time scaling method for diffusion models. Second, we introduce three key innovations: Denoising as Free-Rollout, Guidance Levels as Meta-Actions, and Jumpy Denoising as Fast Simulation. Lastly, we present experimen- tal results demonstrating the effectiveness of MCTD.
Following expansion, simulation estimates the value of a newly added node by sampling a sequence of actions un- til reaching a terminal condition or a predefined rollout depth. The resulting outcome (e.g., cumulative reward in a Markov Decision Process) then updates the node's s estimated value through backpropagation, propagating this informa- tion to all ancestor nodes in the tree. Over multiple iterations, MCTS prunes unpromising branches and refines its value estimates for promising subtrees, guiding the search toward more effective solutions.
2. Preliminaries
2.1. Diffuser:
Diffusion Models for Planning
Diffuser (Janner et al., 2022) addresses long-horizon decision-making by treating entire trajectories as a matrix
80 81
ST
x
Go
a1
or
where 8t and a, denote the state and action at time t, re- spectively. A diffussio process is then trained to iteratively remove noise from samples of x, ultimately producing co- herent trajectories. In practice, this corresponds to reversing a forward noise-injection procedure by learning denoiser (x) over the trajectory space. Since po by itself does not encode reward or other task objectives, Diffuser optionally incorporates a heuristic or learned guidance function Jo (x). This function predicts the return or value of a partially de- noised trajectory, thereby biasing the sampling distribution:
3. Monte Carlo Tree Diffusion
In this section, we first outline the key concepts that en- able MCTD planning: (1) Denoising as Tree-Rollout, (2) Guidance Levels as Meta-Actions, and (3) Jumpy Denoising as Simulation. These innovations form the foundation of MCTD, bridging the gap between traditional tree search methods and diffusion-based planning. We then describe the MCTD process in terms of the four steps: Selection Expansion, Simulation, and Backpropagation.
(x)
(x(x) exp (x)).
(1)
3.1. Denoising as Tree-Rollout
Accordingly, at each denoising step, gradient information from Jg nudges the model toward trajectories that appear both feasible (as learned from the offline data) and promis- ing with respect to returns, in manner akin to classifier guidance in image diffusion (Dhariwal & Nichol, 2021).
Traditional MCTS operates on individual states, leading to deep search trees and significant scalability challenges. Because each node in the tree represents single state, the depth of the tree increases linearly with the planning hori- zon, resulting in an exponential growth of the search space. Furthermore, it lacks the holistic perspective of the entire trajectory that diffusion- planners inherently provide. On the other hand, Diffuser does not offer the tree-like struc-
Diffuusio Forcing (Chen et al., 2024a) extends the above framework by allowing tokenization of x. This tokenization allows each token to be denoised at a different noise level, enabling partial denoising of segments where uncertainty
2

---

Monte Carlo Tree Diffusion
Planning Horizon
Demoising
Xi
Selected node I/
1/0
1/1
1/0
1/
1/0
1/1
Denoising
X
Selection
1/0
X1:
Expanded node
Fast denoising
Value update
x1:2
Expansion
X1.8
Simulation
Evaluate n(x)
(I) Sclection
(2) Expansion
(3) Simulation
(4) Backpropagation
(a) MCTS Perspective
(b) Diffusion Perspective
Figure 1. Two perspectives on Monte Carlo Tree Diffusion (MCTD). (a) MCTS Perspective: The four steps of an MCTD Sunn-SSlectton, Expansion, Simulation, and Backpropagatinn-ar illustrated on partial denoising tree. Fach node corresponds to a partially denoised sub-trajectory, and edges are labeled with binary guidance levels (0 = guidance, I= guided). After new node is expanded, "jumpy" denoising is performed to quickly estimate its value, which is then backpropagated along the path in the tree. (b) Diffusion Perspective: The same process is viewed as partial denoising across both denoising depth (vertical axis) and planning horizon (horizontal axis). Each colored block represents a partially denoised plan at a specific noise level, with darker shades indicating higher noise. Different expansions (0 1) create branches in the forward planning direction, representing alternative trajectory refinements. Notably, the entire row is denoised simultaneously, but with varying denoising levels. The MCTD framework unifies these two perspectives.
ture necessary for intermediate decision- point searches that effectively balance exploitation and exploration.
instance, in of our planning experiments, we use S 5 and N 500.
To address this issue, we first introduce the Denoising as Tree-Rolouu process by leveraging the eemi-autoregressivve denoising process (Chen et al., 2024a). Specifically, we partition the full trajectory x = (x1, (2; ...,...................... xN) into S subplans, x = (2,,Xx,...... xs) such that n,x, 0. Un- like the standard diffuser, where all subplans share the same global denoising schedule, approach allows assigning independent denoising schedules to each subplan. By apply- ing faster denoising to earlier subplans and slower denoising to later subplans, the process effectively becomes causal and semi-autoregressive. This allows the denoising process to determine the future conditioned on the already determined past. Consequently, while preserving the globally consistent and holistic generation advantages of Diffuser, the denois- ing process approximates:
3.2. Guidance Levels Meta-Actions
In MCTS, constructing and searching the tree becomes com- putationally expensive in large action spaces and is fun- damentally impractical for continuous action spaces. To resolve this issue, in MCTD we introduce novel approach by redefining the exploration-exploiaaiin trade-off terms of meta-actions. While meta-actions can be any discrete decision that can control the work we propose implement as guidance levels applied during the denoising process.
in this
For simplicity, consider two guidance levels: GUIDE and NO-GUIDE. In MCTD, we observe that sampling from the prior distribution p(x), i... using the standard diffusion sam- pler represents exploratory behavior as it does not attempt to achieve any goal but only imitates the prior behavior contained in the offline data. Thus, we correspond it to the NO-GUIDE meta-action. Conversely, sampling from goal-directed distribution Py (x). such as through classifier- guided diffusion (Dhariwal & Nichol, 2021), represents exploitative behavior, and thus meta-action, GUIDE, is as- signed. This steers the sampling process toward achieving a specific goal defined by a reward function (x(x).
p(x)= Rd
p(x X1:0-1)
Notably, although this formulation appears autoregressive, it is still executed as a single denoising process by controlling noise levels across subplans (Chen et al., 2024a).
In MCTD, each subplan x., which represents a temporally extended state, is treated as a node in the search tree, rather than using individual states Fn as nodes. This allows the tree to operate at higher level of abstraction, improving efficiency and scalability. Denoising the entire plan x can be viewed as rolling out a sequence of these nodes through a single denoising process Diffuser. Because S & N, the tree depth becomes much smaller than that in MCTS. For
Next, we integrate the concept of meta-actions with the tree-rollout denoising process described in Eqn. (2). To achieve this, we introduce a guidance schedule (g.....,s)) which assigns a guidance meta-action g. E {GUIDE, NO GUIDE} to each corresponding subplan x,, no guidance chosen, the subplan sampled from the
3

---

Monte Carlo Tree Diffusion
explorative tree-rollout prior p(x, (xx](x while subplan with guidance is sampled from (x(X,] X1:s-1) 1) if the GUIDE meta-action is chosen. This allows the guidance levels to be independently assigned to each subplan unlike the standard Diffuser assigning the guidance throughout both the whole trajectory and the whole denoising process.
Algorithm 1 Monte Carlo Tree Diffusion
1:
procedure MCTD(root, iterations) for i 1 to iterations do
2:
3:
node t- root while ASFULLYEXPAADEDD/IN and not ISLEAF(Rode) do node 4- BESTUUCCCUL.D(addt)
4:
5:
6:
By dynamically adjusting the guidance schedule g. this balancing at the level
7:
end while
8: [SEPPANDABLE(rrdd)) then
enables the of subplans within single denoising process. The extended tree-rollout denoising process effectively approximates:
9:
9. + SELECTMETAACTIIO(Oude newSSbblaa + DENOISESUBPLANNnode 9a)
10:
11:
child + CREATEENoDDEhheuuuuyeee
12:
ADDCHILD(nndd,,chil child)
13:
node 4- child
p(x|g) Re II ((x_|xx:e-111
(3)
14:
end if
15:
partial <- G
As a result, this approach enables efficient and scalable
16:
remaining - FASTDENOISING(partic))
17:
full t (portial, remaining) reward + SWAALUATEPLAN(flllPM
planning, even in complex or continuous action spaces,
18:
19:
while mode + null do
3.3. Jumpy Denoising as Fast Simulation
20:
node. mnde.visittoontte +- node. mideevvsittoont +1 node.value node.value + reward
21:
In MCTS, evaluating a node far from a leaf node, where evaluation of the plan is feasible, is critical requirement.
22:
UPDATEMFTAACTTIOSSHEED node, reward)
23:
This is typically addressed in one of two ways: using fast
24:
node &- node.parent end while
25:
forward dynamics model to roll out trajectories to the leaf node, which is computationally expensive, or approximat- ing the node's 's value via bootstrapping, offering faster but less accurate results. However, how to effectively incorpo- rate these simulation strategies into the diffuser framework remains an open question.
26:
end for
27: return BESTCHILD(roOt)
28:
end procedure
tree structure. Unlike traditional MCTS, MCTD nodes cor- respond to temporally extended states, enabling higher-level reasoning and reducing tree depth, which improves scalabil- ity. The guidance schedule g is dynamically adjusted during this step by UCB to balance exploration (NO_GUIDE) and exploitation (GUIDE).
In MCTD, we implement this simulation functionality ing the fast jumpy denoising process based on the Denois- ing Diffusion Implicit Model (DDIM) (Song et al., 2021a). Specifically, when the tree-rollout denoising process has progressed up to the s-th subplan, the remaining steps are denoised quickly by skipping every C step:
Expansion. Once a leaf or partially expanded node is se- lected, the expansion step generates new child nodes by extending the current partially denoised trajectory. Each child node corresponds to a new subplan generated us- ing the diffusion model. Depending on the meta-action Ja. the subplan either sampled from the exploratory prior distribution p(x,/x1:s>> or the goal-seeking distribution Pg(X, (x, (1:s-1). The meta-action controls which behavior is used for expansion, and the newly generated node is added to the tree as an extension of the current trajec- tory. Importantly, the guidance levels are not restricted to binary choices. For instance, it is possible to general- ize the meta-actions to multiple guidance levels, such as {ZERO, LOW, MEDIUM, HIGH), offering finer control over the balance between exploration and exploitation during the expansion process.
(4) This produces full trajectory x (X1:s: 13)1: which is then evaluated using the reward function r(x). While this fast denoising process may introduce larger approximation errors, it is highly computationally efficient, making well- suited for the simulation step in MCTD.
X++1:8 p(Xx+1.s|X1.s g).
3.4. The Four Steps of an MCTD Round
Building on the description above, we detail how the four Expansion, Simula-
traditional steps of tion, and Backpropagation--ar adapted and implemented in MCTD. This process is illustrated in Figure 1.
Selection. In MCTD, the selection process involves travers- ing the tree from the root node to leaf or partially ex- panded node. At each step, child node is chosen based on a selection criterion such as Upper Confidence Bound (UCB). Importantly, this step does not require computation- ally expensive denoising; it simply traverses the existing
Simulation. Simulation in MCTD is implemented via the fast jumpy denoising. When a node is expanded, the remain- der of the trajectory is quickly completed by using the fast denoising as described in Section 3.3. The resulting plan
4

---

Monte Carlo Tree Diffusion
is then given to the plan evaluator r rxx) This approach main- tains computational efficiency while providing sufficient estimate of the plan's 's quality.
search with diffusion models for full trajectory generation, bridging structured search with generative planning.
Backpropagation. After the simulation step, the reward obtained from evaluating the complete plan is backprop- agated through the tree to update the value estimates of all parent nodes along the path to the root. In MCTD, this backpropagation process also updates the iitla--ctton--bssee guidance schedules, enabling the tree to dynamically adjust balance for future iterations. This ensures that promising trajectories, as indicated by their rewards, are prioritized, while sufficient exploration is maintained to prevent premature convergence.
5. Experiments
We evaluate the proposed approach, MCTD, on a suite of tasks from the Offline Goal-conditioned RL Benchmark (OGBench) (Park et al., 2025), which spans diverse do- mains such maze navigation with multiple robot mor- phologies (e.g., point-mass or ant) and robot-arm manip- ulation. Our chosen poiks-point and antmaze navigation, multi-cube manipulation, and a newly introduced visual pointmaze-joine assess a planner's ss ability to handle long- horizon planning, sequential manipulation, and partial vi- sual observability. In the visual pointmaze, an agent per- ceives RGB image observations of the 3D environment, thereby testing each method's resilience to partial observ- ability and the ability to handle image- based planning. De- tailed experimental settings are provided in Appendix A.
the eeploration-exploiaaeee
4. Related Works
Diffusion models (Sohl-Dickstein et al., 2015) have recently shown significant promise for long-horizon trajectory plan- ning, particularly in settings with sparse rewards, by learn- ing to generate entire sequences rather than stepwise ac- tions (Janner et al., 2022; Chen et al., 2024a; Ajay et al., 2023; Zhou et al., 2024; Chen al., 2024c;b; 2025). Various enhancements have been proposed to extend their capabil- ities. For instance, Chen et (2024)))) incorporate causal noise scheduling for semi- rumiiaauorreressivvee planning, and other works introduce hierarchical structures (Chen et al., 2024c;b; 2025), low-level value learning policies (Chen et al., 2024b), or classifier-free guidance (Ajay et al., 2023; Zhou et al., 2024). Despite these developments, the explicit interplay between exploration and exploitation in diffusion sampling has received relatively little attention.
5.1. Baselines
We compare MCTD against several baselines, each employ- ing distinct strategy for test-time computation. First, we in- clude the standard single-shot Diffuser (Janner et al., 2022), which plans only at the start of an episode and executes that plan without further adjustments. Next, we consider two methods that allocate additional test-time compute dif- ferently. MiffuserrRRelaaningg periodically replans at fixed intervals, allowing partial corrections during the episode. This enables us to gauge the benefits of iterative replanning without a full tree search. et al., 2024) generates multiple random trajectorics in par- allel and selects the highest-scoring one according to the same reward used by MCTD. While this increases test-time sampling and applies multiple guidance levels for different samples. it lacks the systematic expansion and pruning of tree-search algorithm. We also compare our method to Diffusion Forcing (Chen et al., 2024a), which introduces a causal denoising schedule that enables semi-uuttreerresivve trajectory generation. This comparison helps isolate the benefits of autooeroorerrssivvee denoising from our full tree- structured approach. Further implementation details about all baselines can be found in Appendix A.1.
-andoom Search (Zhou
Before Diffusion models, the long-hou morizooi planning prob- lem has been also addressed through the sub-goal based tree search (Jurgenson et al., 2020) and environment knowledge based high-level tree search (de Waard et al., 2016).
Chen et al. (2024a) further propose Monte Carlo Guidance (MCG) to leverage multiple sub-plans and average their guidance signals, thereby biasing denoising toward higher- reward outcomes. While this can encourage planning toward an expected reward over multiple rollouts, it does not imple- ment an explicit search mechanism. The detailed discussion Appendix C. Similarly, Ye et al. (2025) apply a discrete diffusion denoising process to tasks such as chess, implicitly modeling MCTS without explicit tree expansion.
5.2. Maze Navigation with Point-Mass and Ant Robots
MCTS (Coulom, 2007) has achieved impressive results across various maniinn problems, particularly when combined with learned policies or value networks (Silver et al., 2016; Schrittwieser et al., 2020). It has also been applied to System reasoning in Large Language Models (LLMs) to enhance structured problem-solving (Xiang et al., 2025; Yao et al., 2024; Zhang al., 2024b). However, to the best of our knowledge, this work is the first to integrate tree
We begin with the pointmaze and antmaze tasks from OG- Bench, featuring mazes of varying sizes (medium, large, and giant). The agent is rewarded for reaching designated goal region, and the relevant dataset (navigate) comprises long- horizon trajectories that challenge each method's 's capacity for exploration. As in prior work (Janner et al., 2022), we employ a heuristic controller for the pointmaze environment, whereas the antmaze environment uses value-learning pol-
5

---

Monte Carlo Tree Diffusion
Table LLng-HHrrzzon Maze Results. Success rates (%) on pointmaze antmaze environments medium, large, and giant mazes for the navigate datasets. Each cell shows the mean + standard deviation.
Diffuser
Environment
Dataset
Diffussion Forcing
MCTD
Base
Replanning
Random Search
medium-navigate--0
+
60 0
60 +
16
0
pointmaze
large-navigate-vO
44 +
40 +
34 + 13
74 +
98 6
giant navigatev)
0 +
0
4 8
50 = 10
100 0
mediun--navigate-v0
36 + 15
40 18
48 4 10
10
100 0
antmaze
14 + 16
26 26
20 1 0
57 1 6
98
gjant-navigate-vO
010
010
418
24 12
94
Diffuser
Diffusion Forcing
MCTD
0, Rewardi-0
Plans
0/0, Reward=0
0/1, Reward=0.26
Executed Trajector
0/0/O, Reward=0
0/1/0, Reward=0.31 0/1/1, Roward=0.43
Figure 2. Comparison of generated plans and actual rollouts three planners- Wlanees--Diffuser, Diffusion Forcing, and MCTD. While Diffuser and Diffusion Forcing fail to produce successful trajectory plans, MCTD succeeds refining its plan adaptively.
0M/0/11 Reward-0.31
Figure 3. Visualization of the MCTD tree-search process with binary guidance set {NO. GUIDE, GUIDE] on a pointmaze- medium task. Each node corresponds to partially denoised tra- jectory, where the left image shows noisy partial plan and the right image shows the plan after fast denoising. The search ex- pands child nodes by selecting either NO.GUIDE or GUIDE, evaluates each newly generated plan, and ultimately converges the highlighted leaf as the solution.
icy (Chen al., 2024b) for low-level control. Appendix A.5 provides more details this task evaluation.
Results Table presents success rates across medium, large, and giant mazes for both point-mass and ant robots. MCTD consistently surpasses other methods by large margin. Notably,
Search, despite using roughly the number of denoising steps MCTD, demonstrates improvement over one-shot Diffuser, derscoring the importance of tree-based systematic branch- ing over random sampling. Diffuser-Repph Replanning similarly fails to outperform the base Diffuser, whereas Diffusion Forcing exhibits higher success rates in larger mazes, in- dicating benefits from a semi-auttreeressive schedule for long horizons. However, only MCTD achieves near-perfect performance the most challenging (large, giant) mazes, owing to its ability to refine partial trajectories and avoid dead-ends or inconsistent plan segments. This process is visualized in Figure The performance on the giant map impressive, because the map contains numerous detour paths to reach the goal as shown in Figure 2. The comparison with Diffuser-Randont Search further highlights MCTD's efficiency, since it offers substantially better outcomes under the same computational budget.
principal reason for the degradation in the baselines is that the value learning policy (Wang al., 2023) struggles in nigh-dimensional control tasks. Despite performance degra- dations due inaccurate execution, MCTD consistently outperforms baselines and the performance gap widens as the maze size increases.
5.3. Robot Arm Cube Manipulation
For multi-cube manipulation tasks from OGBench, robot arm must move one to four cubes to specific table loca- tions as shown in Figure 4a. Increasing the number of cubes grows both the planning horizon and the combinatorial com- plexity. While the planner (MCTD or a baseline) issues high- level positions for the cubes. value-learning policy (Wang et al., 2023) executes local actions. We augment MCTD with object-wise guidance by selecting which eube is ma- nipulated at each step, avoiding simultaneous movements of multiple cubes. We note that the object-wise guidance is particularly suited to MCTD because the tree-expansion mechanism naturally supports discrete imeta-actions" at each node, such as selecting which object to move next.
In the antmaze domain, MCTD continues to perform nearly perfectly, while baseline performance generally degrades more significantly than in pointmaze, except for the medium-sine sized settings under Diffusion Forcing. A
6

---

Monte Carlo Tree Diffusion
Table Robot Cube Manipulation Results. Success rates (%) for single, double, triple, and quadruple cube tasks OGBench,
Parenthetical values in the MCTD MCTD-Replanning column denote performance when using DQL performer trained only on the single-cube
dataset.
Dataset
Diffuser
Diffuser- Diffuser-Replanning Replanning
Diffusion Forcing
MCTD
MCTD-Replannin
single-play-v0
78 i 23
92 1 13
100 :
98 6
100 0
double-play-v0
12 i: 10
12 +: 13
18 + 11
22 : 11
50 1 16 (78 +: 11)
trppe--paa--00
8 =
1 8
16 + 8 0
0
6 (40 + 21)
quadruple-play-v0
0
0
0+
0 (24 +
Table 3. Visual Pointmaze Results. Success rates (%) partially observable, image-based mazes of medium and large sizes.
Dataset
Diffuser
Repluai-R Diffusion Forcing
MCTD
MCTD-Replanning
mediunt--avvggaee--
8 = 13
8 = 10
66 32
82 = 18
90 1
nangge-aviggte-v
0 0
0=0
20 + 21
Initial Observation
Goal
Diffision
Forcing
(a) Stacking Task
(b) Holistic Plan
MCTD
steps
0
50
100
150
200
250
300
330
Figure 5. Task and interaction examples in the visual pointmaze. The agent operates under partial observability, receiving pixel- based initial and goal images. In the top row, Diffusion Forcing struggles to reach the goal due focus on exploitation, failing sufficiently explore the environment. In contrast, MCTD (bottom row) effectively balances exploration and exploitation, successfully navigating to the goal.
(c) Object-wiss wise Plan
Figure (a) Illustration the robot arm cube manipulation task. (b) A holistic plan for moving two cubes, where entanglement between their partial trajectories leads an inaccurate move- ment (blue cube). (c) An object-wise plan produced by MCTD- Replanning, which avoids entanglement and achieves more reliable control over each cube's movement.
from 22% to 50%, and also benefits tasks with more cubes. Another challenge arises from the value-learning policy, which generalizes poorly as the number of cubes increases. Even if MCTD generates a suitable global plan, flawed local control undermines the final outcome. Interestingly, training the policy solely on single-cube data still allows MCTD-R CCCTD-Replanning solve portion (24%) of four-cube scenarios, illustrating the potential of object-wise guidance and iterative replanning.
Single-pass diffusion methods, by contrast, attempt to pro- duce single, holistic trajectory in single pass, leaving no opportunity to decide object-by-object. This guidance and additional task-specific reward shaping are described in Appendix A.6.
5.4. Visual Pointmaze
Results Table 2 shows success rates for one four cubes.
To evaluate our method on image-based planning, partic- ularly under partial observability, we introduce visual pointmaze environment. In this environment, the agent re- ceives top-down pixel images of the initial and goal states, as shown in Figure Unlike the visual antmaze in OGBench, which involves complex locomotion, this environment al- lows us focus how well the planner operates with raw visual inputs without relying on ground-truth states.
All methods perform comparably on single-cube tasks, but performance drops significantly for multi-cube scenarios, especially for standard diffusion baselines that struggle to decide cube ordering and handle subtasks such as stacking. However, HCTTD-Replanne mance gaps from other models for multi-cube scenarios. MCTD exhibits moderate gains (e.g., 22% two cubes), though it suffers from holistic plan entanglements when mul- tiple objects are involved as shown in Figure 4b. To address this, MCTD- CepT--Replnnnnng periodically replans, effectively separating each cube's 's movements as shown in Figure 4c. This improves the success rate on the two-cube problem
shows significant perfor-
We pre-encode observations with a Variational Autoencoder (VAE) (Kingma & Welling, 2014) to generate latent states suitable for diffusion based planning, 2; = VAB(z;). Be- cause Diffuser's s default heuristic controller inapplica-

---

Monte Carlo Tree Diffusion
Large
Giant
expanding tree- to a larger runtime as the bud- get grows. For instance, Diffuser-Random Search shows moderate runtime growth because multiple samples can be denoised in parallel batches, which is not directly feasible in MCTD's tree-structured approach. Nonetheless, MCTD can terminate early whenever a successful plan is found, thereby limiting unnecessary expansions. As a result, on the large maze, MCTD's runtime quickly saturates once the task can be solved with smaller search depth, rather than fully consuming the available budget. A similar pattern emerges in the giant maze: as success rates increase under higher budgets, the runtime growth for MCTD remains sub- linear, indicating that many expansions terminate early once suitable trajectories are located.
Success Rale
40
400
Runtia
300
200
100
3000
9000
15000
17500
Max. number denoising
Max. number of denoising
Diffuser-Random Search Figure 6. Success rates (%) and wall-clock runtime (Sec.) as the maximum number of denoising steps (test-time budget) increases for large and giant pointmaze tasks
MCTD
Diffuser
ble to purely visual inputs, we design an inverse dynamics model that maps consecutive latent observations actions, à fim (2 2t; 2e+1). For guidance, employ posi- tion estimator trained on ground-truth coordinates; however, this estimator does not update the underlying VAE, ensuring that the planner remains in a partially observable regime. Further details on this task are provided Appendix A.7.
5.6. Ablation Studies
In this section, we conduct empirical ablation studies to analyze key components of MCTD, including the statistical nature of tree search, exploration-exploitatio balance, tree depth, and the effectiveness of causal denoising and tree search. These ablation studies are primarily conducted on the PointMaze environment using the navigate datasets.
Results Table 3 compares success rates on medium and large versions of the visual pointmaze. MCTD and MCTD- Replanning show better performance than baselines on image-based, partially observable tasks. MCTD out- performs Diffuser baselines and Diffusion Forcing in the medium maze, presumably due to its tree-structured explo- ration, which better balances exploration and exploitation under perceptual partial observability. In the large maze, single-pass planners fail almost entirely, highlighting the dif- ficulty of long-horizon visual planning. Replannng improves performance by reevaluating partial plans mid- trajectory, though success remains modest overall. This re- sult underscores the inherent challenge of scaling diffusion- based methods extended horizons with visual data under partial observability and suggests that future work may need more sophisticated visual encoders or hierarchical guidance strategies.
5.6.1. GREEDY TREE SEARCH
To evaluate the effectiveness of statistical tree search in balancing exploration and exploitation, we implemented a greedy tree search baseline. This approach searches candi- dates at each branching point and selects the best sample us- ing the same reward function employed in iiffuser-Randow Search and MCTD. For comprehensive evaluation, we var- ied the number of children for this approach from 5 to 20 (MCTD uses 5 children by default).
Table 4. Greedy Tree Search Performances (%) on PointMaze.
Children#
10
15
20
Medium
62:1-6
58::11
60:10
60:10
Large
40-10
40+0
40-00
40:00
Giant
04:0
0::0
0.:0
04:0
The limited performance shown in Table 4 can be attributed to the method's lack of deep search across multiple tra- jectory samples and the absence of backtracking mecha- nisms inherent in proper tree structures.
5.5. Inference- Timee Scalability & Time Complexity
We examine how each planner leverages additional test-time computation by varying the maximum number of denoising steps and measuring both success rate and runtime on the large and giant pointmaze tasks (Figure 6). As the inference- time budget grows, MCTD consistently achieves high success rates-uutimately nearing perfection in giant mazes. In contrast, simply increasing denoising steps (Dif- fuser) or drawing more random samples (Diffuser-Randon) Search) yields modest returns, suggesting limited test-time scalability.
5.6.2. META-ACTION SET INVESTIGATION
A core mechanism of MCTD is balaneing exploration and exploitation to generate globally optimal plans. In this sec- tion, we study this balance by analyzing different meta- action set configurations. We evaluate five meta-action sets ranging from exploration- biased to exploitation biased con- figurations. Set 1 is exploration- biased, consisting of small guidance levels [0, 0.02, 0.05, 0.07, 0.1), while Set repre-
MCTD incurs additional overhead from maintaining and
8

---

Monte Carlo Tree Diffusion
sents our default configuration [0, 0.1, 0.5, 1, 1, Sets 3 and 4 are increasingly exploitation-baased with [0.5, 3, and [4, 5, 6, 7, 8], respectively. Finally, Set 5 includes ex- treme guidance levels [0.1, 0.1, 10, 100] to test whether MCTD can effectively select appropriate guidance levels while ignoring extreme values.
Table 7. Performances Causal Denoising and Tree Search Abla- tion PointMaze.
Method
DFwoCD
DF
MCTDwoCD
MCTD
Medium
44::22
40+:21
17+
100+0
Large
50++1
55+-20
78+6
98++6
Giant
81:00
34::16
18211
10010
Table Performances of Diverse Meta-Action Sets PointMaze.
tially to MCTD's performance. Tree search consistently improves performance across all maze sizes. Meanwhile, causal denoising provides particularly substantial gains for long-horizer planning, with performance improve- ments most pronounced in the giant maze.
Sets
Set
Set
Set 3
Set
Set 5
Medium
86:
100:0
100+0
98:66
100+0
Large
45::5
98:6
90=:
80=:0
100+0
Giant
54
100+00
100+00
74
100+0
The results in Table 5 demonstrate that performance de- grades when meta-actions are heavily biased toward either exploration (Set 1) or exploitation (Set 4). Notably, MCTD maintains robust performance with balanced meta-action sets, even when some values are extreme outliers (Sets 2, 3, and 5). This indicates that maintaining diverse guid- ance options is more critical than their specific values, provided both exploration and exploitation are adequately represented.
6. Limitations & Discussion
While MCTD enhances inference-time scalability by in- tegrating search-based planning with diffusion models, it remains computationally expensive due to its System 2-style deliberative reasoning. Determining when to engage in struc- tured planning versus relying on fast System 1 (model-free) planning is an open challenge, as expensive tree search may not always necessary. Adaptive compute allocation based on task complexity or uncertainty could improve efficiency.
5.6.3. SUBPLAN LENGTII
Another limitation is inefficiency in large-scale search spaces, where evaluating multiple trajectory hypotheses remains computationally demanding despite using low- dimensional meta-actions. A promising direction is amor- tized search, where the system meta-learns from inference- time search to improve exploration efficiency over time. Instead of treating initial exploration random, MCTD could incorporate inference-tine learning mechanisms re- fine its search dynamically. Additionally, self-supervised re- ward shaping could improve trajectory evaluation in sparse- reward settings. Optimizing parallelized denoising, differ- entiable tree search, and model-based rollouts could further enhance efficiency.
Another key configurable component in MCTD tree depth, which can be controlled through the subplan count S. While S is adaptively chosen based on the environment's episode length, in this section we investigate the impact of varying tree depths by manually controlling S.
Table 6. Empirical Results for Subplan Counts on PointMaze- Large.
S
1
3
5
20
Success Rate
80:10
901110
10010
92::10
Run Time (sec.)
11:1
1914
65111
131-110
Scarch #
111+31
64+550
190+:37
500+:0
The results in Table reveal an important trade-off be- tween computational efficiency and performance. With longer subplans (smaller S), the search space is reduced, leading to faster execution times but potentially compromis- ing performance quality. Conversely, very short subplans (e.g., S = 20) require more search iterations, which may exhaust the computational budget before finding optimal solutions.
7. Conclusion
We introduced Monte Carlo Tree Diffusion (MCTD), a framework designed to combine the best of both worlds: the structured search of Monte Carlo Tree Search and the generative flexibility of diffusion planning to enhance the inference-ting scalability of System planning. MCTD leverages meta-actions for adaptive exploration-exploitttion tree-structured denoising for efficient diffusion- expan- sion, and fast jumpy denoising for rapid simulation. Experi- mental results demonstrate that MCTD outperforms existing approaches in various planning tasks, achieving superior scalability and solution quality. Future work will explore adaptive compute allocation, learned meta-action selection, and reward shaping further enhance performance, paving the way for more scalable and flexible System 2 planning.
5.6.4. CAUSAL DENOISING AND TREE SEARCH
To investigate the effectiveness of causal denoising and tree search, we compared four variants: Diffusion Forcing without Causal Denoising (DFwoCD), Diffusion Forcing (DF), MCTD without Causal Denoising (MCTDwoCD), and MCTD.
The results reveal that both components contribute substan-
9

---

Monte Carlo Tree Diffusion
Acknowledgements
via hierarchical multiscale diffusion. arXiv preprint arXiv:2503.201022 2025.
This research was supported by GRDC (Global Re- search Development Center) Cooperative Hub Program (RS-2024-004361655) and Brain Pool Plus Program (No. 2021HID3A2AA0310360 through the National Research Foundation of Korea (NRF) funded by the Ministry of Sci- ence and ICT.
Coulom, R. Efficient selectivity and backup operators in monte-carlo tree search. In van den Herik, H. J., Ciancar- ini, P., and Donkers, H. H. L. M. J. (eds.), Computers and Games, pp. 22-83, Berlin, Heidelberg, 2007. Springer Berlin Heidelberg. ISBN 978-3-540-75538-8
Impact Statement
de Waard, M., Roijers, D. M., and Bakkes, S. C. Monte carlo tree search with options for general video game playing. In IEEE Conference on Computational Intelligence and Games, 2016.
This work introduces Monte Carlo Tree Diffusion (MCTD), framework that integrates diffusion models with Monte Carlo Tree Search (MCTS) for scalable, loni-horizon plan- ning. As a ceneral-purpose planning framework, MCTD itself does not pose direct risks; however, care must be taken when applying it to safety-critical domains, where decision- making impacts human well-being or critical infrastructure. Depending on its implementation, MCTD could influence high-stakes decisions in areas such as autonomous systems, healthcare, finance, necessitating robust oversight and alignment with ethical guidelines. Additionally, its compu- tational demands raise considerations around energy effi- ciency and sustainability, highlighting the need for responsi- ble AI deployment.
Dhariwal, P. and Nichol, A. Diffusion models beat gans on image synthesis. In Advances in Neural Information Processing Systems, 2021.
Guan, X., Zhang, L.. L.., Liu, Y., Shang, N., Sun, Y., Zhu, Y., Yang, F, and Yang, M. rstar-math: Small Ilms can master math reasoning with self-evolved deep thinking. arXiv preprint arXiv:2501.04511 2025.
Hafner, D., Lee, K.-H., Fischer, I., and Abbeel, Deep hierarchical planning from pixels. In Advances in Neural Information Processing Systems, 2022.
References
Hamed, H., Kim, S., Kim, D., Yoon, J., and Ahn, S. Dr. strat- egy: Model-based generalist agents with strategic dream- ing. In International Conference Machine Learning, 2024.
Ajay, A., Du, Y., Gupta, A., Tenenbaum, J. B., Jaakkola, T. S., and Agrawal, P. Is conditional generative model- ing all you need for decision making? In International Conference Learning Representations, 2023.
Janner, M., Du, Y., Tenenbaum, J., and Levine, S. Plan- ning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, 2022.
Brandfonbrener, D., Henniger, S., Raja, S., Prasad, T., Loughridge, C. R., Cassano, F., Hu, S. R., Yang, J. Byrd, W. E., Zinkov, R., and Amin, N. VerMCTS: Synthesiz- ing multi-step programs using a verifier, a large language model, and tree search. The 4th Workshop on Mathe- matical Reasoning and Al at NeurlPS, 2024.
Jurgenson, T., Avner, O., Groshev, E., and Tamar, A. Sub- goal trees: a framework for goal-based reinforcement learning. In International Conference on Machine Learn- ing. 2020.
Chen, B., Monsó, D. M., Du, Y., Simchowitz, M., Tedrake, R., and Sitzmann, V. Diffusson forcing: Next-token pre- diction meets full-sequence diffusion. In Advances on Neural Information Processing Systems, 2024a.
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion based generative models. Ad- vances in Neural Information Processing Systems, 2022.
Chen, C., Back, J., Deng, F., Kawaguchi, K., Gulcehre, C., and Ahn, PlanDQ: Hierarchical plan orchestration via d-conductor and q-performer. In International Confer- ence on Machine Learning, 2024b.
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. In International conference on Learning Repre- sentations, 2014.
Kocsis, L. and Szepesvári, C. Bandit based monte-carlo planning. In European Conference on Machine Learning, 2006.
Chen, C., Deng, F., Kawaguchi, K., Gulcehre, C., and Ahn, S. Simple hierarchical planning with diffusion. In Inter- national Conference on Learning Representations, 2024c.
Park, S., Frans, K., Eysenbach, B., and Levine, S. Ogbench: Benchmarking offline goal-conditioned rl. In Interna- tional conference on Learning Representations, 2025.
Chen, C., Hamed, H., Back, D., Kang. T., Bengio, Y., and Ahn, S. Extendable long-horizon planning
10

---

Monte Carlo Tree Diffusion
Schrittieser, J., Antonoglou, I,, Hubert, T., Simonyan, K., Sifre, L.,, Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839): 604-609, 2020.
Zhou, G., Swaminathan, S., Raju, R. V., Guntupalli, J. S., Lehrach, W., Ortiz, J., Dedieu, A., Lázaro-Gredilla M., and Murphy, K. Diffusson model predictive control. arXiv preprint arX0v:2410005364 2024.
Silver, D., Huang, A., Maddison, C. J., Guez, A., Sifre, L., Van Den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershelvam, V., Lanctot, M., et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484-48)) 2016.
Sohl-Dickstein, Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International Conference Machine Learning, 2015.
Song. J., Meng, C., and Ermon, S. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021a.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-hased generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021b.
Wang, Z., Hunt, J., and Zhou, M. Diffusion policies as an expressive policy class for offline reinforcement learning. In International Conference on Learning Representations, 2023.
Xiang, V., Snell, C., Gandhi, K., Albalak, A., Singh, A., Blagden, C.,, Phung, D., Rafailov, R., Lile, N., Mahan, D., et al. Towards system reasoning in llms: Learning how to think with meta chain-of- ofain-of-though arXiv preprint arXiv:2501. 046622, 2025.
Yao, S., Yu, D., Zhao, J.. Shafran, I., Griffiths, T., Cao, Y., and Narasimhan, K. Tree of thoughts: Deliberate problem solving with large language models. Advances in Neural Information Processing Systems, 2024.
Ye, J., Wu, Z.,, Gao, J., Wu, Z., Jiang, X., Li, Z., and Kong, L. Implicit search via discrete diffusion: A study on chess. In International Conference on Learning Representations, 2025.
Zhang, D., Huang, X., Zhou, D., Li, Y., and Ouyang, W. Accessing gpt-4 level mathematical olympiad solutions via monte carlo tree self-refine with Ilama-3 8b. arXiv preprint arXiv:2406.07399, 2024a.
Zhang, D., Zhoubian, S., Hu. Z., Yue, Y., Dong, Y., and Tang, J. Rest-mets*; Llm self-training via process reward guided tree search. In Advances in Neural Information Processing Systems, 2024b.

---

Monte Carlo Tree Diffusion
A. Experiment Details
A.1. Baselines
We compare MCTD against several baselines, each highlighting different ways to leverage inference-tine scaling on diffusion models or broader planning strategies:
Diffuser (Janner et al., 2022): We use the standard single-shot Diffuser as our primary diffusion baseline. It generates a full trajectory by denoising from maximum to zero noise in one pass, guided by return- return-predictnng function Jg. While it can produce coherent plans, it lacks any mechanism to adapt or refine these plans at test time.
To partially address the lack of iterative search, we evaluated a replanning variant of Diffuser. After an initial denoised plan is sampled, the policy is periodically re-invoked at predefined intervals (e.g., every 10 time steps), using the agent's 's new state as the "start" for a new diffusion process. This strategy leverages additional test-time compute by attempting incremental corrections but does not share information across different replans.
Rnffuser-Reelanningg
We consider a random-searee search variant of Diffuser that generates multiple trajectories from random noise and selects the best candidate according heuristic or learned value function. This approach can be seen "Sample-Scort Score-Ranke method proposed in (Zhou et al., 2024), increasing test-time compute by drawing more samples but lacking the systematic exploration or adaptation found in tree-search paradigms.
Search:
Diffusion Forcing: Diffusion Forcing (Chen et al., 2024a) extends Diffuser by introducing a tokenized, causal noise schedule. Sub-plans of the trajectory can be denoised at different rates, allowing partial re-sampling of specific time segments without re-optimizing the entire plan. In our experiments, we evaluate Diffusion Forcing with the Transformer backbone provided in their official repository, which performed better in our preliminary tests. As the Diffusion Forcing baseline is inherently designed for replanning, we follow that evaluation setting.
A.2. Heuristic Controller, Value-Leearing Policy, and Inverse Dynamics Model
A long-standing challenge in diffusion-based planning is preserving global trajectory coherence while managing local, high- -igh-dimensional state-action control (Chen et al., 2024a;b). For example, PlanDQ (Chen et al., 2024b) couples high-level diffusion planner with a learned low-level policy. Similarly, our approach focuses the diffusion planner on lower-dimensional, representative state information (e.g.. object or agent positions), while delegating detailed action inference to a low-level action module (i.e., a heuristic controller, a value-learning policy, or an inverse dynamics model). The details of this integration are discussed in Algorithm 10. Such hierarchical design allows MCTD and baseline planners to concentrate on broader strategic decisions while leaving the finer-grat mmer-grained execution details this action module.
A.3. Model Hyperparameters
For reproducibility, detail the hyperparameters used in our experiments. These settings were selected based on prior work and empirical tuning ensure stable training and evaluation. Nearly identical hyperparameters were applied consistently across all tasks, except where task-specific configurations were necessary, which are discussed in their respective sections.
1. DIFFUSER
Our Diffuser implementation is based on the official repository (https:/ www janneen/jeeem///mmm reteeirmiiiiiii /////iffuuee ffuser). While retaining the core architecture and training procedure, we introduce the following modifications to align with our experimental setup:
Guidance Function: Whereas the original Diffuser applied goal-inpainting guidance for point-maze tasks, employ distance-based reward function to ensure fair comparison across all baselines.
Replanning Variant: To evaluate the impact of inference-time iterative refinement, we implemented replanning strategy that periodically re-samples trajectories based on the agent's 's updated state.
To assess the effect of scaling inference-time computation, we implement the "Sample- Score-Rank" variant of Diffuser from (Zhou et al., 2024). Furthermore, to incorporate varied guidance, we assign a guidance scale to each sample, drawn uniformly from the set (0.01, 0.05, 0.1, 0.2,
Random Search Variant:
12

---

Monte Carlo Tree Diffusion
Table 8. Diffuser Hyperparameters
Hyperparameter
Value
Learning Rate
2e
EMA Decay
0.995
Precision in Training/Inference
32
Batch Size
32
Max Training Steps
20000
Planning Horizon
Different for each task, please check in the task sections 50 for Diffuser- Reaiussn-.....nnng. otherwise planning horizon
Open Loop Horizon
Guidance Scale
0.1
Beta Schedule
Cosine
Diffusion Model Objective
Zo To-prediction
U-Net Depth
Kernel Size
5
The Number of Channels
32, 128, 256
A.3.2. DIFFUSION FORCING
Our implementation of Diffusion Forcing is based on the officially released source code (https::/github. com/ buoyancy 99/d1 /t///fffuuu ffusion fornnoorcnnn.... The following modifications were made:
Replanning Frequency: Since Diffusion Forcing natively supports inference-time replanning, we evaluated under comparable conditions by applying periodic plan refinement, similar to our Diffuser-Replanning variant.
To maintain consistency across models, we utilized the Transformer- backbone provided in the Diffusion Forcing repository.
Transformer Backbone:
A.3.3. MONTE CARLO TREE DIFFUSION (MCTD)
Our MCTD model inherits most hyperparameters from Diffusion Forcing. However, we introduce specific modifications and add several new hyperparameters related to the tree search component.
A.3.4. HALUE-LEARNING POLICY
To address the challenges of complex action spaces in specific tasks, we integrate value-leariing policy adapted from the PlanDQ source code (Chen al., 2024b) (https: /ttpthubttttttt//// com/ http:/chhcccopplaancc.eccoo
A.4. The Evaluation Details
We followed the OGBench (Park et al., 2025) task configurations, 5 tasks per each environment. For each task, we evaluated 10 random seeds per each model, and reported the averaged success rates and their standard deviations.
A.5. Maze Navigation with Point-Mass and Ant Robots
A. MODIFICATIONS FROM OGBENCH
Although the original OGBench datasets include random noise in the start and goal positions, we eliminated this noise isolate model performance from environmental randomness. This modification facilitates clearer comparisons across methods and ensures that performance differences can be attributed to planning capabilities rather than uncontrolled stochasticity. Additionally, we incorporate velocity into the state representation, which the original benchmark does not provide, to apply the heuristic controller for Point-Maze tasks implemented in (Janner et al., 2022). For the Ant-Maze Giant task, we found cases where well-generated plans failed due to inaccurate execution. To prevent this, we increased the episode length from 1000 to 1500.
13

---

Monte Carlo Tree Diffusion
Table Diffusion Foreing Hyperparameters
Hyperparameter
Value
Learning Rate
Weight Decay
Warmup Steps
10000
Precision in Training
16-mixed
Precision in Inference
32
Batch Size
1024
Max Training Steps
200005
The Number of Frame Stack
10
Planning Horizon
Different for each task, please check in the task sections
Open Loop Horizon
50
Causal Mask
Not Used
Guidance Scale
3 for medium mazes and 2 for other mazes
Scheduling Matrix
pyramid
Stabilization Level
10
Beta Schedule
Linear
Diffusion Model Objective
To zo-prediction
DDIM Sampling eta
0.0
Network Size
128
The Number of Layers
12
The Number of Attention Heads
4
The Feedforward Network Dimension
512
A.5.2. GUIDANCE FUNCTION
To plan long-horizon trajectories while maintaining the effectiveness to reach the goal quickly possible, we applied guidance function to minimize the distance between each state and the goal, represented as >. || gll2 (Chen et al., 2024a). We note that for fair comparison, we applied the same guidance style across all baselines. For example, Diffuser originally used goal-inpainting guidance for Point-Maze, but in this evaluation, we used the distance- guidance instead.
A.5.3. GUIDANCE SCALE
We applied different guidance scales to MCTD for different tasks. For Point-Maze Medium and Large, a scale set of [0, 0.1, 0.5, 1, 2] was used. For Point-Maze Giant, applied [0.5, 1, 4]. For the Ant-Maze tasks, of 1, 5] was used.
A.5.4. PLANNING HORIZON
For the Medium and Large mazes, we applied a planning horizon of 500, and for the Giant maze, a horizon 1000. The trajectory horizons in the original datasets are 1000 for Medium and Large, and 2000 for Giant. We used a shorter planning horizon than the dataset horizon to generate more data through the sliding window technique.
A.5.5. HECRISTIC CONTROLLER AND VALUE-GUIDED POLICY
For Point-Maze, adopted the heuristic controller designed in (Janner et al., 2022). For Ant-Maze, we used value-guided policy (Wang et al., 2023), as done in (Chen et al., 2024b). However, unlike (Chen et al., 2024b), only provide the planned position 10 steps ahead as the subgoal.
A. MCTD REWARD FUNCTION
We designed heuristic reward function. One component of this function penalizes large, physically impossible position changes between consecutive states. Another component provides reward for reaching the goal. Reaching the goal earlier
14

---

Monte Carlo Tree Diffusion
Table 10. MCTD Hyperparameters
Hyperparameter
Value
Learning Rate
5e
Weight Decay
1e
Warmup Steps
10000
Precision in Training
16-mixed
Precision in Inference
32
Batch Size
1024
Max Training Steps
200005
The Number Frame Stack
Planning Horizon
Different for each task, please check in the task sections Same to planning horizon for MCTD, 50 for MCTD-Replanning
Open Loop Horizon
Causal Mask
Not Used
Scheduling Matrix
pyramid
The Maximum Number of Search Guidance Set
500
Different for each task, please check task sections
The number of Partial Denoising
20
The Jumpy Denoising Interval Stabilization Level
10
Beta Schedule
Linear
Diffusion Model Objective
To xo-prediction
DDIM Sampling eta
0.0
Network Size
128
The Number Layers
12
The Number of Attention Heads
The Feedforward Network Dimension
512
yields larger reward, calculated as r (H - ()/H, where H is the total horizon length and / is the current timestep.
A.5.7. VALUE-LEARNING POLICY VARIANCE
For the AntMaze tasks, we employed a value-learning policy following (Chen et al., 2024b). We observed instances where plans were generated successfully, but the policy subsequently failed to execute them effectively. To account for this execution variance, we conducted three trials per task and seed combination for all AntMaze experiments, reporting the maximum success rate.
A. SHORT-HORIZE (STITCH) RESULTS
To examine performance when only short-horizon trajectories are available during training, we evaluated our method on the stitch datasets, with results summarized in Table 12. The trajectory horizon in these datasets is significantly shorter than the number of steps required to reach the goal (e.g., a horizon of 100 steps in the dataset versus over 200 required steps for medium map tasks). Consequently, we applied replanning to Diffuser, Diffusion Forcing, and MCTD. contrast to the long-horizon dataset experiments, the Diffuser demonstrates superior performance over Diffusion Foreing, suggesting that its implicit stitching capability is more effective. Meanwhile, MCTD outperforms the baselines the medium-sized map due its multi-sample exploration. However, its search depth constrained by the limited length of trajectories in the training data, resulting in suboptimal performance the large map.
A.5.9. COMPARISON WITH GOAL-INPAINNIIG DIFFUSER
The Diffuser (Janner et al., 2022) typically employs a learned guidance function to bias generated trajectories toward high- return outcomes. For PointMaze tasks, however, an alternative approach-gode geallinpainting guidance was introduced (Janner al., 2022), which improves single-pass planning by "imagining" an intermediate trajectory whose final state is the
15

---

Monte Carlo Tree Diffusion
Table 11. Value-I.caming Policy Hyperparameters
Hyperparameter
Value
Learning Rate
3e
Learning Eta
1.0
Max Backup
False
Reward Tune
eql cql_antmaze
The Number of Training Epochs
2000
Gradient Clipping
7.0
Top-k
Target Steps
10
Randomness on Data Sampling (p)
0.2
Table 12. Short-Horizon (Stitch) Maze Results. Success rates (%) on caduced-hhrizon stitch variants of the pointmaze environment.
Dataset
Miluuser-Replanning
Diffusion Forcing
MCTD-Replanning
pointnaze-mediun--tiih----
76 + 12
53 + 16
90 + 14
pointmazz-llrge--siihh---
34 16
20 0
20 1 0
designated goal. In the context of Diffuser- Rertlani--R,,,,ningg, adapt goal-inpainting by providing the experienced trajectory contextual information rather than direct denoising target.
As shown in Table 13, goal-inpainting outperforms the standard Diffuser with heuristic guidance across medium, large, and giant pointmazes. By treating the goal as the terminal state to inpaint, the Diffuser model effectively samples trajectories from both ends, leveraging local convolutional mechanisms to reconcile the start and goal conditions. Although this approach yields marked performance gains, MCTD remains superior in long-horizon scenarios, outperforming both the base and replanning variants of the goal-inpainting Diffuser. The key distinction lies in MCTD's tree-structured partial denoising and adaptive branching; while goal-inpainting enhances single-pass generation, it lacks the iterative entrration-xploitaatin loop of an explicit tree search. Moreover, the advantage of two-ended trajectory imagination diminishes as the maze size increases (e.g., giant mazes). Consequently, MCTD uncovers robust solutions even where goal-inpainting shows diminishing returns, underscoring the benefits of tree-based refinement in complex, long-range planning tasks.
A.6. Robot Arm Cube Manipulation
A. OBJECT- WISE GUIDANCE
Given that the robot arm can only manipulate single cube at a time, we adopt an object-wise guidance mechanism in MCTD. Rather than applying single, holistic diffusion schedule to all cubes simultaneously. MCTD treats the selection of which cube to move as a meta-action. This design partially mitigates issues arising from plans that attempt to move multiple cubes concurrently, although multi-object manipulation still requires careful sequencing to avoid suboptimal interleaving.
A. GUIDANCE FUNCTION
The guidance function the same as that used for the tasks (Appendix A.5.2). For MCTD and MCTTD-Replannin the guidance is applied on object-specific basis, aiming to reduce the distance between an object's current state and its corresponding goal state.
A.6.3. GUIDANCE SET
For the cube manipulation tasks, guidance set of [1, 2, 4] is applied for each object. For instance, a task with two cubes, the total guidance set size comprising set of three values for each of the two objects.
16

---

Monte Carlo Tree Diffusion
Table Comparison Results with Goal-Inpainting Diffuser. Success rates (%) on pointmaze and antmaae environments with medium. large, and giant mazes for the navigate datasets.
Diffuser
Goal-Inpainting Diffuer
Environment
Dataset
MCTD
Base
Replanning
Base
Replanning
medium-navigate-vO
58 +
60 + 0
84 + 8
80 + 9
100 1
pointmaze
large-navigate-v0
44 8
40 0
94 + 9
84 + 17
98 + 6
giant-navigate-v0
0
0 0
30 - 16
34 21
100 -
medium
degiium-navvgate--Ve
36 =
40 18
100 0
94
100 -
antmaze
large-navigate-v0
14 + 16
26 13
86 1 9
66 + 13
98 +
giant-navigate-v0
0 0
0 0
12 10
1
94 1
A.6.4. REWARD FUNCTION
To avoid generating unrealistic plans, we augmented the reward function from the maze environment with additional constraints. A plan is deemed unrealistic and assigned reward of 0 if it violates any of the following rules:
No Simultaneous Movements: While minor noise in the plan permissible, significant concurrent movement of multiple objects is forbidden.
Stable Final Placement: A cube's final position must be stable, either on the floor or on top of another cube, suspended in mid-air.
No Collisions: A cube cannot be placed in location already occupied by another cube.
Clearance for Manipulation: A plan cannot move cube if another cube is positioned top of it.
A.6.5. PLANNING HORIZON
The planning horizon was set to 200 for single-cube tasks, matching the episode length. For tasks involving multiple cubes, the horizon was extended to 500.
A.6.6. VALUE-LEARNING POLICY
Similar the antmaze tasks, we employed value-learnng policy (Wang al., 2023). This policy is trained achieve subgoals suggested by the planner, which correspond to planned states 10 steps into the future.
A.6.7. HANDLING REDUNDANT OPERATIONS IN OBJECT-WET WISE PLANNING
To accelerate object-wise planning, we integrated common, repetitive operations into the framework. These operations include automatically dropping an object when reaches its goal and pre-posinn pre-positioning the robot arm before executing the plan for each object. Embedding these routines directly mitigates inconsistencies, such as planning a move beyond the arm's reach or attempting grasp an object when the gripper already occupied.
A.7. Visual Pointmaze
A.7.1. SETUP
Since diffusion- based planners typically operate on state-action representations, we first learn an image encoder using a Variational Autoencoder (VAE) (Kingma & Welling, 2014). The resulting latent states serve as the input our diffusion model. Notably, these latent encodings do not contain explicit positional information; consequently, the planner must infer the underlying geometry navigate the maze indirectly. We design inverse dynamics model-a a compact Multi-Layee Perceptron (MLP)-to to map consecutive latent observations to actions. This model replaces the heuristic controller from the original Diffuser architecture, which is not directly applicable in this partially observable setting.
A.7.2. GUIDANCE
To guide the diffusion process, we introduce a position estimator trained on ground-truth coordinates. Although this estimator provides an approximate distance to the goal, its learning signal is not used update the image encoder. The estimator does
17

---

Monte Carlo Tree Diffusion
supply the planner with the full underlying state, thereby ensuring the image encoder processes only pixel-level data. This design preserves partial observability in the core planning procedure and highlights the open question of how best incorporate task-relevant signals in visual domains. We apply the same guidance function as in the maze task described in Appendix A.5.2.
A.7.3. GUIDANCE SCALES
For this task, we use a set of guidance scales E {0,0.1, 0.5, 1, for both MCTD and MCTD- MCTD-Reelanniig.
A.7.4. DATASET GENERATION AND PLANNING HORIZON
We follow the dataset generation script from OGBench (Park et al., 2025) for the PointMaze environment, setting trajectory horizons to 1000 for the Medium and Large mazes. To generate more training data, we reduced planning horizon of 500 for these mazes and employ sliding- window technique, consistent with our earlier pointmaze experiments.
A.7.5. VARIATIONAL AUTOENCODER
We pretrain a VAE (Kingma Welling. 2014) to encode 64 x 64 x 3 RGB observations into an 8-dimensional latent space. This compresses the high-dimensional visual inputs into a structured representation suitable for downstream plan- ning. The VAE models the latent variable 2 as a Gaussian distribution with mean fr and standard deviation o, using the reparameterization trick:
μ + 6, where 1).
(5)
We use the sampled latent embedding as the input representation for all downstream tasks.
A.7.6. INVERSE DYNAMICS AND POSITION ESTIMATOR MODELS
We trained the inverse dynamics and position estimator models on an offline dataset latent states encoded by the pretrained VAE. Both models are implemented as MI.Ps.
Inverse Dynamics. This model, fine predicts the action å required transition from state 2N to the subsequent state 2,+1- To account for velocity partially observable context, the model is conditioned on both the previous and current latent states, 2-1 and 24:
àe fin (21-1, Re: 26+1).
(6)
For the initial timestep (t = 0), where 2 is unavailable, we set x = 20-
Position Estimator. The position estimator predicts the agent's (zr, y) coordinates from the latent state 24, It provides a weak positional signal to aid the planner for guidance while maintaining the overall partial observability of the task.
A.7.7. THE NECESSITY OF REPLANNING COMPLEX ENVIRONMENTS
In this section, we analyze the failure modes of diffusion bastud planners that do not employ replanning, such as Diffusion Forcing, in the visual pointmaze task. As illustrated in Figure 7, the absence of iterative correction leads trajectory collapse, revealing significant challenges in pixel-based planning. These findings underscore the necessity of replanning for generating robust trajectories in complex and partially observable environments.
B. MCTD Architecture Implementation
The Monte Carlo Tree Diffusion (MCTD) framework relies on three key elements within its tree-search loop: sub-plan noise scheduling, causal semi-autoregreesivee denoising, and partial, jumpy denoising. First, rather than assigning a single noise schedule to the entire trajectory, each sub-trajectory segment receives its own noise-level schedule. This design allows different segments to transition from high to low noise at distinct rates, thereby enabling focused refinement of earlier parts of the plan while preserving flexibility in later segments.
Next, we adopt a causal noise schedule (Chen et al., 2024a), which renders the denoising process semi-autoregressivee earlier plan segments undergo more thorough denoising, while subsequent ones shift more gradually from higher to lower noise. This preserves the globally consistent trajectory generation characteristic of diffusion models but also ensures that
18

---

Monte Carlo Tree Diffusion
UF
Figure 7. Trajectory collapse in Diffusion Forcing without replanning the visual pointmaze. The left panel shows the estimated position of the planned trajectory (yellow points) halting at a certain area. The right panels (top left to bottom right) visualize the planning progression over time, suddenly teleporting to an unknown area This underscores the importance of replanning overcoming single-past pass planning limitations.
critical decisions near the beginning of the plan receive fidherrti refinements. Finally, both planning and simulation employ partial and jumpy denoising via DDIM sampling (Song et al., 2021a), allowing the system to skip directly from noise level to t A. This skip-level update accelerates the evaluation of candidate trajectories during tree expansion and supports partial denoising steps essential for exploration.
We implement these strategies with a Transformer- based model that processes an expandable token sequence, thus avoiding any architectural constraint on the planning horizon. Each sub-plan annotated with a distinct noise-level index, and the Transformer predicts the fully denoised data To from the corresponding noisy tokens. Concretely, employ an To paramtterizztion by training the network to minimize the squared error between the predicted in and the ground-truth To:
Liaan Ez =oll"].
(7)
Because each sub-plan has its own noise schedule, the Transformer naturally accommodates varying horizon lengths and partial denoising intervals without requiring further modifications to the model architecture.
By integrating causal scheduling with skip-level DDIM updates, MCTD achieves a autreeruoorgrssiivee, partial denoising process within tree-search loop. Early plan segments are refined immediately, while future segments remain at higher noise until exploration heuristic signals indicate their importance. This arrangement allows for fast approximate rollouts by skipping multiple noise levels when simulating plan. Consequently, MCTD integrates the holistic generative capacity of diffusion models with the adaptive, branch-and-boume nature of Monte Carlo Tree Search, thereby enabling iterative trajectory refinement in tasks with long horizons and sparse rewards.
C. Detailed Comparison between Monte Carlo Guidance (MCG) and Monte Carlo Tree Diffusion (MCTD)
While MCG (Monte Carlo Guidance) utilizes Monte Carlo sampling to adjust guidance levels during denoising, it lacks the structured search capabilities of MCTD (Monte Carlo Tree Diffusion) and does not explicitly refine trajectories over time. MCG focuses on adjusting guidance strengths rather than performing an explicit trajectory search, making it less effective for long-horizon planning where iterative improvement is crucial. Additionally, MCG does not implement the four standard MCTS steps (Selection, Expansion, Simulation, and Backpropagation) limiting its ability to efficiently allocate inference-time scaling for progressive refinement. Without structured search process, MCG cannot revisit and optimize subplans, making it less scalable for complex decision -akiing tasks. Furthermore, MCG lacks meta-actions for adaptive
passively adjusting guidance instead of dynamically balancing search strategies. Finally, does not incorporate fast jumpy denoising for simulation, making trajectory evaluation less efficient compared to MCTD, which efficiently estimates plan quality without costly forward model rollouts.
19

---

Monte Carlo Tree Diffusion
D. Additional Experiment Results
D.I. UCT Hyperparameter Ablation Study
In addition to the meta-action ablation study (Section 5.6.2), we empirically analyze the exploration-exllitttioo balance by varying the UCT (Koesis & Szepesvári, 2006) hyperparameter in the formula: VUCT V2 + W where v; the node's estimated value, and N and n4 are the visit counts of the parent node and node itself, respectively.
Table 14. UCT Hyperparameter Ablation Study Results on PointMaze.
W
0 (Greedy)
v2 (Default)
5
10
Success Rate (%)
88.113
100.10
100.10
98
98.16
Medium
Run Time (sec.)
15+1
31=4
63+10
76
92::55
The number Search
105+84
77=299
94+88
120=27
134+:23
Success Rate (%)
90::10
98:16
100.10
98:
100:10
Large
Run Time (sec.)
16410
744:34
90+:10
102=14
1044:18
The number of Search
117478
174190
211+26
257=41
2651:43
Success Rate (%)
824:14
100+0
98+66
100+-0
100+E0
Giant
Run Time (sec.)
25+1
215++23
216+112
225-22
230+
The number of Search
228::69
394:1152
442:113
464=13
49317
As shown in Table 14, with W = 0 (pure greedy search), MCTD achieves faster inference times and requires fewer searches, but at the cost of reduced performance. This occurs because greedy search explores the tree depth- -iith-wiic rapidly, reducing computational overhead jump denoising but sacrificing thorough exploration. The default value (W = 1.141) strikes an effective balance, while higher values increase computational costs.
D.2. Jumpiness Scale Ablation Study
Fast denoising (Song et al., 2021a) is a key concept in MCTD that enables tree search within the denoising process. The jumpiness scale C' determines how many denoising steps are skipped during this process. Larger values result in fewer denoising steps, trading accuracy for computational speed. We investigated how different jumpiness scales affect performance and efficiency, including one-shot variant that denoises trajectories in a single step. The results demonstrate that proper selection of C significantly impacts both model performance and efficiency.
Table 15. Jumpiness Scale Ablation Study Results PointMaze.
5
10 (Default)
20
50
One-shot
Success Rate (%)
10010
100:10
98:16
100=0
100:1
100:10
Medium
Run Time (sec.)
84:126
391:16
34:113
29:112
27::13
42:1
The number of Scarch
86-+30
744:32
77+-29
73+:30
65:129
143+:46
Success Rate (%)
100.0
100.0
100-0
98:16
94.113
63:139
Large
Run Time (sec.)
160.169
911335
74.134
65.132
68.149
355.1183
The number of Search
176+-85
188+81
174+900
167--89
174+112
301+134
Success Rate (%)
100.10
98.16
100-1
100-0
88.116
10.121
Giant
Run Time (sec.)
632.1239
258.1115
231-1103
185.88
174-176
164.12
The number Search
390=1188
37941588
394::152
363+154
387+119
500::0
For relatively simple tasks, even aggressive jumpiness (large C') or one-shot denoising maintains high performance. However, as task complexity increases, these approaches significantly degrade performance because value function estimation becomes less accurate under highly jumpy denoising, leading to suboptimal tree search decisions. The default setting (C 10) achieves an excellent balance, maintaining high success rates while reducing computational time compared smaller tested one-shot decoding by directly applying predictions from diffusion model trained to
jump values. Additionally, predict fully denoised data. While this represents the most computationally efficient method, it exhibits less effective search behavior in medium-sized maps and clear performance degradation in larger maps. These results demonstrate that jumpy denoising provides an optimal balance between effectiveness and efficiency.
20

---

Monte Carlo Tree Diffusion
E. Algorithms
Algorithm Monte Carlo Tree Diffusion (MCTD)
procedure MCTD(root, iterations)
2: for i
iterations do
3: node SIL.ECTRROMISIINNOODEOroo))]]
4: if IISEXPANDABLEE(ood) then
5:
node + EXPANDNODE(node)
6: end
7: reward + SIMULATE(node))
8: BACKPROPAGAATERuude, rewarr)
9: end for
10: return BESTCHILDD(rof)
11:
end procedure
Algorithm 3 Selection in MCTD
1: procedure SELECCTRROMSSINGNOOOEEEddd
2: while (SFULLLEEXPANDEEDnnddd))) and not IsLEAF(node) do
3: node <- BESTUCTCHHI.((uud))
D Use UCB for exploration-exploitatio
4: end while
return node
6: end procedure
Algorithm 4 Expansion in MCTD
1: procedure EXPANDNoDD(node)
2: 9x f SELECCMETAACTTON(nddd))
Determine guidance level D. Generate new subplan using diffusion
3: child <- DENOOSEESURRRAANNOOdU...
4: ADDCHILD(mode,, child)
5: return child
6: end procedure
21

---

Monte Carlo Tree Diffusion
Algorithm 5 Simulation in MCTD (Jumpy Denoising)
1: procedure SIMULATE(node)
2: fullPlon <
3: return EVALUATEPLLA((fullPllm
4: end procedure
Algorithm Backpropagation in MCTD
procedure BACKPROPAGATERuode, reward)
2: while node + reull do
3: node.visile condee node.visilCounll +
4: nole. node.value + reward
5: UTIATTEMETTAATTOOSSCHDDLLLEEE,,D
Adjust guidance levels
6: node <- node.parenl
7: end while
8: end procedure
Algorithm 7 Meta-Action and Guidance Selection
1: procedure SELECCTEETAACTIO(nndl
2: return UCBSELECCTION((NO, Low, MEDIUM, HIGH))
3: end procedure
4: procedure DENOISESSUBPLAN(Wod. gs)
5: if g, = NO then
6: return SAMPLEPRIO(G((x/XI.cc.
Exploration
7: else
8: return SAMPLEGUUIDED//S)////xCxx.
Exploitation
9: end if
10: end procedure
Algorithm 8 Jumpy Denoising for Fast Simulation
1: procedure FAsTIUMPYDEENNNN//prrraaaaa
2: return SAMPLE(y((sHHss/XIss.)
P X-shot fast decoding
3: end procedure
22

---

Monte Carlo Tree Diffusion
Algorithm 9 Partial and Jumpy Denoising with DDIM
Require: Initial plan Xto at noise level lo. noise schedule A² = [lo, 41,... 4x]. diffusion model fo(-), guidance scale Ensure: Partially denoised plan X
for k = O, 1. .... 1-1d do
I/ Predict noise
at level t.
2:
t fo ( tx)
DDIM e-prediction or score function
// Apply guide using gradients of return e(guided) ~ NN (E + aLVJ, 2²)
3:
DDIM update from ts to fk+1
Cit e (euided)
+ v atu+1 (guided)
4: X(y-1
Xte
oix
5: end for
6: return X(x
Algorithm 10 Additional ControllerPoliccy/nnvrre Dynamics Model Integration
Require: Long-term planner, Pg. environment E and controller T,, (heuristic controller or value learning policy or inverse dynamics model), planning horizon H
E with initial observation and goal s, g
2: while not done do
// Get a plan from planner Po
3:
Pı: = g)
4:
i= k
>> k steps after plan state is the subgoal for Controller
5:
g' 6'=Px
6:
while < H do
// Compute low-level action with Controller
7:
5 (5,88 // Interact with environment
8:
(o, done) +- e.step(a)
9:
10:
if s ~ g then
11:
i=itk k g = Pk end if
14:
if done then
15:
break
16:
end if
17:
end while
18:
end while
23


Reviews:
[{"review_id": "66M8OkbWmd", "reviewer": "Reviewer_Fb4y", "summary": "This paper proposes a new algorithm, MCTD, which involves MCTS in the diffusion model. MCTD uses whole trajectories as its states and introduces a meta-action to generate child nodes based on whether a guidance function is used. By leveraging the MCTS process, MCTD achieves better success rates and runtime performance in most test settings compared to baselines such as Diffuser and Diffusion Forcing.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> I am not sure why we need a **no-guidance action**.\n> \n\nDiffusion planners often rely on strong guidance, but under unseen goals or complex long-horizon tasks, this can be detrimental. Allowing “no-guidance” (or weak guidance) provides essential exploration. We demonstrate this by **showing the performance degradation when MCTD only uses high-guidance values.**\n\nDefault (Guidance Levels: [0, 0.1, 0.5, 1, 2]) / Only High Guidance (Guidance Levels: [2, 2, 2, 2, 2]) Performances:\n\n- Medium: 100 ± 0 / 98 ± 6\n- Large: 98 ± 6 / 80 ± 0\n- Giant: 100 ± 0 / 78 ± 11\n\n> Code only has README, and README is empty.\n> \n\nWe are working on the code release. We will make sure to release the code after the review process.\n\n> Why is there an exp which is different from the diffuser paper?\n> \n\nThis reflects a probabilistic interpretation common in RL, where optimality can be linked to exponentiating the reward, $\\exp(r(\\tau))$. In the original Diffuser paper, a similar concept was denoted by a heuristic term $h(\\tau)$. We write $\\exp(\\mathcal{J}(\\tau))=h(\\tau)$.\n\n> What is $\\mathbf{x}_i$ in $\\mathbf{x}$? What is $S$ ?\n> \n\nAs written in the page 3, line 136, $\\mathbf{x}_i$﻿ is subplan in full trajectory $\\mathbf{x}$﻿, and $\\bigcap_S \\mathbf{x}_S=\\emptyset$  means there is no shared trajectories between subplans. $S$ is the number of subplans in the full trajectory $\\mathbf{x}$.\n\n> What is “subplan” and “denoising schedule”?\n> \n\nAs written in the page 3, line 136, subplan is the part of trajectories. Denoising schedule is the assignments of different noise levels per each subplan. For instance, in MCTD, near future subplans are considered as low noised, while further future subplans are as highly noised as shown in Figure 1 (b).\n\n> “Because S ≪ N.” What is N?\n> \n\nN is the total length of trajectories and S is the number of subplans.\n\n> I assume x_s is an agent path trajectory but with noise (middle result of the diffusion model). ... so the tree height depends on the diffusion model iteration steps.\n> \n\n$\\mathbf{x}_s$ denotes a **subplan** (a partial trajectory segment) at some intermediate stage of denoising, not the full trajectory. In MCTD, the **tree height is determined by the number of subplans** $S$, rather than the total denoising steps.\n\n> What are \"meta-actions\"?\n> \n\nIn standard MCTS, each node expands by choosing from all possible actions. Because we want to keep the branching factor manageable (especially in continuous spaces), we introduce **meta-actions** to represent **guidance levels**.\n\n> “In MCTD, we observe… in the offline data.” I don’t understand, and there may be some grammar problems.\n> \n\nAlthough we checked that its grammar is not wrong, we agree that it could be clearer with a slight change as follows:\n\n*In MCTD, we observe that sampling from the prior distribution p(x), i.e., using the standard diffusion sampler, represents exploratory behavior as it does not attempt to achieve any goal. Instead, it only imitates the prior behavior contained in the offline data.*\n\nIn the part, we discussed the sampling from the prior $p(\\mathbf{x})$ represents the exploratory behavior by not attempting to achieve any goal.\n\n> could we have **more child nodes**, or how could we determine the **optimal number of child nodes**?\n> \n\nWe show a simple illustration in the Figure 1 with two children. In practice, we can use **multiple guidance levels** and indeed we do in our experiments.\n\nCurrently, the number of child nodes (meta-actions) is a **hyperparameter**. Exploring an adaptive selection of the number of meta-actions for each task is an interesting direction for future work.\n\n> The author should ensure the notations are consistent or explain new notations such as g in equation 4 and g_s in equation 3.\n> \n\n$\\mathbf{g}$ and $g_s$ are introduced in page 3 as the guidance schedule and guidance, but we will revise to reintroduce them right before each equation.\n\n> why do we need a no_guidance action?\n>\n\nWe replied in the no-guidance action investigation.\n\n> ... I think there should be a test using **the same maze map** and robot stacking test settings as in **the diffuser paper**.\n> \n\nWe appreciate the suggestion. While we could not evaluate MCTD on the robot stacking due to time limitation, we **did** evaluate on the D4RL multi-2D maze tasks which are evaluated in Diffuser paper. **The results demonstrate MCTD outperforms Diffuser.**\n\nDiffuser / MCTD Performances:\n\n- Umaze: 128.9 ± 1.8 / 239.7 ± 46.2\n- Medium: 127.2 ± 3.4 / 430.6 ± 104.8\n- Large: 132.1 ± 5.8 / 548.9 ± 195.8\n\nWe measured the performance with the accumulated rewards from the environment as done in Diffuser. The performance is averaged from 100 samples. We analyzed the performance gaps due to the searchablity on MCTD and different guidance types (Diffuser - inpainting / MCTD - reducing distances between states and goal). The inpainting guidance can lead inefficient planning when the start and goal are nearer than the plan length.", "claims_and_evidence": "in detailed comments", "methods_and_evaluation_criteria": "in detailed comments", "theoretical_claims": "NA", "experimental_designs_or_analyses": "in detailed comments", "relation_to_broader_scientific_literature": "in detailed comments", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strengths:\n\nThe idea of MCTD is interesting.\n\nMCTD performance surpasses baselines in most test cases, both in success rate and runtime.\n\n\nWeaknesses:\n\nThe meta-action is not very convincing; I am not sure why we need a no-guidance action.\n\nThe writing should be largely improved.", "other_comments_or_suggestions": "Code only has README, and README is empty.\n\nPreliminaries  \n\nDiffuser equation 1: Why is there an exp which is different from the diffuser paper?  \n\n\n\nMonte Carlo Tree Diffusion  \n\n3.1  \n\n“Specifically, we partition the full trajectory… x_s!=empty.” What is x_i in x? What is S?  \n\nWhat is “subplan” and “denoising schedule”? I am totally lost in this paragraph.  \n\n“Because S ≪ N.” What is N?  \n\nI assume x_s is an agent path trajectory but with noise (middle result of the diffusion model). So basically, this part wants to say MCTD uses x_s as a tree node, so the tree height depends on the diffusion model iteration steps. If we only need 5 steps to recover the whole trajectory from noise, the tree height will be 5.  \n\n3.2  \n\n“While meta-actions… during the denoising process.” So what are “meta-actions”? I still have no idea about them.  \n\n“In MCTD, we observe… in the offline data.” I don’t understand, and there may be some grammar problems.  \n\nTo my understanding, the meta-action is just two sample models: one samples from p, and the other samples from p_g. According to figure 1, it seems MCTD alternates between using two kinds of probabilities. In the tree, each node only has two child nodes. So the question is, could we have more child nodes, or how could we determine the optimal number of child nodes?  \n\n3.3  \n\nThe author should ensure the notations are consistent or explain new notations such as g in equation 4 and g_s in equation 3.  \n\nMy question is, if we already have a reward function and also a guidance function, why do we need a no_guidance action? I think, just like in the diffuser model, encoding this information in the inference part should be enough.  \n\n\n\nExperiments  \n\nSince the baselines are mainly diffuser models, I think there should be a test using the same maze map and robot stacking test settings as in the diffuser paper.", "questions_for_authors": "In the tree, each node only has two child nodes. So the question is, could we have more child nodes, or how could we determine the optimal number of child nodes?  \n\nIf we already have a reward function and also a guidance function, why do we need a no_guidance action? I think, just like in the diffuser model, encoding this information in the inference part should be enough.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mxXVyH6fxO", "reviewer": "Reviewer_iXeu", "summary": "Standard diffusion-based planners don't improve with additional test-time computation. This limits their effectiveness on complex long-horizon tasks.The authors introduce Monte Carlo Tree Diffusion (MCTD) which formulates the diffusion process in such a way that each denoising process can effectively \"branch out\" - i.e. each denoising step can theoretically happen in different goal directions an with different denoising rates. Notably the \"branching out\" of diffusion step gives rise to a tree structure when different branches are explored. This in turn allows for it to search over these branches and build on the monte carlo tree search literature to efficiently explore different branches to find better overall plans at the cost of additional test-time computation. \n\nThe authors show that MCTD significantly outperforms standard diffusion planners and reinforcement learning baselines on challenging long-horizon tasks with sparse rewards albeit at the cost of large test-time computation.\n\n## Post Rebuttal\nI acknowledge the authors' rebuttal and maintain my original assessment.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The claim of \"System 2 planning\" is bit of a stretch as the system does not implictly have a world model it can search all variations upon.\n> \n\nWe acknowledge that there may be differing perspectives on what constitutes System 2 planning in the context of ML. Drawing on Kahneman [1], we define System 2 planning as a deliberate, sequential search across multiple hypotheses at inference time, which aligns with our approach.\n\nRegarding world models, we argue that diffusion models can qualify. Although they may not strictly meet the action-conditioned forward model definition, we believe a narrow view of “world model” is unnecessary for System 2 classification. Since diffusion models implicitly encode structural knowledge, we question whether excluding them is truly warranted.\n\nOur method uses a semi-autoregressive rollout guided by meta-actions, similar to autoregressive, action-conditioned forward models but operating in a more temporally abstract space. We believe that this abstraction retains the hypothesis-driven nature of System 2 planning and supports our contribution to broader discussions on such systems.\n\n[1] Kahneman, Daniel, 1934-2024, author. Thinking, Fast and Slow. New York :Farrar, Straus and Giroux, 2011.\n\n> The claim that MCTD \"overcome limitations\" of both MCTS and diffusion models is also a bit of a stretch.\n> \n\nWe acknowledge that our phrasing could be more precise. Rather than \"overcoming limitations,\" MCTD brings complementary strengths from each approach while addressing certain weaknesses. We will revise this claim more moderately in the final version.\n\n> It is unfortunate that the random search using repeated diffusion was not scaled to **match the test time compute** used by MCTD.\n> \n\nMajor computation overhead on Diffusion planner is on denoising process, and **we evaluated some baselines with matched budget in our paper.** Diffuser-Random Search baseline already uses the same number of denoising steps as MCTD, as discussed in Section 5.2. Additionally, Section 5.5 compares both Diffuser and Diffuser-Random Search against MCTD using equivalent computational budgets.\n\n> theoretical results will be very hard to inherit as the \"fast diffusion\" does not use any model whose sub-optimality can be quantified.\n> \n\nWe agree on this. Thank you pointing out this!\n\n> How about a baseline where the tree search is simply conducted **greedily at every diffusion step**?\n> \n\nFollowing your recommendation, we implemented and evaluated one-step greedy search at every diffusion step on Diffuser for the PointMaze. The approach can search diverse candidates at each step and select the best using the same reward function as MCTD. We scaled this approach to use even more children than MCTD (MCTD uses 5 children). **The results demonstrate that TTC scaling is indeed limited with this approach.** Specifically, we did: \n\nChild# = 5 / 10 / 15 / 20\n\nPerformances:\n\n- Medium: 62±6 / 58±11 / 60±0 / 60±0\n- Large: 40±0 / 40±0 / 40±0 / 40±0\n- Giant: 0±0 / 0±0 / 0±0 / 0±0\n\nThe results are primarily due to the lack of deep searching across multiple samples and the absence of backtracking over tree structure.\n\n> Missing analysis of how **performance varies with subplan count** and **meta-action choices**.\n> \n\nThanks to your feedback, we **additionally conducted comprehensive ablation studies for both subplan length and meta-action choices on PointMaze**. They are discussed in the rebuttal for reviewers 6pZp and TyXv due to the space limitation.\n\n> Searching all different paths the diffusion process could have taken and exploring through them has close **connections with tree of thought** literature from other generative models like LLMs.\n> \n\nThank you for pointing out this. We agree that our work shares important connections with Tree of Thought approaches. We will strengthen these references in our revised paper.\n\n> There could be simple but strong baselines ... greedy search on the MCDS tree structure, **repeated runs of random diffuser-random search**.\n> \n\nWe addressed this in the evaluation for one-step greedy search.\n\n> The work does not provide detailed study on the effects of sub-plan length and meta action choices in a controlled setting.\n> \n\nWe addressed this in the performance analysis for subplan count and meta-action choices.\n\n> The work also does not show what is **the trade-off between using the scale of \"jumpiness\" in the jumpy denoising and overall plan quality**.\n> \n\nWe appreciate this insightful suggestion. We **additionally investigated the effects of the \"jumpiness\" scale** through additional ablation study, which is discussed in the rebuttal for 6pZp's review.\n\n> Did the authors visualize a particular example where **UCB style approach had clear benefits** over just greedy search ?\n> \n\nWe appreciate this insightful suggestion. We **additionally conducted an ablation study controlling the weight C for visit count in the UCB value formulation** which is discussed in the rebuttal for TyXv's review.", "claims_and_evidence": "Key Claim: MCTD formulation of diffusion allows to exploit test-time compute to improve the overall quality of long horizon plans. \nEvidence: Quantitative results across maze navigation tasks show significant improvements over baselines.\nMoreover the integration of MCTS principles into this diffusion framework is sound and the implementation details are detailed. \n\nMinor comment 1: The claim of \"System 2 planning\" is bit of a stretch as the system does not implictly have a world model it can search all variations upon. At best its a very informed heuristic search over the data distribution and does not have any sense of planning in abstract space and grounding it to low level actions. \n\nMinor comment 2:  The claim that MCTD \"overcome limitations\" of both MCTS and diffusion models is also a bit of a stretch. It is definitely the case that MCTD allows us to combine the theory of MCTS with the diffusion process. But it does it at the cost of test time compute and is a different dimension to doing the diffusion process and not really \"overcoming limitations\" in a pure sense. This does not \"magically\" solve the limitation of MCTS being slow and needing a model, or diffusion process being suboptimal when being used in its more standard forms.", "methods_and_evaluation_criteria": "The evaluation criteria is appropriate. \n- The maze navigation tasks (point-mass and ant) provide clear tests of long-horizon planning capabilities with varying difficulty levels\n- The multi-cube manipulation evaluates compositional planning and transfer in a physical domain.\n- Visual maze navigation tests performance under partial observability\n- The metrics (success rates and rewards) directly measure planning effectiveness\n\nIt is unfortunate that the random search using repeated diffusion was not scaled to match the test time compute used by MCTD. It would have been very useful to show that random search alone scales very poorly with more test time compute.", "theoretical_claims": "While relationship with MCTS has been established and is clearly motivated, theoretical results will be very hard to inherit as the \"fast diffusion\" does not use any model whose sub-optimality can be quantified. Moreover the action spaces defined in the diffusion process is technically continuous with a non trivial dynamics model which has not been studied for th eproblem.", "experimental_designs_or_analyses": "- The choice of OGBench's maze environments provides a good testbed for long-horizon planning with sparse rewards\n- Their ablation of \"Diffuser with diffused action\" highlights the failure mode of baseline approaches\n\n\nconcerns:\n- Lack of direct comparison to other test-time optimization methods which can scale with compute. Random diffusion process can be easily scaled to any amount of compute for fair comparison. How about a baseline where the tree search is simply conducted greedily at every diffusion step? It could very well be the case that one step lookahead at every step would also have been enough. \n\n- Missing analysis of how performance varies with subplan count and meta-action choices. IT would have been wonderful to see the effectiveness of the approach being highlighted for varying number of sub-plan count and meta action choices.", "relation_to_broader_scientific_literature": "Understanding the nuances of diffusion process and leveraging it for faster generation of plans has been explored by some related works where they directly predict a candidate first and then do a limited number of diffusion steps. This work explores the opposite direction where test time compute is being exploited. \n\nleveraging test time compute has always been grounded in search in one form or another. Searching all different paths the diffusion process could have taken and exploring through them has close connections with tree of thought literature from other generative models like LLMs.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Weakneses:\n- There could be simple but strong baselines that leverage test time compute in trivial ways to match the compute used by MCDS. For example, greedy search on the MCDS tree structure, repeated runs of random diffuser-random search. \n- The work does not provide detailed study on the effects of sub-plan length and meta action choices in a controlled setting. \n- The work also does not show what is the trade-off between using the scale of \"jumpiness\" in the jumpy denoising and overall plan quality.", "other_comments_or_suggestions": "- Whle analogies are clear with MCTS, its is a bit hard to follow with references to \"nodes\" and \"meta actions\". May be it would have been simpler to simply define a diffusion MDP where each state is a set of partially denoised sub plans and action can denoise one subplan in different directions. This way may be it would have been easier to define and contrast directly with the MCTS formulation. \n\n- It would be very interesting to reformulate the test time compute objective to have \"diverse set of solutions\" rather than just optimizing for the best solution.", "questions_for_authors": "- Did the authors visualize a particular example where UCB style approach had clear benefits over just greedy search ?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7MhLkkIrqE", "reviewer": "Reviewer_TyXv", "summary": "The paper \"Monte Carlo Tree Diffusion for System 2 Planning\" introduces Monte Carlo Tree Diffusion (MCTD), a novel algorithm that merges diffusion models with Monte Carlo Tree Search (MCTS) to enhance test-time compute scalability in planning with diffusion models. It proposes three key innovations: restructuring denoising as a tree-rollout process for semi-autoregressive trajectory refinement, using guidance levels as meta-actions to balance exploration and exploitation, and employing jumpy denoising for planning. MCTD adapts the four MCTS steps—Selection, Expansion, Simulation, and Backpropagation—into a diffusion context, enabling iterative plan improvement. Evaluated on the Offline Goal-conditioned RL Benchmark, MCTD outperforms baselines like Diffuser and Diffusion Forcing in long-horizon tasks such as maze navigation , cube manipulation, and visual pointmaze. The paper claims MCTD is the first to explicitly integrate MCTS with diffusion planning, offering a scalable System 2 reasoning approach.\n\n**After Rebuttal Summary**\nThe reviewers provided a comprehensive ablation that eased my main concerns that this was simply performing a form of standard sampling. I think this ablation improves the strength of the paper and the takeaways regarding the role that MCTS plays in performance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Relation To Broader Scientific Literatures**\n\nWe thank you for acknowledging our work's connection to the broader scientific literatures. We will ensure these connections are clearly articulated in the revised version.\n\n> Lack of ablations investigating the parameters of MCTS search. In particularly investigating the **balancing of the exploitation-exploration trade off.**\n> \n\nWe appreciate this valuable suggestion. Following your suggestion, ****we have conducted comprehensive ablation studies focusing on the MCTS hyperparameters that control exploration-exploitation balance in MCTD on PointMaze. **The results demonstrate the exploration-exploitation balance is one of the key factors to optimize the performance efficiently.** Below are key results from these studies:\n\n**Meta-action choice ablation study**\n\nChoice 1: [0, 0.02, 0.05, 0.07, 0.1]\n\nChoice 2: [0, 0.05, 0.1, 0.5, 1]\n\nChoice 3: [0, 0.1, 0.5, 1, 2]\n\nChoice 4: [0.5, 1, 2, 3, 4]\n\nChoice 5: [2, 3, 4, 5, 6]\n\nChoice 6: [4, 5, 6, 7, 8]\n\nChoice 7: [0, 0.1, 1, 10, 100]\n\nChoices: C1 / C2 / C3 / C4 / C5 / C6 / C7 Performances:\n\n- Medium: 88 ± 16 / 98 ± 6 / 100 ± 0 / 100 ± 0 / 98 ± 6 / 98 ± 6 / 100 ± 0\n- Large: 44 ± 15 / 100 ± 0 / 98 ± 6 / 90 ± 10 / 80 ± 0 / 80 ± 0 / 100 ± 0\n- Giant: 54 ± 23 / 100 ± 0 / 100 ± 0 / 100 ± 0 / 88 ± 16 / 74 ± 18 / 100 ± 0\n\nThe results reveal that when meta-actions are heavily biased toward either exploration (C1) or exploitation (C6), performance degrades significantly. Interestingly, MCTD demonstrates robust performance with balanced meta-action sets, even when some values are extreme outliers (e.g., C7). This suggests that maintaining diverse guidance options is more important than their specific values, as long as both exploration and exploitation are adequately represented.\n\n**UCB hyperparameter ablation study**\n\nWe investigated the influence of the UCB exploration weight parameter $C$ in the formula: $v_{\\text{UCB}} = v_i + C\\sqrt{\\frac{\\ln N}{n_i}}$, where $v_i$ is the node's estimated value, and $N$ and $n_i$ are the visit counts of the parent node and the node itself, respectively.\n\nC=0(Greedy) / 1.141 (Default) / 3 / 5 / 10\n\nMedium:\n\n- Performance: 88 ± 13 / 100 ± 0 / 100 ± 0 / 98 ± 6 / 98 ± 6\n- Run time: 15 ± 1 / 31 ± 4 / 63 ± 10 / 76 ± 23 / 92 ± 55\n- Search#: 105 ± 84 / 77 ± 29 / 94 ± 8 / 120 ± 27 / 134 ± 23\n\nLarge:\n\n- Performance: 90 ± 10 / 98 ± 6 / 100 ± 0 / 98 ± 6 / 100 ± 0\n- Run time: 16 ± 0 / 74 ± 34 / 90 ± 10 / 102 ± 14 / 104 ± 18\n- Search#: 117 ± 78 / 174 ± 90 / 211 ± 26 / 257 ± 41 / 265 ± 43\n\nGiant:\n\n- Performance: 82 ± 14 / 100 ± 0 / 98 ± 6 / 100 ± 0 / 100 ± 0\n- Run time: 25 ± 1 / 215 ± 23 / 216 ± 12 / 225 ± 22 / 230 ± 10\n- Search#: 228 ± 69 / 394 ± 152 / 442 ± 14 / 464 ± 13 / 493 ± 7\n\nWith C=0 (pure greedy search), MCTD shows faster inference time and conducts fewer searches, but performance degrades. This occurs because greedy search tends to explore the tree depth-wise quickly, reducing the overhead from jumpy denoising but compromising thorough exploration. The default value (C=1.141) achieves a good balance, while higher values show diminishing returns in performance despite increased computational costs.\n\n**Other ablation studies**\n\nWe also studied other ablations for **subplan length**, **causal denoising**, **tree search**, and **the scale of jumpiness** which are discussed in the rebuttal for 6pZp's review.\n\n> Could your method be more **generally presented as planning** at a higher level of abstraction and then open to a variety of planning algorithms, or is it inherently tied to an MCTS based approach?\n> \n\nWhile we built our approach on an MCTS backbone, the core idea—treating partial denoising steps as higher-level \"states\" and using meta-actions (guidance choices) to expand them—is not inherently tied to MCTS. In principle, **one could substitute different search algorithms** (e.g., best-first search or A*) to explore and refine these partial trajectories. However, MCTS's ability to balance exploration-exploitation through principled tree search is uniquely suited for navigating the complex, high-dimensional trajectory spaces in diffusion planning. \n\n> The **readout policy** used was best child in MCTS, did you explore any other readout policies?\n> \n\nThank you for suggesting this interesting idea. We agree that we may consider most visited child as an alternative. However, **it seems not be fitting well to our settings** as it pursues finding sequences of a complete high-quality plan instead of inferring the next best action. In future work, we plan to explore other readout strategies and their effects on planning in different domains.\n\n[1] Silver, David, et al. \"Mastering the game of Go with deep neural networks and tree search.\" *nature* 529.7587 (2016): 484-489.", "claims_and_evidence": "Claims made are both supported and convincing.", "methods_and_evaluation_criteria": "Evaluation datasets are appropriate.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Yes sound.", "relation_to_broader_scientific_literature": "The contributions of \"Monte Carlo Tree Diffusion for System 2 Planning\" bridge diffusion-based planning and Monte Carlo Tree Search (MCTS) literature. It builds on diffusion works like Diffuser (Janner et al., 2022) and Diffusion Forcing (Chen et al., 2024a) by adding tree-structured search to improve test-time compute scalability, addressing limitations noted in Karras et al. (2022), and extends guidance concepts from Dhariwal & Nichol (2021) into meta-actions. In the broader RL and MCTS context, it advances Coulom (2006) and applications like AlphaGo (Silver et al., 2016) by planning at a higher abstraction level with subplans as nodes, reducing reliance on error-prone forward models (Hafner et al., 2022) and aligning with hierarchical planning trends (Chen et al., 2024c), offering a novel, efficient alternative for long-horizon tasks and System 2 reasoning.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths\n- Novel integration of MCTS with diffusion that enables planning at a high level of abstraction, which is very promising for making test time planning methods more efficient.\n\nWeaknesses\n- Lack of ablations investigating the parameters of MCTS search. In particularly investigating the balancing of the exploitation-exploration trade off. This would give more insight as to whether the planning is more sampling based or really finding promising branches and exploiting them. Willing to raise my score if such ablations are added.", "other_comments_or_suggestions": "No", "questions_for_authors": "- Could your method be more generally presented as planning at a higher level of abstraction and then open to a variety of planning algorithms, or is it inherently tied to an MCTS based approach?\n\n- The readout policy used was best child in MCTS, did you explore any other readout policies?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Hv4CpmK1lq", "reviewer": "Reviewer_6pZp", "summary": "The paper introduces Monte Carlo Tree Diffusion (MCTD) for long-horizon planning problems. MCTD combines the benefits of generative Diffusion models with tree-based planning in Monte Carlo Tree Search, without the requirement of a forward dynamics model. The primary aim is to utilise Diffusion models for improved long-horizon planning. Additionally, it allows MCTD to scale at test time with additional compute.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Essential References Not Discussed**\n\nWe thank you for recommending traditional long-horizon planning literatures. We will incorporate these valuable references in our revised version.\n\n> ... the lack of details provided in the Experimental Setup. ... lacking in the main manuscript.\n> \n\nThank you for suggesting more experimental details in the main manuscript. In our revised version, we will include the experimental setup in the main text.\n\n> Typo: page 2 “aiming to balencely”, Page 4, should define DDIM.\n> \n\nThank you for identifying these issues! We will correct them.\n\n> When partitioning trajectories into x_1, …, x_S, is S predetermined? How do you select it for each environment? and what effect does **altering S have on the performance**?\n> \n\nThe subplan count $S$ is indeed a predetermined hyperparameter that we adapt to each environment's episode length. Following your insightful question, we conducted comprehensive ablation studies on PointMaze to analyze its effect. **The results demonstrate the suitable $S$ choice leads best effectiveness and efficiency.**\n\n**Subplan length ablation study**\n\nS=1 / 3 / 5 (Default) / 10 / 20\n\nMedium:\n\n- Performance: 80 ± 13 / 96 ± 8 / 100 ± 0 / 96 ± 8 / 98 ± 6\n- Run Time: 11 ± 1 / 14 ± 3 / 31 ± 4 / 103 ± 5 / 91 ± 3\n- Search#: 81 ± 63 / 29 ± 40 / 74 ± 9 / 475 ± 28 / 500 ± 0\n\nLarge:\n\n- Performance: 80 ± 0 / 90 ± 10 / 100 ± 0 / 100 ± 0 / 92 ± 10\n- Run Time: 11 ± 0 / 19 ± 4 / 65 ± 11 / 118 ± 6 / 131 ± 10\n- Search#: 111 ± 31 / 64 ± 50 / 190 ± 37 / 500 ± 0 / 500 ± 0\n\nS=1 / 2 / 4 / 8 (Default) / 15 / 30\n\nGiant:\n\n- Performance: 72 ± 19 / 88 ± 16 / 100 ± 0 / 100 ± 0 / 100 ± 0 / 100 ± 0\n- Run Time: 13 ± 1 / 18 ± 4 / 40 ± 22 / 215 ± 23 / 213 ± 9 / 178 ± 7\n- Search#: 181 ± 78 / 75 ± 104 / 46 ± 41 / 374 ± 48 / 500 ± 0 / 500 ± 0\n\nThese results reveal an important trade-off: With longer subplans (smaller S), the search space is reduced, leading to faster running times but potentially lower performance. Conversely, very short subplans (e.g., S=20) can require more search iterations, potentially exhausting the computational budget before finding optimal solutions.\n\n> Since the Diffusion Forcing paper introduces a casual Denoising schedule, is it feasible to suggest that **the additional performance gains are due to the tree-based planning?** A more indepth discussion on this would be interesting.\n> \n\nThis is an excellent question about disentangling the contributions of tree search (TS) and causal denoising (CD). To address this directly, we conducted an ablation study comparing four configurations: Diffusion Forcing (DF) without CD, standard DF, MCTD without CD, and full MCTD. **The results demonstrate the huge performance gains from TS.**\n\n**Causal Denoising and Tree Search ablation studies**\n\nDF wo CD / DF / MCTD wo CD / MCTD Performances:\n\n- Medium: 44 ± 22 / 40 ± 21 / 87 ± 16 / 100 ± 0\n- Large: 50 ± 14 / 55 ± 20 / 78 ± 6 / 98 ± 6\n- Giant: 8 ± 10 / 34 ± 16 / 18 ± 11 / 100 ± 0\n\nThe results reveal that both components contribute substantially to MCTD's performance. TS consistently improves performance across all maze sizes. Meanwhile, CD provides especially dramatic gains for long-horizon planning in the giant maze.\n\n> What is C in the fast Denoising process? How does **this affect the backpropagated values**?\n> \n\nThe parameter $C$ in fast denoising determines how many levels are skipped during the jumpy denoising process. Larger $C$ values mean fewer denoising steps are performed, trading accuracy for speed. We investigated how different jumpiness scales affect performance and efficiency with the version which de-noises the trajectory in one-shot. **The results demonstrate proper value choice for $C$ affects the model performance and efficiency.**\n\n**Jumpiness scale ablation study**\n\nJump=1 / 5 / 10 (Default) / 20 / 50 / One-shot\n\nMedium:\n\n- Performance: 100 ± 0 / 100 ± 0 / 98 ± 6 / 100 ± 0 / 100 ± 0 / 100 ± 0\n- Run time: 84 ± 26 / 39 ± 16 / 34 ± 13 / 29 ± 12 / 27 ± 13 / 42 ± 12\n- Search#: 86 ± 30 / 74 ± 32 / 77 ± 29 / 73 ± 30 / 65 ± 29 / 143 ± 46\n\nLarge:\n\n- Performance: 100 ± 0 / 100 ± 0 / 100 ± 0 / 98 ± 6 / 94 ± 13 / 63 ± 39\n- Run time: 160 ± 69 / 91 ± 35 / 74 ± 34 / 65 ± 32 / 68 ± 49 / 355 ± 183\n- Search#: 176 ± 85 / 188 ± 81 / 174 ± 90 / 167 ± 89 / 174 ± 112 / 301 ± 134\n\nGiant:\n\n- Performance: 100 ± 0 / 98 ± 6 / 100 ± 0 / 100 ± 0 / 88 ± 16 / 10 ± 21\n- Run time: 632 ± 239 / 258 ± 115 / 231 ± 103 / 185 ± 88 / 174 ± 76 / 164 ± 2\n- Search#: 390 ± 148 / 379 ± 158 / 394 ± 152 / 363 ± 154 / 387 ± 149 / 500 ± 0\n\nFor relatively simple tasks, even aggressive (large $C$) jumpiness or one-shot maintains high performance. However, as task complexity increases, they significantly degrades performance, as the value function estimation becomes less accurate under highly jumpy denoising, leading to suboptimal tree search decisions. The default (Jump=10) achieves an excellent balance, maintaining high success rates while reducing computational time compared to smaller jump.", "claims_and_evidence": "The authors make two claims. \n\nOne that MCTD outperforms standard Diffusion-based methods for long-horizon trajectory planning, which they demonstrate on three experiments (Maze Nav, Cube Manipulation, and Vizual Pointmase) and compare against various Diffusion-based baselines such as Diffuser (with various sampling regimes), and Diffusion Forcing.\n\nSecond, their application of MCTS with Diffusion models allows their Diffusion models to improve with additional Test Time Scaling which is demonstrated with one experiment on the Point Maze task.", "methods_and_evaluation_criteria": "The methods and evaluation criteria make sense for the problem introduced.", "theoretical_claims": "None", "experimental_designs_or_analyses": "The experimental design and analysis appear sound and valid.", "relation_to_broader_scientific_literature": "The paper situates itself relative to the Diffusion literature, and blends concepts from MCTS and tree-based planning to improve the behaviour of standard Diffusion models specifically for long-horizon planning.", "essential_references_not_discussed": "The paper focuses on the Diffusion literature. It does not cite traditional planning literature that attempt to solve long-horizon planning problems via sub-plans [1], or Options [2]. Two suggestions would be: \n\n[1] Jurgenson, T., Avner, O., Groshev, E. and Tamar, A., 2020, November. Sub-goal trees a framework for goal-based reinforcement learning. In International conference on machine learning (pp. 5020-5030). PMLR.\n\n[2] M. de Waard, D. M. Roijers and S. C. J. Bakkes, \"Monte Carlo Tree Search with options for general video game playing,\" 2016 IEEE Conference on Computational Intelligence and Games (CIG), Santorini, Greece, 2016, pp. 1-8,", "other_strengths_and_weaknesses": "**Clarity**: A major strength of the work is the extremely well written manuscript. The Experiments and Results look promising, and the paper presents a new way of combining the benefits of Diffusion with long-horizon planning. \n\nOne weakness for me is the lack of details provided in the Experimental Setup. I don’t believe I could re-implement the proposed method from the manuscript alone. Whilst I believe most of the relevant information is provided in the Appendix (and code will be open-sourced), as a reader, I would benefit from a section in the main paper introducing topics such as: the Diffusion model architecture, the value networks, and training hyperparameters. Plus more details on the experiments themselves, like the action space, the training and compute budgets. This is all lacking in the main manuscript. Adding these details would improve the paper. \n\n**Novelty and Significance**\nTo the best of my knowledge, this is the first application of Diffusion models to a tree-search framework such as MCTS, and the results are promising. The paper could be of interest to many researchers in the field.", "other_comments_or_suggestions": "Typo: page 2 “aiming to balencely” \n\nPage 4, should define DDIM.", "questions_for_authors": "1) When partitioning trajectories into x_1, …, x_S, is S predetermined? How do you select it for each environment? and what effect does altering S have on the performance? \n\n2) Since the Diffusion Forcing paper introduces a casual Denoising schedule, is it feasible to suggest that the additional performance gains are due to the tree-based planning? A more indepth discussion on this would be interesting. \n\n3) What is C in the fast Denoising process? How does this affect the backpropagated values?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper combines Monte Carlo tree search (MCTS) with diffusion for the purpose of planning. By augmenting diffusion planning with MCTS, it is possible to leverage additional test time compute to improve the quality of the plans, a capability that vanilla diffusion doesn't have.

All reviewers with a substantial input agree that this is a strong paper that deserves publication. Although not a magic bullet (the increase in computational cost is significant), it shows a non-trivial way to leverage additional compute in diffusion planning, a topic of significant focus for the ICML community. Experiments are well-designed and convincingly show the improved efficiency, including long-horizon tasks, compositional planning, partial observability, etc. The rebuttal also shows the benefits of the proposed approach wrt one-step greedy search (a more trivial integration of MCTS and diffusion). Overall, convincing and relevant work with solid results.
Decision: Accept (spotlight poster)