StableDRL

Stabilizing Reinforcement Learning for Diffusion Language Models

Jianyuan Zhong^*, Kaibo Wang^*, Ding Ding^*, Zijin Feng, Haoli Bai, Yang Xiang, Jiacheng Sun, Qiang Xu

^* Equal contribution

Official implementation of the ICML 2026 paper
"Stabilizing Reinforcement Learning for Diffusion Language Models."

Overview

Group Relative Policy Optimization (GRPO) is effective for autoregressive language models, but applying it directly to diffusion large language models (dLLMs) can cause reward collapse. StableDRL identifies a self-reinforcing instability loop: noisy estimated importance ratios produce gradient spikes and policy drift, which further amplify the variance of subsequent ratios.

StableDRL breaks this loop with:

• Unconditional clipping, which prevents noisy importance-ratio outliers from producing unbounded updates.
• Self-normalization, which constrains each group update to the convex hull of its per-sample advantage-weighted directions.
• Staircase attention, which extends the method to block diffusion models with leakage-free likelihood estimation.

Note

This code release currently provides the full-attention LLaDA-8B training and evaluation path. The paper additionally evaluates StableDRL on the SDAR-8B block diffusion model.

Results

Figure 1. StableDRL enables stable full-parameter RL training on full-attention and block dLLMs. It improves reasoning accuracy over prior methods across MATH500, GSM8K, Countdown, Sudoku, and AIME'24.

Method

Figure 2. Noisy estimated importance ratios create a spike-drift feedback loop. StableDRL combines unconditional upper clipping and self-normalization to stabilize the policy update.

Installation

Clone the repository and create a Python environment:

git clone https://github.com/JianyuanZhong/StableDRL.git
cd StableDRL

conda create -n stabledrl python=3.10 -y
conda activate stabledrl

Install a CUDA-compatible version of PyTorch, followed by the project and runtime dependencies:

pip install -r requirements.txt
pip install accelerate click pyyaml numpy tqdm psutil requests regex math-verify fire

The released configurations initialize from GSAI-ML/LLaDA-8B-Instruct.

Training

Four task configurations are included:

Task	Configuration	Data source
MATH	configs/math_spg_snis.yaml	qwedsacf/competition_math
GSM8K	configs/gsm8k_math_spg_snis.yaml	openai/gsm8k
Countdown	configs/countdown_spg_snis.yaml	User-provided local JSONL
Sudoku	configs/sudoku_spg_snis.yaml	User-provided local CSV

Before launching a run, update the following fields in the selected configuration:

1. network_kwargs.pretrained_model_name_or_path and tokenizer_kwargs.pretrained_model_name_or_path
2. training_args.run_dir
3. data_loader_kwargs.local_path for Countdown or Sudoku

For example, launch the MATH experiment with 8 GPUs and Fully Sharded Data Parallel (FSDP):

accelerate launch \
  --num_processes 8 \
  --config_file configs/accelerate/fsdp.yaml \
  train.py \
  --config configs/math_spg_snis.yaml

Checkpoints and training logs are written below the configured training_args.run_dir. To resume an interrupted run:

accelerate launch \
  --num_processes 8 \
  --config_file configs/accelerate/fsdp.yaml \
  train.py \
  --config configs/math_spg_snis.yaml \
  --resume-from /path/to/existing/run

Evaluation

Set CKPT_PATH in the corresponding script, then run:

# GSM8K by default
bash scripts/eval_math.sh

# MBPP by default
bash scripts/eval_code.sh

The evaluation entry points expose additional options for dataset, generation length, diffusion steps, and block length:

python math_metrics.py --help
python code_metrics.py --help

Caution

Code benchmark evaluation executes model-generated programs. Run it only inside an appropriately isolated environment.

Inference

import torch
from transformers import AutoTokenizer

from networks.llada_svpo import LLaDASVPO, generate_spg

model_path = "/path/to/checkpoint"  # or GSAI-ML/LLaDA-8B-Instruct

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = LLaDASVPO.from_pretrained(
    model_path,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).cuda().eval()

inputs = tokenizer("Solve: 24 = 6 * 4", return_tensors="pt").to("cuda")

with torch.inference_mode():
    output = generate_spg(
        model,
        inputs.input_ids,
        steps=128,
        gen_length=128,
        block_length=32,
        temperature=0.0,
        prompt_mask=inputs.attention_mask,
    )

generated = output[:, inputs.input_ids.shape[1]:]
print(tokenizer.decode(generated[0], skip_special_tokens=True))

Repository Structure

StableDRL/
├── configs/        # Task and Accelerate/FSDP configurations
├── dataloaders/    # MATH, GSM8K, Countdown, Sudoku, and code datasets
├── networks/       # LLaDA model, generation, and StableDRL objective
├── training/       # Distributed training loop and optimization utilities
├── evaluate/       # Answer parsing and benchmark graders
├── scripts/        # Evaluation launch scripts
└── train.py        # Training entry point

Citation

If you find this work useful, please cite:

@inproceedings{zhong2026stabilizing,
  title     = {Stabilizing Reinforcement Learning for Diffusion Language Models},
  author    = {Jianyuan Zhong and Kaibo Wang and Ding Ding and Zijin Feng and
               Haoli Bai and Yang Xiang and Jiacheng Sun and Qiang Xu},
  booktitle = {Forty-third International Conference on Machine Learning},
  year      = {2026},
  url       = {https://openreview.net/forum?id=YvGtYoFyud}
}

Contact

For questions or problems with the code, please open a GitHub issue.