Expected Reasoning-Step Return Unifies On-Policy Learning from Rewards and Teachers
Authors: Qiangqiang He, Jin Li
Organizations: State Key Laboratory for Novel Software Technology, Nanjing University, Nanjing, China · College of Software Engineering, Southeast University, Nanjing, China
On-policy reasoning models can learn from task rewards or teacher signals, but these sources differ in form and can favor conflicting updates, leaving unclear which should guide a given reasoning action. We introduce Expected Reasoning-Step Return (ERSR), which treats semantic reasoning steps as macro-actions and uses Monte Carlo student-policy rollouts to estimate the expected final task reward of student-generated and teacher-proposed actions in a common return space for step-level comparison. ERSR analysis reveals an outcome-dependent asymmetry: student actions are more beneficial than teacher replacements on successful trajectories, whereas teacher replacements become more beneficial on failed trajectories. We further show that student answer-probe gains track student-step ERSR utility and distinguish beneficial from harmful reasoning steps. Based on these findings, we propose \textbf{Return-Referenced On-Policy Learning (R2OPL)}, which reinforces student reasoning on successful trajectories and distills teacher signals on failed ones, while using group success rate for difficulty scaling and student-probe gains for step-level modulation. Experiments across reasoning benchmarks and teacher--student configurations show that R2OPL consistently outperforms strong baselines. ERSR training dynamics further show that R2OPL jointly exploits substantial utility from both reward- and teacher-side signals, whereas existing hybrids often leave substantial residual utility in one branch.
Figures & tables
Figure 1: ERSR evaluates a reasoning step by estimating its student-step advantage Ak and teacher-replacement advantage AkTR , with all subsequent continuations sampled from the student policy.
Outcome
ERSR Utility
Mean
Positive (%)
Zero (%)
Negative (%)
Δsign (%)
R=1
A
+0.032
51.6
15.9
32.6
+19.0
ATR
+0.007
44.8
14.0
41.2
+3.7
R=0
A
−0.021
30.2
26.8
43.0
−12.9
ATR
+0.064
49.8
24.2
26.0
+23.7
Table 1: Outcome-dependent ERSR utility. Positive, zero, and negative denote the percentages of steps by utility sign; Δsign is the positive–negative percentage-point gap.
A
ATR
Outcome
ΔPS
ΔPT
eDTS
eLT
ΔPS
ΔPT
ΔPT−S
PT−S
eDTS
eLT
R=1
0.246
0.178
−0.092
−0.104
0.003
0.007
0.005
0.038
−0.071
−0.048
R=0
0.211
0.131
+0.109
+0.089
0.008
−0.013
−0.022
0.016
−0.012
−0.002
Table 2: Pearson correlations between step-level signals and ERSR utility. ΔPT−S=ΔPT−ΔPS and PT−S=PT−PS . The highest absolute correlation per utility group is bolded ; shaded cells mark negative teacher-side correlations on successful trajectories.
Figure 2: Student-probe ranking of step utility. (a) Student-step utility across probe quintiles on successful trajectories. (b) Student-step utility across probe quintiles on failed trajectories.
Figure 3: Efficient answer-probe computation. (a) Packed probes attend only to their reasoning prefixes. (b) Probe position IDs align with the corresponding reasoning steps.
Figure 4: ERSR dynamics during training. AMC23 Avg@16 and the remaining ERSR utilities of student reinforcement and teacher distillation are evaluated across training checkpoints.
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
Qwen3-1.7B
Gemma-4-E2B-it
Explicit-step compliance
1020/1024 (99.6%)
772/1024 (75.4%)
Appendix
Table 6: Compliance with Explicit-Step CoT on 1,024 sampled trajectories.
Rule
Boundary signal
Safeguard
Explicit steps
Top-level consecutive Step N headings
Used whenever a valid sequence is present
Terminal answer
Independent Final Answer heading or terminal \boxed{...}
Separate an independent answer section; otherwise retain the answer in the final reasoning step
Semantic fallback
Repeated top-level reasoning headings
Requires at least two complete reasoning sections and excludes non-reasoning headings
Conservative fallback
No reliable structural signal
Keep the remaining response as a single step
Appendix
Table 7: Reasoning-step segmentation hierarchy.
Figure 5: Representative semantic-heading fallback case. Repeated top-level headings provide semantic boundaries when no explicit Step N sequence is present, yielding four segments with the terminal answer retained in the final segment.
Teacher
R=1
R=0
Qwen3-4B-Instruct-2507
+0.007 [ +0.000 , +0.013 ]
+0.0638 [ +0.054 , +0.075 ]
Qwen3-4B
+0.0146 [ +0.009 , +0.020 ]
+0.0218 [ +0.017 , +0.027 ]
Appendix
Table 9: Teacher-replacement ERSR utility across teachers. Values report mean ATR with confidence intervals under successful and failed student trajectories.
Figure 6: ERSR utility across problem difficulty. Student-step utility A is evaluated on successful trajectories, while teacher-replacement utility ATR is evaluated on failed trajectories for the two teachers.
Estimation Budget
A
ATR
Steps
MC@N
Typical Dev.
95% Bound
Typical Dev.
95% Bound
250
4
0.0146
0.0358
0.0226
0.0554
500
2
0.0133
0.0326
0.0177
0.0435
500
8
0.0084
0.0207
0.0150
0.0368
1000
8
0.0060
0.0147
0.0106
0.0260
Appendix
Table 10: Low-cost approximation of aggregate ERSR. Deviations are measured against the full MC@128 estimates ( A=+0.032 , ATR=+0.064 ). The shaded row denotes the configuration used for training-time tracking.
Dataset
Domain
Task Format
Original Size
Evaluated
AMC23
Math
Open-ended short answer
40
40
AIME24
Math
Numerical fill-in
30
30
AIME25
Math
Numerical fill-in
30
30
MATH-500
Math
Open-ended short answer
500
200
GPQA-Diamond
Science
4-choice MCQ
198
198
SciBench
Science
Open-ended numerical
580
200
Appendix
Table 12: Evaluation-set composition across mathematical, scientific, and logical benchmarks.
Hyperparameter
Value
Training framework
verl
Numerical precision
bfloat16
Optimizer
AdamW
Learning rate
1×10−6
Learning-rate schedule
Cosine
Warmup
None
Appendix
Table 13: Shared training and evaluation configurations. Unless otherwise specified, all methods use the same optimization, rollout, and evaluation settings.
Figure 7: AMC23 training dynamics under OPD, GRPO, and R 2 OPL with equal branch scaling. Curves report Avg@16 across training checkpoints.
Figure 8: Branch-scale imbalance in R 2 OPL under equal weighting ( μ=λ=1 ). We compare the gradient norms and advantage magnitudes of the successful-trajectory RL branch and failed-trajectory OPD branch during training.
Figure 9: ERSR utility dynamics under equal branch scaling ( μ=λ=1 ). We compare the remaining student-step utility A and teacher-replacement utility ATR on AMC23 at steps 0 and 60 for OPD and R 2 OPL.
Figure 10: Effect of RL–OPD branch scaling in R 2 OPL. For each (μ,λ) configuration, we report the best AMC23 Avg@16 over training and the remaining student-step and teacher-replacement ERSR utilities at the corresponding best checkpoint.
On-policy distillation (OPD) combines student-generated rollouts with dense token-level supervision from a teacher, but providing such supervision for every rollout requires substantial teacher computation. We introduce Success-Referenced On-Policy Distillation (SR-OPD), which reduces this cost by selecting which prompts and rollouts receive teacher supervision. When the student produces both successful and failed rollouts for the same prompt, a successful rollout can serve as a natural reference for selecting failed rollouts. SR-OPD therefore focuses on such prompts and prioritizes failed rollouts whose hidden-state trajectories show sustained divergence from a successful reference, while accounting for estimated teacher-input cost. Across three teacher-student pairs and six mathematical reasoning benchmarks, SR-OPD uses only 3.46-5.02% of the teacher-input tokens required by Vanilla OPD in the one-pass setting while maintaining comparable reasoning performance. Under a controlled setting matched to 5% of Vanilla OPD's teacher-input budget, further experiments support both key design choices: focusing supervision on prompts with both successful and failed rollouts, and using successful rollouts to guide failure selection. These results indicate that a student's own successful behavior can serve as a useful reference for allocating teacher supervision under a fixed teacher-input budget.
Xiang Chen, Futao Su, Kong Wang +2
Tongji University · Xiaohongshu · Shanghai Jiao Tong University +1
Reinforcement learning can substantially improve a reasoning teacher, but it is unclear which of those improvements survive when the teacher supervises a smaller on-policy student. We study this question in mathematical reasoning by comparing teacher lineages before and after GRPO, multiple student scales, direct GRPO, and several on-policy distillation objectives. The central finding is that transfer is structured rather than scalar: teacher strength alone does not make dense distillation competitive, while an RL-improved teacher creates useful but metric-dependent student gains. This motivates Transfer-Stratified On-Policy Distillation (TS-OPD), which screens training problems by the joint sampled success of the student and teacher, routes acquisition problems to gated forward KL, routes consolidation problems to gated reverse KL, and adds an entropy brake to protect sampled coverage. Across the main comparison, TS-OPD is the strongest student objective for macro average correctness with the GRPO-improved teacher, while pass@K remains more mixed. Ablations show that the gains come from routing and token gating rather than skipping problems. These results support a transfer-aware view of OPD: stronger teachers help when the supervision direction and token budget match the student's observed ability, not merely because the teacher endpoint is stronger.
Xiaoyu Chen, Bo Shao, Tiangang Zhu +5
Institute of Software, Chinese Academy of Sciences · Microsoft · Work done during an internship at Microsoft +1
Reinforcement learning with verifiable rewards (RLVR) supervises mathematical reasoning through final-answer correctness, but provides little guidance on individual tokens. On-policy distillation (OPD) supplies dense feedback on student-generated responses, yet teacher preference need not reflect correctness. Recent hybrids combine OPD and verifier-derived advantages or reweight task credit using teacher ratios. However, teacher guidance enters after verifier-based group normalization, and token reweighting need not preserve the total task credit assigned to each response. We introduce Unified Entropy-Calibrated Credit Redistribution for GRPO (UECR-GRPO), which integrates verifier and teacher signals within a single GRPO-style update at both the response and token levels. Path-Utility Unification (PUU) combines verifier reward and a teacher-to-anchor path log-ratio in a single KL-regularized objective. Its on-policy implementation uses a length-normalized teacher score and combines both rewards before group normalization and PPO clipping, allowing teacher evidence to influence the response ranking. Entropy-Calibrated Redistribution (ECR) then uses the signed teacher--old-policy token gap to redistribute the verifier-derived component. Full-vocabulary teacher entropy attenuates uncertain guidance, while a response-wise zero-sum projection preserves the total task credit and its token-wise sign before clipping. Across five mathematical reasoning benchmarks, UECR-GRPO achieves average Avg@12 accuracies of 17.21% and 65.09% with Qwen3-1.7B and Qwen3-4B students, respectively, exceeding the strongest baseline at each scale by 0.89 and 0.56 percentage points.
Jie Zhang, Jingxiao Yang, Zhehao Huang +2
Shanghai Jiao Tong University · Zhejiang University