💻 computer science

Compositional Visual Planning via Inference-Time Diffusion Scaling

This paper proposes a training-free framework for long-horizon robot planning that achieves stable compositional visual generation by enforcing boundary agreement on Tweedie estimates within a factor graph, rather than on noisy intermediate states, thereby enabling robust generalization to unseen start-goal combinations without additional training.

Original authors: Yixin Zhang, Yunhao Luo, Utkarsh Aashu Mishra, Woo Chul Shin, Yongxin Chen, Danfei Xu

Published 2026-03-04

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yixin Zhang, Yunhao Luo, Utkarsh Aashu Mishra, Woo Chul Shin, Yongxin Chen, Danfei Xu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to clean an entire messy house. You have a video of the robot cleaning the kitchen, and another video of it cleaning the living room. But you don't have a video of it cleaning the whole house from start to finish.

The Problem: The "Jigsaw Puzzle" Glitch
Previous methods tried to solve this by taking the kitchen video and the living room video and trying to tape them together. They would look at the blurry, noisy edges where the two videos meet and try to force them to match.

Think of it like trying to glue two pieces of a puzzle together while they are still covered in fog. Because the "fog" (the noise in the AI's generation process) makes the edges look different, the pieces don't fit perfectly. When you try to stitch them, the robot might suddenly teleport, the walls might disappear, or the robot might try to walk through a table. The plan falls apart because the connection points are shaky.

The Solution: The "Crystal Clear" Blueprint
This paper proposes a smarter way to stitch these plans together. Instead of trying to glue the blurry, foggy edges, the authors say: "Let's wait until the fog clears, then glue the clear picture."

Here is how their method works, using a few analogies:

1. The "Tweedie Estimate" (The Crystal Clear Blueprint)

In the world of Diffusion AI (the tech that generates images and videos), the AI starts with a random cloud of static (noise) and slowly removes the noise to reveal a picture.

Old Way: The AI tries to agree on the plan while it's still a blurry cloud of static.
New Way: The authors wait until the AI has "guessed" what the final, clean picture looks like (called the Tweedie estimate). It's like waiting until the fog lifts and you can clearly see the edge of the puzzle piece before you try to connect it. This makes the connection much more stable.

2. The "Chain of Messengers" (Message Passing)

Imagine you have a long line of people passing a message down a hallway.

The Old Way: Everyone shouts their part of the message at the same time, but they are all shouting over the noise. The message gets garbled by the time it reaches the end.
The New Way: The authors use a "Synchronous and Asynchronous" system.
- Synchronous: Everyone checks the message at the exact same time to make sure the neighbors agree.
- Asynchronous: People pass the message forward and backward, correcting errors as they go, like a relay race where runners double-check the baton handoff.
- The Result: The message (the robot's plan) stays consistent from the very first step to the very last, even if the robot has to do something it has never done before.

3. The "Training-Free" Magic Trick

Usually, if you want a robot to do a new, long task, you have to retrain it for weeks with thousands of examples.

This Paper's Trick: They treat the robot's brain like a Lego set. They already have the small Lego bricks (short videos of simple tasks). They don't need to build new bricks; they just need a better instruction manual on how to snap the existing bricks together.
They do this at the moment of thinking (inference time), not during training. It's like having a master builder who can look at a pile of bricks and instantly figure out how to build a castle, a bridge, or a tower without needing to learn how to make new bricks first.

Why Does This Matter?

Generalization: If you show the robot how to open a drawer and how to push a cup, this method allows it to figure out how to "Open the drawer, then push the cup" even if it never saw that specific combination before.
Stability: The robot doesn't hallucinate or glitch out in the middle of a long task. The transition between "opening the drawer" and "pushing the cup" is smooth and logical.
Efficiency: It saves time and money because you don't need to retrain the AI for every new long-term task.

In a Nutshell:
This paper is about teaching an AI to plan long, complex robot movements by acting like a skilled editor. Instead of trying to edit a blurry, noisy movie, the AI waits to see the clear picture, then uses a smart system of checks and balances to stitch short clips together into a perfect, long movie—without ever needing to learn a new skill.

1. Problem Statement

Long-Horizon Visual Planning: While diffusion models excel at short-horizon robot planning and video generation, scaling them to long-horizon tasks is challenging. Existing methods face two main issues:

Computational Constraints: Training models on long sequences is memory-intensive and data-hungry.
Compositional Instability: Existing approaches attempt to stitch together short, overlapping video segments generated by pretrained models. However, they typically enforce consistency by averaging noisy intermediate states (score averaging). The authors argue that the factorization assumption (the idea that joint distributions can be approximated by products of local factors) breaks down in the noisy diffusion space ( $t > 0$ ), leading to inconsistent global plans, boundary gaps, and "brittle" behavior.

2. Methodology

The paper proposes a training-free framework that performs long-horizon planning by treating the problem as inference over a chain-structured factor graph operating on Tweedie estimates (estimated clean data) rather than noisy states.

Core Insight: The Noisy-Bethe Gap

The authors prove a Noisy-Bethe Gap Theorem, demonstrating that the standard Bethe approximation (used in methods like DiffCollage) is inaccurate when applied to noisy diffusion states ( $t > 0$ ). The gap arises because forward diffusion introduces shared, heteroscedastic distortions that create correlations between local factors, violating the independence assumptions required for simple score averaging.

Solution: Instead of composing on noisy states, the method enforces boundary agreement on the Tweedie estimates ( $x_{0|t}$ ), which represent the model's best guess of the clean data at any timestep.

Technical Framework

Factor Graph Formulation:
- A long-horizon plan is decomposed into $n$ overlapping video chunks (factors).
- The graph is a linear chain where neighboring factors share transition boundary variables.
- Boundary Conditions: The first factor is anchored to the start state ( $s$ ), the last to the goal state ( $g$ ), and intermediate factors must agree on shared transition frames.
Joint Synchronous and Asynchronous Message Passing:
To enforce boundary agreement on the Tweedie estimates, the authors introduce a novel message-passing scheme:
- Synchronous Message Passing: Treats the chain as a Gaussian linear system. It optimizes a global loss ( $L_{sync}$ ) that penalizes deviations from a consistent linear system ( $\Sigma^{-1}x = \eta$ ). This ensures global consistency but can be numerically stiff.
- Asynchronous Message Passing: Uses a Temporal-Difference (TD) style approach with stop-gradient operators. It propagates constraints forward and backward using bootstrapped targets, offering faster convergence and stability at the cost of slight bias.
- Joint Optimization: The method combines both losses ( $L = L_{sync} + L_{async}$ ) to balance global consistency with optimization stability.
Diffusion-Sphere Guided Sampling:
- The method utilizes Diffusion Sphere Guidance (DSG) to steer the sampling process without retraining.
- It interpolates between the unconditional sampling direction (diversity) and the gradient descent direction derived from the message-passing losses (alignment).
- This ensures the generated trajectory satisfies boundary constraints while maintaining sample diversity and local quality.
Inference-Time Scaling:
- The short-horizon diffusion backbone is trained once and frozen.
- At test time, the model generates long-horizon plans by chaining overlapping factors and performing message passing on the latent space (using the Cosmos tokenizer) to reduce dimensionality.
- The final visual trajectory is converted to robot actions via a separate inverse dynamics model.

3. Key Contributions

Theoretical Insight: Identification and proof of the "Noisy-Bethe Gap," showing that factorization assumptions fail in noisy diffusion spaces, necessitating composition on clean (Tweedie) estimates.
Novel Inference Framework: A training-free method for long-horizon planning using a chain-structured factor graph with joint synchronous and asynchronous message passing on Tweedie estimates.
Diffusion-Sphere Guidance: A mechanism to balance constraint enforcement (boundary agreement) with sample diversity during the denoising process.
Generalization: The ability to generalize to unseen start-goal combinations (Out-of-Distribution) by composing known short-horizon fragments, without task-specific fine-tuning.

4. Experimental Results

The method was evaluated on a Compositional Planning Benchmark based on the ManiSkill simulator, covering 100 tasks across 4 scenes (Tool-Use, Drawer, Cube, Puzzle) with varying difficulties.

Video Quality: The proposed method significantly outperformed baselines (DiffCollage, GSC) in Dynamic Quality (Motion Smoothness, Background Consistency) and Static Quality (Imaging), producing coherent, artifact-free trajectories.
Task Success Rates:
- In-Distribution (IND): Achieved high success rates (e.g., 97% on Tool-Use), comparable to strong policy learning baselines.
- Out-of-Distribution (OOD): Demonstrated superior generalization. While baselines like DiffCollage failed almost entirely (0% success) on unseen start-goal pairs, the proposed method maintained high success rates (e.g., 96% on Tool-Use, 52-65% on other complex scenes).
- Real-Robot Deployment: Successfully deployed on a Franka Emika Panda robot, achieving 70-100% success on real-world tasks, significantly outperforming DiffCollage (0-10%).
Ablation Studies: Confirmed that combining synchronous and asynchronous message passing yields better results than using either alone, and that performance scales with increased sampling steps.

5. Significance

This work represents a paradigm shift in compositional generative planning. By moving the composition logic from the noisy diffusion space to the clean data estimate space, the authors solve the fundamental instability plaguing previous stitching methods.

Training-Free: It eliminates the need for retraining or fine-tuning large models for new long-horizon tasks, making it highly efficient and scalable.
Generalization: It enables robots to solve complex, unseen long-horizon tasks by recombining learned short-horizon skills, a critical step toward general-purpose robotics.
Broader Applicability: While focused on robotics, the framework is applicable to other domains requiring long-sequence generation, such as panorama image synthesis and long-form text-to-video generation.