Scalable Maximum Entropy Reinforcement Learning for Diffusion Policies via Adjoint Matching

2026-06-21

AI Summary

This work addresses the challenge of efficiently training diffusion policies in online reinforcement learning without access to real data, where standard approaches such as score matching cannot be directly applied. The authors propose a simulation-agnostic training framework grounded in stochastic optimal control theory, driven by adjoint matching. This method circumvents the need for explicit likelihood estimation or costly backpropagation through the diffusion process, and instead integrates maximum-entropy reinforcement learning with several stability-enhancing strategies. By doing so, it achieves competitive performance while substantially reducing computational overhead, thereby significantly improving the practicality and robustness of diffusion policies in online reinforcement learning settings.

Abstract

Diffusion policies have recently emerged as a powerful paradigm for representing complex action distributions in reinforcement learning (RL). However, their application to online RL remains limited by the challenge of scalable training in the absence of ground-truth data, where standard optimization techniques such as score matching are not directly applicable. In this work, we introduce a highly efficient algorithm for optimizing diffusion policies by leveraging recent advances in stochastic optimal control. Our approach is based on adjoint matching, which enables simulation-free training and circumvents the need for explicit likelihood estimation or costly backpropagation through the diffusion process. Furthermore, we propose several extensions that improve the robustness and stability of the method in practical settings. Empirical results demonstrate that our approach achieves competitive performance while significantly reducing computational overhead, making diffusion policies more viable for online RL scenarios.

Problem

• diffusion policies
• online reinforcement learning
• scalable training
• score matching
• stochastic optimal control

Innovation

• adjoint matching
• diffusion policies
• scalable reinforcement learning
• stochastic optimal control
• simulation-free training