Interactive language agents improved by tree-structured training method
TSBPO: Tree-Structured Branching for Policy Optimization in Interactive Language Agents
Artificial IntelligenceComputation and LanguageMachine Learning
Summary
Training language-based computer agents that interact with environments is challenging because it's hard to know which actions lead to success and how to assign credit to each step. The authors introduce TSBPO, a new way to train these agents by organizing possible action sequences in a tree and focusing updates on parts that lead to different outcomes. Their approach helps the agents learn better from complete interactions and improves success in tasks like manipulating objects in simulations and shopping goal completion. This method adjusts how it treats different parts of responses during training to avoid bias from varying response lengths.
What this means in practice
- •For interactive ai developers: Train more effective language agents that complete tasks in simulated environments using tree-structured sampling to improve learning from delayed outcomes.
- •For automated web service teams: Improve language-based agents that perform goal-directed tasks on websites by more accurately assigning credit for multi-step actions during training.
Authors
Rui Yu, Chong Mu, Jinchuan Zhang
Abstract
Training interactive language agents from terminal outcomes poses three coupled problems: which task and state continuations to sample, how to assign delayed feedback to decisions, and how to aggregate token losses for multi-token actions. We present TSBPO, a tree-structured sampling-to-update paradigm that addresses these problems within one training loop. A likelihood-fitted Root/Branch allocator uses retained task and state evidence to target continuations likely to produce contrasting terminal outcomes under fixed rollout budgets. On the resulting replayed physical trees, trajectory advantages compare immediate sibling segments, while observation-group step advantages compare outcomes at matched task/observation states. The actor then sums clipped token surrogates over each complete response and normalizes across active turns, avoiding response-length averaging in the policy-gradient term. With a Qwen2.5-1.5B actor, fixed checkpoints obtain 97.79% seen and 96.64% unseen stochastic success on ALFWorld and an 81.07% per-episode success rate over six stochastic episodes for each of the 500 WebShop test goals. A WebShop step-term ablation yields 75.93% at the same update-100 evaluation; an ALFWorld aggregation comparison shows higher observed command validity with complete-turn weighting. These results support the evaluated training configuration on two interactive benchmarks.