Beyond On-Policy Exploration: Integrating External Policy Rollouts for Reinforcement Learning in Diffusion Language Models
2026-08-03 • Machine Learning
Machine Learning
AI summaryⓘ
The authors studied how to improve training of diffusion large language models (dLLMs) using reinforcement learning. They found that relying only on examples generated by the model itself (on-policy rollouts) can limit learning when good examples are rare. To fix this, they proposed a method called ERILS, which mixes examples from a stronger external model while controlling example length and treating rewards from different sources separately. Their experiments on tasks like Sudoku showed big improvements using ERILS compared to existing methods. This work highlights the importance of carefully managing example generation and reward processing in training dLLMs.
diffusion large language modelsreinforcement learningon-policy rolloutsexternal policyreward processingrollout length controlzero-shot evaluationSudoku completionmulti-sample performancetraining stability
Authors
Wonseok Lee, Jimyeong Kim, Jungmin Ko, Wonjong Rhee
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.