Graph-Enhanced Policy Optimization in LLM Agent Training
About
Multi-step LLM agents in interactive environments represent a crucial step toward long-horizon decision-making. To train such agents, group-based reinforcement learning is widely adopted, which reinforces trajectories with higher relative performance within the group. However, in most existing methods, every step within a trajectory and every trajectory with the same terminal reward receive identical credit, regardless of their actual contributions. Since different states play different structural roles in an online state-transition graph built from sampled trajectories, their impacts should be differentiated and converted into task-aware credit at both the step and trajectory levels. We therefore present Graph-Enhanced Policy Optimization (GEPO), a framework for dual-level structural credit assignment in multi-step LLM agent training. Specifically, GEPO derives a state-level Task-Conditioned Criticality score that combines topological betweenness on the state-transition graph with semantic similarity to the task prompt. Based on this score, trajectory-level credit is reshaped through a state-adaptive discount, while step-level credit is scaled by the criticality of its successor state. Experimental results show that GEPO outperforms the strongest baselines by 1.1\% in success rate on ALFWorld, 3.2\% on WebShop, and 3.8\% on average across search-augmented QA tasks at the 7B scale. Compared with flat group-based methods, GEPO reduces across-seed variance and concentrates gradient signals on the most critical steps.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Question Answering | Search-QA | Average Score51 | 177 | |
| Interactive Decision-making | ALFWorld (test) | Success Rate94.9 | 81 | |
| Interactive Decision-making | WebShop (test) | Success Rate80.5 | 37 | |
| Single-hop Question Answering | PopQA out-of-domain | Accuracy52.4 | 24 | |
| Multi-hop Question Answering | HotpotQA in-domain | Accuracy45.1 | 22 | |
| Single-hop Question Answering | NQ (in-domain) | Accuracy48.2 | 19 | |
| Single-hop Question Answering | TriviaQA (out-of-domain) | Accuracy68 | 19 | |
| Multi-Hop Search-Augmented QA | 2WIKI (out-of-domain) | Accuracy47.6 | 12 | |
| Multi-Hop Search-Augmented QA | MuSiQue (out-of-domain) | Accuracy24 | 12 | |
| Multi-Hop Search-Augmented QA | Bamboogle (out-of-domain) | Accuracy72 | 12 |