PaperlayerView exact revision on arXiv

Source-linked AI summary

A Good Self-Teacher Meets the Student Where They Are: Joint On-Policy Learning and Teaching

Randy Ardywibowo, Arnav Dalal, Jiantao Jiao

arXiv:2610.10447v1cs.LGcs.AI

TL;DR

Outcome-reward RL gives sparse supervision on difficult, long-horizon tasks, and privileged self-distillation does not ensure guidance that improves the student. The paper characterizes reward-consistent teaching and proposes JOLT, jointly training a privileged teacher and unprivileged student; it reports performance and efficiency gains across reasoning and agentic tasks.

  • Problem

    Outcome rewards provide sparse supervision, especially on difficult, long-horizon tasks, while privileged conditioning alone does not ensure that a teacher's distillation update improves the student.

    Source: §1 · p. 1, §2.2 · p. 4

  • Method

    JOLT jointly trains a shared policy as a privileged teacher with reward and student-directed KL regularization and as an unprivileged student receiving dense on-policy distillation.

    Source: §1 · p. 2

  • Results

    JOLT+ improves over GRPO across reasoning and agentic benchmarks, with a 4.2-percentage-point gain on MATH.

    Source: §1 · p. 3

  • Takeaways & Limitations

    JOLT reaches competitive performance without direct student rewards and uses fewer training tokens, while adding student rewards further strengthens the approach.

    Source: §5 · p. 10

    The gradient-proportionality result applies to a fixed student policy under the ideal teacher; joint updates to shared parameters fall outside that guarantee.

    Source: §2.3 · p. 5

Abstract

from arXiv · show

Reinforcement Learning (RL) from outcome rewards suffers from sparse supervision, particularly on difficult, long-horizon tasks where successful trajectories are rare and costly to generate. On-Policy Distillation (OPD) offers an attractive alternative by providing dense token-level supervision from a stronger teacher along the student's own generations. Self-distillation methods further remove the need for a separate teacher model by conditioning the same policy on privileged information to serve as its own teacher. However, privileged conditioning alone does not guarantee that the resulting distillation update improves the student. Indeed, privileged information can lead the teacher to solve tasks through shortcuts unavailable to the student, producing supervision poorly matched to the student's current behavior. Consequently, even a higher-performing teacher can provide guidance that degrades student performance. To address this, we analyze how the choice of privileged teacher affects the student's update. We derive a necessary and sufficient condition for the teacher's local distillation update to be a positive multiple of the student's reward gradient. Our analysis suggests that the teacher should not only perform well on the task, but also provide guidance suited to the student's current capabilities. This characterization motivates a practical teacher-training surrogate that combines outcome rewards with token-level Kullback-Leibler (KL) regularization toward the student. Based on this result, we propose Joint On-Policy Learning and Teaching (JOLT), which jointly trains a single policy in two roles: a privileged teacher using a KL-regularized objective, and an unprivileged student using dense on-policy distillation. Across mathematical reasoning, coding, tool use, and terminal use, JOLT improves training efficiency and performance, with further gains from student rewards.

1 INTRODUCTION

Outcome-reward RL provides sparse supervision on difficult, long-horizon tasks, while privileged self-distillation can produce guidance mismatched to the student. The paper motivates JOLT, which jointly trains a privileged teacher and unprivileged student, and reports gains across reasoning and agentic tasks.

1 INTRODUCTION

  • Outcome rewards arrive only after a trajectory ends, making supervision sparse on difficult, long-horizon tasks where successful trajectories may be rare and costly.On-Policy Distillation instead evaluates student-sampled actions under a teacher, providing a dense signal at each visited state.Source: §1 · p. 1
  • Privileged conditioning can produce shortcuts unavailable to the student, so a higher-performing teacher may still induce a distillation update that locally lowers student reward.The resulting guidance may also be poorly matched to the student's behavior or unreliable after an erroneous student-generated prefix.Source: §1 · p. 2
  • A useful teacher is characterized by its effect on the student: the paper derives a necessary and sufficient condition for its local distillation update to align positively with the student's reward gradient.The proposed teacher favors high-advantage actions under student continuations while remaining close to the student's policy, motivating outcome rewards with a teacher–student KL penalty.Source: §1 · p. 2
  • JOLT jointly trains one shared policy as a privileged teacher and unprivileged student, with each role learning on-policy from its own generations.The teacher uses outcome rewards and token-level KL regularization toward the student; the student receives dense teacher distillation, and JOLT+ can add direct student outcome rewards.Source: §1 · p. 2
  • JOLT+ gains 4.2 percentage points over GRPO on MATH, 3.0 on LCBv6, 9.9 on AppWorld, and 4.5 in Terminal-Bench pass@32.On MATH, JOLT matches GRPO accuracy using 13× fewer completion tokens and 6.5× fewer processed training tokens, without direct rewards on student trajectories.Source: §1 · p. 3

2 METHODOLOGY

The section connects on-policy distillation to reward learning, then characterizes when a privileged teacher’s update aligns with the student’s reward gradient. JOLT turns that characterization into joint teacher–student training with shared parameters.

2.1 BACKGROUND

  • A state-dependent baseline centers action values into advantages, which recover the expected-return policy gradient while reducing variance without bias.Because action values depend on future trajectories, sparse terminal rewards can make useful credit estimates costly to obtain.Source: §2.1 · p. 3
  • OPD gives dense supervision by scoring student-sampled actions under a stronger teacher; OPSD uses the same policy with privileged context instead.OPD’s teacher–student log-probability ratio supervises each sampled action, while OPSD retains dense supervision without an additional teacher model.Source: §2.1 · pp. 3–4

2.2 A REWARD-CONSISTENT TEACHER

  • A reward-consistent teacher is one whose local OPD update matches the student’s expected-return gradient up to a positive scaling.The section derives a necessary and sufficient condition for this proportionality.Source: §2.2 · p. 4
  • An exponential, KL-regularized policy-improvement update is a sufficient construction for a teacher whose local OPD update is reward-consistent.The construction motivates training the privileged teacher toward a target distribution based on the student’s advantages.Source: §2.2 · p. 4

2.3 JOINT ON-POLICY LEARNING AND TEACHING (JOLT)

  • The proportionality result applies to an ideal teacher for a fixed student; joint teacher updates to shared parameters fall outside that guarantee.JOLT’s formulation is agnostic to how privileged context is constructed and supports learning from original tasks alongside easier assisted versions.Source: §2.3 · p. 5
  • JOLT combines a privileged teacher’s reward-based, token-level KL-regularized training with student on-policy distillation, updating shared parameters from each role’s own trajectories.Since student continuation advantages are generally unavailable, its practical teacher surrogate uses returns from teacher trajectories and a separate token-level KL loss.Source: §2.3 · pp. 4–5

3 RELATED WORK

Related work spans trajectory-based distillation, privileged self-distillation, and efforts to improve teacher reliability or train teachers for student learning. The paper distinguishes its update-based teacher characterization from these approaches.

Related work

  • Sequence-level distillation trains on teacher-generated prefixes, while OPD queries the teacher along student-generated trajectories to provide dense reverse-KL supervision.The different prefix distributions motivate on-policy supervision.Source: §3 · pp. 5–6
  • Figure 2 compares training accuracy, teacher–student KL, and final evaluations across three benchmarks; Figure 3 reports AppWorld training performance and JOLT+’s update-efficiency comparison with GRPO.The captions report higher mean evaluation accuracy for JOLT than distillation baselines in Figure 2 and about 1.8× fewer updates for JOLT+ to reach GRPO’s final training performance in Figure 3.Source: Figure 2 · p. 6, Figure 3 · p. 7
  • Privileged self-distillation conditions the student model on added information, providing dense guidance without a separately trained teacher model.Prior work uses sources such as verified solutions, demonstrations, feedback, system prompts, or prior experience.Source: §3 · p. 6
  • Prior studies report unreliable guidance on student-generated prefixes from support mismatch, high entropy, privileged-context leakage, or reasoning-style differences.Other work assesses or modifies the student update; this paper instead characterizes teacher distributions whose OPD update follows the student’s reward gradient.Source: §3 · p. 6
  • Teacher-learning approaches optimize explanations, problem generation, or task success with learnability signals; π-Distill jointly trains shared teacher and student roles but uses off-policy student updates.The paper connects its exponential teacher reweighting to KL-regularized policy improvement.Source: §3 · pp. 6–7

4 EXPERIMENTS

JOLT and JOLT+ improve on distillation baselines across the evaluated reasoning and agentic tasks, with student rewards adding gains. The experiments also compare training-token efficiency, teacher–student agreement, and student-gradient error patterns.

Cross-task results

  • JOLT+ has the highest reported mean accuracy on MATH, LCBv6, and AppWorld; JOLT also exceeds GRPO on MATH and LCBv6.On AppWorld, JOLT remains below π-Distill and GRPO. Table 1 reports MATH/LCBv6 accuracy and AppWorld state-test pass rate.Source: §4 · pp. 7–8, Table 1 · p. 7

Training-token efficiency

  • On MATH, JOLT reaches high accuracy with fewer training tokens than GRPO and surpasses OPSD’s early plateau under both token counters.Figure 4 compares completion tokens and processed training tokens.Source: Figure 4 · p. 8

Terminal-Bench

  • On Terminal-Bench, JOLT improves over the base model and both distillation baselines at every reported pass@k, while GRPO scores higher.Table 2 reports pass@k after TMax training and uses one training run per method.Source: §4 · p. 8, Table 2 · p. 9

Student–teacher alignment

  • In the illustrated GSM8K response, JOLT’s teacher–student log-probability differences are smaller and concentrated at fewer token positions than OPSD’s.Figure 5 compares the methods on the same problem; the passage describes this as an example.Source: Figure 5 · p. 8, §4 · p. 9

Gradient bias–variance analysis

  • Along the evaluated training trajectories, GRPO’s gradient error is dominated by variance, while OPSD develops increasing empirical bias and JOLT variants retain low variance with smaller bias.The bias and MSE references use Monte Carlo reward gradients; the analysis notes that these measurements do not isolate teacher training’s effect or establish reward consistency of the complete joint update.Source: Figure 6 · p. 9, §4 · p. 9

5 CONCLUSION

JOLT jointly trains one policy as a privileged teacher and an unprivileged student, with the teacher’s objective balancing task performance and proximity to the student. The paper reports improved effectiveness and efficiency from joint training, with further gains when student rewards are included.

Joint training

  • JOLT jointly trains a single policy in privileged-teacher and unprivileged-student roles, with each learning from its own trajectories.The formulation accommodates different sources of privileged information.Source: §5 · pp. 9–10

Teacher objective

  • The analysis links teacher supervision to the student’s reward objective and motivates balancing teacher task performance with proximity to the current student.The paper reports that JOLT can reach competitive performance without direct student reward and achieve high accuracy with fewer training tokens.Source: §5 · pp. 9–10

Student rewards

  • Adding student rewards further strengthens the approach by combining dense teacher supervision with direct feedback on student generations.The paper presents this as a way to supplement sparse rewards with supervision that adapts to the model’s current capabilities, without a separate teacher.Source: §5 · p. 10

A EXPERIMENTAL DETAILS AND ADDITIONAL RESULTS

The appendix specifies run counts, uncertainty conventions, hyperparameters, checkpoint selection, and additional evaluation and sensitivity results.

Experimental settings

  • Most experiments use five independent training seeds and report normal-approximation 95% confidence intervals, with stated exceptions for Terminal-Bench and gradient analysis.The gradient analysis uses batch-bootstrap intervals; the appendix also lists task-specific run counts and evaluation-band conventions.Source: A EXPERIMENTAL DETAILS AND ADDITIONAL… · p. 14

Hyperparameters and checkpoints

  • Hyperparameters are selected separately by method using training curves and final training performance, then fixed before repeated-seed comparisons.The appendix identifies B as problems per update and G as group size per active role; JOLT/JOLT+ use four teacher and four student rollouts per problem on GSM8K and MATH.Source: A EXPERIMENTAL DETAILS AND ADDITIONAL… · p. 14, Table 3 · p. 14

Additional evaluation results

  • Figure 7 presents evaluation curves for GSM8K, MATH, LCBv6, and AppWorld, with training summaries for the first three tasks and final AppWorld evaluation.The plotted bands and error bars show standard deviations across runs; LCBv6 evaluates complete test suites and AppWorld reports development-set state-test pass rate.Source: Figure 7 · p. 15

Hyperparameter sensitivity

  • Figure 8 examines GSM8K sensitivity across JOLT configurations and learning rates across methods.Source: Figure 8 · p. 16

A.1 GSM8K

The appendix specifies training and evaluation protocols across GSM8K, MATH, LiveCodeBench v6, AppWorld, and Terminal-Bench, including how privileged context differs between teacher and student roles.

A.1 GSM8K

  • On GSM8K, the strongest JOLT configurations use learning rates near 10−6 and coupling coefficients below one; unit coefficients are unstable and higher rates degrade all methods.Training uses the official 7,473-problem training split; evaluation uses 1,319 test problems with one greedy response per problem.Source: §A.1 · p. 16

A.2 MATH

  • MATH training uses 7,500 problems and gives the teacher a worked reference solution, while evaluation uses student prompts without that solution.The full test split contains 5,000 problems; training and evaluation use temperature 1.0 and a 16,384-token completion limit.Source: §A.2 · p. 17

A.3 LIVECODEBENCH-V6

  • LiveCodeBench v6 uses 131 problems, with approximately half of each problem’s private tests visible during training and the complete 4,852-test suite used for evaluation.Four unprivileged student responses per problem are evaluated against the expanded suite, including withheld inputs.Source: §A.3 · pp. 17–18

A.4 APPWORLD

  • AppWorld training uses all 90 public training tasks, while evaluation uses 57 development tasks and reports state-test pass rate averaged across tasks.Teacher guidance summarizes prior attempts into strategies, mistakes, and workflows; each run starts with an empty task-local skill bank.Source: §A.4 · p. 19

A.5 TERMINAL-BENCH

  • Terminal-Bench gives the teacher private verifier files that are omitted from student context and evaluation, which grades explicit submissions against hidden verifiers.Evaluation covers the complete task set and samples 128 trajectories per task, but uses one training run per method, so variability across training seeds is unassessed.Source: §A.5 · p. 20

A.6 CONSISTENCY ACROSS SETTINGS

The section assesses whether JOLT+ consistently exceeds GRPO across the evaluated settings by comparing final-checkpoint runs and estimating uncertainty in their ordering.

A.6 CONSISTENCY ACROSS SETTINGS

  • The probability of improvement averages pairwise final-checkpoint comparisons with equal weight for GSM8K, MATH, LCBv6, and AppWorld.Uncertainty is estimated by stratified percentile bootstrap, resampling training runs within each method and setting.Source: §A.6 · pp. 21–22
  • With only three runs in most settings, bootstrap intervals may underestimate uncertainty, particularly when all observed comparisons favor one method.Terminal-Bench is reported separately because it has one training run per method.Source: §A.6 · p. 22
  • 94.7% average probability of improvement, with a nominal 95% bootstrap interval of [83.3%, 100.0%], indicates JOLT+ exceeded GRPO in nearly all observed run comparisons.Every observed JOLT+ run exceeded every observed GRPO run on LCBv6 and AppWorld.Source: Table 4 · p. 22, §A.6 · p. 22

B TOKEN EFFICIENCY AND SCALING

This section compares MATH accuracy against training-token expenditure and estimates how much budget each method requires at matched accuracy levels.

B TOKEN EFFICIENCY AND SCALING

  • The analysis combines token-budget sweeps and extended-training checkpoints to trace student accuracy across early learning and larger budgets.Figure 4 presents each method’s Pareto frontier using training-token expenditure and mean evaluation accuracy.Source: B TOKEN EFFICIENCY AND SCALING · p. 22
  • Processed training tokens include prompt and completion positions, teacher-only context, scoring passes, and repeated optimization passes, so token expenditure does not directly measure wall-clock savings.JOLT completion-token counts include both teacher and student completions.Source: B TOKEN EFFICIENCY AND SCALING · p. 22
  • The fitted accuracy-versus-budget relationship is sigmoidal in log training-token budget, with separate parameters estimated for each method and token counter.The fit uses bounded nonlinear least squares with multiple initializations and reports RMSE.Source: B TOKEN EFFICIENCY AND SCALING · pp. 22–23
  • At tested MATH accuracy targets, JOLT requires approximately 9.5–14.4× fewer completion tokens and 4.3–7.4× fewer processed training tokens than GRPO.The ratios use fitted budgets for both methods and are point estimates without confidence intervals.Source: Table 7 · p. 23, B TOKEN EFFICIENCY AND SCALING · p. 23

C TEACHER-OBJECTIVE ABLATIONS

The ablations compare teacher-objective components and student rewards on MATH, testing their associations with student accuracy while noting that teacher and student updates share parameters.

C TEACHER-OBJECTIVE ABLATIONS

  • The ablation retains student distillation while separately testing teacher KL-only and teacher outcome-reward-only objectives, with configurations tuned independently.The student-reward control uses JOLT+’s student learning configuration but applies no explicit teacher optimization.Source: C TEACHER-OBJECTIVE ABLATIONS · p. 24
  • 73.80 ± 0.74% held-out accuracy for JOLT+ exceeds 66.52 ± 0.73% for OPSD + Student Reward at step 75, supporting a benefit from explicit teacher optimization on MATH.The control retains student rewards and dense privileged self-distillation but disables both teacher-objective terms.Source: C TEACHER-OBJECTIVE ABLATIONS · p. 24, Table 8 · p. 24
  • Both teacher-training variants improve over OPSD as training progresses, and the teacher outcome-only variant exceeds the KL-only variant in held-out accuracy.Because the roles share parameters, the comparisons cannot distinguish improved distillation from direct transfer through teacher updates.Source: C TEACHER-OBJECTIVE ABLATIONS · p. 24, Figure 10 · p. 25

D TOKEN-LEVEL ALIGNMENT

At final checkpoints, JOLT and JOLT+ show smaller teacher–student token-level log-probability differences than OPSD and π-Distill. Figure 11 compares both the direction and magnitude of those differences on student-generated responses.

D TOKEN-LEVEL ALIGNMENT

  • Red marks tokens with higher teacher log probability, while blue marks tokens with lower teacher log probability.Color intensity reflects the magnitude of the difference, with symmetric limits around zero shared across methods.Source: D TOKEN-LEVEL ALIGNMENT · p. 25
  • The visualizations use each method’s final checkpoint and a randomly selected problem and responses.Teacher and student probabilities are compared at each student-generated token using the same response prefix.Source: D TOKEN-LEVEL ALIGNMENT · p. 25
  • JOLT and JOLT+ exhibit smaller teacher–student log-probability differences throughout generations than OPSD and π-Distill.The comparison extends Figure 5 to include π-Distill and JOLT+.Source: D TOKEN-LEVEL ALIGNMENT · p. 25, Figure 11 · p. 26

E GRADIENT BIAS–VARIANCE ANALYSIS

The analysis compares student-gradient estimators with a Monte Carlo reward-gradient reference, both along each method’s training trajectory and at fixed checkpoints. Distillation-only and combined estimates show stronger directional agreement than reward-only estimates at the evaluated checkpoints, while several measurement and scope caveats constrain interpretation.

Reference and measurements

  • The reference is an empirical target built from eight independent REINFORCE estimates, each using 512 completions per problem, rather than an exact expected-reward gradient.The measurements use 32 independent estimates per batch, each with eight student completions per problem, at specified MATH checkpoints.Source: E GRADIENT BIAS–VARIANCE ANALYSIS · p. 26

Measurement limitations

  • Lower measured variance or Euclidean discrepancy alone does not establish stronger directional alignment because positive rescaling changes those statistics without changing direction.The reference mean’s estimated variance relative to its squared norm ranges from 0.074 to 0.661 across evaluated batches, checkpoints, and trajectories.Source: E GRADIENT BIAS–VARIANCE ANALYSIS · pp. 26–27
  • These measurements reflect both the policies reached and their estimators, and reference uncertainty contributes to discrepancy, especially when the reference gradient is small.The analysis excludes teacher gradients and does not establish alignment of the combined shared-parameter update.Source: E GRADIENT BIAS–VARIANCE ANALYSIS · p. 27

Comparisons during training

  • Along evaluated training trajectories, JOLT and JOLT+ maintain substantially smaller disagreement with the reference than OPSD while retaining low variance.GRPO’s error is predominantly variance, whereas OPSD shows low variance but increasing disagreement.Source: E GRADIENT BIAS–VARIANCE ANALYSIS · p. 27

Directional comparisons at fixed checkpoints

  • At evaluated fixed checkpoints, distillation-only and combined estimates have lower directional dispersion and higher mean reference alignment than reward-only estimates.The pattern appears across all four trained policies and persists after gradient-magnitude normalization.Source: E GRADIENT BIAS–VARIANCE ANALYSIS · pp. 27–28, Figure 12 · p. 28

π-Distill sensitivity

  • π-Distill’s weaker MATH performance persists across tested settings, while its AppWorld development state-test pass rate is near GRPO and above OPSD.The tested MATH changes include learning rates, mixing choices, schedules, warmup, and loss normalization.Source: E GRADIENT BIAS–VARIANCE ANALYSIS · pp. 28–29

Related papers

Loading 2610.10447v1…