Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

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.

Jiazhen Yuan, Zhike Gong, Jinquan Hang, Zhengbiao Bai, Wei Zhao• 2025

Related benchmarks

TaskDatasetResultRank
Question AnsweringSearch-QA
Average Score51
177
Interactive Decision-makingALFWorld (test)
Success Rate94.9
81
Interactive Decision-makingWebShop (test)
Success Rate80.5
37
Single-hop Question AnsweringPopQA out-of-domain
Accuracy52.4
24
Multi-hop Question AnsweringHotpotQA in-domain
Accuracy45.1
22
Single-hop Question AnsweringNQ (in-domain)
Accuracy48.2
19
Single-hop Question AnsweringTriviaQA (out-of-domain)
Accuracy68
19
Multi-Hop Search-Augmented QA2WIKI (out-of-domain)
Accuracy47.6
12
Multi-Hop Search-Augmented QAMuSiQue (out-of-domain)
Accuracy24
12
Multi-Hop Search-Augmented QABamboogle (out-of-domain)
Accuracy72
12
Showing 10 of 10 rows

Other info

Follow for update