MENTOR: Reinforcement Learning via Flexible Teacher-Optimized Rewards for Tool-Use Distillation

Choi, ChangSu; Song, Hoyun; Kim, Dongyeon; Cho, Minkyung; Jung, WooHyeon; Park, SunJin; Bae, NohHyeob; Yu, Seona; Lim, KyungTae

doi:10.18653/2026.emnlp-main.1286

Distilling LLM tool use into small language models with on-policy RL and set-based teacher-alignment rewards

EMNLP 2026

ChangSu Choi^1,*, Hoyun Song^2,1,*, Dongyeon Kim¹, Minkyung Cho¹, WooHyeon Jung¹, Sunjin Park³, NohHyeob Bae³, Seona Yu³, KyungTae Lim^1,†

¹Korea Advanced Institute of Science and Technology (KAIST) ²Department of Artificial Intelligence, Dankook University ³LG CNS

^*Equal contribution ^†Corresponding author

Abstract

Distilling the tool-use capabilities of large language models (LLMs) into small language models (SLMs) is essential for their practical application. The predominant approach, supervised fine-tuning (SFT), is an off-policy distillation method that suffers from poor out-of-domain (OOD) generalization because it rigidly aligns with static teacher trajectories. While reinforcement learning (RL) offers an alternative, the capacity limitations of SLMs pose a severe dilemma: sparse outcome rewards provide insufficient guidance, whereas strict trajectory matching imposes overly restrictive constraints. To bridge this capacity-driven gap, we propose MENTOR, an on-policy distillation framework that introduces a flexible yet process-aware reward structure. Instead of enforcing rigid replication, MENTOR uses the teacher’s reference to guide tool-use behavior, balancing behavioral alignment with downstream performance. Extensive experiments on controlled executable-tool benchmarks demonstrate that MENTOR improves OOD tool-use performance compared to SFT and strict RL baselines. Our findings suggest that within verifiable tool-use environments, flexible tool-use alignment offers a more effective approach than strict trajectory replication for developing adaptable small models.

Keywords Tool-use distillation Knowledge distillation On-policy distillation Small language models (SLMs) Tool-integrated reasoning (TIR) Reinforcement learning (GRPO) Reward design Function calling LLM agents Out-of-domain generalization

Method

MENTOR training loop: the teacher produces a reference trajectory of reasoning, tool-use and observation steps plus an answer; the student generates several rollouts; each rollout is scored by three reward functions (correctness Rc, teacher alignment Ra, tool validity Rv); the rewards go through GRPO group computation to produce advantages that update the student.

Figure 3. Overview of the MENTOR training framework. A trajectory consists of Reasoning (R), Tool-Use (T) and Observation (O) steps and a final Answer (A). Each student rollout is scored against the teacher’s reference trajectory by three reward functions. The first rollout uses the teacher’s tools in a different order and still receives all three rewards; a rollout with a failed tool call loses R_v, one using a different tool set loses R_a, and one that uses a different tool set and gives a wrong answer loses both R_a and R_c.

Small models trained with outcome-only rewards learn to avoid tools, while strict trajectory-matching rewards (as in ToolRL) make them call tools but steadily accumulate invalid calls. MENTOR (Model Enhancement by Non-rigid Teacher-Optimized Rewards) sits between these extremes: it uses the teacher’s trajectory as a reference for which tools to use, not as a sequence to copy.

Teacher references. A teacher (Qwen3-235B-Thinking) solves each training question, producing a trajectory $\tau^{(t)} = \langle (r_1, t_1, o_1), \ldots, (r_{L_\tau}, t_{L_\tau}, o_{L_\tau}) \rangle$ of reasoning, tool-use and observation steps and an answer $\hat{y}^{(t)}$ . Only trajectories whose answer matches the ground truth are kept.
On-policy rollouts. The student generates a group of rollouts $O^{(s)}$ for the same question, executing its Python tool calls in a sandbox.
Composite reward. Each rollout is scored with $R(O^{(s)}, O^{(t)}) = w_c R_c + w_a R_a + w_v R_v$ where R_c = 1 if the final answer matches the teacher’s answer; the teacher-alignment reward R_a = 1 if the student used exactly the same set of tools as the teacher, $\{t \mid t \in \tau^{(s)}\} = \{t \mid t \in \tau^{(t)}\}$ , regardless of order or count; and the tool-validation reward R_v = 1 only if every tool call in the student’s trajectory executes without error.
Policy update. The student is optimized with GRPO: rewards are normalized within each rollout group into advantages that update the student policy.

Training uses mathematical reasoning (AceReason-Math) as a controlled domain. Students are Qwen2.5 (1.5B, 7B) and Qwen3 (1.7B, 8B). Baselines are the vanilla model, SFT on teacher trajectories, outcome-based RL, and Strict RL (a re-implementation of ToolRL’s trajectory-matching reward) in the same RL framework. In-domain evaluation uses six math benchmarks (Math-Forge, Omni-MATH, AIME24, AIME25, AMC23, MinervaMath); out-of-domain evaluation uses tools unseen in training: the BFCL-v4 function-calling benchmark and three retrieval QA benchmarks (Bamboogle, 2WikiMultiHopQA, HotpotQA). An alignment score (AS) based on the Jensen-Shannon divergence measures how close a model’s tool-usage distribution is to the teacher’s.

Main Results

Best in-domain accuracy at every scale. MENTOR reaches overall math accuracy of 48.26% on Qwen3-8B (SFT: 41.47%, Strict RL: 45.77%, vanilla: 39.12%; the 235B teacher: 52.06%), and 9.89% on Qwen2.5-1.5B (SFT: 3.17%, Strict RL: 6.95%).
Better out-of-domain tool use. On retrieval QA with unseen tools, MENTOR scores 20.73 EM on Qwen2.5-7B (SFT: 14.57, Strict RL: 19.10) and 35.44 on Qwen3-8B (SFT: 34.14, Strict RL: 34.66). On BFCL-v4 it reaches 43.39 on Qwen3-8B (SFT: 42.16, Strict RL: 42.95).
SFT barely transfers. SFT gains little out of domain (e.g., BFCL-v4 overall on Qwen2.5-7B: 26.38 → 26.50), while MENTOR reaches 33.52.
Alignment tracks performance. In the reward ablation, the full MENTOR reward (Setting 6) gives the best Math (27.95), BFCL (31.52), RAG (21.08) and alignment score (88.63); across settings, performance and alignment score correlate with Pearson’s r = 0.85.
Fewer invalid calls. Configurations that include the tool-validation reward R_v drive the invalid tool-call rate to near zero during training, while those without it stay unstable.

**Table 1.** Overall scores: in-domain Math accuracy (%, macro-average of six benchmarks) and out-of-domain BFCL-v4 accuracy (%, official weighted average) and RAG exact match (%, macro-average of Bamboogle, 2WikiMultiHopQA, HotpotQA). Mean over 10 runs. Bold marks the best student per model as in the paper. (Subset of Table 1 in the paper: only the Overall columns; per-benchmark scores with standard deviations are in the paper.)
Model	Method	Math (Acc) ↑	BFCL-v4 (Acc) ↑	RAG (EM) ↑
Qwen3-235B (teacher)	Vanilla	52.06	50.13	39.57
Qwen2.5-1.5B	Vanilla	1.70	24.00	1.66
	SFT	3.17	24.59	2.43
	Outcome-based RL	6.69	25.12	7.59
	Strict RL	6.95	24.85	6.39
	MENTOR (ours)	9.89	25.70	8.39
Qwen2.5-7B	Vanilla	21.32	26.38	14.05
	SFT	22.48	26.50	14.57
	Outcome-based RL	26.13	32.94	18.89
	Strict RL	26.64	33.15	19.10
	MENTOR (ours)	28.63	33.52	20.73
Qwen3-1.7B	Vanilla	33.37	29.81	16.53
	SFT	36.34	30.13	17.45
	Outcome-based RL	38.62	30.65	18.18
	Strict RL	38.94	30.95	18.42
	MENTOR (ours)	40.90	31.20	19.23
Qwen3-8B	Vanilla	39.12	41.35	31.83
	SFT	41.47	42.16	34.14
	Outcome-based RL	45.19	42.78	34.81
	Strict RL	45.77	42.95	34.66
	MENTOR (ours)	48.26	43.39	35.44

@inproceedings{choi-etal-2026-mentor, title = "{MENTOR}: Reinforcement Learning via Flexible Teacher-Optimized Rewards for Tool-Use Distillation", author = "Choi, ChangSu and Song, Hoyun and Kim, Dongyeon and Cho, Minkyung and Jung, WooHyeon and Park, SunJin and Bae, NohHyeob and Yu, Seona and Lim, KyungTae", editor = "Lee, Hung-Yi and Dev, Sunipa and Li, Junyi Jessy and Elliott, Desmond", booktitle = "Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing", month = oct, year = "2026", address = "Budapest, Hungary", publisher = "Association for Computational Linguistics", url = "https://aclanthology.org/2026.emnlp-main.1286/", doi = "10.18653/2026.emnlp-main.1286", pages = "29289--29319", ISBN = "979-8-89176-459-0" }

MENTOR: Reinforcement Learning via Flexible Teacher-Optimized Rewards for Tool-Use Distillation

Abstract

Method

Main Results

BibTeX