Diffusion policy fine-tuning improves robot manipulation with stepwise rewards
DIA: Denoising Intermediate Advantage for Diffusion Policy Optimization
Robotics
Summary
Robot policies help machines decide how to move to complete tasks. When these policies learn only from examples, their skills are limited by the quality of those examples. The authors propose a new method called DIA that gives the robot feedback on parts of its decision-making process during action generation. This targeted feedback helps the robot learn faster and perform better, finding smarter ways to complete tasks than earlier methods.
What this means in practice
- •For robotics engineers: Train robot manipulation policies that learn more efficiently by giving feedback on each step of their multi-step action planning process.
- •For industrial automation teams: Improve robotic task success rates and discover efficient task sequences in automated assembly and material handling.
Authors
Arjun Sohal, Yuchi Zhao, Miroslav Bogdanovic, Alan Aspuru-Guzik
Abstract
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 com- bines 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.