RTMC: Step-Level Credit Assignment via Rollout Trees

Advantage estimation challenge

• GRPO assigns one advantage to an entire trajectory, while learned critics add overhead and can struggle with sparse binary rewards.These limitations leave multi-turn episodes without a way to distinguish beneficial steps from harmful ones.Source: §1 · pp. 1–2

Our key insight

• RTMC aggregates return statistics across rollouts that reach a common state to estimate per-step advantages without training an additional model.The method treats each tool call as an action and uses state-action signatures to recognize equivalent states across differing histories.Source: §1 · p. 2

Contributions

• RTMC organizes group rollouts into a tree to compute per-step Q-values and advantages, with prior-based smoothing for sparsely visited nodes.The smoothing is intended to provide informative advantages throughout the tree.Source: §1 · p. 2
• State-action signatures compress interaction histories into compact, comparable representations that make tree construction tractable.For software engineering tasks, the system uses view bucketing, content hashing, and action classification.Source: §1 · p. 2

Empirical validation

• 52.2% pass@1 on SWE-bench Verified, a 3.2 percentage point improvement over GRPO, with additive gains across the progressive ablation.The ablation progresses from the untrained baseline through GRPO and step-level rewards to Rollout-Tree MC.Source: §1 · p. 2

Single-turn RLVR

• GRPO shares a group-level advantage across all response tokens, so it cannot distinguish good steps from bad ones within a multi-step episode.This contrasts with agentic trajectories, where steps contribute unevenly to the outcome.Source: §2 · p. 3

Multi-turn agentic RL as an MDP

• Agentic tasks involve multi-turn interactions in which a tool action's consequences may appear only many steps later.The paper formalizes these tasks as an MDP with problem-solving progress as state, tool calls as actions, environment transitions, sparse terminal rewards, and a discount factor.Source: §2 · p. 3

Credit assignment and advantage estimation

• Prior methods use anchor-state grouping, implicit process rewards, retrospective critics, extra branch rollouts, trajectory purification, or reflective retries, but retain limitations in components, computation, or step-level estimates.The paper also contrasts prefix-tree methods: VinePPO requires extra inference, while TEMPO uses raw token prefixes as state identity.Source: §2 · p. 3
• RTMC decouples state identity from raw context with explicit signatures, enabling cross-rollout tree construction and per-step advantage estimation without extra inference or learned components.The related-work discussion identifies raw token-prefix identity as a poor fit when equivalent agentic progress follows different histories.Source: §2 · p. 3

3 Rollout-Tree Monte Carlo Advantage Estimation

• RTMC preserves critic-free training while seeking step-level discrimination beyond GRPO's uniform trajectory advantage.Rollouts sharing a state and then taking different actions provide comparisons for estimating action quality without additional models or computation.Source: §3 · p. 4

3.1 From Rollouts to Trees

• State-action signatures make states and actions comparable across rollouts; cumulative-history signatures appear at most once per rollout, satisfying the first-visit condition.Each trajectory alternates states, actions, and rewards, with variable-length action spans in the agentic setting.Source: §3.1 · p. 4

3.2 Monte Carlo Advantage Computation

• The first-visit Monte Carlo estimator averages discounted returns across rollouts visiting each state-action pair to estimate Q, then subtracts the state value for the advantage.A step-level advantage is broadcast to the tokens in that action segment.Source: §3.2 · p. 5

3.3 Procedure and Properties

• Two linear passes build the statistics table and assign per-step advantages, with O(N · L) scalar-operation complexity and optional per-problem normalization.The first-visit estimates are unbiased; variance is higher for leaf states, and the tree collapses to trajectory-level GRPO at the coarsest state abstraction.Source: §3.3 · p. 6

3.4 Corrections to the Advantage Estimator

• Single-visit states yield zero advantage, so prior smoothing blends local estimates with the group success rate using n_prior = 2.At one visit, the prior contributes 67%, and its contribution vanishes as visits grow.Source: §3.4 · pp. 6–7
• For coding tasks, validation after code changes receives a small reward to counter discounting that can penalize multi-step verification.The shaping uses β = 0.05, preserves the optimal policy, and may not be needed in other agentic domains.Source: §3.4 · p. 7

4 State-Action Signature Design

• Naïve histories rarely match despite equivalent progress, while overly coarse signatures mix returns and overly fine ones approach trajectory-level estimation.Signatures address both combinatorial histories and functional overlap between tools such as bash and dedicated commands.Source: §4 · p. 7

4.1 Action Signature

• Tool invocations are classified by effect on progress, with equivalent tools mapped to the same category and structured signatures recording scope, target, and applicable results.View, modify, test, and execute actions retain operation-specific distinctions; think and finish are encoded without file targets.Source: §4.1 · p. 7, Table 1 · p. 8

4.2 State Signature

• State signatures record ordered per-file operation sets and non-file flags, abstracting action order while retaining progress distinctions.File views use 100-line buckets, and edits use four-character hashes of concatenated old and new text.Source: §4.2 · p. 8

5.1 Experimental Setup

• The setup evaluates four configurations on SWE-bench Verified: pretrained baseline, GRPO, GRPO with step rewards, and Rollout-Tree MC.Training uses an 8.1K-instance R2E-Gym subset without SWE-bench overlap; RL configurations share 64 trajectories per step and common optimization settings.Source: §5.1 · p. 9

5.2 Main Results

• Each refinement in the ablation improves performance, with stage gains of 1.4–2.2 percentage points and Rollout-Tree MC performing best.Table 3 reports pass@1 on 500 SWE-bench Verified problems, using the best checkpoint and majority vote over three runs.Source: §5.2 · p. 9, Table 3 · p. 9

5.3 Learning Curves

• Across 150 R2E training steps, resolve rates separate progressively from episode-level GRPO through step rewards to tree-based advantages.The metric includes uniform-outcome rollout groups excluded from gradient updates, so it reflects overall problem-solving ability rather than training-signal quality.Source: §5.3 · pp. 9–10, Figure 2 · p. 10

5.4 Ablation Study

• Removing prior smoothing drops pass@1 from 52.2% to 49.7%, below GRPO + Step Reward at 50.4%.The authors attribute the loss to states visited only once, whose zero advantages become informative when value estimates are anchored to the group success rate.Source: §5.4 · p. 10, Table 4 · p. 10

6.1 Advantages of the Approach

• Rollout-Tree MC compares each action’s return with outcomes from rollouts sharing its state, enabling within-trajectory differentiation of good and bad actions.This grounds advantages in relative action quality at each decision point rather than absolute returns subject to problem-difficulty confounds.Source: §6.1 · p. 11
• Tree construction and advantage computation are linear in rollout count and episode length, with no separate value network, gradient computation, or extra inference.The method computes advantages from rollout return statistics alone.Source: §6.1 · p. 11

6.2 Limitations and Future Work

• The absolute gains remain modest, and the authors attribute the current performance ceiling primarily to hand-crafted signatures.The reported improvement is 5.4 pp over the untrained baseline and 1.8 pp over GRPO + Step Reward.Source: §6.2 · p. 11
• Manually defined signatures can merge distinct states when surface heuristics miss action effects or fragment the tree when distinctions are too fine.The signature rules include action classification, view bucketing, and content hashing.Source: §6.2 · p. 11
• With N=8 rollouts per problem, most states are visited once, so raw Monte Carlo estimates yield zero advantage and the prior cannot replace genuine multi-action comparisons.More rollouts or coarser signatures could increase shared states; the validation reward shaping is specific to SWE tasks where verification steps are penalized by γ-discounting.Source: §6.2 · p. 11
• Learned signatures could capture semantic state similarities missed by rules, but would reintroduce dependence on an external model.The tree estimator itself is agnostic to how signatures are produced; the authors suggest gains may widen as signature quality or rollout budgets increase.Source: §6.2 · p. 11
• RL training solves some problems the base checkpoint fails but also regresses on a subset of previously solved instances, leaving mitigation open.The authors identify better credit assignment, curriculum design, and checkpoint selection as possible directions to investigate.Source: §6.2 · pp. 11–12
• Evaluation uses SWE-bench Verified with one base model, so broader generality across benchmarks, task domains, and model families remains unestablished.Source: §6.2 · p. 12

• Rollout-Tree MC organizes group rollouts into a shared tree to enable step-level credit assignment for agentic RL without a critic.Source: §7 · p. 12
• Experiments on SWE-bench Verified show consistent gains over trajectory-level and step-level baselines, with each refinement in credit assignment contributing to the overall improvement.Source: §7 · p. 12
• The approach adds negligible overhead and integrates directly into existing GRPO-based pipelines.Source: §7 · p. 12
• Hand-crafted state-action signatures are the main limitation, and combining rule-based abstraction with learned representations is proposed for scaling to broader agentic domains.Source: §7 · p. 12

• RTMC down-weights frequently visited actions with low estimated advantage and steers search toward branches reaching finish with high success likelihood.Source: A Rollout Tree Visualization · p. 14
• Nodes represent states and edges represent actions; green marks passing finishes, red failing finishes, and gray incomplete trajectories.The root is the episode’s initial state; incomplete trajectories ended by timeout or exceeding maximum context length.Source: A Rollout Tree Visualization · p. 15
• The visualization is a rollout-tree example extracted from training-time rollouts on datalad_final_05ea.Source: Figure 3 · p. 15

B.1 Per-Rollout Advantage Comparison

• GRPO assigns uniformly positive advantages to success rollouts and uniformly negative advantages to failure rollouts, without within-trajectory discrimination.Each rollout is shown across the GRPO, GRPO+Step, and RTMC methods.Source: §B.1 · p. 16, Figure 4 · p. 16
• GRPO+Step improves slightly over GRPO through gamma-discounted cumulative returns but remains coarser than RTMC.Source: §B.1 · p. 16, Figure 4 · p. 16
• RTMC produces step-varying advantages, including sign changes within a rollout, reflecting local action quality.The comparison covers eight rollouts: four successful and four failed.Source: Figure 4 · p. 16, §B.1 · p. 16

B.2 Quadrant Analysis

• The analysis classifies steps into four agreement quadrants using an advantage threshold of |A| > 0.01.The quadrants distinguish both methods agreeing positive or negative from RTMC rescuing or penalizing a step relative to GRPO+Step.Source: §B.2 · p. 17
• 15% of success-rollout steps fall where GRPO+Step assigns positive advantage but RTMC penalizes the action as suboptimal.In failure rollouts, RTMC rescues 4 steps (3%) with positive relative advantage, while 88% fall in the Both− quadrant.Source: Figure 5 · p. 17, §B.2 · p. 17