Technical Summary: Scalable Maximum Entropy Reinforcement Learning for Diffusion Policies via Adjoint Matching

Problem Statement

Diffusion policies have emerged as a powerful paradigm for representing complex, multi-modal 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. Standard training objectives like score matching or bridge matching are not directly applicable because online RL lacks a static dataset of optimal samples.

Existing approaches to training diffusion policies in online settings face three primary limitations:

1. Reverse KL Divergence: Methods minimizing the reverse KL divergence require backpropagation through the entire diffusion process. This incurs massive memory overhead and computational inefficiency, preventing scaling to larger problems.
2. Importance Weighted Matching: Approaches using importance sampling to minimize errors between predicted and target actions struggle to scale to high-dimensional action spaces.
3. Q-Score Methods: Techniques utilizing the gradient of the Q-function ($\nabla_a Q$) often rely on overdamped Langevin dynamics. These processes typically require truncation after finite steps, offering few theoretical guarantees in the finite-time regime, or rely on heuristic sampling techniques that lack guarantees of generating samples from the true target distribution.

Methodology: Adjoint Matching Diffusion Policy (AMDP)

The authors propose Adjoint Matching Diffusion Policy (AMDP), a framework that formulates diffusion policy optimization for Maximum Entropy RL (MaxEnt-RL) as a Stochastic Optimal Control (SOC) problem. The core innovation is the adaptation of Adjoint Matching (AM), specifically Reciprocal Adjoint Matching (RAM), to enable simulation-free training.

1. Theoretical Foundation

The method frames the policy improvement step as finding a drift function $u$ such that the marginal distribution of a diffusion process at terminal time $\tau=1$ matches the optimal MaxEnt policy $\pi^*(a|s) \propto \exp(Q^\pi(s,a)/\alpha)$.

• Schrödinger Bridge: The problem is cast as a Schrödinger bridge problem, minimizing the kinetic cost of the control $u$ subject to boundary conditions.
• Reciprocal Adjoint Matching: Instead of backpropagating through the diffusion path (which requires storing intermediate states), AMDP utilizes the reciprocal projection of the controlled process onto an uncontrolled reference process. This allows the loss to be minimized by sampling intermediate states $X_\tau$ directly from the terminal state $X_1$ using a conditional reciprocal process $\Pi^0_{\tau|1}$.
• Simulation-Free Training: By sampling $X_\tau$ from the terminal state $X_1$ (which is stored in a replay buffer), the algorithm avoids simulating the forward SDE during the gradient update. This circumvents the need for explicit likelihood estimation and costly backpropagation through the diffusion chain.

2. Key Algorithmic Components

• Action Space Adaptation (Error Function Squashing):
  To handle bounded action spaces (e.g., $A = [-1, 1]^d$), the authors introduce a differentiable transformation $f: \mathbb{R}^d \to [-1, 1]^d$. Unlike the standard element-wise $\tanh$ squashing, AMDP employs a scaled error function (erf) squashing:
  $$f_{\text{erf}}(x) = (\text{erf}(k x_i))_{i=1}^d$$
  The scaling factor $k$ is chosen such that the determinant of the Jacobian of the transformation exactly matches the Gaussian marginal density of the reference process. This specific choice causes the log-ratio of the reference density and the Jacobian determinant to become a constant, significantly simplifying the regression objective and improving numerical stability.

• Policy Iteration:

  • Evaluation: A modified "soft" Bellman backup is used. Since the marginal entropy of diffusion models is intractable, the authors propose a tractable entropy lower bound derived from the control cost, which preserves convergence properties.
  • Improvement: The policy is updated by minimizing the AMDP loss, which is a fixed-point iteration. The unique fixed point of this loss corresponds to the optimal control $u^*$.

• Trust-Region Updates:
  To ensure stability in online settings where the target distribution shifts rapidly, a trust-region constraint is added. This penalizes the deviation of the new drift $u$ from the previous iterate $u_{\text{old}}$, preventing aggressive updates that could lead to distribution collapse or Q-function divergence. This is solved via a relaxed Lagrangian formalism.

Key Contributions

1. Scalable Training Algorithm: AMDP introduces a highly efficient algorithm that optimizes diffusion policies without backpropagating through the diffusion process, drastically reducing memory requirements.
2. Simulation-Free Optimization: By leveraging reciprocal adjoint matching, the method allows for optimization using only terminal actions stored in a replay buffer, enabling multiple gradient updates per sample without re-simulating the SDE.
3. Numerically Stable Action Squashing: The proposal of error function squashing provides a stable alternative to $\tanh$, simplifying the training objective by canceling out specific terms in the loss function.
4. Stable Policy Updates: The integration of a trust-region loss ensures stable policy iteration in online settings without altering the theoretical guarantees of the adjoint matching framework.

Experimental Results

The authors evaluated AMDP on 63 environments across MuJoCo Playground, ManiSkill, and HumanoidBench in the on-policy setting, and on 7 high-dimensional dog and humanoid environments (DMC suite) in the off-policy setting.

• On-Policy Performance:

  • AMDP matches or surpasses strong Gaussian baselines (REPPO, PPO, SPO) and state-of-the-art diffusion methods (DIME, DPPO, FPO).
  • In complex, high-dimensional tasks (HumanoidBench), AMDP outperforms all baselines significantly, even without Best-of-N (BoN) sampling.
  • AMDP converges faster than REPPO in sample efficiency on MuJoCo DMC and humanoid locomotion tasks.
  • When combined with BoN sampling (selecting the best of 16 samples), AMDP achieves superior performance across all benchmarks.

• Off-Policy Performance:

  • On DMC dog and humanoid tasks, AMDP matches the performance of DIME (which optimizes reverse KL) and converges slightly faster, demonstrating the competitiveness of the adjoint matching objective in off-policy settings.

• Computational Efficiency:

  • Wall-clock Time: AMDP incurs only a small constant overhead (~10%) compared to highly efficient Gaussian methods like REPPO, regardless of the number of diffusion steps.
  • Scalability: In contrast, methods using reverse KL backpropagation are 10x slower with 16 diffusion steps and up to 72-80x slower with 128 steps.
  • Model Scaling: The authors successfully trained a large diffusion policy (106M parameters, ~300x larger than standard models) using AMDP, demonstrating its ability to scale to complex models where backpropagation-based methods would fail due to memory constraints.

Significance and Claims

The paper claims that AMDP provides a theoretically grounded, computationally efficient, and robust pathway for deploying expressive generative policies in real-world RL applications.

• Overcoming Memory Bottlenecks: By eliminating the need for backpropagation through the diffusion chain, AMDP makes training diffusion policies viable for online RL scenarios where memory and computational resources are constrained.
• Theoretical Guarantees: Unlike heuristic Q-score methods, AMDP is derived from Stochastic Optimal Control with fixed-point guarantees, ensuring that the policy converges to the optimal MaxEnt distribution.
• Practical Viability: The combination of simulation-free training, numerical stability via erf-squashing, and trust-region updates addresses the primary barriers (scalability, stability, and efficiency) that have previously limited the adoption of diffusion policies in online RL.

The authors conclude that their approach significantly reduces computational overhead while maintaining competitive performance, making diffusion policies a more viable alternative to traditional Gaussian policies for complex continuous control tasks.