Prefix tree method improves training of language agents for complex tasks
PTGPO: Prefix Tree Group Policy Optimization
Machine LearningArtificial Intelligence
Summary
Training AI agents to perform long, step-by-step tasks is tough because rewards only come at the end, giving little guidance on intermediate actions. The authors created a new approach called PTGPO that uses a tree structure to better assign credit to each step in a task without needing extra data or a separate judge. This method helps agents learn more effectively by understanding the impact of each action along the way. Tests on several interactive task benchmarks showed it beats existing methods.
What this means in practice
- •For ai software developers: Improve training of AI systems that handle complex text-based interactive tasks by providing more precise feedback on intermediate steps.
- •For robotics engineers: Enhance learning efficiency for robots facing long sequences of decisions by using fine-grained credit assignment without extra environment trials.
Authors
Ruolin Chen, Yang Li, Jihang Wang, Zhuoya Zhao, Yuchen Liu, Jiaming Xu, Qian Zhang
Abstract
Training language agents on long-horizon interactive tasks is challenging because sparse terminal rewards provide weak credit signals for intermediate actions. Existing group-based methods either assign a shared trajectory-level advantage or estimate step-level advantages from state groups, which may conflate the current action with subsequent decisions. We propose Prefix Tree Group Policy Optimization (PTGPO), a critic-free method that constructs prefix trees over trajectory suffixes within anchor-state groups. PTGPO assigns each tree edge a credit given by the difference between its child and parent empirical mean returns, capturing the conditional return change revealed by each additional action. These prefix-tree credits provide fine-grained learning signals without requiring a critic or additional rollouts. Experiments on interactive benchmarks ALFWorld, WebShop and Sokoban show that PTGPO consistently improves over RL baselines, demonstrating its effectiveness.