菜单
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 当前能够提供的文件。
正在显示 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.
基于 SOC 职业分类
