메뉴
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.
기본적으로 소스를 먼저 확인하는 Prompt가 선택됩니다. 직접 명령으로 전환하거나 로컬 사본을 다운로드할 수도 있습니다.
설치 여부를 결정하기 전에 SKILL.md와 SkillsMP에 표시된 보조 파일을 읽어 보세요.
Codex 또는 Claude로 설치 이 Prompt를 복사해 Codex, Claude 또는 다른 어시스턴트에 붙여 넣으면 Skill 페이지를 검토하고 설치를 진행할 수 있습니다.
직접 명령은 검토 Prompt를 거치지 않습니다. 실행하기 전에 소스를 확인하세요.
npx skills add https://github.com/ADu2021/skillXiv --skill tree-search-llm-agent-rl명령은 한 줄로 유지됩니다. 복사하기 전에 가로로 스크롤해 전체 내용을 확인하세요.
로컬 사본을 원하시나요? SkillsMP에서 현재 제공할 수 있는 파일을 다운로드하세요.
SOC 직업 분류 기준
SKILL.md 표시 중
| name | tree-search-llm-agent-rl |
| title | Tree Search for LLM Agent Reinforcement Learning |
| version | 0.0.2 |
| engine | skillxiv-v0.0.2-claude-opus-4.6 |
| license | MIT |
| url | https://arxiv.org/abs/2509.21240 |
| keywords | ["agent RL","tree search","rollout sampling","process supervision","policy optimization","multi-turn tasks","step-level advantage","GRPO"] |
| description | 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. |
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.
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.
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.
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.
Generate M independent trajectories from the initial state without branching. These serve as roots and establish behavioral diversity.
def initialize_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 in range(num_roots):
trajectory = {
'root_id': root_id,
'nodes': [],
'rewards': []
}
state = env.reset()
for step in range(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
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.
def sample_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 in enumerate(trees):
# Filter to non-leaf nodes (not the final step)
non_leaf_indices = list(range(len(tree['nodes']) - 1))
if not 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
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.
def expand_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 = 0
for 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 in range(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 in range(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 += 1
return trees, branch_count
def reconstruct_state_at_node(env, tree, node_idx):
"""
Replay trajectory prefix to reconstruct state at given node.
"""
state = env.reset()
for i in range(node_idx):
node = tree['nodes'][i]
obs, _, _ = env.step(node['action'])
state = obs
return state
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.
def compute_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 in enumerate(trees):
for node_idx, node in enumerate(tree['nodes']):
if not node['children']:
# Leaf node: no siblings to compare
intra_advantages[(tree_id, node_idx)] = 0.0
continue
# 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
def aggregate_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'] else 0.0
return final_reward
Compute baseline advantages across all trajectories to stabilize training. Then perform policy gradient updates using both intra and inter advantages.
def compute_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
def update_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 = 0
for tree_id, tree in enumerate(trees):
for node_idx, node in enumerate(tree['nodes']):
key = (tree_id, node_idx)
if key not in 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 += 1
if 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 > 0 else 0.0
Orchestrate initialization, sampling, expansion, and updates across multiple iterations.
def train_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 in range(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}")
| 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. |
| learning_rate | 1e-5 | 1e-6 to 1e-4 | Smaller LLMs (1.5B-3B) use 1e-5; larger models (7B+) tolerate 5e-6. Start conservative. |
| advantage_normalization | True | - | 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). |
Tree-GRPO excels in scenarios with:
Avoid Tree-GRPO if:
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.
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.
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.