Code and Resources

• Open-source repository: KMnO4-zx/agentic-rl-lab
• Synchronous code: 01-demo-sync.py
• Asynchronous code: 02-demo-async.py
• PyTRIO website and registration: https://pytrio.cn/ (for remote training, sampling, and weight saving)
• SwanLab registration: https://swanlab.cn/login (for tracking training progress and experimental metrics)
• DeepMath-103K dataset: AI-ModelScope/DeepMath-103K
• PyTRIO Skill: SwanHubX/pytrio-skill
• SwanLab Skill: SwanHubX/swanlab-skill

On-policy distillation (OPD) applies knowledge distillation to states visited by an autoregressive student model. Traditional knowledge distillation often trains student models to learn from a fixed dataset or answers pre-generated by a teacher, but during actual inference, the student enters its own generation distribution, which can easily lead to discrepancies between training and inference. The Generalized Knowledge Distillation (GKD) introduced in 2023 explicitly proposed having the student model receive feedback from the teacher based on its own generated sequences. Around the same time, MiniLLM systematically investigated the use of reverse KL divergence, which is better suited for generative models. Today, OPD typically combines these two key ideas: the student generates the trajectory, the teacher only evaluates these trajectories, and reverse KL divergence is then used to guide the student towards the teacher's output.

I used PyTRIO because OPD needs the teacher to rescore exactly the tokens sampled by the student. PyTRIO provides sample() and compute_logprobs(). The latter takes prompt + student completion and returns aligned log probabilities without generating a replacement answer. Local code computes reverse KL and advantages, then calls forward_backward() and optim_step(), while the service handles training and inference. The PyTrio Skill can also help Codex or Claude Code draft a demo from my training logic, which I then review and adapt.

Each OPD step reads prompts, samples completions from the current student and retains each completion token's old-policy log probability. The teacher scores the identical prompt + student completion sequence. Their difference is reverse_kl = student_logprob - teacher_logprob; -reverse_kl supplies the advantage for an importance-sampling update. Exporting a fresh student sampler after the update closes the on-policy loop: student generation → teacher scoring → student update.

OPD in plain language

OPD is not about having students memorize answers prepared in advance by the teacher (which is the logic behind SFT). Instead, it aims to first identify the errors students are most likely to make, and then allows the teacher to correct those errors step-by-step, following the actual path the student takes (the training signal for OPD is more "gentle" compared to SFT).

Comparison Item	Traditional Offline Distillation	On-Policy Distillation
Training Data Source	Fixed dataset or teacher-generated data	Real-time sampling from the current student
Teacher's Role	Generates target answers or provides fixed soft targets	Calculates the log probability of the student's actual trajectories
What the Student Learns	The teacher's behavior within the data distribution	The teacher's preferences for the states the student is currently accessing
Main Advantages	Simple, stable, and data can be prepared in advance	Mitigates distribution shift between training and inference
Main Costs	Requires preparing distillation data	Requires student sampling and teacher scoring in each iteration

Here, "on-policy" refers to trajectories generated by the current student model. The teacher model can be larger than the student, or it can be a similarly sized model that has been specifically trained. The key is not the size of the models, but rather whether the teacher can provide reliable, aligned log probabilities for the token sequences used by the student.

A crucial point is that the student model and the teacher model must share a compatible tokenizer. Since the training signal for OPD is based on the difference in log probabilities at the token level, if the two models use incompatible tokenizers, it will not be possible to directly compare the log probabilities for the same token ID.

OPD's Core Objectives

Let:

• $x$ is the prompt.
• $y$ is the completion generated by student sampling.
• ${\pi }_{\theta }$ represents the student currently being trained.
• ${\pi }_T$ is the frozen teacher.

A common OPD objective minimizes reverse KL from the student distribution to the teacher:

$$D_{\mathrm{KL}}({\pi }_{\theta }\parallel {\pi }_T)={\mathbb{E}}_{y\sim {\pi }_{\theta }(\cdot \mid x)}[\log {\pi }_{\theta }(y\mid x)-\log {\pi }_T(y\mid x)].$$

After breaking down a generated sequence into tokens, the sampling estimate for the token at position $t$ is:

$$r_t^{\mathrm{KL}}=\log {\pi }_{\theta }(y_t\mid x,y_{<t})-\log {\pi }_T(y_t\mid x,y_{<t}).$$

The synchronous demo converts this into an advantage:

$$A_t=-\beta r_t^{\mathrm{KL}},$$

The corresponding code consists of only two lines.

reverse_kl = np.asarray(student_lps) - np.asarray(teacher_lps)
advantages = -args.kl_penalty_coef * reverse_kl

Intuitively:

• When student_logprob > teacher_logprob, the student favors the token more; advantage < 0 discourages it.
• When student_logprob < teacher_logprob, the teacher favors the token more; advantage > 0 encourages it.
• kl_penalty_coef, or $\beta$ in the formula, controls the strength of alignment with the teacher.

Individual tokens can have negative reverse_kl values. KL only guarantees non-negativity after taking the expected value of the complete distribution for the student. Therefore, occasionally seeing negative token-level averages in the logs does not necessarily indicate a formula error, but should be evaluated in conjunction with additional factors such as the number of steps, sequence length, and variance.

What does this OPD training code do?

The current directory provides two implementations.

01-demo-sync.py   同步版，适合阅读和理解完整数据流
02-demo-async.py  异步版，并发执行 batch 内的 student rollout 和 teacher 打分

This article primarily focuses on 01-demo-sync.py. The default configuration is:

• Data: AI-ModelScope/DeepMath-103K from ModelScope;
• Student: Qwen/Qwen3.5-4B, trains a rank 32 LoRA;
• Teacher: Qwen/Qwen3.6-27B, only calculates log probabilities and does not participate in optimization;
• Loss: The loss for importance_sampling, calculated using PyTRIO;
• Logs: SwanLab;
• Training Service: PyTRIO remotely executes forward pass, backward pass, optimizer step, sampling, and weight saving.

DeepMath in this experiment is a prompt-only dataset. The code only reads the question field, and does not use standard answers or calculate mathematical accuracy. The supervisory signal for OPD is entirely derived from the teacher's token-level preferences regarding the student's completion.

Step-by-Step Explanation of the OPD Training Code

Step 1: Prepare the prompt data.

The script will download DeepMath's Parquet shards from ModelScope and then read local files using Hugging Face datasets. By default, it only processes the first shard, shuffling and truncating the data with a fixed seed.

dataset = load_dataset(
    "parquet",
    data_files=shard_paths,
    split="train",
    cache_dir=str(args.dataset_dir / ".datasets_cache"),
)

dataset = dataset.shuffle(seed=args.seed)
if args.sample_size > 0:
    dataset = dataset.select(range(min(args.sample_size, len(dataset))))

build_prompt renders a question with the chat template and asks for the final answer inside \boxed{}:

messages = [{"role": "user", "content": content}]
prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=enable_thinking,
)
prompt_ids = tokenizer.encode(prompt, add_special_tokens=False)

When switching to other tasks, typically only the data reading and build_prompt components need to be modified. OPD does not require the dataset to have an answer field, as the primary supervision comes from the teacher log probabilities.

Step 2: Create 'student' and 'teacher'.

The unified entry point for PyTRIO is ServiceClient.

service_client = trio.ServiceClient()

training_client = service_client.create_lora_training_client(
    base_model=args.base_model,
    rank=args.lora_rank,
    seed=args.seed,
)

teacher_client = service_client.create_sampling_client(
    base_model=args.teacher_base_model or args.base_model,
    model_path=args.teacher_model_path,
)

The responsibilities of the two clients differ.

• training_client holds a trainable LoRA for students, performing forward propagation, backpropagation, optimization, and weight saving.
• teacher_client remains frozen and only calls compute_logprobs for specified token sequences.

If teacher_model_path points to a pre-trained LoRA sampler weight, it can be used to distill a specialized domain teacher; otherwise, it will directly use teacher_base_model.

Step 3: Export the sampler from the current student.

The training client does not directly replace the rollout sampler. At the beginning of each step, the code creates a sampling client based on the current student weights.

if student_sampler is None or step % args.sampler_refresh_steps == 0:
    student_sampler = training_client.save_weights_and_get_sampling_client()

By default, sampler_refresh_steps=1, which means the sampler is refreshed after each parameter update. This aligns most closely with the strict definition of on-policy learning. Increasing the refresh interval can reduce the overhead of weight synchronization, but the resulting trajectories will gradually reflect data generated by an outdated policy.

Step 4: sample student trajectories

For each prompt, the student generates group_size completions.

result = student_sampler.sample(
    prompt=trio.ModelInput.from_ints(prompt_ids),
    num_samples=args.group_size,
    sampling_params=sampling_params,
    return_text=False,
).result()

The most important thing in every sequence is:

• seq.tokens: completion tokens generated by the student.
• seq.logprobs: The logprob value provided by the old student policy when generating these tokens.

The OPD system utilizes group_size primarily to obtain more on-policy trajectories for the same prompt, but it does not calculate relative rewards within groups as GRPO does. Each completion is independently supervised by the teacher on a token-by-token basis.

Step 5: The teacher assigns scores to the same trajectory.

This is the most critical step for OPD. Instead of generating new answers, the teacher reads the tokens that the student has already produced.

all_ids = prompt_ids + completion_ids
all_logprobs = teacher_client.compute_logprobs(
    trio.ModelInput.from_ints(all_ids)
).result()

teacher_lps = all_logprobs[len(prompt_ids):]

The slice begins from len(prompt_ids) because it only requires the teacher logprob for the completion segment. The code also checks:

if len(teacher_lps) != len(completion_ids) or any(v is None for v in teacher_lps):
    raise ValueError("Invalid teacher logprobs for completion tokens")

If the teacher generates a new completion and then compares it to the student's completion, it will be impossible to align the tokens and states one-to-one. Consequently, it will no longer function as the OPD currently implemented.

Step 6: align inputs with next-token targets

Autoregressive training aligns each input position with the next token as its target:

prompt_loss_len = len(prompt_ids) - 1

input_ids = prompt_ids + completion_ids[:-1]
target_ids = [0] * prompt_loss_len + completion_ids
padded_logprobs = [0.0] * prompt_loss_len + student_lps
padded_advantages = [0.0] * prompt_loss_len + list(advantages)

Then, construct the Datum component of PyTRIO.

datum = trio.Datum(
    model_input=trio.ModelInput.from_ints(input_ids),
    loss_fn_inputs={
        "target_tokens": np.asarray(target_ids, dtype=np.int64),
        "logprobs": np.asarray(padded_logprobs, dtype=np.float32),
        "advantages": np.asarray(padded_advantages, dtype=np.float32),
    },
)

At prompt positions, target_tokens, logprobs, and advantages are filled with zeros. This section does not participate in the update for OPD, but the placeholder values ensure that all four sequences have a completely consistent length.

len(model_input)
= len(target_tokens)
= len(logprobs)
= len(advantages)

The completion interval is what truly generates the training signal.

Step 7: Update student.

Once a step has collected all the data, submit the forward/backward pass and optimizer updates.

fwd_bwd = training_client.forward_backward(
    datums,
    loss_fn="importance_sampling",
)
optim = training_client.optim_step(adam)

fwd_bwd_result = fwd_bwd.result()
optim.result()

logprobs stores old-policy log probabilities from rollout. The importance-sampling objective uses the current-to-old probability ratio when weighting advantages; it does not replace those saved old log probabilities. Therefore, the complete logic for each step is as follows:

刷新当前 student sampler
        ↓
student 采样并保存 old logprobs
        ↓
teacher 对同一 completion 计算 logprobs
        ↓
计算 reverse KL advantage
        ↓
importance_sampling + optimizer step

Step 8: Record the metrics and save the weights.

The script will log reverse KL values, the number of completion tokens, step durations, and metrics returned by the PyTRIO trainer. Upon completion of training, it will save the sampler weights, which can be used for inference.

save_result = training_client.save_weights_for_sampler(
    args.save_weights_name
).result()
print(f"Saved weights: {save_result.path}")

The returned path can be further passed to the model_path sampling client for inference or subsequent evaluation.

Run the synchronous demo

Project dependencies are listed in the pyproject.toml file located in the root directory.

uv sync
trio login

First, run a very small smoke test to confirm that authentication, models, and the data link are all functioning correctly.

uv run python 02-opd/general-opd/01-demo-sync.py \
  --steps 1 \
  --batch-size 1 \
  --group-size 1 \
  --max-tokens 128 \
  --sample-size 20 \
  --swanlab-mode disabled

Once you're certain that the model weights are being saved completely, you can then increase the training volume.

uv run python 02-opd/general-opd/01-demo-sync.py \
  --steps 10 \
  --batch-size 4 \
  --group-size 4 \
  --max-tokens 512 \
  --sample-size 100 \
  --swanlab-mode online

The first time you run the program, it will download the necessary DeepMath Parquet shards to the datasets/DeepMath-103K/ directory. The complete dataset consists of 10 shards. You can control the number of shards downloaded and read using --num-shards. We recommend starting with the default setting of 1 shard to verify the training process.

How to Interpret Training Metrics.

Indicator	Meaning
data/datums	The number of completions actually participating in training for this step.
data/completion_tokens_mean	The average number of tokens per completion.
data/completion_tokens_total	The total number of tokens for all completions in this step.
data/completion_tokens_per_second	The total number of tokens divided by the total time spent in the entire OPD step.
opd/reverse_kl_mean	The average difference in log probability between the student and teacher models when sampling tokens.
opd/reverse_kl_std	The degree of dispersion of the reverse KL divergence at the token level.
trainer/*	Training metrics, such as loss, returned by the PyTRIO step.

Here, completion_tokens_per_second refers to full-step training throughput, encompassing sampler refreshes, student sampling, teacher log probabilities, forward/backward passes, and optimizer steps. It is not equivalent to the tokens/second rate of a pure generation service.

When observing experiment OPD, avoid focusing solely on a single step (reverse_kl_mean). It is more valuable to check multiple aspects simultaneously.

1. Is the reverse KL divergence stabilizing overall across multiple steps?
2. Is the completion length experiencing any sudden collapses?
3. Is the trainer loss exhibiting any NaN values or unusual spikes?
4. Are the independent task evaluations, both before and after training, genuinely improving?

KL divergence only indicates that the student is closer to the teacher; it does not automatically equate to an increase in task accuracy.

How to Customize Your Universal OPD Task

1. Replace dataset.

As long as it can return a prompt, it can be used with OPD. For example:

• Math problems: question;
• Coding tasks: problem;
• Dialogue distillation: messages;
• Agent tasks: observations or historical trajectories.

If the data consists of multiple message turns, instruct build_prompt to receive the complete messages directly, without first compressing the role information into a single, generic text string.

2. Replace teacher.

The teacher need not be larger, but should provide stronger guidance on the target skill. Common options include:

• Models with a larger number of parameters.
• Models that have undergone Supervised Fine-Tuning (SFT) or Reinforcement Learning (RL) in a specific domain.
• Specialized LoRA modules trained on the same base model.
• Models incorporating enhanced system prompts or contextual knowledge.

3. Ensure tokenizer compatibility.

The current code uses the student tokenizer to generate token IDs, and then passes these IDs to the teacher. Consequently, the student and teacher must share a compatible vocabulary and a consistent set of special token conventions. If the two tokenizers are incompatible, it will not be possible to directly compare log probabilities for the same IDs; instead, a text-level alignment strategy needs to be redesigned, or a compatible model should be selected.

4. Adjust the sampling distribution.

temperature, top_p and top_k determine which student trajectories are sampled.

• The sampling strategy is overly conservative, resulting in insufficient diversity in the generated sequences. The teacher can only correct a limited range of errors.
• The sampling approach is too random, which leads to a large number of low-quality tokens being generated. This will increase both the training variance and the cost associated with the teacher.
• max_tokens directly impacts the sampling process, teacher scoring, and training costs for each step.

5. Adjust KL intensity.

If kl_penalty_coef is too small, teacher alignment may be weak. If it is too large, individual token signals can dominate and destabilize updates. It's recommended to initially maintain the value of 1.0, observe the reverse KL divergence, loss, generated length, and task evaluation, and then gradually adjust the value.

6. Transitioning from synchronous to asynchronous.

The synchronous version is best suited for understanding and debugging. Once the data flow is confirmed to be correct, you can proceed to run it.

uv run python 02-opd/general-opd/02-demo-async.py \
  --steps 10 \
  --batch-size 4 \
  --group-size 4 \
  --max-tokens 512 \
  --sample-size 100 \
  --swanlab-mode online

The asynchronous version processes prompts within a batch concurrently and submits multiple teacher logprob requests for different completions under the same prompt concurrently. It does not alter the OPD algorithm, but instead improves the concurrent utilization of resources during the remote sampling and scoring phases.

Summary

The minimum closed-loop system for model OPD consists of just five elements.

1. Sample completions from the current student.
2. Save the token log probabilities from student sampling.
3. The teacher calculates the token log probabilities for the same completion.
4. Calculate the reverse-KL advantage using the difference between the two.
5. Update the student using importance sampling and refresh the sampler.

As long as the data can be structured into prompts, the student can return sampling log probabilities, and the teacher can calculate the log probabilities for the same token sequence, this process can be adapted from mathematical problems to code, dialogue, tool calls, and agent trajectory distillation.

References

• On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes / GKD
• MiniLLM: Knowledge Distillation of Large Language Models
• PyTRIO
• SwanLab