Train LLM agents via tree-search rollout sampling and step-wise advantage estimation. Achieve 1.5x more rollouts within fixed token budgets and implicit step-level preference learning through dual-level advantage computation on tree structures.
Install with Codex or Claude Copy this prompt, paste it into Codex, Claude, or another assistant, and let it review the skill page and install it for you.
A direct command skips the review prompt. Inspect the source before running it.
Train LLM agents via tree-search rollout sampling and step-wise advantage estimation. Achieve 1.5x more rollouts within fixed token budgets and implicit step-level preference learning through dual-level advantage computation on tree structures.
Train LLM Agents with Tree-Structured Rollouts and Step-Level Supervision
The Problem: Sparse Feedback in Long-Horizon Agent Tasks
Traditional reinforcement learning for LLM agents faces two critical obstacles. First, outcome rewards alone provide only terminal feedback—the model learns which final trajectory succeeded but not which intermediate steps contributed to success. This sparse supervision severely hampers learning efficiency. Second, rollout budgets are expensive. Each agent query involves multiple tool calls and API interactions, making per-trajectory training prohibitively costly.
Existing methods either resort to dense intermediate rewards (requiring expensive labeling) or accept degraded learning from sparse terminal feedback. Tree-GRPO (Tree-based Group Relative Policy Optimization) solves this by reorganizing how rollouts are sampled and how credit flows through decision points.
Core Concept: Trees Instead of Chains
Standard agent RL generates independent trajectories in parallel—one thought-action-observation sequence per rollout. Tree-GRPO replaces this with tree-structured exploration where nodes share common prefixes.
The key insight: when an agent takes the same action from the same state in two different rollouts, those rollouts can branch from that shared point. This prefix sharing enables more rollouts within the same token budget. Additionally, the tree structure enables backward credit assignment through sibling comparison—you can estimate the value of a step by comparing trajectories that diverge at that point.
The method treats each complete agent step (Thought-Action-Observation triplet) as a tree node, not finer granularities like tokens or sentences. This coarse-grained approach aligns with the discrete decision points where agents actually choose actions.
Architecture Overview
Tree-GRPO operates in three stages:
Initialization: Generate M independent chain trajectories serving as tree roots. These establish diverse starting points for exploration.
Sampling: Randomly select N non-leaf nodes from each existing tree for expansion. This creates a pool of decision points where the agent might have chosen differently.
Expansion: Generate K new branches from each selected node, producing new trajectories. Common prefixes (everything before the selected node) are reused, eliminating redundant computation.
The result is a forest of trees, each rooted in a different chain but sharing internal node prefixes. This structure naturally encodes both within-tree comparisons (siblings branching from the same parent) and across-tree comparisons (rollouts from different roots).
Advantage estimation occurs at two levels:
: The grouped relative advantage within a single tree, computed by comparing trajectories that share history up to a given node. This captures fine-grained step value.
Intra-tree advantage
Inter-tree advantage: The baseline advantage computed across all rollouts, stabilizing training. This prevents variance explosion from per-node estimation.
The final advantage = intra-tree advantage + inter-tree advantage, producing step-wise process supervised signals from outcome rewards alone.
Implementation: Tree-Structured Rollout and Credit Assignment
Stage 1: Initialize Base Trajectories
Generate M independent trajectories from the initial state without branching. These serve as roots and establish behavioral diversity.
definitialize_trees(agent, env, num_roots=4, max_steps=10):
"""
Generate independent trajectories as tree roots.
Args:
agent: LLM agent with act() and observe() methods
env: Agent environment (retrieval, web browsing, etc.)
num_roots: Number of independent chains to generate
max_steps: Maximum steps per trajectory
Returns:
List of trajectory dictionaries with nodes
"""
trees = []
for root_id inrange(num_roots):
trajectory = {
'root_id': root_id,
'nodes': [],
'rewards': []
}
state = env.reset()
for step inrange(max_steps):
# Agent generates thought and action
thought = agent.generate_thought(state)
action = agent.generate_action(thought, state)
# Execute and observe
obs, reward, done = env.step(action)
node = {
'step': step,
'thought': thought,
'action': action,
'observation': obs,
'state': state,
'children': [] # For branches added later
}
trajectory['nodes'].append(node)
state = obs
if done:
break# Final outcome reward at trajectory end
final_reward = env.get_reward()
trajectory['rewards'].append(final_reward)
trees.append(trajectory)
return trees
Stage 2: Sample Expansion Points from Existing Nodes
Randomly select non-leaf nodes across all trees. These are decision points where alternative branches will be explored, enabling comparison of different actions from the same state.
defsample_expansion_nodes(trees, num_per_tree=3):
"""
Select non-leaf nodes for branching.
Args:
trees: List of trajectory trees
num_per_tree: Number of nodes to select from each tree
Returns:
List of (tree_id, node_index) tuples
"""
expansion_targets = []
for tree_id, tree inenumerate(trees):
# Filter to non-leaf nodes (not the final step)
non_leaf_indices = list(range(len(tree['nodes']) - 1))
ifnot non_leaf_indices:
continue# Random sampling with replacement
selected = np.random.choice(
non_leaf_indices,
size=min(num_per_tree, len(non_leaf_indices)),
replace=False
)
for node_idx in selected:
expansion_targets.append((tree_id, node_idx))
return expansion_targets
Stage 3: Branch from Selected Nodes
For each selected node, generate K new trajectories that continue from that point. The prefix (everything up to that node) is shared with the original trajectory, eliminating wasted computation.
defexpand_branches(trees, expansion_targets, agent, env, branches_per_node=2):
"""
Generate new trajectories from selected nodes.
Args:
trees: List of trajectory trees
expansion_targets: List of (tree_id, node_idx) to branch from
agent: LLM agent
env: Agent environment
branches_per_node: Number of branches per selected node
Returns:
Updated trees with new branches attached
"""
branch_count = 0for tree_id, parent_node_idx in expansion_targets:
parent_tree = trees[tree_id]
parent_node = parent_tree['nodes'][parent_node_idx]
# Reconstruct state at parent node by replaying prefix
state = reconstruct_state_at_node(env, parent_tree, parent_node_idx)
for branch_num inrange(branches_per_node):
# Generate alternative action from the same state
thought = agent.generate_thought(state, diversity_boost=True)
action = agent.generate_action(thought, state)
branch_trajectory = {
'parent': (tree_id, parent_node_idx),
'nodes': [],
'rewards': []
}
current_state = state
max_remaining = 10 - parent_node_idx
for step inrange(max_remaining):
obs, reward, done = env.step(action)
branch_node = {
'step': parent_node_idx + step,
'thought': thought,
'action': action,
'observation': obs,
'state': current_state
}
branch_trajectory['nodes'].append(branch_node)
current_state = obs
if done:
break# Continue exploration from branch
thought = agent.generate_thought(obs)
action = agent.generate_action(thought, obs)
final_reward = env.get_reward()
branch_trajectory['rewards'].append(final_reward)
# Attach branch to parent node
parent_node['children'].append(branch_trajectory)
branch_count += 1return trees, branch_count
defreconstruct_state_at_node(env, tree, node_idx):
"""
Replay trajectory prefix to reconstruct state at given node.
"""
state = env.reset()
for i inrange(node_idx):
node = tree['nodes'][i]
obs, _, _ = env.step(node['action'])
state = obs
return state
Stage 4: Compute Intra-Tree Advantages
Within each tree, estimate the value of each node by comparing trajectories that diverge at that node. Trajectories with higher rewards are preferred, yielding step-level relative advantages.
defcompute_intra_tree_advantages(trees):
"""
Estimate step-level advantage within each tree via sibling comparison.
Returns:
Dictionary mapping (tree_id, node_idx) to advantage value
"""
intra_advantages = {}
for tree_id, tree inenumerate(trees):
for node_idx, node inenumerate(tree['nodes']):
ifnot node['children']:
# Leaf node: no siblings to compare
intra_advantages[(tree_id, node_idx)] = 0.0continue# Collect trajectory rewards for all children
child_rewards = []
for child in node['children']:
# Trace complete trajectory through child and branches
total_reward = aggregate_branch_reward(tree, node_idx, child)
child_rewards.append(total_reward)
if child_rewards:
child_rewards = np.array(child_rewards)
# Grouped relative advantage: deviation from mean
mean_reward = child_rewards.mean()
node_advantage = (child_rewards[0] - mean_reward) / (child_rewards.std() + 1e-8)
intra_advantages[(tree_id, node_idx)] = node_advantage
return intra_advantages
defaggregate_branch_reward(tree, node_idx, child_trajectory):
"""
Compute total trajectory reward combining prefix and branch.
"""
prefix_length = node_idx + 1
total_steps = prefix_length + len(child_trajectory['nodes'])
final_reward = child_trajectory['rewards'][0] if child_trajectory['rewards'] else0.0return final_reward
Stage 5: Compute Inter-Tree Advantages and Update Policy
Compute baseline advantages across all trajectories to stabilize training. Then perform policy gradient updates using both intra and inter advantages.
defcompute_advantages(trees, intra_advantages, gamma=0.99):
"""
Compute combined advantages: intra-tree + inter-tree baseline.
Returns:
Dictionary of step-level advantages
"""# Collect all final rewards across trees
all_final_rewards = []
for tree in trees:
if tree['rewards']:
all_final_rewards.append(tree['rewards'][0])
inter_tree_baseline = np.mean(all_final_rewards)
inter_tree_std = np.std(all_final_rewards) + 1e-8
advantages = {}
for (tree_id, node_idx), intra_adv in intra_advantages.items():
# Normalize inter-tree advantage
final_reward = trees[tree_id]['rewards'][0]
inter_adv = (final_reward - inter_tree_baseline) / inter_tree_std
# Combine both levels
combined_adv = intra_adv + inter_adv
advantages[(tree_id, node_idx)] = combined_adv
return advantages
defupdate_policy(model, trees, advantages, learning_rate=1e-5):
"""
Perform policy gradient updates using step-level advantages.
This is equivalent to step-level DPO: prefer high-advantage trajectories.
"""
total_loss = 0.0
update_count = 0for tree_id, tree inenumerate(trees):
for node_idx, node inenumerate(tree['nodes']):
key = (tree_id, node_idx)
if key notin advantages:
continue
advantage = advantages[key]
# Policy gradient: log_prob * advantage
log_prob = model.compute_log_prob(
thought=node['thought'],
action=node['action'],
state=node['state']
)
# Gradient step: increase probability of high-advantage actions
loss = -log_prob * advantage
total_loss += loss
update_count += 1if update_count > 0:
avg_loss = total_loss / update_count
avg_loss.backward()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
optimizer.step()
return avg_loss.item() if update_count > 0else0.0
Training Loop
Orchestrate initialization, sampling, expansion, and updates across multiple iterations.
deftrain_tree_grpo(agent, env, num_iterations=10, roots_per_iter=4,
branches_per_node=2, updates_per_iter=1):
"""
Main training loop for Tree-GRPO.
"""for iteration inrange(num_iterations):
print(f"Iteration {iteration + 1}/{num_iterations}")
# Stage 1: Initialize
trees = initialize_trees(agent, env, num_roots=roots_per_iter)
initial_reward = np.mean([t['rewards'][0] for t in trees])
print(f" Initial reward: {initial_reward:.4f}")
# Stage 2-3: Sample and expand
expansion_targets = sample_expansion_nodes(trees, num_per_tree=3)
trees, branch_count = expand_branches(
trees, expansion_targets, agent, env,
branches_per_node=branches_per_node
)
print(f" Generated {branch_count} branches from {len(expansion_targets)} nodes")
# Stage 4-5: Compute advantages and update
intra_advantages = compute_intra_tree_advantages(trees)
advantages = compute_advantages(trees, intra_advantages)
loss = update_policy(agent.model, trees, advantages)
print(f" Update loss: {loss:.6f}")
# Evaluate on held-out set
eval_reward = evaluate(agent, env, num_trajectories=8)
print(f" Eval reward: {eval_reward:.4f}")
Practical Guidance
Hyperparameter Reference
Parameter
Default
Range
Guidance
roots_per_iteration
4
2-8
Start with 4. Higher values increase diversity but require more tokens. For smaller budgets, use 2-3.
branches_per_node
2
1-4
Determines tree width. 2-3 balances exploration and token cost. Use 1 for severe budget constraints.
nodes_per_tree
3
1-5
How many expansion points per tree. Smaller datasets favor 2-3; larger datasets tolerate 4-5.
Always use. Prevents dominance by outlier trajectories.
gamma (discount factor)
0.99
0.95-0.99
Rarely changed. Use 0.99 for long-horizon tasks (15+ steps).
When to Use Tree-GRPO
Tree-GRPO excels in scenarios with:
Multi-turn agent tasks where credit assignment matters. If your task involves 5+ decision steps, tree structure unlocks per-step learning.
Severe token budgets. If you can afford only 100-200 rollouts per iteration, prefix sharing yields 1.5x more samples at no extra cost.
Sparse outcome rewards. If you only have terminal rewards (success/failure), the method automatically extracts step-level signals via tree structure.
Diverse intermediate states. Tree-GRPO benefits from many unique decision points. Highly repetitive tasks (same action always succeeds) see smaller gains.
Smaller models (1.5B-7B). Tested extensively on Qwen and Llama in this range. Larger models still benefit but may require careful tuning.
When NOT to Use
Avoid Tree-GRPO if:
You have step-level supervision. If you can label intermediate correct actions, use step-level DPO directly—it's simpler and faster. Tree-GRPO's advantage is extracting step signals from outcome rewards; dense supervision makes this unnecessary.
Your task is single-step. A single thought-action-observation with no branching points means zero advantage in tree structure. Use standard policy gradient instead.
Token budget is unlimited. If you can generate 500+ independent chains per iteration, the prefix-sharing optimization becomes marginal. Standard chain-based RL is simpler.
State space is deterministic. If replaying an action always yields the same observation, there's no value in branching from intermediate states. The method assumes stochastic environments.
Outcome rewards are dense. If nearly every trajectory yields distinct outcome values, tree advantages don't provide additional signal beyond outcome. Standard approaches suffice.
You need real-time training. Tree expansion and state reconstruction add latency. For systems requiring immediate adaptation, streaming rollouts are faster.
Critical Implementation Pitfalls
Pitfall 1: Inefficient state reconstruction. Replaying the full prefix to reach an expansion node is expensive. Cache or implement environment-specific shortcuts. For retrieval tasks, avoid re-querying the same documents.
Pitfall 2: Unstable advantage normalization. If a single trajectory dominates rewards, intra-tree advantages explode. Always use grouped relative advantages and normalize by standard deviation. Test stability on held-out trees first.
Pitfall 3: Insufficient branch diversity. If branching only explores minor action variations, the method reduces to noisy independent chains. Push diversity: use temperature > 1.0 or top-k sampling during branch generation. Verify that branches produce meaningfully different observations.
Pitfall 4: Over-expansion. Selecting too many expansion nodes per tree or too many branches per node exhausts budget without improving learning. Start with 2-3 nodes per tree and 2 branches; ablate upward only if you see consistent gains.
Pitfall 5: Misaligned node granularity. Using token-level or sentence-level nodes instead of complete steps (Thought-Action-Observation) breaks the method. Empirical evidence shows coarse-grained nodes outperform fine-grained decomposition for agent tasks.
Tuning Workflow
Set baseline: Run standard chain-based RL for 100 iterations, record final metric.
Enable trees: Use conservative settings (4 roots, 2 branches, 3 nodes). Run 100 iterations, compare metric.
Measure efficiency: Compute cost per unit performance gain. If trees are 1.5x cheaper, proceed.
Increase exploration: Raise branches_per_node to 3 or nodes_per_tree to 4. Track if diminishing returns appear.
Stabilize training: If advantage normalization causes instability, reduce learning rate by 2x. If loss plateaus, increase learning rate by 1.5x.
Validate on test set: Always evaluate on held-out tasks before deploying. Cross-dataset performance often lags in-distribution gains.
Theoretical Grounding
Tree-GRPO's advantage computation is mathematically equivalent to step-level Direct Preference Optimization (DPO). When intra-tree group relative advantages are computed correctly, they implicitly optimize log(π(preferred_trajectory)) - log(π(dispreferred_trajectory)), which is the DPO objective. This equivalence means the method inherits DPO's theoretical guarantees: it learns from preferences without explicit reward models.
The paper proves this equivalence by showing that the intra-tree advantage gradient, when aggregated across the tree, has the same expectation as the step-level DPO gradient. This theoretical insight validates using outcome rewards alone and justifies the two-level advantage structure.
Reference
For the full paper, implementation, code, and experimental details, visit https://arxiv.org/abs/2509.21240. The authors' official code repository is available at https://github.com/AMAP-ML/Tree-GRPO with training scripts for multi-hop QA, single-hop QA, and web agent tasks.
