Tree-Guided Diffusion Planner
Hyeonseong Jeon, Cheolhong Min, Jaesik Park
Abstract
Planning with pretrained diffusion models has emerged as a promising approach for solving test-time guided control problems. Standard gradient guidance typically performs optimally under convex, differentiable reward landscapes. However, it shows substantially reduced effectiveness in real-world scenarios with non-convex objectives, non-differentiable constraints, and multi-reward structures. Furthermore, recent supervised planning approaches require task-specific training or value estimators, which limits test-time flexibility and zero-shot generalization. We propose a Tree-guided Diffusion Planner (TDP), a zero-shot test-time planning framework that balances exploration and exploitation through structured trajectory generation. We frame test-time planning as a tree search problem using a bi-level sampling process: (1) diverse parent trajectories are produced via training-free particle guidance to encourage broad exploration, and (2) sub-trajectories are refined through fast conditional denoising guided by task objectives. TDP addresses the limitations of gradient guidance by exploring diverse trajectory regions and harnessing gradient information across this expanded solution space using only pretrained models and test-time reward signals. We evaluate TDP on three diverse tasks: maze gold-picking, robot arm block manipulation, and AntMaze multi-goal exploration. TDP consistently outperforms state-of-the-art approaches on all tasks. The project page can be found at: https://tree-diffusion-planner.github.io.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext a674c761-b031-44a1-aa74-04a0f275fbf0Builds on23
- Diffusion Models Beat GANs on Image SynthesisPrafulla Dhariwal, Alexander Quinn NicholNeurIPS 2021 · 13,211 citations
- Conservative Q-Learning for Offline Reinforcement LearningAviral Kumar, Aurick Zhou, George Tucker, Sergey LevineNeurIPS 2020 · 2,881 citations
- Offline Reinforcement Learning with Implicit Q-LearningIlya Kostrikov, Ashvin Nair, Sergey LevineICLR 2022 · 1,402 citations
- Planning with Diffusion for Flexible Behavior SynthesisMichael Janner, Yilun Du, Joshua B. Tenenbaum, Sergey LevineICML 2022 · 1,115 citations
- Model Based Reinforcement Learning for AtariLukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski et al.ICLR 2020 · 969 citations
Related papers
- Test-Time Guidance for Flow-Based Generative Models via Parallel Tempering on Source DistributionsShih-Hsin Wang, Joel Keller, Taos Transue, Drake Brown et al.ICML 2026
- Monte Carlo Tree Diffusion for System 2 PlanningJaesik Yoon, Hyeonseo Cho, Doojin Baek, Yoshua Bengio et al.ICML 2025
- Projected Coupled Diffusion for Test-Time Constrained Joint GenerationHao Luan, Yi Xian Goh, See-Kiong Ng, Chun Kai LingICLR 2026 · 7 citations
- Diffusion Tree Sampling: Scalable inference‑time alignment of diffusion modelsVineet Jain, Kusha Sareen, Mohammad Pedramfar, Siamak RavanbakhshNeurIPS 2025 · 41 citations
- Constrained Diffusion with Trust SamplingWilliam Huang, Yifeng Jiang, Tom Van Wouwe, C. Karen LiuNeurIPS 2024 · 13 citations
