📊 statistics

Maximizing Rollout Informativeness under a Fixed Budget: A Submodular View of Tree Search for Tool-Use Agentic Reinforcement Learning

This paper introduces InfoTree, a training-time tree-search framework for tool-use agentic reinforcement learning that formalizes rollout informativeness as a submodular maximization problem to derive an uncertainty-aware selection strategy (UUCB) and an adaptive budget allocator, thereby significantly outperforming existing methods across diverse reasoning and tool-use benchmarks while maintaining robustness and efficiency.

Original authors: Yuelin Hu, Zhenbo Yu, Zhengxue Cheng, Wei Liu, Li Song

Published 2026-05-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yuelin Hu, Zhenbo Yu, Zhengxue Cheng, Wei Liu, Li Song

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to solve complex puzzles (like math problems or coding tasks) by letting it practice over and over again. In the world of AI, this practice is called "rollouts." The robot tries to solve a problem, gets a reward if it's right, and a penalty if it's wrong. The goal is to learn from these attempts.

However, there's a big problem: The "Echo Chamber" Effect.

If you ask the robot to try the same hard puzzle 16 times, it might get the exact same wrong answer 16 times. Or, if it's an easy puzzle, it might get the exact same right answer 16 times. In both cases, the robot learns nothing new because there's no variety. It's like asking a student to take the same multiple-choice test 16 times; if they get it wrong every time, they don't learn why they were wrong, they just get frustrated.

This paper introduces a new method called INFOTREE to fix this. Here is how it works, using simple analogies:

1. The Problem: The "Boring Class"

The authors call this the "Collapse." If the robot's attempts are all identical, the training signal (the lesson) vanishes. They proved mathematically that no matter how many times you let the robot try (even if you give it a huge budget of attempts), if it's a hard problem, it will eventually get stuck in a loop of identical, unhelpful answers. It's like a teacher who only asks students to raise their hands if they already know the answer; the ones who don't know never get a chance to learn.

2. The Solution: The "Curious Explorer" (Submodular Maximization)

Instead of letting the robot pick answers randomly, INFOTREE uses a smart strategy to pick which path to explore next. The authors treat this like a game of "Maximizing Variety."

They use a mathematical concept called Submodularity. Think of it like packing a suitcase:

If you pack a shirt, it adds value.
If you pack a second shirt of the exact same color, it adds very little new value.
But if you pack a different item (like a hat or shoes), it adds a lot of new value.

INFOTREE acts like a smart packer. It looks at the robot's current attempts and asks: "Which next step will give us the most new information?" It doesn't just look for the "best" answer; it looks for the answer that is different from the others.

3. The Three Ingredients of the "Smart Selector"

To decide which path to explore, the system uses a formula (called UUCB) that mixes three ingredients, like a recipe for a good stew:

The "Confidence" Ingredient (Coverage): "Have we tried this path before?" If the robot is confident and has seen this path often, it doesn't need to go there again.
The "Curiosity" Ingredient (Novelty): "Have we ever been to this part of the map?" If a path is new and unexplored, the robot is encouraged to go there.
The "Chaos" Ingredient (Contrast/Entropy): "Are the answers here messy and different?" The system actively looks for places where the robot is confused or where different attempts lead to different results. This "messiness" is actually good news because it means there is a lot to learn.

By balancing these three, the robot avoids the "boring class" and ensures every practice session teaches it something new.

4. The Safety Net: The "Rescue Team" (Adaptive Budget Allocator)

Sometimes, even a smart selector gets stuck. Maybe the robot is so confused that every path it tries leads to a dead end.

The Fix: INFOTREE has a little "Rescue Team" (the Adaptive Budget Allocator). It watches the robot's practice. If it sees the robot is about to waste all its time on a dead end, the Rescue Team says, "Stop! Let's try one wild, crazy guess just to see if we can break the pattern."
The Result: This saves the training session from being wasted, turning a "useless" practice round into a useful one.

5. The Speed Boost: "Speculative Expansion"

Usually, this smart selection process is slow because the computer has to wait for one calculation to finish before starting the next.

The Fix: INFOTREE uses a "Speculative" trick. It lets the computer guess the next step before the previous calculation is fully finished. If the guess is right, great! If it's wrong, it just rolls back and tries again.
The Result: This makes the whole process much faster (cutting the time wasted by over 10%), so the robot can learn more in less time.

The Bottom Line

The paper tested this new method (INFOTREE) on nine different types of challenges, from solving hard math competitions (like the AIME) to helping robots browse the web and write code.

The Results:

Better Learning: The robot learned significantly faster and solved more problems than previous methods.
No More Wasted Time: It stopped the robot from getting stuck in loops of identical answers.
Robust: The system worked well even when the settings were changed slightly, meaning it's not a "fragile" trick that only works in perfect conditions.

In short, INFOTREE is a way to teach AI agents by ensuring they never practice the same mistake twice. It forces them to explore the "messy" and "different" parts of the problem space, turning wasted effort into valuable lessons.

Technical Summary: Maximizing Rollout Informativeness under a Fixed Budget

1. Problem Formulation

The paper addresses a structural bottleneck in Reinforcement Learning (RL) for tool-augmented Large Language Models (LLMs) under sparse terminal rewards (e.g., $R \in \{-1, +1\}$ ). The core issue is Rollout Informativeness under a Fixed Budget (RIFB).

In standard Group Relative Policy Optimization (GRPO), the advantage signal is normalized by the standard deviation of rewards within a group of $B$ rollouts. If all $B$ independent rollouts for a given prompt yield the same outcome (all success or all failure), the standard deviation becomes zero, causing the advantage to collapse to zero. Consequently, no gradient mass is injected into the policy, rendering the compute spent on those rollouts useless.

The authors formalize this as an impossibility result for independent sampling:

Theorem 1: For any budget-agnostic independent sampler, the probability of collapse (zero variance) is $p^B + (1-p)^B$ , where $p$ is the per-prompt success probability.
Implication: As $p \to 0$ or $p \to 1$ (hard prompts), the non-collapse probability grows at most linearly with the budget ( $B \cdot \min(p, 1-p)$ ). Thus, simply increasing the rollout budget $B$ cannot escape the zero-gradient regime for difficult prompts; the collapse rate remains bounded away from zero regardless of budget size.

2. Methodology: INFOTREE

To overcome the limitations of independent sampling, the authors propose INFOTREE, a training-time tree-search framework that optimizes which states to expand rather than how many to sample.

2.1 Submodular Formulation

The central insight is that the problem of selecting which intermediate state to expand to maximize tree informativeness can be cast as a monotone submodular maximization problem. The objective function $F$ is composed of three terms:

Coverage: Tracks the backed-up value $\bar{Q}(s)$ of the leaves.
Novelty: Tracks the exploration term based on visit counts (classical UCB term).
Contrast: Tracks the diversity of outcomes among sibling nodes, derived from token-level entropy.

The authors prove that $F$ is monotone and submodular. This allows the use of a greedy selector with a theoretical $(1 - 1/e)$ approximation guarantee for the optimal expansion schedule.

2.2 The UUCB Selector

The greedy solution to the submodular problem yields the Uncertainty-aware Upper Confidence Bound (UUCB) score, which is not an ad-hoc heuristic but a closed-form marginal gain of the objective:
$\text{UUCB}(v) = \bar{Q}(v) + c\sqrt{\frac{\ln N(p_a(v))}{N(v)+1}} + \lambda_H \tilde{H}(v) - \lambda_C \tilde{C}(v) - \lambda_D \tilde{d}(v)$

$\bar{Q}(v)$ : Coverage gain.
Second term: Novelty gain (standard UCB exploration).
$\lambda_H \tilde{H}(v)$ : Contrast gain, representing depth-stratified token-level entropy (normalized to be comparable across depths).
Penalties: $\lambda_C$ and $\lambda_D$ penalize tool-call costs and expansion depth, respectively, to respect the budget constraint.

2.3 System Components

INFOTREE integrates the UUCB selector with two additional mechanisms to address practical weaknesses:

Adaptive Budget Allocator (ABA): A lightweight MLP trained offline to predict whether a prompt with a uniform-outcome tree (likely to collapse) can be "rescued" by a single high-temperature rollout. If the predicted rescue probability exceeds a threshold, an extra rollout is triggered. This lifts the mixed-outcome ratio from ~58% to ~76% with <5% budget overhead.
Speculative Expansion: To reduce wall-clock latency, workers evaluate UUCB scores using potentially stale Q-values (bounded staleness). A reconciliation step accepts or rolls back expansions based on version tags. This reduces overhead from 14.3% to 4.8%.
Hierarchical Advantage: A gradient signal derived from sibling contrast along the path from root to leaf, providing an alternative credit assignment mechanism that complements the group-normalized GRPO advantage.

3. Key Contributions

Theoretical Formalization: The paper defines RIFB and proves that independent sampling suffers from persistent collapse on hard prompts, regardless of budget size.
Submodular Reduction: It establishes that tree construction under a fixed budget is a submodular maximization problem, providing a theoretical $(1-1/e)$ guarantee for the greedy UUCB selector.
Analytic Derivation: It demonstrates that the entropy bonus in UUCB is an analytic consequence of the submodular objective (Contrast term), rather than an empirical trick.
System Design: The introduction of ABA and Speculative Expansion creates a practical framework that balances informativeness with computational efficiency.
Orthogonality: The authors show that their selection mechanism operates on an axis orthogonal to trajectory re-weighting (e.g., CW-GRPO) and rollout reuse (e.g., Tree-GRPO prefix sharing), allowing for compositional gains.

4. Experimental Results

The method was evaluated on nine benchmarks across mathematical reasoning, web-search, and tool-rich coding/OS agents.

Performance: INFOTREE consistently outperforms flat GRPO, DeepSearch, Tree-GRPO, AT2PO, CW-GRPO, and RC-GRPO.
- On AIME 2025, INFOTREE achieves a Mean@32 of 30.2 (vs. 27.7 for Flat GRPO and 29.1 for Tree-GRPO).
- On GAIA, it achieves 56.0% accuracy (vs. 44.8% for Flat GRPO).
Composability: Combining INFOTREE with Tree-GRPO prefix sharing or CW-GRPO contribution weights yields further gains, confirming the orthogonality of the approaches.
Robustness: A $5 \times 5 \times 5$ hyperparameter grid over the UUCB coefficients ( $\lambda_H, \lambda_C, \lambda_D$ ) shows that >75% of the space lies within 2.1 points of the optimum, indicating the method is not fragile to hyperparameter tuning.
Efficiency: Despite the tree search, the wall-clock overhead is reduced to 4.8% via Speculative Expansion. The "steps-to-target" (steps to reach Mean@32 of 28) drops from 248 (Flat GRPO) to 138 (INFOTREE), a 1.57x improvement in wall-clock time.
Gradient Informativeness: Direct measurement of the squared gradient norm shows INFOTREE produces a 3.9x higher RIFB than flat GRPO.

5. Significance and Claims

The paper claims to reposition the design of rollout strategies from an engineering preference to a formally grounded optimization problem.

From Heuristic to Theory: The work transforms the entropy bonus from an empirical trick into a necessary component of a submodular objective, providing a theoretical justification for its inclusion.
Solving the Collapse: By explicitly targeting the "zero-gradient" regime through submodular selection and adaptive budgeting, INFOTREE addresses a fundamental limitation of current GRPO-based agents on hard prompts.
Complementarity: The results suggest that maximizing rollout informativeness via tree construction is a distinct and complementary axis to existing methods that focus on re-weighting or diversifying independent samples.

The authors note limitations, including the assumption of a fixed policy within a tree-construction episode and the preliminary nature of results on open-ended tasks without verifiable rewards. However, the framework is presented as a robust solution for training tool-augmented agents under sparse rewards.