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

Aggregated IQM learning curves on five benchmarks: MuJoCo Playground DMC, ManiSkill3, IsaacLab, MuJoCo Playground Humanoid, and HumanoidBench. TruDi matches or exceeds all baselines and wins clearly on the humanoid benchmarks.

Aggregated performance on five continuous-control benchmarks. TruDi vs. strong on-policy baselines (REPPO, PPO) and the online diffusion method DIME; curves show the interquartile mean (IQM) of episode return with 95% stratified bootstrap confidence intervals over 10 seeds. (a) MuJoCo Playground DMC (20 tasks); (b) ManiSkill3 (14 tasks); (c) IsaacLab (6 tasks, normalized returns); (d) MuJoCo Playground Humanoid (4 tasks); (e) HumanoidBench (27 tasks). TruDi matches or exceeds strong baselines on standard control (a–c) and significantly outperforms all baselines on difficult, high-dimensional humanoid tasks (d–e).

TL;DR: TruDi enables diffusion policies for on-policy RL with massively parallel simulations by enforcing a trust region over the entire diffusion trajectory, a tractable alternative to the intractable marginal KL constraint. Across 73 tasks from 4 benchmark suites, TruDi matches strong Gaussian on-policy baselines on standard control tasks and shows clear gains on challenging humanoid control.

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.

The Challenge

Trust-region constraints are essential for stable on-policy RL: they prevent premature convergence when the data distribution shifts rapidly between updates. Enforcing them requires the policy's action likelihood, which is intractable for diffusion policies because the executed action is the marginal of the full reverse denoising chain. Prior diffusion-RL methods therefore train off-policy with large replay buffers, which reintroduces the memory bottleneck that massively parallel on-policy training avoids.

Our Approach: TruDi

1. Trajectory-level trust region. TruDi views the denoising process as a latent-variable model. The data-processing inequality then gives a tractable upper bound: the KL between the marginal action distributions is bounded by the KL between the joint distributions over full diffusion trajectories,

\[ D_{\mathrm{KL}}\!\left(\overleftarrow{\pi}^{\,0}_{\mathrm{old}}(a^0 \mid s) \,\big\|\, \overleftarrow{\pi}^{\,0}_{\theta}(a^0 \mid s)\right) \;\le\; D_{\mathrm{KL}}\!\left(\overleftarrow{\pi}^{\,0:N}_{\mathrm{old}}(a^{0:N} \mid s) \,\big\|\, \overleftarrow{\pi}^{\,0:N}_{\theta}(a^{0:N} \mid s)\right). \]

Constraining this bound enforces a trust region over the whole generation process. Together with the same inequality applied to the maximum-entropy objective, this yields the optimization problem

\[ \min_{\theta}\; \mathbb{E}_{\rho^{\pi}}\!\left[ D_{\mathrm{KL}}\!\left(\overleftarrow{\pi}^{\,0:N}_{\theta} \,\big\|\, \overrightarrow{\pi}^{\,0:N}\right) \right] \quad \text{s.t.} \quad \mathbb{E}_{\rho^{\pi}}\!\left[ D_{\mathrm{KL}}\!\left(\overleftarrow{\pi}^{\,0:N}_{\mathrm{old}} \,\big\|\, \overleftarrow{\pi}^{\,0:N}_{\theta}\right) \right] \le \epsilon, \]

where \(\overleftarrow{\pi}\) denotes the denoising (generative) process and \(\overrightarrow{\pi}\) the noising process, whose action marginal is anchored at the Boltzmann distribution \( \exp\!\left(Q^{\pi_{\mathrm{old}}}(s,a^0)/\alpha\right)/\mathcal{Z}(s) \) of the soft Q-function. This is where the reward signal enters the objective. Both KL terms factorize over the diffusion steps and can be approximated from samples.

2. Maximum-entropy actor–critic with dual descent. TruDi follows REPPO-style primal–dual updates: a soft Q-function fitted with TD-λ and HL-Gauss, a split objective that switches to a pure trust-region penalty whenever the constraint is violated, and auto-tuned entropy temperature α and Lagrangian multiplier λ.

3. Probability-flow ODE for evaluation. At evaluation time, actions are generated deterministically by an Euler discretization of the probability-flow ODE, with the score scaled by a factor \(c\):

\[ a^{n-1} = a^{n} + \left( \beta^{n} a^{n} + 2c\,\eta^{2}\beta^{n}\, \nabla \log \overrightarrow{\pi}^{\,n}(a^{n} \mid s) \right)\delta. \]

For \(c = \tfrac{1}{2}\) this is the deterministic counterpart of the reverse SDE and matches its marginals. Scaling the score can be read as tempering the intermediate marginals, since \( c\,\nabla \log \overrightarrow{\pi}^{\,n}(a^{n} \mid s) = \nabla \log \left(\overrightarrow{\pi}^{\,n}(a^{n} \mid s)\right)^{c} \). Larger values of \(c\) therefore sharpen the marginals and bias trajectories toward high-density, high-return regions.

Denoising trajectories under SDE sampling, ODE with c=0.5, and ODE with c=1.0. Higher score scaling sharpens marginal distributions and produces greedier samples.

Scaling the score function induces greedier sampling. Trajectories of the denoising process from the Gaussian prior (right) to the target action distribution (left). (a) stochastic SDE sampling; (b) the probability-flow ODE with score scale c = 0.5 matches the SDE's marginals; (c) c = 1.0 sharpens the intermediate marginals and concentrates trajectories on the modes. In practice this yields higher returns at evaluation time.

Results

The aggregated comparison across all 73 tasks is shown in the teaser figure above. TruDi and REPPO form the top tier on standard control (DMC, ManiSkill3, IsaacLab), while on the humanoid suites TruDi opens a clear gap over every baseline. The largest gains appear on the hardest environments (window cleaning, balancing, hurdling, and cube), where Gaussian baselines frequently plateau early. The sections below look at individual findings in more detail.

Comparison with Diffusion & Flow-Based Policies

We also compare against two methods with expressive generative policies: DPPO, which applies PPO-style clipping to diffusion policies, and FPO, which uses flow matching. On the DMC suite, TruDi achieves comparable or better sample efficiency than FPO and clearly outperforms DPPO. This gap becomes much more pronounced on the high-dimensional Humanoid tasks, where DPPO and FPO make little progress while TruDi learns reliably.

These results suggest that the trajectory-level trust-region formulation, rather than per-step clipping, is what makes expressive policies viable in the massively parallel on-policy setting.

Learning curves comparing TruDi, FPO and DPPO on MuJoCo Playground DMC (left) and Humanoid (right).

Left: MuJoCo Playground DMC. Right: MuJoCo Playground Humanoid, where the gap is most pronounced.

Multimodality on Symmetric Tasks

PushT (rotating a T-block 180°, where clockwise and counter-clockwise are both optimal) and StackCube (where red-on-blue and blue-on-red are both optimal) each have two symmetric optimal solutions. Behavior entropy over 200 deterministic runs measures whether a policy captures both modes (1 = perfectly balanced, 0 = mode collapse). Best results highlighted in orange.

Task	Method	Entropy ↑	Episode Return
PushT	REPPO	0	154.92 ± 5.85
	DIME	0.20 ± 0.20	171.92 ± 1.08
	TruDi	0.32 ± 0.13	167.32 ± 4.42
StackCube	REPPO	0	83.92 ± 0.24
	DIME	0.68 ± 0.15	84.40 ± 0.69
	TruDi	0.87 ± 0.08	84.69 ± 0.10

The Gaussian baseline collapses to a single deterministic solution (entropy 0) on both tasks. TruDi captures both modes without sacrificing return.

PushT: the two rotation directions, one per mode

StackCube: blue-on-red and red-on-blue solutions

Each environment admits two equivalent optimal strategies, creating a bimodal optimization problem.

Computational Efficiency & Wall-Clock Time

Method	Training Time (h)	Episode Return (IQM)
PPO	0.95 ± 0.35	0.1 ± 3.7
REPPO	1.07 ± 0.35	29.6 ± 7.2
DIME	1.38 ± 0.36	15.6 ± 9.4
TruDi (Ours)	1.95 ± 0.46	34.8 ± 3.0

Wall-clock cost and performance after 50M steps on MuJoCo Playground Humanoid. TruDi achieves the highest performance with a manageable increase in training duration.

Left: TruDi outperforms REPPO under a matched wall-clock budget. Right: ablation over diffusion steps T; performance saturates at T=8.

Left: wall-clock comparison ( TruDi, REPPO): under a strictly matched budget (≈ 2.5 h), TruDi still substantially outperforms REPPO, and extending REPPO's training to 150M steps yields no further improvement. Right: diffusion-step ablation for TruDi, T ∈ {1, 4, 8, 16, 32}; line color encodes the step count, from T = 1 (light) to T = 32 (dark). T = 1 fails to learn and performance saturates at T = 8.

TruDi's advantage stems from the expressiveness of the diffusion policy and the stability of the trust-region updates rather than from additional compute time.

Ablation Studies

Sensitivity to the trust-region threshold. We sweep ε from 0.01 to 50 on IsaacLab. When the constraint is too loose, the update behaves like the unconstrained case and final return drops clearly; when it is too strict, learning slows down. A broad intermediate band (ε ≈ 0.1 to 0.4) works best.

Evaluation strategies. The probability-flow ODE with score scale c = 1.0 consistently outperforms SDE sampling and best-of-K selection with the learned Q-function (K = 10, 20). Tracing deterministically to the mode of the policy distribution is more reliable than sampling around it.

Left: sensitivity to the trust-region threshold epsilon on IsaacLab. Right: ODE evaluation beats SDE sampling and best-of-K selection on Humanoid.

Left: sensitivity to the trust-region threshold on IsaacLab. Line color encodes ε on a log scale: ε = 0.01 (strict) → ε ≈ 0.1–0.4 (best) → ε = 50 (loose).
Right: evaluation strategies on MuJoCo Playground Humanoid (legend in the plot). The probability-flow ODE with c = 1.0 performs best.

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