Reason in Chains, Learn in Trees: Self-Rectification and Grafting for Multi-turn Agent Policy Optimization

Yu Li; Sizhe Tang; Tian Lan

cs.AIApr 8, 2026

Reason in Chains, Learn in Trees: Self-Rectification and Grafting for Multi-turn Agent Policy Optimization

Authors: Yu Li, Sizhe Tang, Tian Lan

Organizations: Department of Electrical and Computer Engineering, George Washington University

Abstract

Reinforcement learning for Large Language Model agents is often hindered by sparse rewards in multi-step reasoning tasks. Existing approaches like Group Relative Policy Optimization treat sampled trajectories as independent chains, assigning uniform credit to all steps in each chain and ignoring the existence of critical steps that may disproportionally impact reasoning outcome. In this paper, we propose T-STAR(Tree-structured Self-Taught Agent Rectification), a framework that recovers the latent correlated reward structure across seemingly independent trajectories. Specifically, we consolidate trajectories into a unified Cognitive Tree by identifying and merging functionally similar steps/nodes. It enables an Introspective Valuation mechanism that back-propagates trajectory-level rewards through the tree to obtain a new notion of variance-reduced relative advantage at step-level. Using the Cognitive Tree, we also develop In-Context Thought Grafting to synthesize corrective reasoning by contrasting successful and failed branches at critical divergence points/steps. Our proposed Surgical Policy Optimization then capitalizes on the rich policy gradient information concentrated at these critical points/steps through a Bradley-Terry type of surgical loss. Extensive experiments across embodied, interactive, reasoning, and planning benchmarks demonstrate that T-STAR achieves consistent improvements over strong baselines, with gains most pronounced on tasks requiring extended reasoning chains.

Figures & tables

Figure 1: An illustration using WebShop Operation. Three trajectories share identical prefixes but diverge at the attribute verification step, yielding different outcomes. Existing methods with trajectory-level rewards assign inconsistent credit for the prefixes in each trajectory, while critical decision steps fail to be distinguished and utilized for policy optimization.

Figure 2: Overview of T-STAR . Given $M$ sampled trajectories per task, T-STAR consolidates them into a Cognitive Tree by merging functionally equivalent nodes, then computes structural values via Bellman backup that propagate downstream success rates to intermediate nodes. This enables trajectory stitching where successful reasoning steps receive appropriate credit even within failed rollouts. At divergence points where children exhibit large value differences, the agent performs self-rectification by generating corrective reasoning that transfers successful logic to failed contexts, producing preference pairs for surgical policy optimization.

Figure 3: Thought grafting mechanism. At divergence points identified by large value spreads among children, the agent observes contrasting outcomes between successful and failed branches, then generates rectified reasoning that transfers the successful logic to the failed context. The resulting preference pairs provide step-level supervision for surgical policy optimization.

Figure 4: Training dynamics of T-STAR on WebShop and Sokoban, showing success rate, grafting coverage and anchor reuse, and value spread

\Delta V

at divergence points across training iterations.

Type	Method	ALFWorld (Success Rate %)							WebShop
Type	Method	Pick	Look	Clean	Heat	Cool	Pick2	All	Score	Succ.
Closed-Source Model
Prompting	GPT-4o	86.8	74.4	79.1	90.9	93.6	83.1	84.6	84.2	61.7
Prompting	Gemini-1.5-Pro	81.6	68.1	72.9	84.8	88.6	77.7	79.0	79.3	54.7
Qwen2.5-3B-Instruct
Prompting	ReAct	48.2	40.9	44.8	51.7	55.4	38.6	46.6	50.7	25.3

Table 1: Performance on interactive and embodied tasks. Success rate (%) for ALFWorld subtasks and score/success for WebShop. T-STAR consistently outperforms across all baselines and architectures.

Figure 5: Training dynamics on (a) ALFWorld, (b) WebShop, and (c) Multi-hop QA. T-STAR achieves faster convergence and higher final performance. Shaded regions indicate standard deviation across three seeds.

Type	Method	Single-Hop QA (EM)				Multi-Hop QA (EM)				Avg.
Type	Method	NQ	Trivia	PopQA	Avg.	Hotpot	2Wiki	MusiQ	Bamb	Avg.
Closed-Source Model
Prompting	GPT-4o	59.5	73.0	61.5	64.7	54.0	50.5	47.0	44.5	55.8
Prompting	Gemini-1.5-Pro	56.0	69.5	57.5	61.0	50.0	46.5	43.0	40.5	52.0
Qwen2.5-3B-Instruct
Prompting	Direct	23.5	36.0	25.0	28.2	19.0	16.5	13.0	11.5	21.2

Table 2: Performance on search-augmented QA tasks (Exact Match). T-STAR shows largest gains on multi-hop reasoning. Best performance (highlighted in bold) is always achieved by T-STAR .

Type	Method	Sokoban (Success %)			Blocksworld
Type	Method	Easy	Med.	Hard	Stack	Unstack
Closed-Source Model
Prompting	GPT-4o	78.3	45.4	17.4	84.6	89.8
Prompting	Gemini-1.5-Pro	72.7	40.4	14.4	79.7	85.4
Qwen2.5-3B-Instruct
RL Training	PPO	16.9	4.7	0.7	24.1	32.4

Table 3: Performance on logical planning tasks. T-STAR enables recovery from dead-ends through thought grafting.

Figure 6: Ablation study on ALFWorld. Left: component contribution. Right: hyperparameter sensitivity (

\epsilon_{\text{kl}}

\lambda

\delta

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Component	Time Complexity	Dominant Cost
Standard GRPO	$O(MT\cdot L)$	Forward/backward passes
T-STAR Additional Overhead
KL Estimation	$O(M^{2}TK)$	Forward passes (inference)
Graph Merging	$O(M^{2}T)$	CPU operations
Q-tree Backup	$O(MT)$	Arithmetic operations
Thought Grafting	$O(p_{\text{div}}\cdot MT\cdot l_{z})$	Forward passes (generation)

Appendix

Table 4: Computational complexity breakdown.

M

: trajectories per task,

T

: average trajectory length,

K

: MC samples,

L

: context length,

l_{z}

: thought length,

p_{\text{div}}

: divergence fraction.

T-STAR Breakdown
Method	Time/Iter (s)	Total Time (h)	Relative
GRPO	42.3	1.88	1.00 $\times$
GRPO + T-STAR	51.8	2.30	1.22 $\times$
Tree Construction	5.2	–	–
Q-tree + Divergence	0.8	–	–
Thought Grafting	3.5	–	–

Appendix

Table 5: Runtime comparison on ALFWorld (160 training iterations).

Model Size	GRPO Time (s)	T-STAR Overhead (s)	Relative Overhead
1.5B	28.4	8.2	28.9%
3B	42.3	9.5	22.5%
7B	78.6	11.3	14.4%

Appendix

Table 6: Scaling behavior of T-STAR overhead with model size on ALFWorld.

Category	Dataset	Train	Test	Steps
Embodied	ALFWorld	3,321	140	12.4
Embodied	WebShop	10,587	500	8.6
Single-hop	NQ	79,168	3,610	3.2
	TriviaQA	78,785	11,313	2.8
	PopQA	–	14,267	2.5
Multi-hop	HotpotQA	90,447	7,405	6.4

Appendix

Table 7: Dataset statistics across four task categories.

Param	Val	Param	Val
KL threshold $\epsilon_{\text{kl}}$	0.25	Learning rate	5e-6
MC samples $K$	16	Batch size	32
Trajectories $M$	8	Training steps	160
Discount $\gamma$	0.99	EMA coef. $\alpha$	0.95
Divergence $\delta$	0.3	Max seq. len	4096
Surgical wt. $\lambda$	0.15	Max steps	20

Appendix

Table 8: Hyperparameter settings for T-STAR .

Method	Pk	Lk	Cl	Ht	Co	Pk2	All
GRPO	79.5	67.0	73.5	83.0	87.0	71.0	76.8
+ T-STAR	82.0	70.5	76.0	85.5	89.0	75.5	79.8
DAPO	83.0	71.0	77.0	86.0	90.0	75.0	80.3
+ T-STAR	85.5	74.0	80.5	88.5	92.0	79.0	83.3
GiGPO	89.0	77.0	84.0	92.0	95.0	82.0	86.5
+ T-STAR	91.5	80.5	87.0	94.5	97.0	85.5	89.3

Appendix

Table 9: Per-task results on ALFWorld (Qwen2.5-3B). Pk=Pick, Lk=Look, Cl=Clean, Ht=Heat, Co=Cool.

Dataset	Depth	Nodes	Merge	$\|V_{\text{div}}\|$
ALFWorld	12.4	47.2	0.52	5.8
WebShop	8.6	38.5	0.44	4.2
HotpotQA	6.4	28.3	0.45	3.6
Sokoban-E	12.8	52.1	0.58	6.4
Sokoban-H	24.6	89.4	0.41	8.2

Appendix

Table 10: Cognitive tree statistics. Merge ratio = 1 - (nodes after / nodes before merging).

Dataset	Valid	Error	Success
ALFWorld	91%	83%	72%
WebShop	88%	79%	68%
HotpotQA	85%	76%	65%
Sokoban	82%	71%	58%

Appendix

Table 11: Grafting quality. Valid: semantically correct; Error: identifies mistake; Success: leads to completion.

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

GraphPO: Graph-based Policy Optimization for Reasoning Models
Jun 17, 2026Yuliang Zhan, Xinyu Tang, Jian Li +7Reinforcement Learning LLM Reasoning with Graphs
Gaoling School of Artificial Intelligence, Renmin University of China · Ant Group
Segment-Aligned Policy Optimization for Multi-Modal Reasoning
May 2, 2026Lei Gao, Zhuoming Li, Mengxi Jia +4RL for Language Model Reasoning Policy Optimization
Fudan University · China Telecom Artificial Intelligence Technology (Beijing) Co., Ltd. · Southeast University +1
Contrastive Branch Policy Optimization
Aug 25, 2026Ying Wang, Changlin Qiu, Bang Lin +4RL for Language Model Reasoning Step-Level Credit Assignment in RL
Alibaba Group, Hangzhou, China · Harbin Institute of Technology, Harbin, China

Jun 17, 2026J/K move · Enter open · S save

GraphPO: Graph-based Policy Optimization for Reasoning Models

Yuliang Zhan, Xinyu Tang, Jian Li, Dandan Zheng+6

Gaoling School of Artificial Intelligence, Renmin University of China · Ant Group

Reinforcement Learning LLM Reasoning with Graphs RL for Language Model Reasoning Reinforcement Learning with Verifiable Rewards Process Supervision

Open paper AskPDF arXiv

Reinforcement Learning with Verifiable Rewards (RLVR) has become a standard paradigm for enhancing the capability of large reasoning models. RLVR typically samples responses independently and optimizes the policy using from final answers. This paradigm has two limitations. First, independently responses often contain similar intermediate reasoning steps, causing redundant exploration and wasted computation. Second, sparse final-answer rewards make it hard to identify useful steps. Tree-based methods partly address this problem by sharing prefixes and comparing branches from the same prefix to provide fine-grained signals. However, tree branches are still expanded independently. When different branches reach similar reasoning states, they cannot share information and repeat similar exploration. Moreover, tree-based methods ignore such dispersion and only perform local comparisons within separate branches, which can lead to higher variance in advantage estimation. To address this challenge, we propose GraphPO (Graph-based Policy Optimization), a novel RL framework that represents rollouts as a directed acyclic graph, with reasoning steps as edges and semantic states summarized from the reasoning paths as nodes. GraphPO merges semantically equivalent reasoning paths into equivalence classes, allowing them to share suffixes and reallocating budget away from redundant expansions to diverse exploration. Furthermore, we assign efficiency advantages to incoming edges and correctness advantages to outgoing edges, thereby improving inference efficiency while deriving process supervision from outcome. Theory shows that GraphPO reduces advantage-estimation variance and enhances reasoning efficiency. Experiments on three LLMs across reasoning and agentic search benchmarks show that GraphPO consistently outperforms chain- and tree-based baselines with the same token budgets or response budgets.

Abstract

RL for Language Model Reasoning

Explore similar work

GraphPO: Graph-based Policy Optimization for Reasoning Models

Segment-Aligned Policy Optimization for Multi-Modal Reasoning

Contrastive Branch Policy Optimization