Reason in Chains, Learn in Trees: Self-Rectification and Grafting for Multi-turn Agent Policy Optimization
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
| Type | Method | ALFWorld (Success Rate %) | WebShop | |||||||
| 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 |
| Type | Method | Single-Hop QA (EM) | Multi-Hop QA (EM) | Avg. | ||||||
| NQ | Trivia | PopQA | Avg. | Hotpot | 2Wiki | MusiQ | Bamb | |||
| 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 |
| Type | Method | Sokoban (Success %) | Blocksworld | |||
| 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 |
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
| Component | Time Complexity | Dominant Cost |
|---|---|---|
| Standard GRPO | Forward/backward passes | |
| T-STAR Additional Overhead | ||
| KL Estimation | Forward passes (inference) | |
| Graph Merging | CPU operations | |
| Q-tree Backup | Arithmetic operations | |
| Thought Grafting | Forward passes (generation) | |
| Method | Time/Iter (s) | Total Time (h) | Relative |
|---|---|---|---|
| GRPO | 42.3 | 1.88 | 1.00 |
| GRPO + T-STAR | 51.8 | 2.30 | 1.22 |
| T-STAR Breakdown | |||
| Tree Construction | 5.2 | – | – |
| Q-tree + Divergence | 0.8 | – | – |
| Thought Grafting | 3.5 | – | – |
| 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% |
| Category | Dataset | Train | Test | Steps |
| Embodied | ALFWorld | 3,321 | 140 | 12.4 |
| 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 |
| Param | Val | Param | Val |
|---|---|---|---|
| KL threshold | 0.25 | Learning rate | 5e-6 |
| MC samples | 16 | Batch size | 32 |
| Trajectories | 8 | Training steps | 160 |
| Discount | 0.99 | EMA coef. | 0.95 |
| Divergence | 0.3 | Max seq. len | 4096 |
| Surgical wt. | 0.15 | Max steps | 20 |
| 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 |
| Dataset | Depth | Nodes | Merge | |
|---|---|---|---|---|
| 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 |
| Dataset | Valid | Error | Success |
|---|---|---|---|
| ALFWorld | 91% | 83% | 72% |
| WebShop | 88% | 79% | 68% |
| HotpotQA | 85% | 76% | 65% |
| Sokoban | 82% | 71% | 58% |
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