Trust-Region Diffusion Policies for Massively Parallel On-Policy RL

2026-06-13

AI Summary

This work addresses the challenge of unstable training in large-scale parallel online policy reinforcement learning, where diffusion policies suffer from rapidly shifting data distributions. To tackle this issue, the paper introduces diffusion policies into this setting for the first time and proposes a trust-region optimization method that constrains the entire diffusion trajectory using a KL divergence bound, effectively balancing policy expressiveness and training stability. The approach is evaluated across 73 tasks spanning four standard benchmarks, demonstrating strong empirical performance: it matches or surpasses strong baselines on conventional tasks and achieves substantial improvements on challenging humanoid control problems, establishing new state-of-the-art results.

Abstract

Reinforcement learning with massively parallel simulations has become a standard framework for developing robust, deployable policies; however, most existing approaches still rely on simple Gaussian policy parameterizations. Diffusion models provide a more expressive policy class and have shown strong performance on challenging control problems, yet most diffusion-based RL methods are designed for offline or off-policy training. In this work, we ask whether diffusion policies can be trained effectively in the massively parallel, on-policy regime. To this end, we introduce Trust-region Diffusion Policies (TruDi), which enables diffusion policies for on-policy RL with massively parallel simulations. This setting is particularly challenging because the data distribution changes quickly across updates, making stable training with complex policies difficult. TruDi addresses this by integrating a trust-region optimization rule to enforce a KL-divergence constraint over the entire diffusion trajectory. Empirically, we evaluate TruDi on a diverse set of 4 massively parallel RL benchmarks comprising a total of 73 tasks. Across these tasks, TruDi consistently outperforms or is on-par with strong baselines on standard tasks and achieves clear gains on more challenging humanoid control tasks, establishing a strong new baseline for massively parallel on-policy RL.

Problem

• diffusion policies
• on-policy reinforcement learning
• massively parallel simulation
• trust-region optimization
• policy training stability

Innovation

• diffusion policies
• trust-region optimization
• on-policy reinforcement learning
• massively parallel simulation
• KL-divergence constraint