PR-OPD: Privileged Representation On-policy Self-Distillation for Agentic Reinforcement Learning

AI summary

Overview

Research area: Machine learning — reinforcement learning for language-model agents, specifically privileged on-policy self-distillation and representation-level distillation.

Technical level: Advanced. The paper assumes familiarity with GRPO, on-policy distillation, hidden-state alignment, and linear probing.

Scope: The paper diagnoses why token-probability distillation underuses privileged hindsight skills in agentic RL, and proposes aligning hidden states across all layers instead, evaluated on ALFWorld and WebShop with two Qwen backbones.

What This Paper Is About

Agents trained with reinforcement learning typically receive one reward per episode, so recent work lets the same model act as a "privileged teacher" — re-reading its own responses with a hindsight skill in context — and distills that teacher's token probabilities back into a skill-free student. The authors show this channel is weak: a skill lifts WebShop success from 42.2% to 56.2% yet changes the probabilities of fewer than a quarter of sampled tokens, while changing the hidden states of over 80% of response tokens. Their goal is to transfer the skill's advantage through the channel that actually carries it — the teacher's internal representations — without adding a separate teacher, a skill library, or inference cost.

Key Contributions

1. Two diagnostic phenomena. The paper identifies Invisible Advantage (a skill makes the teacher act better but barely changes its token distribution, especially at action tokens) and Much to Align (the same skill broadly reshapes hidden states, and linear probes on those states recover which skill was used while probes on output features stay near chance).

2. PR-OPD. A method in which, after a GRPO warm start, the policy writes one hindsight skill per completed trajectory, re-reads its own sampled responses with that skill as a stop-gradient teacher, and aligns its projected hidden states to the teacher's at every layer alongside the reward objective. Student, hindsight analyzer, and teacher are three roles of one weight set; the only added module is a per-layer MLP projector, discarded at inference.

3. Best aggregate results in every reported setting. On ALFWorld and WebShop with Qwen2.5-3B-Instruct and Qwen3-1.7B-Instruct, PR-OPD attains the highest overall ALFWorld success and the highest WebShop accuracy and score, improving over GRPO by up to 4.7 points on ALFWorld and 14.0 points in WebShop accuracy (the 3B backbone alone improves by 11.7 points).

4. Layer-wise and structural analysis of where the skill acts. Probing and CKA analyses show skill information peaks in middle layers and persists to the output, motivating alignment at every layer rather than only the last.

Main Findings

• The output channel is nearly silent. Adding a self-generated skill to the policy's context changes the next-token probabilities of fewer than a quarter of response tokens, and at action tokens almost none change — while the skill still raises WebShop success from 42.2% to 56.2%.

• Hidden states change broadly. The same two passes differ in the hidden states of more than 80% of response tokens, including action tokens where the output channel shows nothing. Linear probes distinguish a trajectory's own skill from another trajectory's, while probes on output features stay at chance.

• Main results (step 150). With Qwen2.5-3B-Instruct: ALFWorld overall success 71.1% versus GRPO 66.4%, SDAR 68.8%, and Skill-SD 69.5%; WebShop accuracy 67.2% and score 83.7 versus GRPO 55.5/70.4 and Skill-SD 59.4/75.6. With Qwen3-1.7B-Instruct: ALFWorld 46.9% versus GRPO 43.8%; WebShop accuracy 57.0% and score 73.8 versus GRPO 43.0/66.4 and Skill-SD 54.7/73.7. Against Skill-SD specifically, PR-OPD gains 7.8 points in WebShop accuracy and 8.1 in score on the 3B backbone.

• Token-level distillation underperforms. From the same warm start and the same self-generated skills, distilling token probabilities falls to 57.8% ALFWorld success (8.6 points below GRPO), whereas hidden-state alignment reaches 72.7% (6.3 points above) — a 14.9-point gap between two runs that differ only in how the teacher is read.

• The projector matters more for the smaller model. On the 3B backbone the projector moves success from 72.7% to 71.1%; on the 1.7B backbone, direct alignment without the projector falls to 36.7% (below GRPO's 43.8%), and the projector lifts it to 46.9%.

• Warm-start placement outweighs the number of supervised updates. Starting Stage II at N=120 instead of N=100 raises WebShop accuracy from 49.5% to 57.0% with 20 fewer supervised updates. Paired against GRPO from the same checkpoint over six seeds, alignment from step 120 gains 9.62 points after 30 updates and 8.22 after 10, versus 3.15 from step 100 and 3.63 from step 140.

• Skill effects are specific, not mere prompt-length effects. Hidden-state change grows from early to middle layers and stays high through later layers, exceeding length-matched unrelated text by up to nine times for 3B and seven times for 1.7B, with the trajectory's own skill moving states the most.

• Probes confirm decodability by depth. Against another trajectory's skill, middle-layer probes reach 65.4% (3B) and 79.4% (1.7B); against a token-shuffled skill, 86.2% and 98.3%; against no skill or unrelated text, 88.6% to 99.8%. Output features range from 50.2% to 54.0% on the last two distinctions. For 1.7B, accuracy peaks at 79.4% in the middle layers and drops to 70.1% at the final layer.

• The skill makes targeted structural edits. On the initial 3B policy, same-layer CKA averages 0.996, so overall structure stays largely intact, yet from layer 13 onward the own skill changes structure more than all controls — 1.5 times as much as another trajectory's skill and six times as much as unrelated text, with the debiased estimate giving the same ordering.

Methodology in Plain English

The method has two stages. Stage I is a plain GRPO warm start for N updates (N=100 on ALFWorld, N=120 on WebShop, out of 150 total). This is necessary because the model must be able to solve the task before its own written reflections are worth learning from.

Stage II adds representation supervision to the same reward objective. Each iteration: the policy collects trajectories without any skill; after each episode it serializes the interaction record and outcome and prompts itself for a short list of actionable rules (a "skill") — describing a reusable workflow if the episode succeeded, or the critical mistakes and their fixes if it failed. Skills are regenerated from every new batch rather than retrieved from a library. The same policy then runs two forward passes over the identical sampled response tokens: a student pass with ordinary context and a teacher pass with the skill prepended. Because weights are shared, layer ℓ of the student pairs with layer ℓ of the teacher with no mapping to learn; because the teacher's prompt is longer, states are matched by response token rather than by position.

The student's state at each token and each layer goes through a small per-layer MLP projector (one 512-dimensional hidden layer) and is pulled toward the stop-gradient teacher state by a squared distance on L2-normalized vectors — which reduces to a cosine-distance term. This loss covers every Transformer layer and every response token, including reasoning and action tokens, and is added with weight λ=0.01. The projector exists so the student's own language-model head is not forced to copy states computed under a longer, skill-bearing prompt. The teacher is recomputed at every update because it shares the evolving weights, so the target moves with the policy.

Why This Matters

Impact on research. The paper reframes privileged self-distillation for agents: the question is not how to better weight a per-token scalar, but which channel carries privileged information at all. The finding that skill information peaks in middle layers and degrades at the final layer is a concrete argument against last-layer-only alignment, which prior work adopted for a different reason (depth-pairing mismatch with a separate teacher). The diagnostics — probing output features against hidden states — give the community a reusable test for whether a distillation signal is being wasted.

Real-world applications:

• Multi-step customer-service or web-navigation agents that must learn from sparse outcome signals without shipping extra context at inference.
• Robotics and embodied instruction-following agents trained on multi-turn interaction with terminal-only rewards.
• Enterprise agents where external skill libraries or separate teacher models are impractical due to latency, cost, or data-governance constraints.
• Post-training pipelines that want denser learning signal from the same rollouts without additional human annotation.

Industry relevance. PR-OPD adds no inference overhead, no separate teacher, and no external skill library — only a projector that is discarded after training. That profile fits production RL fine-tuning loops where serving cost and latency are fixed constraints.

Future Directions

• Why the projector is essential for the smaller backbone but nearly free for the larger one is observed but not explained; a capacity- or depth-dependent account would tell practitioners when to include it.
• Choosing the warm-start point currently requires a sweep (N=100, 120, and 140 give markedly different gains). A principled criterion for when a policy's hindsight becomes worth reading is open.
• Extending beyond self-written skills — for example, whether the same representation-level channel transfers privileged information from other sources such as environment state, tool traces, or simulator access.
• Robustness of the alignment target as the policy drifts: the teacher is recomputed every update, so questions remain about stability over longer training than the 150 updates reported here.

Target Audience

Researchers and engineers working on RL post-training for language-model agents, particularly those already using GRPO-style methods or privileged/on-policy self-distillation. It is also useful for practitioners interested in representation-level distillation and for anyone who needs evidence about which layer to align when the teacher and student share weights. Readers without background in RLVR or hidden-state distillation will find the empirical diagnostics accessible but the method section demanding.

Authors’ abstract

Language-model agents are usually trained by reinforcement learning from one reward per episode, and privileged self-distillation enriches it by letting the same policy, given a skill, teach its skill-free self through token probabilities. However, we identify two phenomena that question this channel. Invisible Advantage: a skill in context lifts WebShop success from 42.2% to 56.2%, yet changes the probabilities of fewer than a quarter of the sampled tokens. Much to Align: a skill changes the hidden states of over 80% of response tokens, in a way that linear probes can trace back to the specific skill. To exploit this, we propose Privileged Representation On-policy Self-Distillation (PR-OPD). After a GRPO warm start, the policy writes a hindsight skill for each trajectory, re-reads its own responses with that skill as a stop-gradient teacher, and aligns its projected hidden states to the teacher's at every layer alongside the reward objective, with no external skill library, separate teacher, or inference overhead. On ALFWorld and WebShop with two backbones, PR-OPD achieves the best overall results in every setting, improving over GRPO by up to 4.7 points in ALFWorld success and 14.0 points in WebShop accuracy. Code is available at https://github.com/balibata/PR-OPD.