Diffusion On-Policy Distillation Trainer

Last updated: 08/10/2026

This example distills an OCR-tuned SD3.5-Medium teacher into a fresh SD3.5-Medium student. The student generates images with its own policy, a frozen teacher scores every denoising step of those trajectories, and the student minimizes the KL between its transition and the teacher’s (distill_kl). The OCR reward is only monitored, never optimized — the reward curve shows the student reaching the teacher’s reward level through distillation alone.

Compared with running FlowGRPO again, distillation needs no reward in the loss and converges in far fewer steps. See the algorithm doc for the loss and all config knobs.

Installation

Follow the installation guide to set up the base environment, then install the OCR reward dependency:

pip install Levenshtein

The provided script uses a single node with 3 GPUs: 2 for actor + rollout (the frozen teacher shares them), 1 for the reward model server.

Prepare the dataset

Obtain the raw OCR dataset from the original Flow-GRPO repository:

  • https://github.com/yifan123/flow_grpo/tree/main/dataset/ocr

Place it under $WORKSPACE/data/ocr (where WORKSPACE defaults to $HOME), then preprocess it into parquet files:

python3 examples/flowgrpo_trainer/data_process/sd3_ocr.py \
  --input_dir $WORKSPACE/data/ocr \
  --output_dir $WORKSPACE/data/ocr/sd3

The script reads:

ocr_train_path=$WORKSPACE/data/ocr/sd3/train.parquet
ocr_test_path=$WORKSPACE/data/ocr/sd3/test.parquet

Prepare the teacher

The teacher is a full diffusers checkpoint from the same pipeline family as the student, with the same scheduler. The natural way to get one is the SD3.5 FlowGRPO OCR example: train the LoRA, merge it into the base transformer (peft merge_and_unload), and save the merged pipeline. Any stronger same-family checkpoint works the same way.

Run

TEACHER_PATH=/path/to/merged-teacher \
  bash examples/diffusionopd_trainer/sd35/run_sd35_medium_ocr_distill.sh

What to expect

  • actor/distill_kl_loss starts clearly positive (the teacher’s weights differ from the student’s) and falls by more than an order of magnitude over the first ~40 steps.

  • The validation OCR reward climbs from the base-model level to the teacher’s level within a few tens of steps, even though the loss never sees it.

  • timing_s/teacher reports the once-per-step teacher scoring stage.

To combine distillation with a task reward instead of replacing it, keep diffusion_loss.loss_mode=flow_grpo and set actor.use_distill_loss=True — see the algorithm doc for details.