TreeHCA improves multi-turn language agent training with better credit assignment
TreeHCA: Hindsight-Guided Exploration and Credit Assignment for Multi-Turn LLM Agents
Machine Learning
Summary
Teaching language models to solve problems over multiple steps is hard because it's unclear which decisions helped or hurt the final answer. The authors introduce TreeHCA, a method that revisits earlier decisions to explore and evaluate alternative reasoning paths. This creates a tree of possible outcomes, allowing the model to assign credit more accurately to each step. Their tests showed TreeHCA outperforms previous methods, leading to more precise answers.
What this means in practice
- •For llm developers: Train multi-turn language agents to assign feedback more precisely to each reasoning step, improving overall answer accuracy.
- •For automated customer support teams: Improve multi-step chatbot responses by enabling models to explore alternative dialogue paths and better evaluate which steps lead to successful outcomes.
Authors
Debajoy Mukherjee, Jeremy Carleton, Sushil Vemuri, Srinivas Shakkottai, Dileep Kalathil
Abstract
Outcome-supervised group-relative policy optimization (GRPO) assigns the same advantage to every token in an LLM agent trajectory, leaving useful and unhelpful decisions undifferentiated. Hindsight credit assignment (HCA) can refine this feedback, but its estimates rely on trajectories containing informative alternatives at consequential histories. We introduce TreeHCA, which couples hindsight-guided exploration with hindsight credit assignment. TreeHCA uses changes in reference-answer likelihood to revisit intermediate histories and sample alternative continuations, forming a tree of reasoning and tool interactions. Because this exploration changes the distribution of sampled continuations, we derive the corresponding inverse-hindsight importance correction, including an exact finite-pool identity. We use this structure in a normalized tree backup that propagates terminal outcomes into turn-level advantages. This shared hindsight structure is TreeHCA's central design principle, linking informative exploration with hindsight-guided credit assignment. Across seven question-answering benchmarks and three Qwen backbones, TreeHCA outperforms all baselines in aggregate exact match, validating the effectiveness of paired exploration and credit assignment for multi-turn LLM agent training. Our code is available at github.com/TheBlackCat22/TreeHCA.