GraphOPD: Graph-Augmented On-Policy Distillation for LLM Agents

Bohan Lin; Liyi Chen; Zhuoning Guo; Muyang Li; Qimeng Wang; Yan Gao; Yao Hu; Yudong Zhang

cs.LGOct 6, 2026

GraphOPD: Graph-Augmented On-Policy Distillation for LLM Agents

Authors: Bohan Lin, Liyi Chen, Zhuoning Guo, Muyang Li, Qimeng Wang, Yan Gao, Yao Hu, Yudong Zhang

Organizations: University of Science and Technology of China · Xiaohongshu · The Hong Kong University of Science and Technology (HKUST-GZ)

Abstract

On-policy distillation post-trains large language model agents by supplying dense, step-level guidance from a teacher policy when the reinforcement-learning reward is sparse and arrives only once per trajectory. Existing instantiations allocate this guidance by the size of the teacher-student divergence at each step, on the single-turn intuition that a large disagreement marks a mistake worth correcting. Once decisions chain over many turns, that rule misfires, since an early drift enters every later context both policies condition on, leaving the teacher consistent with the drifted trajectory instead of flagging its cause, while interchangeable steps register large but outcome-irrelevant divergences. We demonstrate this on an agentic benchmark, where distilling the highest-divergence steps brings no consistent benefit over random selection. To this end, we introduce GraphOPD, the first method to bring graph-based structural augmentation into on-policy distillation for agent capabilities. It reads which steps enabled which later ones from the environment's own record of state changes, immune to the drift that corrupts the teacher-student gap, organizes them into a dependency graph, scores each step by a random-walk stationary distribution over it, and fuses that structural credit with the divergence signal into a trajectory-relative mask concentrating supervision on each rollout's highest-aptitude steps. Across three model scales and eleven baselines on ALFWorld, WebShop, and SearchQA, GraphOPD shows competitive performance throughout, improving over the strongest baseline by up to +5.8 pp. An executed-replay audit further shows that this structural credit score tracks true causal impact far above chance, that both fused signals are independently necessary, and that the same signal transfers to out-of-domain tool-integrated reasoning.

Figures & tables

Figure 1: Motivating empirical study. All runs distill only the $k$ highest-KL steps or $k$ randomly chosen steps per trajectory, $k\in\{5,8,10,12\}$ , on ALFWorld with a 7B Qwen2.5-Instruct agent under GRPO for 150 steps, all other conditions identical. ( Left ) Random selection stays ahead at every budget $k$ , by $+1.5$ to $+3.2$ pp in success rate, with the per-budget gaps $\Delta$ annotated above each pair. ( Middle ) At $k=5$ the KL-guided run trails in success rate (solid curves, left axis) throughout and ends at 76.3% against 83.5% for random, while its policy entropy (dashed curves, right axis) contracts as the random run’s expands. ( Right ) The per-step KL of 466 successful and 466 failed steps from 50 rollouts overlaps by 84% between the two outcomes, with medians (dashed) at 0.10 and 0.09, so success and failure are barely separable by divergence.

Figure 2: GraphOPD overview. From a single rollout, two per-step signals are computed in parallel, a structural credit signal from the Footprint-built dependency graph and its random-walk stationary distribution, and a divergence signal from the teacher-student gap normalized within the trajectory. The two are fused into a distillation aptitude $A(t)$ , converted into a trajectory-relative mask $m(t)$ , and the resulting masked distillation loss is added to the unweighted GRPO objective; every stage reads only the rollout itself and needs no extra model calls.

Qwen2.5-3B-Instruct
	ALFWorld							SearchQA								WebShop
Method	Pick	Look	Clean	Heat	Cool	Pick2	Avg	NarQA	Triv	Pop	Hotp	2Wk	MuS	Bam	Avg	Score	Succ.
Vanilla	44.4	11.1	6.2	15.4	28.6	12.5	21.9	24.6	48.1	31.0	26.3	25.3	7.2	59.7	31.7	6.7	0.8
GRPO	91.2	62.5	96.2	61.9	65.0	47.4	75.0	39.3	60.6	41.1	37.4	34.6	15.4	26.4	36.4	79.8	63.3
Skill-Prompt	51.7	66.7	48.4	0.0	4.3	10.0	28.9	23.7	46.2	30.6	24.4	22.1	7.5	12.5	23.9	0.2	0.8
OPSD	48.8	41.7	16.7	0.0	15.8	16.7	28.1	0.1	0.1	0.1	0.0	0.0	0.0	0.0	0.0	11.3	3.1

Table 1: Main results. Success rate (ALFWorld), accuracy (SearchQA), and success rate (WebShop) across three model scales. Bold and underline marks the best and the second best result per column within each model block.

	Qwen3-1.7B			Qwen2.5-3B			Qwen2.5-7B
Method	WebShop	SearchQA	ALFWorld	WebShop	SearchQA	ALFWorld	WebShop	SearchQA	ALFWorld
GraphOPD	73.4	42.9	63.3	78.9	46.1	88.6	83.6	47.3	91.9
Soft-Pruning	71.1	37.4	62.5	77.3	44.2	85.9	82.0	45.3	90.6
Divergence-only	64.8	32.4	59.4	74.2	40.8	84.4	79.6	42.3	87.5
Structural-only	69.5	36.4	60.2	75.0	43.3	84.4	81.3	43.6	89.1
Shuffled Centrality	68.8	40.1	59.4	72.7	41.2	82.8	78.1	41.9	85.9

Table 2: Ablation study. Success rate (WebShop), accuracy (SearchQA), and success rate (ALFWorld). Variant definitions are given in the text above, and the GraphOPD row reproduces the three-run means of the full method, while each ablation variant is a single-run point estimate. Bold marks the best result per column within each model block.

$\alpha$	0	0.2	0.5	0.7	1.0
ALFWorld	89.1	90.6	91.9	87.5	87.5
WebShop	81.3	79.7	83.6	76.6	79.6

Table 3: ALFWorld and WebShop success rate in percent for Qwen2.5-7B-Instruct across the signal-mixing coefficient

\alpha

Figure 3: Analysis. ( Left ) Average graph nodes, equivalently trajectory length, average edges, and average density of the dependency graphs built on ALFWorld at 7B over the 150 training steps; nodes and edges fall from about 47 and 39 to a plateau of 25 to 30 while density rises from about 0.03 to between 0.10 and 0.14, so trajectories grow shorter and denser. ( Middle ) Accuracy radar on the four out-of-domain tool-integrated-reasoning benchmarks, AIME24 and AIME25 under avg@32 and LiveCodeBench-v5 and GPQA-Diamond under avg@8, plus their average, for GraphOPD against Vanilla, GRPO, OPSD, and SDAR; the GraphOPD polygon dominates every axis, averaging 29.0% against 20.0% for SDAR. Figure 3 (right) makes the mechanism concrete on a fragment of one real successful ALFWorld trajectory, where raw KL peaks at a cabinet revisit that never decides the outcome while $\ell_{\text{graph}}$ peaks at the decisive object-placement step within this fragment.

Signal	$\rho$	Hit@10%
$\ell_{\text{graph}}$	0.58	0.48
$\ell_{\text{model}}$	0.34	0.41
Random	$-0.06$	0.13

Table 4: Credit-content audit on ALFWorld at 7B. Metrics defined in the text.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Environment	Key action types	Footprint key format
ALFWorld	go to, take, put, heat/cool/clean, open/close	location:<loc> , visible_at:<loc> , holding:<obj> , placed:<obj>@<loc> , object_state:<obj>:<prop> , container_state:<container>
WebShop	search[q], click[product_id], click[attr], click[buy]	query_token:<word> , one per query word, selected_product:<id> , selected_attr:<attr> , purchase_action
SearchQA	<search> , <think> , <answer>	query_token:<word> , retrieved_passages , reasoning_state , answer_candidate , passage_entity:<entity>

Appendix

Table 5: Footprint extractor specifications by environment.

Benchmark	Graph construction (s/step)	Total step time (s)	Share
ALFWorld	0.6659	342.0936	0.19%
WebShop	0.4535	144.2885	0.31%
SearchQA	0.4144	322.2142	0.13%

Appendix

Table 6: Wall-clock cost of graph construction relative to the total per-training-step time.

	Qwen3-1.7B			Qwen2.5-3B			Qwen2.5-7B
Method	ALFWorld	SearchQA	WebShop	ALFWorld	SearchQA	WebShop	ALFWorld	SearchQA	WebShop
GraphOPD	63.3 $\pm\,1.18$	42.9 $\pm\,0.05$	73.4 $\pm\,1.60$	88.6 $\pm\,0.57$	46.1 $\pm\,0.02$	78.9 $\pm\,0.43$	91.9 $\pm\,0.14$	47.3 $\pm\,0.40$	83.6 $\pm\,0.43$
SDAR	51.3 $\pm\,5.09$	40.7 $\pm\,0.72$	59.9 $\pm\,0.93$	81.8 $\pm\,3.81$	43.4 $\pm\,0.03$	69.3 $\pm\,0.91$	84.4 $\pm\,1.60$	47.5 $\pm\,1.21$	82.0 $\pm\,0.38$
OPID	58.2 $\pm\,1.85$	39.6 $\pm\,0.78$	66.4 $\pm\,2.99$	82.8 $\pm\,1.50$	44.6 $\pm\,0.11$	74.5 $\pm\,0.14$	89.1 $\pm\,0.48$	48.2 $\pm\,0.74$	80.0 $\pm\,1.00$
GEPO	58.1 $\pm\,2.58$	40.1 $\pm\,0.38$	67.9 $\pm\,1.20$	80.7 $\pm\,3.84$	44.6 $\pm\,1.91$	75.3 $\pm\,1.35$	85.9 $\pm\,5.02$	45.5 $\pm\,2.85$	80.0 $\pm\,1.85$

Appendix

Table 7: Mean and run-to-run standard deviation (pp) over three independent runs for the four multi-run methods of the main results, on ALFWorld and WebShop success rate and SearchQA accuracy.

Source within the turn	Footprint key format
Code block and reasoning text	symbol:<identifier> , one per name bound by an assignment, a function definition, or a comprehension, call:<function> , one per interpreter function invoked, import:<module> , one per module imported, answer_candidate:<value> , for the value placed inside \boxed{} or submitted as the final program
Interpreter response	result:<value> , one per value printed or returned, test_result:<example_id> , one per public example checked on LiveCodeBench-v5

Appendix

Table 8: Footprint extractor specification for the four out-of-domain tool-integrated-reasoning benchmarks.

Method	AIME24	AIME25	LCB-v5	GPQA-D	Avg
Vanilla	11.9	8.8	14.2	21.7	14.2
GRPO	15.1	13.3	15.1	28.3	18.0
OPSD	19.2	12.2	16.4	31.5	19.8
SDAR	21.4	16.3	15.8	26.3	20.0
GraphOPD	29.8	21.4	28.2	36.4	29.0

Appendix

Table 9: Task accuracy (%) on four out-of-domain tool-integrated-reasoning benchmarks, with the scoring protocol per benchmark given in the text above.

Steps	Late checkpoint, step 150 (succeeds)	Early checkpoint, step 20 (fails)
0–2	Action: go to sinkbasin 1 Action: go to countertop 1 Action: take soapbottle 1 from countertop 1	Action: go to cabinet 1 Action: open cabinet 1 Action: examine cabinet 2
3–5 / 3–49	Action: go to cabinet 1 Action: open cabinet 1 Action: move soapbottle 1 to cabinet 1 Result: reward $=10$ , episode ends at step 5	Actions (47 steps): examine cabinet $\times 25$ , look $\times 6$ , open cabinet $\times 5$ , go to cabinet $\times 3$ , plus 8 other revisit actions, never mentioning the soapbottle Raw KL: ranges 0.16 to 0.65, mean 0.32 Result: reward $=0$ , times out at step 49
	The step-150 checkpoint carries the soapbottle from step 2 onward, so every later action follows directly from the object it already holds.	The step-20 checkpoint repeatedly opens and examines cabinets without ever picking up the soapbottle, so its failure is an entity it never engaged with rather than a late misstep.

Appendix

Table 10: A side-by-side comparison of two rollouts in ALFWorld. Both checkpoints attempt the same task and environment seed (”put some soapbottle on cabinet”, seed=1000 ) from the same training run. Their strategies differ from the first action, so the early checkpoint never engages the target object at all.

Step	Action	Raw KL	$\ell_{\text{graph}}(t)$
3	examine countertop 1	0.33	$-0.90$
4	go to cabinet 1	0.36	$-0.28$
5	open cabinet 1	0.29	$-0.57$
6	go to countertop 1	0.26	$-0.90$
7	examine countertop 1	0.17	$-0.28$
8	go to cabinet 1	0.29	0.31

Appendix

Table 11: Qualitative case study on a window (steps 3 to 10) of a real successful ALFWorld trajectory (”put some soapbottle on cabinet”). Raw KL peaks at step 4, a cabinet-revisit step that does not decide the outcome, while

\ell_{\text{graph}}(t)

peaks at step 10, the object-placement action decisive within this window. Bold marks each column’s peak.

Step	Action	Raw KL	$\ell_{\text{graph}}(t)$
0	go to fridge 1	0.31	$-0.75$
1	examine fridge 1	0.44	$-1.12$
2	go to countertop 1	0.28	$-0.68$
3	take potato 1 from countertop 1	0.22	$-0.82$
4	go to microwave 1	0.24	0.95
5	open microwave 1	0.26	$-0.04$

Appendix

Table 12: Qualitative case study on a complete representative successful Heat-type episode (twelve steps; actions and observations follow the ALFWorld grammar and Footprint conventions of Appendix A ). The structural credit peaks at step 7, the processing-location revisit four steps before the terminal placement, with the earlier processing-location approach (step 4) and the heating action (step 6) also well above neutral; the terminal placement (step 11) stays near neutral. Raw teacher-student KL is reported for contrast: its maximum (step 8) coincides with the structural-credit minimum. Bold marks each column’s peak.

Explore similar work

Jul 6, 2026cs.LG

Multi-Turn On-Policy Distillation with Prefix Replay

We study on-policy distillation (OPD) for agentic tasks, where an LLM agent interacts with an environment over multiple turns and a student imitates a teacher over these multi-turn interaction histories. Fully online OPD is costly because each update requires fresh student rollouts through the environment and teacher queries at visited histories. We propose Replayed-Prefix On-Policy Distillation (ReOPD), an off-environment alternative that reuses pre-collected teacher trajectories as replayed prefixes: the student acts at selected steps, while the teacher provides dense per-step supervision without executing new environment interactions. We show that multi-turn OPD introduces a prefix trap: making histories more student-on-policy improves relevance to the student, but can query the teacher on histories where its target is unreliable. This creates a two-sided distribution shift between student occupancy and teacher reliability. ReOPD addresses this by treating multi-turn OPD as a reliability-aware prefix distribution design and implements it with a simple step-decaying sampling schedule that emphasizes early, lower-shift prefixes. Across mathematical reasoning with Python and search environments over multiple teacher and student model scales, ReOPD preserves or improves OPD-level accuracy, uses zero tool calls during student training, and is at least 4

\times

faster per rollout than OPD. ReOPD therefore turns expensive agent-environment interaction into a reusable offline resource, enabling scalable distillation across tools, tasks, and environments.

Baohao Liao, Hanze Dong, Christof Monz +3

1Microsoft Research · University of Amsterdam

Efficient On-Policy Distillation Teacher Interaction History

Ask PDF

May 26, 2026cs.AI

StepOPSD: Step-Aware Online Preference Distillation for Agent Reinforcement Learning

Reinforcement learning for multi-turn agents suffers from a credit-assignment mismatch: rewards are sparse and trajectory-level, while success often hinges on a few local decisions. Existing online policy distillation (OPD) provides denser token-level supervision, but typically treats heterogeneous agent trajectories as monolithic strings rather than causal interaction units. We present StepOPSD, a post-rollout preference self-distillation framework that takes the agent step as the unit of credit redistribution. StepOPSD decomposes trajectories into action-centered step segments, rescoring them under hindsight-enriched teacher contexts and converting token-level log-probability gaps into sign-preserving advantage shaping with a normalized per-step credit budget before the GRPO update. Across ALFWorld and Search-QA with Qwen3-1.7B and Qwen2.5-3B-Instruct, StepOPSD attains best or second-best results on subsets most sensitive to local causal errors, including first-place performance on ALFWorld Heat (79.1%), PickTwo (95.0%), Search-QA TriviaQA (61.6%), and tied-best performance on HotpotQA (40.4%). The results further reveal a consistent two-knob law: smaller α_clip acts as a broadly stabilizing local trust region, whereas the optimal global mixing strength λ_mix remains task-dependent. These findings suggest that step-aware distillation is most useful when trajectory-level rewards are weakly aligned with the local action that determines downstream success.

Yanfei Zhang, Xu Lin, Chenglin Wu

Independent Researcher · Tencent · DeepWisdom

Online-Policy Distillation Rp-Opsd Dataset Distillation Preference Alignment Learning

Ask PDF

May 31, 2026cs.LG

Trust Region On-Policy Distillation

On-Policy Distillation (OPD) is a fundamental technique for efficient post-training of large language models (LLMs), with broad applications in agent learning, multi-task enhancement, and model compression. However, OPD training becomes unstable when the teacher and student distributions differ substantially, as teacher supervision on student-generated tokens may yield unreliable policy gradients and even cause optimization failure. This work addresses reliable on-policy token-level supervision through credit assignment strategies, and proposes Trust Region On-Policy Distillation, TrOPD. It features the following characteristics: 1) Trust-Region On-Policy Learning: TrOPD performs OPD only in regions where the teacher provides reliable supervision, mitigating the optimization difficulty of the K1 reverse-KL estimator under distribution mismatch. 2) Outlier Estimation: For outlier regions, we explore gradient clipping, masking, and forward-KL estimation to reduce the adverse effects of unreliable supervision. 3) Off-Policy Guidance: The student continues generation from teacher prefixes and uses forward KL to imitate off-policy guidance, encouraging on-policy exploration toward reliable regions. Experiments show that TrOPD consistently outperforms SoTA OPD baselines, including OPD, EOPD, and REOPOLD, across mathematical reasoning, code generation, and general-domain benchmarks.

Xingrun Xing, Haoqing Wang, Boyan Gao +2

Samsung Research, Beijing, China · University of Oxford · Peking University

Online-Policy Distillation Token-Level Supervision Kullback-Leibler Divergence

Ask PDF