Distillation Models are Good Samplers for Diffusion Reinforcement Learning

We present DMSampler, a framework that accelerates online diffusion reinforcement learning by replacing expensive training-time policy rollouts with a co-evolving few-step distilled sampler. Instead of repeatedly sampling the policy model for roughly 50 denoising steps, DMSampler generates reward-evaluation samples in only 4--8 steps while periodically re-distilling the sampler from the updated policy, yielding an order-of-magnitude reduction in rollout cost. The framework alternates between two stages: an RL phase that optimizes the policy using hybrid samples from the old policy and distilled sampler, and a distillation phase that realigns the few-step sampler to the improved policy. Intuitively, the distilled sampler acts as a fast proxy for the current policy during RL, and is refreshed whenever the policy improves so that sampling remains both efficient and aligned. Two designs make this loop stable and effective: hybrid distillation sampling preserves on-policy structure during rollout, and reward-aware distillation reuses high-reward trajectories to reduce forgetting during compression. Experiments on text-to-image and text-to-video generation show that DMSampler improves OCR, GenEval, and VBench performance while substantially reducing GPU hours, and that the same idea can be combined with multiple diffusion RL optimizers. Our code will be available at: https://github.com/HiDream-ai/DMSampler.

Diffusion models can generate high-quality images and videos, but improving them with online reinforcement learning is very expensive because the model must repeatedly generate many samples during training. Each sample often requires dozens of denoising steps, so most of the training time is spent on sampling rather than learning. We propose DMSampler, a method that speeds up training-time sampling by using the current policy for the most important early denoising steps and a much faster few-step distilled sampler for the remaining steps. To keep this fast sampler reliable, we update it together with the main model: the main model learns from rewards, and the distilled sampler is then refreshed to match the improved model. We also reuse high-reward samples during distillation so that the fast sampler preserves the behaviors that reinforcement learning has improved. This makes diffusion reinforcement learning much cheaper while maintaining or improving generation quality. In experiments on text-to-image and text-to-video tasks, DMSampler improves text rendering, image alignment, and video generation benchmarks while substantially reducing training cost.