Skill-Guided On-Policy Distillation

Why People Wanted It

A tool agent can finish a task successfully or fail, but that final outcome does not say which intermediate action deserved credit. When an interaction contains a bad tool result, the useful lesson may be a reflective choice—checking, changing course, or not repeating an unhelpful action—buried among many tokens. The real question is how to give the current policy denser guidance on those choices without replacing the outcome signal or requiring extra environment rollouts.

The Core Move

Roll out the ordinary policy on a mixture of clean and perturbed interactions, then let that same policy re-score its sampled action tokens with an environment-matched skill prompt; use the skill-induced score shift as detached, confidence-weighted guidance while retaining the ordinary reinforcement-learning objective.

For a sampled action token $a$ in ordinary history $h$, let $\ell_s=\log\pi_\theta(a\mid h)$ and $\ell_t=\log\pi_\theta(a\mid H(h,\kappa))$, where $\kappa$ is the environment’s skill card. The teacher shift and soft confidence are

$$\Delta=\operatorname{sg}(\ell_t-\ell_s),\qquad \gamma=\sigma(\beta_{\mathrm{opd}}\Delta).$$

The shift is stop-gradient: it weights token guidance from the skill-conditioned view, while policy gradients update the ordinary-context branch.

Quiz

In reflection on-policy distillation, what does the skill-conditioned branch do?

It generates a separate environment trajectory with the skill prompt.

It replaces verifier returns with token labels.

It supplies the skill prompt at deployment.

It re-scores already sampled action tokens under a skill-augmented history.

$ check

Why Noisy Agentic Reinforcement Learning Keeps Failures in the Mix

Here is the part that gets technical: the RL signal is computed over a group that deliberately mixes clean and perturbed environment rollouts, rather than treating noisy failures as a separate contrast set. For task $q$, trajectories receive verifier returns $R_{q,i}$ and group-relative advantages

$$A_{q,i}=\frac{R_{q,i}-\mu_q}{\sigma_q}.$$

Here $\mu_q$ and $\sigma_q$ are the return mean and standard deviation over the mixed group. The choice matters: a failed perturbed trajectory remains part of the on-policy distribution and contributes through its outcome, but it is not normalized against a noisy-only subgroup. The RL loss applies a clipped policy-ratio objective to valid agent-generated tokens; tool returns and simulated user messages are masked out. Its job is still trajectory-level outcome learning, not to label a particular reflection token as good.

Why the Skill Branch Scores the Same Tokens

Now the version for readers who want more of the machinery: the skill-conditioned branch does not generate a second trajectory. The policy first samples a trajectory without a skill card. Given the known environment, the corresponding skill is loaded directly, and a deterministic history construction $H(h,\kappa)$ adds it only for rescoring those already sampled action tokens. This is on-policy distillation because the student tokens came from the current behavior policy, not from a separate expert demonstration.

For each valid action token, the ordinary branch supplies $\ell_s$ and the skill branch supplies $\ell_t$. The OPD loss uses the detached difference $\Delta=\operatorname{sg}(\ell_t-\ell_s)$ and soft weight $\gamma=\sigma(\beta_{\mathrm{opd}}\Delta)$. Minimizing the token-level term $-\gamma\ell_s$ raises the ordinary-context score of sampled tokens that the skill-conditioned view favors. Detaching the teacher score prevents the policy from changing the target by manipulating the skill branch; only the ordinary branch receives this guidance gradient. A weak or negative shift is down-weighted, not categorically discarded.

Why ORPO Combines Two Different Scales of Feedback

Here is the deeper objective-level point: outcome learning and reflective distillation answer different questions. The group-relative RL term asks whether a whole interaction did well compared with its peers; OPD asks which sampled action tokens receive a higher score when the environment skill is present. The joint objective is

$$\mathcal{L}_{\mathrm{ORPO}}=\mathcal{L}_{\mathrm{RL}}+\mathcal{L}_{\mathrm{OPD}}.$$

The sum keeps every trajectory in the outcome term, while the confidence-weighted token term supplies denser, skill-guided pressure on the sampled actions. In the described setup the skill branch is a forward-pass-only teacher: it does not query the environment, and its gradients are detached. What remains open is whether the balance between these two signals stays useful when the skill-conditioned score shift is weak or misleading.

Minimal Causal Code

The loop below isolates the causal idea: sampled tokens come from ordinary rollouts; the skill context only re-scores them; the detached shift weights an update to ordinary-context log-probabilities. It omits the clipped RL term and batching details.

import torch
import torch.nn.functional as F

# Each batch item contains an already sampled action-token sequence.
# ordinary_logp and skilled_logp are aligned per-token scores from the
# current policy, evaluated on ordinary and skill-augmented histories.
ordinary_logp = policy.log_prob(sampled_tokens, ordinary_histories)
skilled_logp = policy.log_prob(sampled_tokens, skilled_histories)

shift = (skilled_logp - ordinary_logp).detach()
confidence = torch.sigmoid(beta_opd * shift)

# valid_agent_tokens excludes tool returns and simulated user messages.
weights = confidence * valid_agent_tokens
opd_loss = -(weights * ordinary_logp).sum() / valid_agent_tokens.sum().clamp_min(1)

# RL loss is computed separately from mixed-group verifier returns.
loss = rl_loss + opd_loss
loss.backward()
optimizer.step()

What You Gain, and What You Lose

You gain token-level guidance tied to the current policy’s own actions, while preserving verifier-based outcome learning and avoiding a second environment rollout for the skill branch. The confidence gate makes the guidance soft: tokens with a larger positive skill-induced shift receive more weight, while weak or negative shifts receive less.

You also accept limits. The teacher is the same policy under a privileged training-time prompt, not an independent expert; its score shift is a guidance signal, not proof that a token is correct. Skill coverage and the usefulness of the matched skill therefore matter. The method adds a skill-conditioned forward pass and requires the skill library during training, even though the skill prompt is removed at deployment. Finally, group-normalized outcomes remain coarse: OPD helps assign token-level pressure, but does not turn every reflection choice into a verified label.