PaperlayerView exact revision on arXiv

Source-linked AI summary

Back to the Definition: Estimating Step-Level Advantages via Trajectory Graphs for Agentic Reinforcement Learning

Xincheng Yao, Haobo Fu, Weiming Liu, Chongyang Zhang

arXiv:2609.28963v1cs.AIstat.ML

TL;DR

Group-based advantage estimation is reliable at the response level, but becomes biased at the step level when trajectories at the same index represent different states. GRAFT combines rollout trajectories into a graph, estimates node values with Bellman iteration, and derives step advantages from value differences, with Graph GAE to reduce value-estimation bias. The paper reports consistent gains over GRPO and superior performance to recent agentic RL methods across multi-turn benchmarks.

  • Problem

    GRPO-style group estimates are reliable for responses but biased at the step level because trajectories at the same step index can occupy different states.

    Source: §1 · p. 2, §2 · p. 3

  • Method

    GRAFT merges rollout trajectories into a graph, estimates node values by Bellman iteration, and assigns edge advantages from node-value differences, with Graph GAE to reduce value-estimation bias.

    Source: §1 · p. 2, §3.2 · p. 5, §3.3 · p. 6

  • Results

    GRAFT reports consistent gains over GRPO and superior performance to recent agentic RL algorithms across multi-turn benchmarks, with gains also reported on searchQA tasks.

    Source: §4.2 · p. 8, §6 · p. 10

  • Takeaways & Limitations

    The graph-based estimates are presented as step-level advantages aligned with the standard RL definition, using outcome rewards rather than PRM-assigned step rewards.

    Source: §1 · p. 2, §6 · p. 10

Abstract

from arXiv · show

Group-based reinforcement learning (RL) methods, such as GRPO and its variants, have become a leading paradigm for training reasoning and agentic large language models (LLMs). While their group-normalized advantage estimation is reliable at the response level, it becomes systematically biased at the step level, since coarse-grained trajectory-level advantages are hard to accurately reflect the contribution of individual steps (i.e, failed trajectories may contain valuable steps). Revisiting the foundational RL definition, we notice that GRPO's success on single-turn tasks stems from its advantage estimation strategy, which adheres to the basic definition: the mean reward of multiple actions sampled from the same state constitutes a credible state-value estimate. Extending the faithful estimation to step-level would in principle demand sampling multiple actions from each intermediate state, which is too costly on a per-state basis. To mitigate this issue, we propose a Graph-based Faithful sTep-level credit-assignment framework (GRAFT) that grafts all rollout trajectories into a trajectory graph, recovering node state-values via Bellman iteration on the graph, and assigning credit to each edge by the node value difference. Theoretically, the estimated step-level advantage faithfully adheres to the basic advantage definition in RL. To further ensure the reliability of step-level advantage estimation, we further propose Graph GAE, which extends GAE to the trajectory graph for reducing the impact of state-value estimation bias. Experiments across a range of multi-turn agentic benchmarks show consistent gains over GRPO and superior performance compared to recent agentic RL algorithms. Code will be available at https://github.com/xcyao00/GRAFT.

1 INTRODUCTION

GRPO’s group-based advantage estimate is well-grounded when a response is treated as one action, but trajectory-level advantages are too coarse for multi-turn step credit. GRAFT combines trajectories in a graph to estimate step-level values and advantages from shared states.

1 INTRODUCTION

  • GRPO’s group mean estimates the prompt’s state value when each response is treated as a single action, aligning its advantage with the standard RL definition.GRPO and its variants have gained traction for their critic-free design.Source: §1 · pp. 1–2
  • In multi-turn tasks, trajectory-level advantages can over-credit redundant or erroneous steps and obscure valuable steps in failed trajectories.Per-step group normalization can also be biased because trajectories at the same step index may occupy different states; resampling continuations from each state is costly.Source: §1 · p. 2
  • GRAFT merges semantically identical states across rollout trajectories into shared graph nodes, aggregating the actions taken from each state into on-state action groups.State matching uses exact matching or embedding-based similarity matching.Source: §1 · p. 2
  • GRAFT estimates node values with Bellman iteration and assigns each edge a step advantage from the difference between its endpoint values.It also extends GAE to the trajectory graph to reduce the impact of state-value estimation bias.Source: §1 · p. 2

2 PRELIMINARY ANALYSIS

The preliminary analysis explains why group-based estimation is credible for a shared initial prompt but systematically biased when applied across distinct intermediate states. It then motivates faithful step-level estimation through same-state action samples.

2 PRELIMINARY ANALYSIS

  • In standard RL, a state’s value is the expected action value, and an advantage measures an action’s return relative to that state value.For LLMs, a prompt is the initial state and a response can be treated as a single action with its reward as the Q-value.Source: §2 · p. 3
  • GRPO’s group reward mean estimates the prompt value when responses are treated as single actions, but assigning that response advantage uniformly to tokens mismatches the action granularity.The response-level estimator aligns with the standard advantage definition, with GRPO adding normalization.Source: §2 · p. 3
  • At later steps, trajectories may occupy different states, so rewards from actions taken in other states bias the estimate of the current state’s value.This mismatch can be worst when every first-step action differs and therefore every second-step state is distinct.Source: §2 · p. 3
  • Sampling N −1 extra continuations from the same intermediate state alongside the observed action enables an unbiased advantage estimate, but requires additional rollouts.The first step remains aligned with GRPO because the prompt is shared across responses.Source: §2 · p. 3

3 PROPOSED METHOD

GRAFT combines rollouts in a trajectory graph to estimate intermediate-state values with Bellman iteration, then assigns step-level credit from value differences. Graph GAE extends this approach with multi-step estimates to reduce the impact of value-estimation bias.

3.3 STEP-LEVEL ADVANTAGE ESTIMATION

  • For each transition, GRAFT assigns the discounted successor-node value minus the current-node value as its step advantage, which matches the standard advantage definition when step rewards are zero.Positive values indicate movement toward higher-value states, while negative values indicate movement toward lower-value states.Source: Figure 1 · p. 4, §3.3 · p. 5

3 PROPOSED METHOD

  • Aggregating similar states across trajectories forms shared action groups, enabling more faithful step-level estimation without additional rollout cost despite the O(N^2T) cost of sampling multiple actions per intermediate state.Trajectory-level advantages can obscure valuable steps in failed trajectories and give unwarranted credit to redundant or erroneous steps in successful ones.Source: §3 · p. 4

3.1 GRAPH CONSTRUCTION

  • The task-specific directed graph represents canonicalized states as nodes and rollout transitions as edges.GRAFT reorganizes the N rollouts into a unified graph so actions from the same state can be aggregated.Source: §3.1 · p. 4, §3 · p. 4

3.2 STATE-VALUE ESTIMATION VIA BELLMAN ITERATION

  • Bellman iteration propagates terminal rewards backward; each non-sink node’s value aggregates expected next-state values across its distinct actions using empirical transition probabilities.The action-level update approximates the standard value function under sparse rewards; reverse-BFS ordering accelerates convergence, and nodes unable to reach a sink retain value zero.Source: §3.2 · p. 5

3.3 STEP-LEVEL ADVANTAGE ESTIMATION

  • At the population level, graph-based credit has lower estimation error than the smaller of the GiGPO and GRPO credit errors, with the inequality strict when the specified state-action set has positive visitation probability.Graph Bellman backups aggregate different continuations at shared states into a lower-error, state-specific signal.Source: §3.3 · p. 6
  • Graph GAE combines subsequent mean k-hop TD residuals with λ-weighting to reduce the impact of state-value estimation bias.At λ = 0 it reduces to the one-step estimate, with low variance but high bootstrap bias; increasing λ expands the horizon and trades lower bootstrap bias for higher variance.Source: §3.3 · p. 6
  • GRAFT normalizes advantages within each node’s outgoing-edge group; for a one-edge group it sets the mean to zero and uses the absolute value as the standard deviation.Source: §3.3 · p. 7

4 EXPERIMENTS

The experiments compare GRAFT with established agentic RL baselines across multi-turn benchmarks and test its components, task efficiency, and added training cost. GRAFT performs strongly across the reported benchmarks, while its added runtime is small relative to total training time.

4.1 EXPERIMENTAL SETUP

  • Evaluation covers ALFWorld, WebShop, and SearchQA using Qwen2.5-Instruct models, with training and evaluation configurations based on GiGPO.For SearchQA, NQ and HotpotQA are in-domain evaluation datasets; the remaining datasets assess out-of-domain generalization.Source: §4.1 · pp. 7–8

4.2 MAIN RESULTS

  • GRAFT improves over GRPO and outperforms GiGPO and GraphGPO on ALFWorld and WebShop across the reported model sizes.Reported average success-rate gains over GRPO are 24.6% and 20.8% on ALFWorld, and 25.5% and 16.7% on WebShop, for the 1.5B and 7B models, respectively.Source: §4.2 · p. 8
  • On SearchQA, GRAFT reaches a 48.6% average success rate at 7B and outperforms Search-R1, StepSearch, GiGPO, and GraphGPO.The reported gains span single-hop and multi-hop reasoning datasets; Table 2 marks in-domain and out-of-domain datasets with † and ⋆.Source: §4.2 · p. 8, Table 2 · p. 8

4.3 ABLATION STUDIES & FURTHER ANALYSIS

  • The complete GRAFT framework achieves the highest ablation performance, while removing advantage normalization substantially reduces results on both benchmarks.The reported examples are an ALFWorld “All” decline from 96.10 to 71.33 and a WebShop success rate of 72.93 without normalization; GRPO-C and Graph GAE also contribute gains.Source: §4.3 · p. 9, Table 3 · p. 9
  • GRAFT-trained models require fewer interaction turns than GRPO- or GiGPO-trained models on both WebShop and ALFWorld.The analysis links fewer turns with lower inference costs and fewer environment API calls.Source: §4.3 · p. 9, Figure 3 · p. 9
  • GRAFT adds no GPU-memory overhead, and its graph construction and graph-based advantage estimation occupy 0.43% of total RL training time.The added work occurs during advantage estimation and is reported as negligible compared with rollout and policy update.Source: §4.3 · pp. 9–10, Figure 4 · p. 9

5 RELATED WORK

Related work traces LLM reinforcement learning from preference alignment toward online outcome-based training, while noting persistent step-level credit-assignment challenges. It reviews learned process rewards and alternatives that derive finer-grained guidance from sampling or state transitions.

Reinforcement Learning for LLMs

  • Early PPO-based RLHF methods were computationally burdensome due to separate reward and value models, prompting a shift toward DPO and variants; offline DPO methods are inherently performance-bounded compared with online methods.Source: §5 · p. 10
  • Online outcome-based methods such as GRPO, DAPO, and GSPO train with verifiable rewards without SFT cold-start or process supervision, but retain step-level credit-assignment challenges.Process Reward Models offer finer-grained guidance, but the section notes annotation costs and susceptibility to reward hacking.Source: §5 · p. 10

Agentic Reinforcement Learning

  • Agentic RL work addresses sparse rewards through process rewards, entropy-based exploration, state-based grouping, hindsight criticism, and graph-based reward shaping.The section distinguishes GraphGPO’s heuristic reward shaping from the paper’s approach and points to Appendix A.2 for a detailed comparison.Source: §5 · p. 10

6 CONCLUSION

The conclusion identifies step-level bias in GRPO-style group advantages and presents GRAFT as a graph-based alternative grounded in standard RL definitions. It reports consistent benchmark gains and theoretical convergence under graph saturation.

6 CONCLUSION

  • GRPO-style group advantage estimation is reliable at response level but systematically biased when extended to individual steps.Source: §6 · p. 10
  • GRAFT merges rollout trajectories into a unified graph, estimates node values with Bellman iteration, and assigns edge credit by node-value differences.The conclusion states that the resulting step-level advantage adheres to the standard RL definition and converges to the true target as the graph becomes sufficiently saturated.Source: §6 · p. 10
  • Across diverse multi-turn agentic benchmarks, GRAFT consistently outperforms GRPO and recent state-of-the-art agentic RL algorithms while remaining critic-free and PRM-free.Source: §6 · p. 10

A.1 ADVANTAGES OVER GIGPO

GRAFT estimates step-level advantages by propagating value information across trajectories in a graph, rather than relying on returns from steps grouped only by state. The authors report advantages over GiGPO, GraphGPO, SALT, and RTMC, and describe GRAFT’s overall graph-based training procedure.

A.1 ADVANTAGES OVER GIGPO

  • GRAFT’s graph-based value estimates use downstream rewards beyond trajectories that visit a state, whereas GiGPO’s state-grouped estimates are high-variance and sample-inefficient for infrequently visited states.The step advantage γV(s′) − V(s) matches the standard advantage under deterministic transitions; GiGPO’s Monte Carlo returns have high variance from stochastic futures.Source: §A.1 · p. 15
  • Bellman propagation lets GRAFT’s state values aggregate reward signals from downstream trajectories, giving it a higher information-propagation degree than GiGPO throughout training.For a branching state whose successors are visited by K trajectories, its value incorporates rewards from all K, while GiGPO uses only trajectories that directly visit the state.Source: §A.1 · p. 15

A.2 DISCUSSION WITH GRAPHGPO

  • GRAFT treats the trajectory graph as an empirical MDP and applies Bellman iteration to estimate node values, then derives step advantages from value differences instead of GraphGPO’s shortest-path reward heuristic.GraphGPO assumes beneficial actions reduce distance to a successful state; GRAFT also uses a calibrated GRPO objective to address misalignment between step advantages and optimization.Source: §A.2 · p. 16
  • GRAFT consistently outperforms GraphGPO across multiple multi-turn benchmarks, achieves better results than GiGPO and SALT, and is theoretically guaranteed to converge to true advantages as the graph saturates.The authors attribute its performance to advantage estimates grounded in standard RL definitions rather than GraphGPO’s heuristic approach.Source: §A.1 · p. 15, §A.2 · p. 16, §A.3 · p. 16

A.3 DISCUSSION WITH SALT

  • Unlike SALT’s averaging of inherited trajectory-level advantages for repeated transitions, GRAFT uses graph Bellman bootstrapping to estimate advantages even for actions appearing only once.The authors report significantly better results than SALT in Table 1.Source: §A.3 · p. 16

A.4 DISCUSSION WITH RTMC

  • RTMC’s return averaging over state-action pairs is described as fundamentally equivalent to GiGPO’s state-grouped step advantage estimation, so the authors say GRAFT’s advantages over GiGPO apply equally to RTMC.RTMC averages returns for each (s, a) pair before state-group normalization, but the authors say this does not introduce an essential difference.Source: §A.4 · p. 17

B THE OVERALL ALGORITHMIC PROCEDURE FOR GRAFT

  • GRAFT’s procedure builds a graph from rollouts, estimates node values with Bellman iteration, computes step advantages from value differences, and updates the policy using calibrated optimization.It also extends GAE to the graph to obtain more robust advantage estimates.Source: B THE OVERALL ALGORITHMIC PROCEDURE… · p. 17

C PROOF OF PROPOSITION 1

The proof compares graph-based, state-grouped, and trajectory-level credit against the true state-dependent advantage. Under stated conditions, it establishes lower estimation error for graph-based credit and explains how its empirical accuracy depends on value estimates.

C PROOF OF PROPOSITION 1

  • The target at each step is the state-dependent advantage Aπ(St, At), whereas trajectory-level credit assigns the same signal to every step in a trajectory.For graph-based TD credit, the proof uses deterministic transitions so each state-action pair has a unique successor.Source: C PROOF OF PROPOSITION 1 · pp. 18–20
  • The population graph-based credit has lower estimation error than state-grouped and trajectory-level credit, with strict inequality when Ω has positive visitation probability under π.The proof’s comparison uses nonnegative population errors for the two baselines and establishes strictness under the stated visitation condition.Source: C PROOF OF PROPOSITION 1 · pp. 18–19
  • The empirical graph-based estimator approaches the zero-error population target as its graph-based state-value estimates become more accurate.With correct state aggregation and adequate coverage, additional independent rollouts can improve both empirical baselines and transition statistics, which act in complementary directions.Source: C PROOF OF PROPOSITION 1 · pp. 20–21
  • E(XG) ≤ min{E(XS), E(XE)} under the proposition’s rollout and environment assumptions, with both comparisons strict when condition 40 is strict.The result compares mean-squared advantage estimation error for GRAFT’s graph-based credit, GiGPO’s state-grouped credit, and GRPO’s trajectory-level credit.Source: C PROOF OF PROPOSITION 1 · p. 20

D.1 HYPERPARAMETER ABLATION STUDIES

The appendix examines Graph GAE’s discount and weighting factors, then analyzes how states recur across trajectories during training. The authors report robust performance across a reasonable hyperparameter range and substantial sharing of state groups.

D.1 HYPERPARAMETER ABLATION STUDIES

  • Larger γ yields better performance, while λ is typically tuned per benchmark because the optimal lookback depth varies with task complexity.The authors report robust performance across a reasonable range of hyperparameters.Source: §D.1 · p. 22, Table 4 · p. 22

D.2 NODE GROUP SIZE

  • WebShop and ALFWorld average node group sizes are around 4 and 5, respectively, and more than 60% of steps in both benchmarks belong to groups of at least two.Figure 5 shows WebShop group size initially rises and then falls, while ALFWorld’s rises amid repetitive task actions; the authors link shared state groups to more faithful step-level estimation.Source: §D.2 · p. 22, Figure 5 · p. 22, Figure 6 · p. 22

D.3 INFORMATION PROPAGATION DEGREE

GRAFT propagates reward information across more trajectories than GiGPO during training. The appendix reports that this reduces the share of steps receiving information from only one trajectory.

D.3 INFORMATION PROPAGATION DEGREE

  • Up to 60% of GiGPO steps receive reward information from only one trajectory, compared with 20% in GRAFT; GRAFT also has a higher propagation degree than GiGPO throughout training.The authors define propagation degree as the number of distinct trajectories contributing reward information to a node’s value estimate.Source: §D.3 · pp. 22–23, Figure 7 · p. 23, Figure 8 · p. 23

D.4 ADDITIONAL RESULTS

The additional evaluation extends GRAFT to Qwen3 models and reports stronger performance than GiGPO and GraphGPO on both Qwen3-4B and Qwen3-8B.

D.4 ADDITIONAL RESULTS

  • The evaluation reproduces GiGPO and GraphGPO for comparison on the Qwen3 model family.It extends the Qwen2.5-Instruct experiments to Qwen3-4B and Qwen3-8B.Source: §D.4 · p. 23
  • The authors conclude that GRAFT transfers beyond a single base model.The reported Qwen3 results support generalizability across the evaluated model families.Source: §D.4 · p. 23
  • GRAFT significantly outperforms GiGPO and GraphGPO on both Qwen3-4B and Qwen3-8B.The table covers ALFWorld and WebShop.Source: Table 5 · p. 23, §D.4 · p. 23

D.5 CASE STUDY

The ALFWorld case studies illustrate how graph state-values distinguish action quality in both all-success and mixed-outcome trajectories. The section also describes the benchmarks used for evaluation.

D.5 CASE STUDY

  • Graph state-values rise near successful terminal states, remain low where success is unreachable, and sit between these extremes for mixed-outcome states.Node color represents state-value; Figure 9 contains only successful trajectories, while Figure 10 mixes successful and failing ones.Source: §D.5 · p. 23, Figure 9 · p. 24
  • In the cloth task, directly putting the uncleaned cloth away receives negative advantage, while going to the sink supports the required cleaning sequence.The trajectory later succeeds after corrective actions, but the graph still distinguishes the earlier erroneous step.Source: §D.5 · pp. 23–24
  • In the kettle task, moving toward the second kettle raises downstream values, whereas returning to an already-visited location lowers them.The example includes both successful and failing trajectories and highlights decision points tied to progress toward the goal.Source: §D.5 · p. 24

E.1 DESCRIPTIONS OF BENCHMARKS

  • The evaluation benchmarks cover long-horizon embodied tasks, interactive online shopping, and single- and multi-hop search-augmented question answering.The named environments are ALFWorld, WebShop, and SearchQA.Source: §E.1 · p. 24

E.2 TRAINING DETAILS

Training follows prior GiGPO configurations for comparison, with benchmark-specific rollout, reward, and state-matching settings. Prompts provide task context and structured reasoning or action formats.

E.2 TRAINING DETAILS

  • Training and evaluation configurations follow GiGPO across benchmarks to support direct, fair comparison.The appendix then gives benchmark-specific settings.Source: §E.2 · p. 24
  • ALFWorld uses exact state matching, up to 50 environment steps, and a group size of 8 across group-based RL methods.Its rule-based rewards are +10 for success, 0 for failure, and −0.1 for invalid actions.Source: §E.2 · p. 26
  • WebShop uses exact state matching, episodes capped at 15 steps, and a group size of 8 for group-based RL methods.Its rule-based rewards are +10 for success, 0 for failure, and −0.1 for invalid actions.Source: §E.2 · p. 26
  • SearchQA uses similarity-based state matching with threshold τ = 0.9, a four-turn limit, and a group size of 5.Its rule-based rewards are +1 for success, 0 for failure, and −0.01 for invalid actions.Source: §E.2 · p. 26

E.3 TRAINING PROMPTS

  • Prompts include recent history, step-by-step reasoning, and explicit action or search-and-answer tags tailored to each benchmark.History length is 2 for ALFWorld and WebShop, while SearchQA retains full history.Source: §E.3 · p. 26

Related papers

Loading 2609.28963v1…