Research map / Image generation / Fast diffusion sampling
Resisting Stochastic Risks in Diffusion Planners with the Trajectory Aggregation Tree
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Trajectory Aggregation Tree (TAT), a training-free tree that merges states across multiple sampled diffusion-planner trajectories to filter out infeasible (artifact) states and improve planning reliability and speed.
Problem
Diffusion planners stochastically generate infeasible or unreliable trajectories, which harms decision reliability and stability, and extra denoising steps needed for quality slow planning.
Approach
Aggregate states from many historical and current sampled trajectories into a dynamic tree (trajectories as branches, states as nodes), weight nodes by cumulative decayed visitation, and act on the highest-weight child node, so unreliable states are marginalized without retraining or extra sampling.
Key results
- Diffuser with TAT improves average Maze2D single-task score from 119.5 to 126.2 and multi-task from 129.4 to 136.2.
- On D4RL locomotion, Diffuser with TAT reaches a 81.3 average vs 77.5 for Diffuser, and RGG with TAT reaches 84.0 vs 81.2, with no task-level drops.
- With warm-start sampling at 6 steps, Diffuser with TAT scores 79.6 vs 74.3 baseline at 20 steps, with planning latency of 0.55s vs 1.68s (more than threefold faster).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add TAT to Diffuser | Diffuser · Maze2D Large | score: 123.0 -> 133.4 | Helped |
| Add TAT to Diffuser | Diffuser · Multi2D Large | score: 132.1 -> 143.8 | Helped |
| Add TAT to Decision Diffuser | Decision Diffuser · Kuka block stacking conditional | score: 54.0 -> 62.3 ±3.5 | Helped |
| Add TAT to Diffuser | Diffuser · MuJoCo locomotion average | normalized average return: 77.5 -> 81.3 | Helped |
| Add TAT to RGG | RGG · MuJoCo locomotion average | normalized average return: 81.2 -> 84.0 | Helped |
| Increase the trajectory batch size from 1 to 32 (ablation) | Diffuser^TAT · Hopper Medium | score: 74.0 ±2.2 -> 82.6 ±1.9 | Helped |
| Increase the trajectory batch size from 32 to 512 (ablation) | Diffuser^TAT · Hopper Medium | score: 82.6 ±1.9 -> 86.9 ±1.9 | Helped |
| Reduce sampling to 6 steps and use TAT-weighted warm-start planning (ablation) | Diffuser^TAT · Hopper Medium | score: 74.3 -> 79.6 ±2.0 | Helped |
| Reduce sampling to 4 steps with TAT-weighted warm-start planning (ablation) | Diffuser^TAT · Hopper Medium | score: 74.3 -> 64.8 ±1.4 | Hurt |
| Use TAT-weighted warm-start instead of vanilla warm-start at 8 sampling steps (ablation) | Diffuser · Hopper Medium | score: 77.2 ±2.3 -> 82.4 ±1.7 | Helped |
| Use TAT-weighted warm-start instead of vanilla warm-start at 4 sampling steps (ablation) | Diffuser · Hopper Medium | score: 58.9 ±1.2 -> 64.8 ±1.4 | Helped |
Surprising findings
- Performance of Diffuser with TAT keeps increasing with more aggregated trajectories (batch size up to 512), with gains becoming marginal at large sizes.
- Tree overhead is small: under 5% of total planning time and about 16.8 MiB of memory on Hopper Medium.
Limitations
- Theoretical analysis assumes independent artifact generation across trajectories and worst-case artifact agreement, with a binomial/CLT approximation.
- Requires sampling many trajectories per step (e.g., 32-64), which increases computation relative to a single trajectory.
- Evaluated only on offline control tasks; the paper does not test other generative domains.
Experimental setup
Datasets: Maze2D Multi2D Kuka block stacking D4RL MuJoCo locomotion
Benchmarks: Maze2D Multi2D Kuka block stacking D4RL MuJoCo locomotion
Metrics: normalized average return score planning time memory usage
Baselines: MPPI CQL IQL BCQ BC DT TT MOPO MOReL MBOP Diffuser Decision Diffuser
Builds on: Diffuser Decision Diffuser Restoration Gap Guidance classifier-free guidance warm-start planning
Closest papers by meaning
- Monte Carlo Tree Diffusion for System 2 Planning (ICML 2025 · similarity 0.89)
- Compositional Monte Carlo Tree Diffusion for Extendable Planning (NeurIPS 2025 · similarity 0.89)
- DiffuserLite: Towards Real-time Diffusion Planning (NeurIPS 2024 · similarity 0.89)
- Tree-Guided Diffusion Planner (NeurIPS 2025 · similarity 0.88)
- Fast Monte Carlo Tree Diffusion: 100× Speedup via Parallel and Sparse Planning (NeurIPS 2025 · similarity 0.88)
- DiRecT: Safe Diffusion-Based Planning via Receding-Horizon Denoising (NeurIPS 2026 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.