Expected Reasoning-Step Return Unifies On-Policy Learning from Rewards and Teachers

Qiangqiang He; Jin Li

cs.AISep 26, 2026

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

Abstract

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 (R $^2$ OPL)}, 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 R $^2$ OPL consistently outperforms strong baselines. ERSR training dynamics further show that R $^2$ OPL 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 $\widehat{A}_{k}$ and teacher-replacement advantage $\widehat{A}_{k}^{\mathrm{TR}}$ , with all subsequent continuations sampled from the student policy.

Outcome	ERSR Utility	Mean	Positive (%)	Zero (%)	Negative (%)	$\bm{\Delta_{\mathrm{sign}}}$ (%)
$R=1$	$\widehat{A}$	$+0.032$	51.6	15.9	32.6	$+19.0$
$R=1$	$\widehat{A}^{\mathrm{TR}}$	$+0.007$	44.8	14.0	41.2	$+3.7$
$R=0$	$\widehat{A}$	$-0.021$	30.2	26.8	43.0	$-12.9$
$R=0$	$\widehat{A}^{\mathrm{TR}}$	$+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;

\Delta_{\mathrm{sign}}

is the positive–negative percentage-point gap.

Outcome	$\Delta P^{S}$	$\Delta P^{T}$	$e^{D_{\mathrm{TS}}}$	$e^{L_{T}}$	$\Delta P^{S}$	$\Delta P^{T}$	$\Delta P^{T-S}$	$P^{T-S}$	$e^{D_{\mathrm{TS}}}$	$e^{L_{T}}$
	$\bm{\widehat{A}}$				$\bm{\widehat{A}^{\mathrm{TR}}}$
$R=1$	0.246	0.178	$-0.092$	$-0.104$	0.003	0.007	0.005	0.038	$\mathbf{-0.071}$	$-0.048$
$R=0$	0.211	0.131	$+0.109$	$+0.089$	0.008	$-0.013$	$\mathbf{-0.022}$	0.016	$-0.012$	$-0.002$

Table 2: Pearson correlations between step-level signals and ERSR utility.

\Delta P^{T-S}=\Delta P^{T}-\Delta P^{S}

and

P^{T-S}=P^{T}-P^{S}

. 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.

Panel of Figure 5 — 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	$\bm{R=1}$	$\bm{R=0}$
Qwen3-4B-Instruct-2507	$+0.007$ [ $+0.000$ , $+0.013$ ]	$\mathbf{+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

\widehat{A}^{\mathrm{TR}}

with confidence intervals under successful and failed student trajectories.

Figure 6: ERSR utility across problem difficulty. Student-step utility

\widehat{A}

is evaluated on successful trajectories, while teacher-replacement utility

\widehat{A}^{\mathrm{TR}}

is evaluated on failed trajectories for the two teachers.

Estimation Budget		$\bm{\widehat{A}}$		$\bm{\widehat{A}^{\mathrm{TR}}}$
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 (

\widehat{A}=+0.032

\widehat{A}^{\mathrm{TR}}=+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\times 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 (

\mu=\lambda=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 (

\mu=\lambda=1

). We compare the remaining student-step utility

\widehat{A}

and teacher-replacement utility

\widehat{A}^{\mathrm{TR}}

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

(\mu,\lambda)

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.

RL for Language Model Reasoning

Reinforcement learning for language model reasoning uses reward-driven training to improve multi-step problem solving, tool use, and reliable answers.

See all 1,094 papers

Spend Teacher Tokens Where They Matter: Success-Referenced On-Policy Distillation
Oct 2, 2026Xiang Chen, Futao Su, Kong Wang +2Language Model Distillation On-Policy Distillation
Tongji University · Xiaohongshu · Shanghai Jiao Tong University +1
Transfer-Stratified On-Policy Distillation for RL-Improved Reasoning Teachers
Oct 5, 2026Xiaoyu Chen, Bo Shao, Tiangang Zhu +5RL for Language Model Reasoning Selective Knowledge Distillation
Institute of Software, Chinese Academy of Sciences · Microsoft · Work done during an internship at Microsoft +1
When and Where to Trust the Teacher: Unifying On-Policy Distillation and GRPO through Entropy-Calibrated Credit Assignment
Sep 23, 2026Jie Zhang, Jingxiao Yang, Zhehao Huang +2RL for Language Model Reasoning Group Relative Policy Optimization
Shanghai Jiao Tong University · Zhejiang University

Oct 2, 2026J/K move · Enter open · S save

Spend Teacher Tokens Where They Matter: Success-Referenced On-Policy Distillation

Xiang Chen, Futao Su, Kong Wang, Jiayi Chen+1

Tongji University · Xiaohongshu · Shanghai Jiao Tong University+1

Language Model Distillation On-Policy Distillation Knowledge Distillation Selective Knowledge Distillation

Open paper AskPDF arXiv

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.

Abstract

RL for Language Model Reasoning

Explore similar work

Spend Teacher Tokens Where They Matter: Success-Referenced On-Policy Distillation

Transfer-Stratified On-Policy Distillation for RL-Improved Reasoning Teachers

When and Where to Trust the Teacher: Unifying On-Policy Distillation and GRPO through Entropy-Calibrated Credit Assignment