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

TL;DR

FORGE-Tree introduces a modular control-layer that augments frozen Vision-Language-Action models with a Diffusion-Forcing head and Monte Carlo Tree Diffusion to plan long-horizon manipulation as segment-level refinements. By aligning denoising to stage structure, performing budgeted, segment-wise edits, and guiding search with a differentiable scene graph, it preserves executed prefixes while allocating compute where geometry is tight. The approach, trained on EMMA-X and evaluated on LIBERO and ManiSkill, yields substantial improvements over native VLA baselines across OpenVLA and Octo backbones, particularly on long-horizon tasks. Its plug-and-play nature, interpretability via Relational scene graphs, and compute-aware planning offer a practical path to robust, scalable long-horizon manipulation in real-world settings.

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/

Figures (5)

• Figure 1: System overview. Pretrained VLAs encode the instruction and observations, builds a scene graph $\mathcal{G}$, our diffusion head predicts noise, the partial-denoising sampler edits a selected future segment with meta-actions $a=(k,m,s,w,\tau)$, and MCTD evaluates candidates with geometry-aware rewards before executing the first segment and replanning.
• Figure 2: Diffusion Forcing (DF) with partial denoising. Left: Noise as masking—subtasks $\mathcal{S}_j$ share per-subtask timesteps $t[i]{=}t_j$ (darker cells indicate larger $t$). Middle (training): conditioned on $c\!=\!f_{\text{VLA}}(o,u)$, the diffusion head predicts $\epsilon_\theta(x_t[i],c,t[i])$, recovers $\hat{x}_0[i]$, and is supervised by the DF loss (noise MSE + end-of-trajectory and stage-end geometric terms + smoothness) chi2023diffusionchen2024diffusion. Right (inference): we partially denoise a selected segment $S_{k,m}$ using jumpy DDIM song2020denoising with geometry guidance $w\nabla_{x_t}U(\hat{x}_0;\mathcal{G})$; only the segment evolves while the complement is frozen, and the first segment is committed before replanning.
• Figure 3: Stage-aware MCTD. For each stage we run an MCTS whose edges are meta-actions $a{=}(k,m,s,w,\tau,t)$. Selection uses P-UCT with scene-graph priors, silver2017mastering, and unvisited edges are ordered by a fast heuristic $r_{\text{fast}}(n,a)$. Simulation performs partial denoising on $S_{k,m}$ (jumpy DDIM, guidance strength $w$) to produce a candidate $\hat{x}_0$, which is scored by the geometry-aware return $R(\hat{x}_0;\mathcal{G})$. Returns are backed up to update $(N,W,Q)$; the best child is executed for one segment, observations update $\mathcal{G}$, and the process recedes to the next stage.
• Figure 4: Top: OpenVLA collides (red dashed boxes). Bottom: FORGE-Tree edits only the upcoming segment with geometry guidance and places actions inside the goal region with proper clearance.
• Figure 5: Left→right: snapshots from a LIBERO task (“put the white mug on the left plate and put the yellow mug on the right plate”) with FORGE-Tree. In each panel, the top diagram shows the current scene graph $\mathcal{G}$ with nodes (Gripper, Mug1–3, Plate1–2, Table) and edges labeled by relations (in, on); the bottom image shows the corresponding RGB frame with color-matched boxes. As the episode advances, relations update (e.g., in(Gripper, Mug), on(Mug, Plate)), marking subgoal completion. FORGE-Tree uses this graph both to bias expansion (priors over meta-actions) and to score candidates via geometry relation predicates, enabling stage-wise planning that satisfies the relational goals.