Beyond On-Policy Exploration: Integrating External Policy Rollouts for Reinforcement Learning in Diffusion Language Models

2026-08-03

AI Summary

Diffusion-based language models in reinforcement learning suffer from inefficient training due to their reliance on sparse, self-generated successful trajectories, which yield insufficient reward signals. This work proposes ERILS, the first method to systematically address the challenges of length mismatch and unstable reward handling when jointly training on external high-reward trajectories and agent-generated ones. By integrating length-controllable external trajectories and a source-specific reward mechanism, ERILS enables efficient and stable policy optimization. The approach achieves substantial performance gains on Sudoku, Countdown, and MATH500 benchmarks, notably attaining a 98.4% four-choice accuracy on Sudoku—compared to 40.3% for the baseline—while maintaining approximately 90% single-shot deterministic accuracy.

Abstract

Recent reinforcement learning methods for diffusion large language models (dLLMs) commonly rely on on-policy rollouts generated by the target dLLM itself. When successful on-policy rollouts are scarce, however, on-policy training may receive little positive reward and make only limited progress. To mitigate this problem, we explore incorporating higher-reward rollouts generated by a stronger external policy alongside on-policy rollouts from the target dLLM. However, directly incorporating these external rollouts introduces two practical challenges: differences in rollout length and instability when jointly processing rewards from on-policy and external rollouts. To address these challenges, we propose External Rollout Integration with Length Control and Source-Specific Processing (ERILS), which controls external-rollout length and processes the rewards of on-policy and external rollouts separately. Experiments on Sudoku, Countdown, and MATH500 under zero-shot evaluation show that ERILS improves multi-sample performance across all three tasks, with the largest gains on Sudoku. On Sudoku, ERILS achieves 98.4% best-of-4 completion accuracy, compared with 40.3% for the strongest baseline. ERILS also maintains approximately 90% deterministic single-completion accuracy on Sudoku across generation lengths of 128, 256, and 512 tokens. Our component analysis further shows that length-controlled external rollouts are more effective than uncontrolled external rollouts, and that source-specific reward processing avoids the training collapse observed with joint reward processing. These results show that rollout construction and reward processing are important design dimensions when integrating external rollouts into dLLM reinforcement learning.

Problem

• reinforcement learning
• diffusion language models
• on-policy exploration
• external policy rollouts
• reward sparsity

Innovation

• diffusion language models
• reinforcement learning
• external policy rollouts
• length control
• source-specific reward processing