Information Gain-based Rollout Policy Optimization: An Adaptive Tree-Structured Rollout Approach for Multi-Turn LLM Agents

🕒 Published (v1): 2026-07-07 12:47 UTC · Source: Arxiv · link

Why this paper was selected

Tree-structured information-gain rollout for RL on multi-turn LLM agents; novel optimization signal

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

IGRPO is a reinforcement learning framework for multi-turn search-augmented LLM agents that adaptively allocates rollout budget to intermediate states by their information gain — the marginal improvement in ground-truth answer likelihood — rather than uniformly. This induces a theoretically grounded "teacher distribution" over trajectories that serves as a principled policy optimization target, unifying adaptive exploration and policy learning in one framework.

Problem

Existing RL methods for long-horizon search agents (both chain-based and tree-based) allocate rollout computation without explicitly estimating the utility of intermediate states. Chain-based methods waste full trajectories on already-unpromising prefixes; tree-based methods use heuristic or entropy-based branching rules that may still expand low-value nodes. Neither explicitly measures whether an intermediate state contains evidence useful for answering the question, leading to wasted budget on misleading or uninformative branches.

Method

IGRPO defines node-level information gain as the marginal improvement in the policy's normalized log-likelihood of the ground-truth answer \(a\) upon transitioning from a parent node to a child node:

\[\text{IG}(h) = s_\theta(h; a) - s_\theta(\text{parent}(h); a), \quad s_\theta(h; a) = \exp\!\left(\frac{1}{L}\sum_{j=1}^L \log \pi_\theta(a_j \mid h, a_{<j})\right)\]

At each expansion stage \(i\), active nodes \(h \in \mathcal{A}_i\) are selected with softmax-proportional probability based on a composite informativeness score \(\text{val}(h) = \frac{s_\theta(h;a) + \text{IG}(h)}{2}\):

\[p_i(h) = \frac{\exp(\gamma \cdot \text{val}(h))}{\sum_{h' \in \mathcal{A}_i} \exp(\gamma \cdot \text{val}(h'))}\]

A stage budget \(B_i\) nodes are independently sampled and expanded, biasing the tree toward informative branches without hard thresholds.

Theoretical grounding: Theorem 1 shows that as \(B_t \to \infty\), the empirical trajectory distribution converges to \(\mu_\theta(o) \propto \pi_\theta(o) \exp(\gamma V(o))\), where \(V(o) = \sum_t \text{val}(h_t)\) is the cumulative value. Theorem 2 shows this is the unique optimizer of \(\max_\mu \mathbb{E}_{o \sim \mu}[V(o)] - \frac{1}{\gamma} D_\text{KL}(\mu \| \pi_\theta)\). Since \(V(o)\) correlates strongly with final reward \(R(o)\), \(\mu_\theta\) acts as a reward-aligned teacher distribution. The IGRPO policy objective is identical in form to GRPO (Eq. 9) but samples trajectories from \(\mu_{\theta_\text{old}}\) rather than \(\pi_{\theta_\text{old}}\).

Key Contributions

• Information gain-based tree rollout: Soft, probabilistic budget allocation toward informative nodes, replacing heuristic hard-threshold branching.
• Theoretical characterization: Proof that IG-based rollout induces a limiting distribution \(\mu_\theta(o) \propto \pi_\theta(o)\exp(\gamma V(o))\), which is the unique solution to a KL-regularized value-maximization variational problem.
• Unified framework: The induced \(\mu_\theta\) directly yields a policy optimization target, connecting adaptive exploration and policy learning without separate design choices.
• Empirical validation: Consistent improvement over strong baselines across seven search-augmented QA benchmarks under matched rollout budgets.

Results

• 3B backbone (Qwen2.5-3B-Instruct): IGRPO average EM = 45.4% vs. best baseline GiGPO 42.3% — a +3.1% gain.
• 7B backbone (Qwen2.5-7B-Instruct): IGRPO average EM = 48.2% vs. best baseline StepSearch 47.3% — a +0.9% gain.
• Per-dataset highlights (3B): PopQA 47.5% vs. Tree-GRPO 43.7%; Bamboogle 65.7% vs. GiGPO 64.1%; 2Wiki 40.2% vs. GiGPO 37.0%; HotpotQA 38.9% vs. GiGPO 36.9%.
• Per-dataset highlights (7B): Bamboogle 71.8% vs. StepSearch 68.9%; HotpotQA 44.6% vs. IGPO 42.9%; NQ 48.0% vs. StepSearch 46.4%.
• IGRPO outperforms both chain-based (Search-R1, ZeroSearch, StepSearch, GiGPO, IGPO) and tree-based (Tree-GRPO) baselines on average across both scales.

Limitations

• The limiting distribution theorem (Theorem 1) requires \(B_t \to \infty\) per layer; finite budgets used in practice only approximate \(\mu_\theta\).
• Information gain computation requires access to ground-truth answer tokens at rollout time, which is standard for RL training but precludes use in reward-free or preference-only settings.
• The correlation between \(V(o)\) and \(R(o)\) is empirically observed during IGPO training; the theoretical bridge to the reward objective assumes this correlation holds throughout IGRPO training.
• Evaluation is restricted to search-augmented QA; generalization to other multi-turn agentic tasks (e.g., code execution, tool-use beyond retrieval) is not demonstrated.
• ZeroSearch results are omitted for several benchmarks in the 3B setting, making some comparisons incomplete.

Relevance to Agentic AI / LLM Agents

IGRPO directly addresses a core efficiency problem in training multi-turn LLM agents: how to allocate a fixed computation budget across a branching trajectory space during rollout collection. The information gain scoring provides a computationally cheap (no extra sampling required) intermediate signal that can replace heuristic branching rules in tree-based RL. The theoretical connection between adaptive rollout and a principled teacher distribution is a meaningful contribution to the line of work on credit assignment in long-horizon agent training (alongside GiGPO, IGPO, Tree-GRPO). For practitioners building RL pipelines for tool-using agents, this suggests that soft informativeness-weighted expansion is preferable to uniform or entropy-threshold-based branching.