FORGE-Tree: Diffusion-Forcing Tree Search for Long-Horizon Robot Manipulation

TL;DR

FORGE-Tree introduces a novel diffusion-based tree search method for long-horizon robot manipulation, improving success rates by local trajectory refinement and selective search, addressing challenges of error accumulation and compute allocation.

Contribution

It presents FORGE-Tree, a new plug-in control layer combining diffusion forcing and Monte Carlo tree diffusion for enhanced long-horizon manipulation.

Findings

Improves success rate by up to 17.2 percentage points over baselines.

Scales performance with search budget while preserving executed prefix.

Maintains gains under limited compute budgets, especially on long-horizon tasks.

Abstract

Long-horizon robot manipulation tasks remain challenging for Vision-Language-Action (VLA) policies due to drift and exposure bias, often denoise the entire trajectory with fixed hyperparameters, causing small geometric errors to compound across stages and offering no mechanism to allocate extra test-time compute where clearances are tight. To address these challenges, we introduce FORGE-Tree, a plug-in control layer that couples a stage-aligned Diffusion Forcing (DF) head with test-time Monte Carlo Tree Diffusion (MCTD). With a frozen VLA encoder, DF aligns timesteps to subtask stages; during inference we partially denoise only a target segment while keeping other tokens frozen, turning trajectory refinement into a sequence of local edits. We then apply Monte Carlo Tree Diffusion to select the next segment to refine. A scene graph supplies priors for expansion and geometry relation-aware scoring for rollouts, yielding tree-structured denoising whose performance scales with search budget while preserving the executed prefix. Evaluation on LIBERO, FORGE-Tree improves success rate by 13.4 to 17.2 pp over the native VLA baselines with both OpenVLA and Octo-Base. Gains remain consistent under comparable compute budgets, especially on long-horizon variants. Videos available at: https://taco-group.github.io/FORGE-Tree/