Technical Summary: Diffusion ReRoll

Problem Statement

Sequential prediction is fundamental to robotics, encompassing tasks such as video prediction, long-horizon planning, policy learning, and behavior synthesis. While diffusion models have become a standard approach for these tasks due to their ability to capture high-dimensional, temporally correlated, and multimodal outputs, they face a critical limitation in revisability.

Standard full-sequence diffusion models perform a monotonic denoising process where the entire sequence shares a single noise level. This often leads to premature commitment: early segments of a sequence may reach a low-noise (deterministic) state before terminal constraints or future context are fully incorporated. If an early prediction is erroneous, it becomes difficult to correct later in the process because the "prefix" is already locked in.

Prior work, such as Diffusion Forcing, introduced per-token noise levels to enable causal denoising (resolving earlier tokens before later ones). However, even with causal structures, once a token is denoised, it typically remains fixed. Existing revision methods (e.g., inpainting, masked re-prediction) often require separate modules or do not directly address how to decide which parts of a sequence should be corrected during the generative process itself. The core challenge is how to enable iterative cross-horizon revision—allowing earlier and later segments to mutually refine one another—without discarding the sequential structure or requiring a new model architecture.

Methodology: Diffusion ReRoll

The authors propose Diffusion ReRoll, a framework that treats the denoising process itself as a revision interface. Instead of a single monotonic trajectory, Diffusion ReRoll organizes the generation process using a schedule matrix that defines piecewise-linear noise-level patterns, termed linear chunks.

Core Mechanisms

1. Schedule Matrix and Linear Chunks:
   The framework assigns a specific noise level to each token at each denoising step via a schedule matrix $S$. Unlike full-sequence diffusion (shared noise) or standard causal denoising (monotonic per-token noise), Diffusion ReRoll constructs virtual linear chunks. These chunks represent regions of the horizon that are denoised together.

2. ReRoll Events (Selective Re-noising):
   The key innovation is the ReRoll event. When the leading edge of a linear chunk reaches a predefined reset noise level ($k_{reset}$), the corresponding token region is re-noised (injected with higher noise, up to the maximum noise level $K$).

   • This allows the re-noised region to be refined again using updated context from the rest of the sequence.
   • While some regions continue denoising, others return to a high-noise state, preventing premature commitment to erroneous predictions.
   • This enables iterative cross-horizon revision, where earlier and later segments can revise one another before the final sequence is produced.

3. Training with Randomized Linear Chunks:
   To support this structured denoising schedule, the model is trained with randomized linear chunks. Instead of sampling token-wise noise levels independently (as in Diffusion Forcing), the training process samples:

   • The number of chunks ($N_{chunk}$).
   • The slope of the noise profile within each chunk.
   • The starting noise level for each chunk.
   • Masking strategies for bidirectional training (e.g., masking the meeting point in bidirectional schedules).
     This exposes the model to structured, piecewise-linear stochasticity, aligning the training distribution with the deployment schedule.

4. Bidirectional Schedules:
   The framework supports bidirectional schedule matrices where denoising waves originate from both the start and end of the horizon. Two variants are proposed:

   • Meeting Matrix: Forward and backward waves meet and continue revising from both directions.
   • Taking Over Matrix: Backward waves provide early context but are eventually overwritten by forward causal waves, preserving a causal final structure.

Implementation Details

• Architecture: Diffusion ReRoll does not require a new backbone. It utilizes standard transformer or U-Net architectures with per-token noise-level conditioning. Each token receives its own noise-level embedding.
• Inference: The denoising process follows the schedule matrix, using DDIM updates for downward transitions (denoising) and explicit forward noising for ReRoll transitions (resetting noise).

Key Contributions

1. Diffusion ReRoll Framework: A schedule matrix-based framework for revisable robotic sequence generation that utilizes virtual linear chunks and selective re-noising.
2. Programmable Revision Interface: The schedule matrix serves as a controllable interface to induce specific revision and information-flow patterns (e.g., forward, backward, or bidirectional) without altering the underlying model architecture.
3. Empirical Validation: Demonstrated improvements across three distinct domains: long-horizon planning, policy learning, and unified video-action modeling.

Experimental Results

The authors evaluated Diffusion ReRoll against full-sequence diffusion (Diffuser), causal denoising (Diffusion Forcing), and standard diffusion policies (Diffusion Policy) on several benchmarks.

1. Long-Horizon Planning (OGBench PointMaze & AntMaze)

• Guidance-based Replanning: Diffusion ReRoll achieved an average success rate of 88.1%, a 21% relative gain over Diffusion Forcing (72.8%).
• Goal Inpainting: Achieved 90.8% success, a 23% relative gain over Diffuser (73.8%).
• Observation: ReRoll effectively recovered from early planning errors that caused full-sequence diffusion to commit to invalid paths.

2. Policy Learning (LIBERO-10 & RoboCasa)

• LIBERO-10 Multi-task: Across various observation histories and prediction horizons, Diffusion ReRoll improved average success from 40.1% (Diffusion Policy) to 62.7%, a 56.5% relative gain.
• RoboCasa Single-task: Improvements were observed but less consistent than in multi-task settings, suggesting that the benefits of cross-horizon revision are most pronounced in tasks with diverse temporal contexts.

3. Unified Video-Action Modeling (UWM)

• Performance: On LIBERO-10, ReRoll improved OOD joint-policy success from 57.2% to 71.0%.
• Consistency: It achieved the lowest zero-action error (0.301m vs. 0.324m for UWM), indicating superior consistency between generated actions and visual dynamics, particularly under out-of-distribution (OOD) conditions.

Significance and Claims

The paper claims that structured re-noising is an effective mechanism for revisable robotic sequence generation. By treating denoising not just as noise removal but as a process capable of iterative cross-horizon revision, Diffusion ReRoll allows models to correct early errors using future context.

The authors emphasize that:

• The gains are driven by the selective re-noising mechanism, which prevents premature commitment to erroneous prefixes.
• The schedule matrix acts as a flexible design interface, allowing practitioners to control information flow (causal vs. non-causal) without retraining the model architecture.
• The approach is particularly effective in long-horizon and multi-task scenarios where temporal dependencies are complex and errors in early segments can propagate catastrophically.

The work is modest in its scope, noting that all experiments are currently in simulation and that the method requires schedule tuning (slope, reset level, number of events) for new tasks. However, it establishes a new paradigm for using diffusion schedules as a controllable mechanism for robust robotic planning and policy learning.