PaperlayerView exact revision on arXiv

Source-linked AI summary

DIA: Denoising Intermediate Advantage for Diffusion Policy Optimization

Arjun Sohal, Yuchi Zhao, Miroslav Bogdanovic, Alan Aspuru-Guzik

arXiv:2609.12245v2cs.RO

TL;DR

Diffusion-policy fine-tuning faces difficult credit assignment across stochastic denoising decisions. DIA learns values over partially denoised actions to provide per-step credit alongside environment-level feedback. Across four benchmarks, it improves final performance, with larger gains on challenging long-horizon problems.

  • Problem

    Stochastic denoising sequences make credit assignment and stable policy-gradient optimization challenging during reinforcement-learning fine-tuning.

    Source: §I · p. 1

  • Method

    DIA learns values over partially denoised actions, uses them to form denoising-level advantages, and combines that inner signal with the environment-level estimator.

    Source: §I · p. 2

  • Results

    Across Robomimic, FurnitureBench, Franka Kitchen, and D3IL, DIA improves final performance over baselines, with larger gains on long-horizon problems; on bimanual TRANSPORT, it achieves 23% higher return than DPPO at nearly the same success rate.

    Source: §VI · p. 6, A. Results · p. 6

  • Takeaways & Limitations

    DIA reaches successful states through shorter, more direct executions and discovers successful task-level sequences when complete solutions are weakly represented or absent from demonstrations.

    Source: §VI · p. 7

    Using the terminal-Q target rather than sampled returns reduces variance at the cost of potential bias from the learned Q-ensemble.

    Source: §B.2 · p. 11

Abstract

from arXiv · show

Diffusion-based robot policies have become widely used in robotic manipulation, where they are typically trained with behavior cloning. However, policies trained purely from demonstrations are limited by the quality and coverage of the available data. Reinforcement learning can further improve the performance of these pretrained policies through interaction. A common approach is to use policy-gradient methods that formulate diffusion-policy fine-tuning as an outer environment MDP together with an inner denoising MDP. However, existing methods typically assign the same environment-level credit to all denoising steps used to construct an action chunk, without distinguishing which intermediate decisions contributed most to the final return. We introduce Denoising Intermediate Advantage (DIA), a policy-gradient method that learns a value function over partially denoised actions and uses it to construct a denoising level advantage for each step of the generative process. DIA combines this inner credit signal with the standard environment-level PPO advantage, providing state-dependent credit throughout the denoising chain. Across Robomimic, FurnitureBench, Franka Kitchen, and D3IL, DIA consistently improves final performance over existing diffusion-policy fine-tuning methods. Beyond final reward, DIA reaches successful states more efficiently and can shift farther from the pretrained behavior distribution, enabling it to discover more effective and efficient task-level strategies and subtask sequences that baseline methods fail to reach.

I. INTRODUCTION

Diffusion-policy fine-tuning uses environment feedback to improve demonstration-trained policies, but existing credit assignment does not account for the realized intermediate denoising states. DIA adds denoising-state values to assign state-dependent credit, and the paper situates this approach among policy-gradient and Q-based methods for generative policies.

I. INTRODUCTION

  • DIA addresses the gap between shared environment-level credit and the changing intermediate states of denoising by learning values over partially denoised actions.Its denoising-level advantage is combined with the standard environment-level estimator; reported gains include more efficient successful executions and discovery of task-level sequences.Source: §I · pp. 1–2
  • Across Robomimic, FurnitureBench, Franka Kitchen, and D3IL, DIA improves task performance over existing diffusion-policy fine-tuning methods, with larger gains on challenging long-horizon tasks.The reported benefits extend to state- and image-based policies, more efficient successful states, and successful task-level sequences that baseline methods fail to discover.Source: §I · pp. 1–2

A. Policy gradients through the generative chain

  • Generative-chain policy-gradient methods optimize denoising or integration decisions, while DPPO nests denoising inside the environment MDP and uses a shared advantage with fixed depth weighting.Other approaches address flow policies, residual actions, or image-generation chains with a shared terminal reward.Source: A. Policy gradients through the… · p. 2

B. Q-based fine-tuning of diffusion policies

  • Q-based fine-tuning methods guide actions using learned values, but typically treat the generative process as a fixed sampler rather than assigning credit to individual generative decisions.Q-Flow propagates terminal value to intermediate flow states, whereas the described Q-based methods do not attribute credit to individual generative decisions.Source: B. Q-based fine-tuning of diffusion… · p. 2

A. Policy-gradient reinforcement learning

Policy-gradient methods use value-based advantages to weight decision scores, while diffusion policies generate actions through stochastic denoising chains. DIA treats those denoising decisions as the process being optimized, with the completed action producing the environment reward and next state.

A. Policy-gradient reinforcement learning

  • Actor–critic methods weight each policy score term by the advantage rather than the return, and GAE trades critic bias against rollout variance through λ.The value baseline reduces variance in the score-term weighting.Source: A. Policy-gradient reinforcement… · p. 2

B. Diffusion policies

  • A diffusion policy conditions a stochastic denoising chain on an observation to produce an action or short action chunk, representing multimodal behavior through its endpoint distribution.Each Gaussian transition has a closed-form log-density even though the endpoint marginal does not.Source: B. Diffusion policies · pp. 2–3

IV. DENOISING INTERMEDIATE ADVANTAGE

  • DIA learns values along the denoising chain, forms per-decision advantages bootstrapped from the completed action, and combines them with the environment-return estimator.The method is presented as a policy-gradient approach for diffusion policies.Source: §IV · p. 3

A. The denoising process is the primary MDP

  • Each environment action is generated through K stochastic denoising decisions, with the initial state as noise and the fully denoised endpoint executed as the action.The flattened process state includes the environment state, current latent, and denoising index; reward arrives when the final denoising step produces the action, and the next chain begins from fresh noise.Source: A. The denoising process is the… · p. 3

B. Exact policy gradient on the flattened MDP

The section develops an exact policy-gradient formulation for the flattened denoising process. It proceeds from trajectory-level expressions to score terms for denoising decisions and a value-function baseline.

B. Exact policy gradient on the flattened MDP

  • The derivation starts from the discounted-return objective and trajectory probability, then differentiates the trajectory expression and substitutes denoising-step score terms.The text states that the score term depends on the internal denoising path and refers to an appendix proof.Source: B. Exact policy gradient on the… · p. 3
  • The resulting formulation expresses the policy gradient using per-decision terms across the denoising chain rather than only a trajectory-level expression.The supplied derivation presents successive policy-gradient forms after substituting the denoising score term.Source: B. Exact policy gradient on the… · p. 3
  • Subtracting a value-function baseline transforms the policy-gradient expression into a form weighted by the difference between action value and the denoising-state value.The figure contrasts shared, depth-weighted DPPO credit with DIA’s state-dependent inner advantage combined with the outer advantage.Source: B. Exact policy gradient on the… · p. 3, Figure 2 · p. 4

C. Value and advantage on partially denoised actions

DIA assigns values and advantages to partially denoised actions, allowing credit to vary across a denoising chain. The completed action’s Q-value supplies the terminal value for this inner process.

C. Value and advantage on partially denoised actions

D. Estimating the inner advantage

The inner advantage is estimated within the denoising chain using GAE and a learned value estimate of the environment Q-value. DIA combines this signal with the outer environment advantage, matching their scales per batch and using α to control the inner contribution.

D. Estimating the inner advantage

  • GAE estimates the inner denoising-chain advantage, while the value function is trained against the environment Q-value for the current state and fully denoised action.The value-function loss is mean-square error.Source: D. Estimating the inner advantage · p. 4
  • The environment Q-value is estimated with TD-learning using the mean of an ensemble of independently initialized Q functions to reduce noise.Source: D. Estimating the inner advantage · p. 4

E. Advantage mixing and scale matching

V. EXPERIMENTS

The experiments evaluate DIA against diffusion-policy fine-tuning baselines across four benchmarks and report consistent performance gains, particularly on challenging long-horizon tasks. They also examine execution efficiency, end-effector routes, and task sequences that are weakly represented or absent in demonstrations.

V. EXPERIMENTS

  • Evaluation covers robomimic, FurnitureBench, Franka Kitchen, and D3IL across manipulation, assembly, kitchen, and navigation settings; methods generally share pretrained checkpoints and other settings.The comparison exception is Q-chunking, run in its own framework with its own flow policy and without the shared pretrained checkpoint.Source: §V · p. 5

A. Results

  • On robomimic, DIA achieves the highest final reward among compared methods and raises TRANSPORT return 23% over DPPO at nearly the same success rate.The reported return difference is small on shorter-horizon LIFT, CAN, and SQUARE tasks.Source: A. Results · pp. 5–6
  • On TRANSPORT, DIA reduces median time to task completion by 21% versus DPPO at similar success rates.Figure 4 measures time to first success in environment steps over 200 evaluation episodes per seed.Source: Figure 4 · p. 6, A. Results · p. 6
  • On FurnitureBench, DIA improves ONE-LEG return by 12% and success by 3%, and LAMP-MED return by 20% and success by 12% over DPPO.The larger gains are reported for the more challenging LAMP-MED task.Source: A. Results · p. 6
  • On D3IL AVOID, DIA finds shorter, more direct end-effector routes across three pretrained policies, while DPPO remains closer to curved behavior-cloning trajectories.Figure 5 plots 50 trajectories per panel for each policy and task variant.Source: Figure 5 · p. 6, A. Results · p. 6

VI. CONCLUSION

  • Across four benchmarks, DIA improves final performance over baselines, especially on long-horizon tasks, while producing efficient executions and discovering solutions to missing task sequences.The conclusion also reports shorter, more direct executions and successful task-level sequences when complete solutions are weakly represented or absent from behavior-cloning demonstrations.Source: §VI · pp. 6–7

A. Results

  • In Kitchen PARTIAL, DIA improves completed subtasks by approximately 15% and full-task success by 54%; in MIXED, it completes approximately 24% more subtasks and reaches 69% success versus 0% for BC and DPPO.Figure 6 shows ordered subtask-completion flows, with path width representing the fraction of episodes taking each path.Source: A. Results · p. 6, Figure 6 · p. 7

VI. CONCLUSION

  • The evaluation is simulation-based; extending DIA to physical robots remains future work, with real-world on-policy rollouts potentially costly.The authors identify sample efficiency and sim-to-real transfer as next steps.Source: §VI · p. 7

APPENDIX A POLICY GRADIENT DERIVATION

The derivation uses baselines independent of the sampled denoising decision to preserve the policy-gradient expectation. It expresses each inner advantage as a one-step change in value along the denoising chain.

APPENDIX A POLICY GRADIENT DERIVATION

  • A baseline that does not depend on the sampled next denoising choice can be subtracted from its score-weighted return without changing the expectation.The derivation applies this property to a baseline conditioned on the environment state and current partially denoised action.Source: Appendix A · pp. 9–10
  • Each inner advantage is the difference between the next denoising-state value and the current denoising-state value.Summing these one-step changes along a sampled chain cancels intermediate values and connects the chain to the completed-action value.Source: Appendix A · p. 9
  • Later denoising-step terms vanish from the score-weighted expression, leaving the current decision’s contribution.The derivation uses the zero conditional mean of later decisions’ score terms given their chain history, and the downstream return’s independence from the internal path once state and completed action are fixed.Source: Appendix A · pp. 9–10

B.1 Scope of the inner MDP: within-trajectory advantages

DIA keeps environment-level reward estimation separate from within-chain denoising credit. The section contrasts this two-level design and its action-Q inner-value target with a flattened MDP and alternative targets.

B.1 Scope of the inner MDP: within-trajectory advantages

  • A flat critic must predict future denoising noise across later action chains, making its target high-variance and its adjacent-step advantages noisy and poorly identified.The authors report that collapsing the levels destabilizes training even on simple robomimic tasks.Source: §B.1 · p. 10
  • Flattening also expands a T-step environment trajectory to TK steps, whereas DIA keeps environment reward estimation at length T and assigns separate within-chain credit.The section states that discounting and bootstrapping across intervening denoising steps can attenuate the outer reward signal.Source: §B.1 · p. 10

B.2 Inner-value target: regressing to the action’s Q

  • DIA trains each intermediate-state value toward the Q-value of the action that the remaining denoising chain will produce.The inner value is fit to that action-Q target at every denoising index.Source: §B.2 · p. 10

B.1 Scope of the inner MDP: within-trajectory advantages

  • On TRANSPORT, the flat combined-chain GAE reaches train reward 235, versus 375 for DIA and 318 for DPPO.The ablation averages five seeds; Figure 7(a) compares the flat alternative with DIA’s two-level design.Source: §B.1 · p. 10, Figure 7 · p. 11

B.2 Inner-value target: regressing to the action’s Q

  • On TRANSPORT, training the inner value on the environment GAE return reaches train reward 332, against 375 for DIA over five seeds.Figure 7(b) compares this alternative target with DIA.Source: Figure 7 · p. 11, §B.2 · p. 11
  • Replacing the sampled return with its conditional expectation given the completed action preserves the target mean while removing the environment-return variance conditional on that action.The terminal-Q target retains variance from which action the remaining chain produces, but not from the environment response after execution.Source: §B.2 · p. 11
  • Because implementation uses learned critics rather than exact V and Q, the action-Q target trades sampled-return variance for potential Q-ensemble bias.The sampled return stays close to realized rewards at high variance; the terminal-Q target removes that variance at the cost of Q-ensemble bias.Source: §B.2 · p. 11
  • The Gaussian-delta control reaches train reward 299 on TRANSPORT, below DPPO’s 318 and DIA’s 375, recovering none of DIA’s gain.The control replaces inner TD residuals with mean-zero Gaussian draws before forming the inner advantage.Source: §B.2 · pp. 11–12

B.3 Scale matching the inner advantage

DIA rescales its inner advantage to keep its contribution interpretable relative to the outer advantage. The raw-scale ablation shows that without rescaling, the inner signal can contribute little to the mixture.

B.3 Scale matching the inner advantage

  • DIA rescales the inner advantage to match the outer advantage’s per-iteration standard deviation before combining them.Both standard deviations are computed globally over the batch.Source: §B.3 · p. 12
  • After rescaling, α fixes the inner contribution’s standard deviation as a fraction of the outer signal, independent of the Q scale at that iteration.The inner contribution’s standard deviation is exactly α times the outer advantage’s.Source: §B.3 · p. 12
  • On LIFT, the raw-scale mixture ends at 262.1, versus 266.4 for DPPO and 271.0 for DIA, in a single-seed ablation.The inner advantage averaged 3.6% of the outer advantage, supplying 0.7%–3.5% of the mixed advantage at α=0.5.Source: §B.3 · p. 12

APPENDIX C PROBING THE INNER VALUE OVER THE DENOISING CHAIN

The appendix probes whether the learned inner value varies along denoising chains and relates those predictions to changing action choices. It also reports per-iteration runtime and evaluation setup details.

APPENDIX C PROBING THE INNER VALUE OVER THE DENOISING CHAIN

  • The probe evaluates bVk at every chain position for sampled chains at 12 timesteps spanning pre-reward and post-reward states.It uses one successful SQUARE episode with T=100 and first reward at t=26.Source: Appendix C · p. 12, §C.1 · p. 12

C.1 Within-trajectory bVk predictions

  • Before the first reward, the inner value changes by an average magnitude of 11.9 across the chain and peaks at +30.66; from t=24 onward, its change is at most 1.18.The first reward arrives at t=26, and the smaller later change indicates the action choice no longer influences reward in this episode.Source: §C.1 · p. 12
  • With the initial latent fixed, inner-value predictions separate increasingly along the chain between samples grouped by the emitted action chunk’s lateral commitment.The environment value is 10.00 for every sample by construction, while the separation grows as the chain resolves toward the route.Source: §C.1 · pp. 12–13
  • On AVOID, the pretrained policy always passes right, while by iteration 9 the fine-tuned policy always passes left; iteration 7 is the crossover.Figure 8 contrasts DIA’s changing paths with DPPO on the same states.Source: §C.1 · p. 12, Figure 8 · p. 13

APPENDIX D WALL-CLOCK OVERHEAD

  • DIA requires 2–5% more per-iteration wall time than DPPO on the robomimic tasks.The measurements were taken sequentially on the same RTX 3090 at nenv=50.Source: Appendix D · p. 13, Table III · p. 13

APPENDIX E ADDITIONAL TASK AND TRAINING DETAILS

  • The evaluation spans four benchmarks with low- and high-dimensional observations and horizons from 100 to 1000 environment steps.Within each task, compared methods share actor architecture, pretrained behavior-cloning checkpoint, observation normalization, and environment settings, except QC, which uses its own framework and flow policy.Source: Appendix E · p. 13

E.1 Robomimic

The evaluation covers four manipulation benchmarks, using task-specific demonstrations, observations, actions, and training settings. The Robomimic setup includes diverse multi-human demonstrations, while the benchmark suite also spans assembly, kitchen subtasks, and obstacle avoidance.

E.1 Robomimic

  • Robomimic evaluates LIFT, CAN, SQUARE, and TRANSPORT with binary per-step success rewards and multi-human demonstrations from 300 teleoperated episodes per task.The multi-human split includes varied skill, hesitations, retries, and suboptimal grasps; TRANSPORT is bimanual and multi-stage.Source: §E.1 · pp. 13–14

E.2 FurnitureBench

  • ONE-LEG and LAMP are long-horizon FurnitureBench assembly tasks at medium randomization, with 700- and 1000-step horizons, respectively.ONE-LEG uses 50 teleoperated demonstrations totaling 24,091 transitions.Source: §E.2 · p. 14

E.3 Franka Kitchen

  • The Franka Kitchen evaluation uses COMPLETE, PARTIAL, and MIXED datasets, with returns counting solved subtasks out of a maximum of four.COMPLETE targets four subtasks in order; PARTIAL and MIXED differ in targeted subtasks and demonstration data.Source: §E.3 · p. 15

E.4 D3IL Avoid

  • D3IL AVOID tests planar obstacle reaching with three variants that differ in demonstration data and a reward penalizing obstacle contact or leaving the arena.The demonstrations cover routes passing on either side of each obstacle, and the variants use progressively larger route subsets.Source: §E.4 · p. 15

E.5 Listed training hyperparameters

  • The appendix lists shared and per-task DIA training hyperparameters, including network sizes, while DPPO uses matching values for shared hyperparameters.Table VI defines [n]d as d hidden layers of width n.Source: §E.5 · p. 15, Table VI · p. 16

Related papers

Loading 2609.12245v2…