💻 computer science

DIA: Denoising Intermediate Advantage for Diffusion Policy Optimization

The paper introduces Denoising Intermediate Advantage (DIA), a policy-gradient method that enhances diffusion-based robot policies by assigning state-dependent credit to intermediate denoising steps, thereby enabling more efficient exploration and superior task performance compared to existing fine-tuning approaches.

Original authors: Arjun Sohal, Yuchi Zhao, Miroslav Bogdanovic, Alan Aspuru-Guzik

Published 2026-09-15

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Arjun Sohal, Yuchi Zhao, Miroslav Bogdanovic, Alan Aspuru-Guzik

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

Robots that can manipulate objects with human-like dexterity have long been a goal of artificial intelligence, but teaching them to do so is a delicate balancing act. For years, the most effective method has been "behavior cloning," where a robot learns by watching videos of humans performing tasks and then mimicking those movements. This approach has given rise to a powerful new class of robot controllers based on a technology called diffusion. Much like how a sculptor might start with a rough block of stone and gradually chip away material to reveal a statue, these diffusion models start with random noise and slowly refine it into a precise sequence of actions. While this method is excellent at capturing the variety and nuance of human demonstrations, it has a fundamental limitation: the robot is strictly bound by the quality and variety of the data it was shown. If the training videos never showed a robot solving a problem in a new way, the robot will likely fail when faced with that problem in the real world. To overcome this, researchers have begun combining these diffusion models with reinforcement learning, a technique where an agent learns by trial and error, receiving rewards for good outcomes and penalties for mistakes. The challenge, however, is that diffusion models do not make a single decision at once; they generate an action through a long, step-by-step process of refinement. Determining exactly which step in that long chain was responsible for success or failure has proven to be a difficult puzzle for existing methods.

A team of researchers at the University of Toronto and the Vector Institute has proposed a new solution to this problem called Denoising Intermediate Advantage, or DIA. Their work addresses a specific flaw in how current systems give credit to a robot's learning process. In standard approaches, when a robot successfully completes a task after a long sequence of refinement steps, the system assigns the same amount of credit to every single step in that sequence. It is as if a team of painters were working together to finish a mural, and the manager praised every brushstroke equally, regardless of whether a particular stroke was a minor detail or the crucial line that defined the entire image. The researchers argue that this is inefficient because the intermediate steps in the generation process contain valuable information about where the robot is heading. By treating the entire refinement chain as a single block, previous methods missed the opportunity to learn from the specific decisions made at each stage of the process.

To fix this, the DIA method introduces a way to evaluate the robot's progress at every single step of the refinement process, not just at the end. The system learns to predict the value of the action being formed as it takes shape, step by step. This allows the robot to understand that some intermediate decisions are more critical to the final outcome than others. Instead of waiting until the very end to see if the task was completed, the system provides feedback throughout the generation process, guiding the robot to make better choices earlier on. This creates a more nuanced learning signal that helps the robot distinguish between a lucky success and a well-executed strategy. The researchers tested this approach across four different sets of robotic tasks, ranging from simple object manipulation to complex, multi-step assembly jobs that require a robot to move its arms in coordination over a long period.

The results of these tests were consistent and significant. On a standard set of benchmarks known as Robomimic, which includes tasks like lifting objects, moving cans, and drawing squares, the DIA method consistently outperformed existing techniques. In the most difficult task, a bimanual transport challenge where a robot must move an object using two arms, the new method achieved a 23 percent higher reward score than the previous best approach, while maintaining the same high rate of success. This improvement was not just about getting the job done more often; it was about doing it better. The robots trained with DIA reached the successful state more quickly, taking fewer steps to complete the task. In one specific test, the time it took for the robot to succeed was reduced by 21 percent. This suggests that the robot was not merely stumbling upon a solution but was learning a more efficient path to the goal.

The power of this method became even more apparent when the researchers tested it on tasks where the training data was incomplete or missing the full solution. In a kitchen simulation where a robot had to perform a sequence of subtasks like turning on a stove or opening a cabinet, the standard training data often showed only parts of the full sequence. Previous methods struggled here, often getting stuck in loops or repeating irrelevant actions because they were too reliant on the exact patterns seen in the training videos. The DIA method, however, was able to recombine the partial demonstrations it had seen to create entirely new, successful sequences of actions. In the most challenging version of this test, where no single demonstration showed the complete task, the baseline methods failed completely, achieving zero percent success. The DIA method, by contrast, succeeded 69 percent of the time. It managed to piece together the necessary steps to complete the task, effectively learning a strategy that was not explicitly present in any of the original training examples.

These findings suggest that by paying attention to the intermediate steps of decision-making, robots can break free from the limitations of their training data. The method allows them to explore new strategies and discover more efficient ways to solve problems, rather than simply copying what they have seen before. While the experiments were conducted in simulation, the results point toward a future where robots can adapt more flexibly to complex, real-world environments. The researchers note that the next step will be to test these ideas on physical robots, a transition that will require overcoming the practical challenges of collecting data in the real world. For now, the work demonstrates that giving credit to the right moments in a robot's thought process can lead to significantly smarter and more capable machines.

Technical Summary: Denoising Intermediate Advantage (DIA) for Diffusion Policy Optimization

Problem Statement
Diffusion-based policies have become a standard for behavior cloning (BC) in robotic manipulation due to their ability to model multi-modal action distributions. However, policies trained solely on demonstrations are limited by the quality and coverage of the data. While Reinforcement Learning (RL) can fine-tune these policies through interaction, credit assignment remains a significant challenge. Unlike single-step policies, diffusion policies generate actions through a sequence of stochastic denoising steps. Existing fine-tuning methods, such as DPPO, formulate this as a two-level Markov Decision Process (MDP): an outer MDP over environment states and an inner MDP over the denoising process. A critical limitation of current approaches is that they assign the same environment-level credit (advantage) to all denoising steps used to construct an action chunk. They fail to distinguish which intermediate denoising decisions contributed most to the final return, often relying on fixed heuristics (e.g., depth-based weighting) rather than state-dependent signals. This leaves the structure of the denoising trajectory itself underutilized for optimization.

Methodology
The authors propose Denoising Intermediate Advantage (DIA), a policy-gradient method that introduces a value function over partially denoised actions to enable fine-grained credit assignment within the generative chain.

Two-Level MDP Formulation: The method treats the denoising process as a sequential decision process. At each environment time step $t$ , the policy generates an action $a_t$ through $K$ denoising steps, transitioning from initial noise $x_{t,0}$ to the final action $x_{t,K}$ . The environment reward is only received at the final step of the denoising chain ( $k=K-1$ ), while intermediate steps yield zero reward.
Inner Value and Advantage: DIA defines a value function $V_k(s_t, x_{t,k})$ over the partially denoised state at step $k$ of the chain. This value represents the expected return of the environment given the current state and the partially denoised action. The method derives an inner advantage $A_k$ for each denoising step, calculated as the difference between the value of the next state and the current state ( $V_{k+1} - V_k$ ).
Estimation via GAE: To estimate the inner advantage, the authors employ Generalized Advantage Estimation (GAE) within the denoising chain. They train a critic to predict the value of partially denoised actions, bootstrapping from the estimated Q-value of the fully denoised action ( $\hat{Q}(s_t, a_t)$ ).
Advantage Mixing: The final policy update combines the standard environment-level advantage ( $A^{out}$ ) with the newly derived inner denoising-level advantage ( $A^{in}$ ). The combined advantage is formulated as $\hat{A}^{comb}_{t,k} = \hat{A}^{out}_t + \alpha c_B \hat{A}^{in}_{t,k}$ , where $\alpha$ controls the relative scale and $c_B$ is a scaling factor to match the variances of the two estimators.

Key Contributions

DIA Algorithm: Introduction of a policy-gradient method that learns values over intermediate denoising states, enabling state-dependent credit assignment for individual decisions within the generative process.
Theoretical Derivation: Derivation of the denoising-level advantage and its integration with standard environment-level policy-gradient estimators for diffusion-policy fine-tuning.
Empirical Validation: Demonstration across four benchmarks (Robomimic, FurnitureBench, Franka Kitchen, and D3IL) that DIA consistently improves final performance, particularly on challenging long-horizon tasks, compared to existing fine-tuning methods like DPPO, QSM, and IDQL.

Results

Performance Gains: DIA consistently achieves higher final returns and success rates than baseline methods. On the Robomimic TRANSPORT task (a difficult bimanual task), DIA achieved a 23% higher return than DPPO at a similar success rate. On FurnitureBench, gains were observed in both return (up to 20%) and success rate (up to 12%) for long-horizon assembly tasks.
Efficiency: Beyond raw reward, DIA reaches successful states more efficiently. In the TRANSPORT task, DIA reduced the median time to task completion by 21% compared to DPPO. In D3IL AVOID, DIA trajectories were shorter and more direct, deviating less from the optimal path than those of DPPO.
Behavioral Discovery: In the Franka Kitchen benchmark, specifically in the MIXED setting where no demonstration contained the full task sequence, DIA achieved a 69% full-task success rate, whereas both the baseline BC policy and DPPO achieved 0%. DIA successfully recombined separately demonstrated subtasks into effective task-level sequences that baseline methods failed to discover.

Significance and Claims
The paper claims that DIA addresses a fundamental gap in diffusion-policy RL by utilizing the internal structure of the generative process for credit assignment. The authors argue that treating all intermediate decisions with a single environment-level signal ignores the information contained in the evolving denoising trajectory. By incorporating intermediate denoising-state values, DIA allows the policy to:

Shift farther from the pre-trained behavior distribution.
Discover more effective and efficient task-level strategies.
Compose subtask behaviors into successful sequences even when such sequences are weakly represented or absent in the original demonstrations.

The authors conclude that these results demonstrate the effectiveness of incorporating intermediate denoising-state values into policy-gradient fine-tuning, particularly for complex, long-horizon robotic manipulation tasks. Future work is noted as extending these methods to physical robots and improving sample efficiency for real-world deployment.

Technical Summary: Denoising Intermediate Advantage (DIA) for Diffusion Policy Optimization

More like this