On-Policy Distillation

Reading depth

Choose A Pass

Every period in machine learning has a few phrases that start appearing everywhere. At first they sound like names for very specific techniques. After a while, they begin to sound like names for a wider anxiety. On-policy distillation is one of those phrases.

The anxiety is easy to understand. We now have models that can do impressive things when they are large, slow, expensive, or allowed to think for longer. We also have products that need models to be cheaper, faster, more stable, and easier to control. Somewhere between these two facts sits the hope behind distillation: maybe the expensive behavior can be taught to the practical model.

But there is a catch. The behavior we want to teach is not always visible in a clean dataset. Sometimes it appears only after the model begins to act, make partial mistakes, call tools, recover, refuse, continue, or drift. In other words, the model creates the very situations we need to train on.

That is where the word on-policy begins to matter.

The name is compact, but it contains several different ideas. To understand what it means, we first need to separate these ideas from one another.

In this post we will use four basic notions:

a policy
on-policy and off-policy data
online reinforcement learning
distillation

After these notions are clear, the phrase “on-policy distillation” becomes much less mysterious.

The important part is where the examples come from.

Policy

In reinforcement learning, a policy is the rule by which an agent chooses an action. In language modeling, we can use the same word in a slightly informal but useful way. A model receives a context and produces a continuation. This continuation might be a final answer, a tool call, a refusal, a plan, or a question to the user.

For the purposes of this post, we will call this behavior the model’s policy.

This means that the policy is not just the model’s weights. It is the behavior produced by those weights under some decoding setup, product context, tools, system prompt, and input distribution.

This distinction matters because on-policy data is defined with respect to the policy that generated it.

Notice the small shift here. We are no longer speaking only about a model as an object stored on disk. We are speaking about the model as something that acts inside a setting. The setting matters. The tools matter. The sampling setup matters. The product wrapper matters. A policy is the model as behavior, not merely the model as a file.

On-Policy And Off-Policy Data

Suppose we have a model we want to improve. We ask it to solve a set of tasks and save the responses. These responses were produced by the current model. In that case, we will call the resulting data on-policy data.

Now suppose instead that our data comes from another source. It may come from an older model, from human-written examples, from a benchmark, or from a stronger model that generated answers without first looking at the current student’s mistakes. In that case, we will call the data off-policy data.

Let us consider a simple example.

We have a coding assistant that often calls the wrong file-search command. If we train it on a clean dataset of expert-written solutions, that dataset may show the correct final answer. But it may not show the assistant’s actual bad intermediate state: it searched the wrong directory, got irrelevant results, and then tried to recover.

If we sample traces from the assistant itself, we will see those states: the bad search, the irrelevant results, the attempted recovery. That is the value of on-policy data. It is not necessarily cleaner than off-policy data, but it is closer to the world the current model actually creates.

Online Reinforcement Learning

Online reinforcement learning is another related idea. In an online RL loop, the policy acts, receives a signal, and is then updated. After the update, the policy acts again. The new policy may now visit different states, which means the next data is also different.

So the training process is not a single pass over a fixed dataset. The model and the data-generating process affect each other.

On-policy distillation borrows this sampling shape, but changes the feedback. Instead of waiting for a single final reward, it uses a teacher distribution to give token-level scores.

Think of a junior developer trying to fix a bug. Plain RL is like running the test suite at the end and only hearing pass or fail. Off-policy distillation is like reading clean pull requests from a senior engineer; the examples are good, but they may not show the strange branch you personally walked into. On-policy distillation is closer to a senior engineer reviewing your own diff line by line: this line is fine, this line is suspicious, this line is where the bug probably entered.

Distillation

Distillation is the transfer of behavior from a teacher to a student. The teacher is usually more expensive, more capable, or more controlled. The student is usually cheaper, smaller, faster, or easier to deploy.

In the simplest case, the teacher writes an answer and the student learns to imitate that answer.

For example, a large reasoning model may solve a math problem with a long internal process. We may then train a smaller model to produce the final answer, or to produce a shorter version of the reasoning trace. The smaller model will not become identical to the larger model, but it may inherit some useful behavior.

There is one point to be careful about. Distillation is not always about copying final answers. In language models, the teacher also defines a distribution over tokens. On-policy distillation uses that distribution directly, on the student’s own sampled text.

Definition

We can now combine the pieces.

On-policy distillation is a training loop in which the current student policy generates trajectories, a teacher model scores the tokens in those trajectories, and the student is updated using those dense scores.

Method	Sampling	Reward signal
Supervised fine-tuning	off-policy	dense
Reinforcement learning	on-policy	usually sparse
On-policy distillation	on-policy	dense

In a minimal form:

Choose a task distribution.
Sample trajectories from the current student policy.
Ask the teacher for log probabilities on the sampled student tokens.
Compare student and teacher probabilities, often with a reverse-KL-style signal.
Use the token-level signal as a dense reward or advantage.
Update the student.
Evaluate the updated student.
Repeat only if the update was useful.

The method is mostly contained in steps 2 and 5: collect from the current student, then score densely instead of only at the end.

Why The Distribution Matters

It is tempting to describe distillation in a very clean way:

teacher produces target, student imitates target, student gets better.

This picture is useful for off-policy distillation, but it hides a problem. The student does not live inside the teacher’s dataset. Once deployed, it produces actions. Those actions create the next state.

If the model makes a bad tool call, it must now deal with the bad tool result. If it gives an uncertain answer confidently, the user may ask a different follow-up. If it refuses too early, the conversation enters a different path. The model is not only answering prompts. It is shaping the future prompts and states it will see.

This is why on-policy data matters. It is data from the world induced by the current policy.

This is also why the method is different from ordinary dataset curation. Instead of starting with a fixed set of clean examples, the loop first asks what the model actually did and then decides which parts of that behavior should become training signal.

Target Distribution

The first design choice is the task distribution. This answers the question: where do we want the student to improve?

For a coding assistant, the distribution might include bug fixes, refactors, failing tests, and code review requests. For an agent, it might include multi-step tasks with tool calls. For a support assistant, it might include real support conversations and policy-sensitive questions.

If the distribution is too easy, the loop will mostly produce examples the student already handles. If the distribution is too strange, the student may become better at edge cases while ordinary behavior gets worse.

So the task distribution should be chosen deliberately. It should represent the behavior we actually want to move.

This point sounds modest, but it is often where the whole method succeeds or fails. If we choose the wrong place to look, even a very good teacher will give us beautiful answers to the wrong questions.

Current Policy

The current policy is the student as it exists now. This sounds obvious, but it is easy to lose track of.

If we update the student, then the next round of samples comes from a different policy. If the model becomes more verbose, future traces become more verbose. If it starts calling tools more often, future traces include more tool calls. If it learns a bad shortcut, that shortcut may appear repeatedly in the next training data.

This currentness is both the strength and the danger of the method.

There is a kind of honesty in this. The current policy may not show us the model we wish we had. It shows us the model we actually have. That can be uncomfortable, but it is useful.

Teacher Logprobs

The teacher in this setup is not mainly writing a better answer. For each position in the student’s trajectory, we ask how likely the teacher would be to produce the same next token given the same prefix.

Suppose the student is solving a math problem. The first few steps may be reasonable. Then one token begins a wrong operation. A final-answer reward may only say “wrong” after the whole solution is finished. Teacher logprobs can point closer to the first bad step.

In practice, the teacher is often treated with too much mystery. We say “teacher model” and move on. But the teacher is also a system with prompts, tools, context, sampling choices, and blind spots. If the student is not magic, neither is the teacher.

Research Note

The language model as reward model

The elegant trick is that we do not need a separate learned reward model for this part of the loop. The teacher language model itself gives a reward-like signal: tokens that are probable under the teacher are treated as better, and tokens that are improbable under the teacher are treated as worse.

This does not make the method magically correct. The teacher can still be biased, weak, or poorly matched to the student’s states. But it removes one familiar failure point: there is no separate proxy reward model for the student to exploit. In that limited sense, reverse KL is harder to hack, because low penalty directly means high probability under the teacher’s own distribution. This is also why the method sits near ideas like inverse RL and DPO: a model distribution is being turned into a reward-like training signal.

Dense Token Rewards

The main advantage over ordinary RL is density.

In a simple RL setup, the student may generate a long solution and receive one reward at the end. The reward tells the model whether the rollout worked, but not where the reasoning went wrong.

On-policy distillation can provide a signal at every token. If the student token has high probability under the teacher, the penalty is small. If the student token has low probability under the teacher, the penalty is larger.

One way to write this is reverse KL between the student and teacher distributions on tokens sampled from the student. To make the notation less mysterious, let $\pi_\theta$ be the student policy and let $\pi_T$ be the teacher policy. Let $x$ be the prompt, and let the student produce a rollout:

y_{1:T} \sim \pi_\theta(\cdot \mid x)

At time t, the state is not only the original prompt. It is the prompt together with everything the student has already written:

s_t = (x, y_{<t})

Now we look at the token the student actually picked, y_t. The teacher is asked a very specific question: in this same state s_t, how likely was this token?

\begin{aligned} \text{student score} &= \log \pi_\theta(y_t \mid s_t) \\ \text{teacher score} &= \log \pi_T(y_t \mid s_t) \end{aligned}

The per-token reverse KL penalty has the following shape:

\mathrm{KL}_t = \log \pi_\theta(y_t \mid s_t) - \log \pi_T(y_t \mid s_t)

If the teacher also finds the token natural, the second term is not too small and the penalty stays mild. If the teacher thinks the token is strange in that prefix, $\log \pi_T(y_t \mid s_t)$ becomes very low and the penalty grows.

This is the small but important point. The teacher is not judging a sentence from outside. It is judging the student’s own next step inside the student’s own path.

In training code, the sampled trajectory usually comes from the old copy of the student, the one that produced the rollout. So the practical signal is often written with $\pi_{\text{old}}$ :

\begin{aligned} \mathrm{rkl}_t &= \log \pi_{\text{old}}(y_t \mid s_t) - \log \pi_T(y_t \mid s_t) \\ A_t &= -\mathrm{rkl}_t \\ &= \log \pi_T(y_t \mid s_t) - \log \pi_{\text{old}}(y_t \mid s_t) \end{aligned}

Then the update uses this $A_t$ in a policy-gradient-style loss. A positive advantage says: make this sampled token more likely. A negative advantage says: move away from this token in this state.

This is also why $\pi_{\text{old}}$ appears in the notation. The rollout was sampled before the current gradient update, so the data came from a behavior policy. When the current parameters $\pi_\theta$ are updated using that same sampled trajectory, the small policy mismatch is handled with an importance-sampling-style loss. The advantage is treated as a fixed coefficient; the gradient flows through the current policy log probability, not through the teacher score.

There is a detail worth not skipping. Reverse KL is not symmetric:

D_{\mathrm{KL}}(\pi_\theta \parallel \pi_T) \ne D_{\mathrm{KL}}(\pi_T \parallel \pi_\theta)

This direction matters because the samples come from the student. The teacher is being queried on the states the student actually visits, not only on clean teacher-written states.

Distillation Update

After computing the dense signal, we update the student.

At this point the method starts to look less like ordinary supervised learning and more like a small policy-improvement loop. The student has produced behavior. The teacher has turned that behavior into token-level preferences. The update asks the current student to keep the choices that looked teacher-like and reduce the choices that did not.

Then set the token-level advantage from the teacher signal. In plain language, tokens that look unlikely under the teacher receive worse advantages. Tokens that look teacher-like receive less penalty.

This matters for compute too. The teacher does not have to sample a full new answer for every prompt; it only scores tokens that already exist. That is a small change in wording, but a large change in engineering shape. Scoring is usually easier to batch, easier to cache, and easier to attach to partial rollouts than asking the teacher to generate a new completion every time.

There are still choices. We can train on complete rollouts or partial rollouts. We can combine the dense teacher signal with a final environment reward. We can mix on-policy distillation with older off-policy data to preserve broad behavior. These choices do not change the central shape of the update: generate from the student, score under the teacher, and move the student with a loss that respects where the sample came from.

A minimal Python version of the token signal looks like this:

from dataclasses import dataclass


@dataclass(frozen=True)
class TokenScore:
    token: str
    student_logprob: float
    teacher_logprob: float


def reverse_kl_penalty(score: TokenScore) -> float:
    return score.student_logprob - score.teacher_logprob


def token_advantage(score: TokenScore) -> float:
    return -reverse_kl_penalty(score)


def token_advantages(scores: list[TokenScore]) -> list[float]:
    return [token_advantage(score) for score in scores]

This is not the full training loop. It only shows the shape of the token-level signal. The real update still needs an RL-style loss around it, usually with an importance-sampling correction when the rollout was collected from a previous copy of the student.

Evaluation Gate

The evaluation gate decides whether the new student is actually better.

This cannot be checked only on the sampled failures. The sampled failures should improve; otherwise the run did not do its job. But the model may also change around those failures.

It may become more verbose. It may become more confident. It may call tools more often. It may become better on hard prompts and worse on simple prompts. It may copy the teacher’s style in places where that style is not wanted.

So we need both local and global checks.

Without this gate, the loop can produce new checkpoints without producing a better model. Motion is not the same thing as progress.

Deployment Feedback

Deployment feedback is not only explicit user feedback. It includes the trace of the model in the product.

Which prompts cause retries? Which answers get edited? Which tool calls fail? Where does the user abandon the flow? Which evals pass alone but fail inside the workflow?

These are signals about the states created by the current policy. But logs by themselves are not yet a dataset; each example needs enough metadata to explain where it came from and what system produced it.

Failure Modes

On-policy distillation is not automatically better than offline distillation. It is more targeted, but that also means it can be more sensitive to mistakes in the loop.

Many of the failures below come from the same root. Reverse KL is mode-seeking: it can help the student commit to teacher-like behavior instead of spreading probability across weak alternatives. In the wrong setting, that same property can narrow the model too much or make teacher preferences look like truth.

When It Is Useful

On-policy distillation is useful when the model’s failures are distributional. That means the failures appear because of the model’s own behavior, not merely because the original dataset lacked a clean example.

It is a good fit for:

agent tool-use failures
reasoning models that need cheaper deployment
refusal or policy-boundary calibration
recovery after partial mistakes
product workflows where the model creates the next state
domains where static datasets become stale quickly

It is a weaker fit for:

simple formatting fixes
tasks where high-quality offline data already matches deployment
cases where the teacher uses hidden capabilities the student cannot access
situations where there is no reliable evaluation gate

Conclusion

On-policy distillation is not simply teacher teaches student. It is useful because it trains on the states the model actually enters.

It is also why the method can go wrong. If the loop samples the wrong states, uses a poorly matched teacher, or ignores product constraints, it can turn small habits into larger ones.

The practical takeaway is simple. Before asking how to distill, ask where the examples come from. In on-policy distillation, that question is not a detail. It is the method.

Perhaps this is why the phrase has become interesting. It names a technical loop, but also a change in attitude: let the model reveal its own mistakes, then decide carefully which of those mistakes are worth turning into training data.

AI audio summary