Source-linked AI summary
SPA-RL: Reinforcing LLM Agents via Stepwise Progress Attribution
Hanlin Wang, Chak Tou Leong, Jiashuo Wang, Jian Wang, Wenjie Li
TL;DR
LLM agents face sparse, delayed rewards that make credit assignment across long trajectories difficult. SPA redistributes final rewards into stepwise progress contributions and combines them with grounding signals, achieving consistent improvements across agent benchmarks in success rate and grounding accuracy.
Problem
Sparse rewards supplied only at task completion make it difficult to assign feedback to earlier actions in long-horizon LLM-agent tasks.
Method
SPA trains a progress estimator to decompose final rewards into stepwise contributions and combines them with grounding signals as intermediate RL rewards.
Results
SPA consistently achieves state-of-the-art performance across Webshop, ALFWorld, and VirtualHome in success rate and grounding accuracy.
Takeaways & Limitations
SPA provides more informative intermediate feedback for credit assignment and practical training of LLM agents on complex, long-horizon tasks.
Takeaways & Limitations
With PPO and GAE, early actions receive exponentially vanishing weights when rewards are sparse and delayed, limiting reward propagation across long trajectories.
Abstract
from arXiv · showhide
Reinforcement learning (RL) holds significant promise for training LLM agents to handle complex, goal-oriented tasks that require multi-step interactions with external environments. However, a critical challenge when applying RL to these agentic tasks arises from delayed rewards: feedback signals are typically available only after the entire task is completed. This makes it non-trivial to assign delayed rewards to earlier actions, providing insufficient guidance regarding environmental constraints and hindering agent training. In this work, we draw on the insight that the ultimate completion of a task emerges from the cumulative progress an agent makes across individual steps. We propose Stepwise Progress Attribution (SPA), a general reward redistribution framework that decomposes the final reward into stepwise contributions, each reflecting its incremental progress toward overall task completion. To achieve this, we train a progress estimator that accumulates stepwise contributions over a trajectory to match the task completion. During policy optimization, we combine the estimated per-step contribution with a grounding signal for actions executed in the environment as the fine-grained, intermediate reward for effective agent training. Extensive experiments on common agent benchmarks (including Webshop, ALFWorld, and VirtualHome) demonstrate that SPA consistently outperforms the state-of-the-art method in both success rate (+2.5\% on average) and grounding accuracy (+1.9\% on average). Further analyses demonstrate that our method remarkably provides more effective intermediate rewards for RL training. Our code is available at https://github.com/WangHanLinHenry/SPA-RL-Agent.
1 Introduction
SPA addresses delayed rewards in long-horizon LLM-agent tasks by redistributing final outcomes into stepwise progress signals and combining them with action-grounding feedback. The framework is reported to improve training and benchmark performance.
- Delayed end-of-trajectory rewards make it difficult to determine which earlier actions helped or harmed task completion.
- Existing process-supervision approaches emphasize short-term local optimization, which can overlook alignment with long-term task achievement.
- SPA decomposes delayed rewards into incremental stepwise contributions aligned with cumulative progress toward final task completion.
- SPA uses predicted stepwise progress as dense intermediate rewards and combines it with grounding signals for executable actions.
- SPA achieves state-of-the-art performance across Webshop, ALFWorld, and VirtualHome, improving success rate and grounding accuracy over existing methods.
2 Preliminaries
The paper formulates LLM-agent interaction as a POMDP and explains why PPO struggles when rewards are sparse and delayed. With GAE, credit assigned to early actions can decay exponentially, motivating intermediate supervisory signals.
- Agentic tasks are modeled as POMDPs with instruction, state, action, observation, transition, and reward spaces.
- An LLM agent repeatedly selects actions from its interaction history, receives observations, and stops when the task is completed or the step limit is reached.
- Sparse rewards are typically zero before the final step, where the terminal reward indicates task success or failure.
- PPO optimizes expected cumulative reward using clipped policy updates, advantage estimates, and a separately trained value network.
- When rewards occur only at termination, TD errors vanish for most steps, leaving early actions with exponentially vanishing weights under GAE.
- These limitations make vanilla PPO ineffective for long-horizon sparse-reward tasks and motivate SPA's intermediate supervisory signals.
3 Method
SPA constructs dense, goal-oriented intermediate rewards by estimating each step’s contribution to task completion and combining it with action executability. These fused rewards replace sparse terminal feedback during PPO-based policy optimization.
- 3.1 Behavior Cloning: SPA first gives LLM agents basic planning ability through supervised fine-tuning on successful expert trajectories in a ReAct-style format.Thought–action pairs are treated as single units for behavior cloning.
- 3.2 Reward Redistribution via Stepwise Progress Attribution: The progress estimator assigns each trajectory action a contribution score and constrains their sum to match the observed final task reward.It is trained with mean squared error between predicted task completion and the observed final reward.
- 3.2 Reward Redistribution via Stepwise Progress Attribution: SPA applies the trained estimator to exploration trajectories to produce per-step contribution scores that redistribute delayed rewards into dense feedback.The exploration data are collected by broadly sampling state–action sequences with the base agent.
- 3.3 Reinforcement Learning with Intermediate Rewards: The fused immediate reward combines progress contributions with a binary grounding signal indicating whether each action can be successfully executed.Positive hyperparameters α and β balance contribution against grounding.
- 3.3 Reinforcement Learning with Intermediate Rewards: The RL objective replaces sparse terminal rewards with fused rewards throughout return and advantage estimation while retaining PPO’s clipped surrogate form.This produces a dense reward stream that reflects both stepwise progress and action executability.
- 3.3 Reinforcement Learning with Intermediate Rewards: The fused reward improves credit assignment and policy optimization by combining task-success and environment-grounded learning signals.This distinguishes the approach from vanilla PPO, which relies solely on sparse end-of-episode feedback.
4 Experiments
The experiments evaluate SPA on WebShop, ALFWorld, and VirtualHome against supervised, rejection-sampling, PPO, and agent-specific RL baselines. SPA achieves the strongest reported ALFWorld performance, while WebShop grounding accuracy is omitted because its predefined action set yields 100% grounding.
- 4.1 Experimental Setup: The evaluation covers WebShop for web navigation and ALFWorld and VirtualHome for embodied household tasks with delayed final rewards.Each environment is modeled as a partially observable Markov decision process.
- 4.1 Experimental Setup: WebShop grounding accuracy is not reported because its predefined per-step action set consistently produces 100% grounding accuracy.The figure uses WS and VH to denote WebShop and VirtualHome.
- 4.1 Experimental Setup: The study compares SPA with SFT, RFT, PPO, ArCHer, and StepAgent baseline methods.These baselines include behavior cloning, rejection-sampling fine-tuning, direct final-reward PPO, and hierarchical or step-level RL approaches.
- 4.2 Overall Results: In unseen ALFWorld tasks, SPA reaches 79.1% task success, 3.6% above StepAgent, and achieves 91.7% grounding accuracy.It improves success rates in five of six task categories, including PICK at 95.8% and CLEAN at 90.3%.
- 4.2 Overall Results: Figure 3 compares five intermediate-reward variants in ALFWorld, with SPA outperforming all alternatives and the strongest heuristic, MC.Random and Mean rewards perform similarly to or worse than standard PPO.
5 Analyses and Discussions
SPA improves credit assignment by providing meaningful stepwise rewards, especially for long-horizon tasks, while fine-grained step-level attribution balances granularity and efficiency.
- Progress estimator effectiveness: SPA’s progress-based rewards outperform alternative reward variants, while random or uniformly distributed rewards provide little benefit over PPO.SPA achieves the highest success rate, whereas Random and Mean perform similarly to or worse than standard PPO.
- Credit assignment: Credit assignment identifies which earlier actions contributed to outcomes when rewards arrive later.This is the central problem addressed by redistributing delayed rewards across intermediate steps.
- Credit assignment: Trajectory-level credit assignment is coarse, whereas token-level methods perform better but incur optimization inefficiencies due to higher complexity.The analysis presents step-level credit assignment as a balance between granularity and efficiency.
6 Related Work
Related work on LLM agent learning spans supervised fine-tuning and reinforcement learning, while process supervision adds intermediate signals to address delayed feedback.
- LLM Agent Learning: Supervised fine-tuning methods generally use expert trajectories to improve agent planning capabilities.Some approaches also construct reflection trajectories through manual exploration strategies.
- LLM Agent Learning: LLM agent learning approaches are broadly divided into supervised fine-tuning and reinforcement-learning-based methods.These categories differ in how they improve planning and learn from environmental interaction.
- Process Supervision in RL: Process supervision addresses the difficulty of evaluating intermediate actions when agents receive rewards only after completing multi-step interactions.Intermediate signals can enhance reinforcement learning, although acquiring trajectory evaluations may require extensive environment interactions.
7 Conclusion
The paper proposes SPA to address sparse and delayed rewards by redistributing final rewards into stepwise signals. Experiments across agent benchmarks report improved success rates and grounding accuracy.
- Conclusion: SPA redistributes final rewards into fine-grained stepwise signals aligned with overall task progress.This provides more precise credit assignment for intermediate actions.
- Conclusion: SPA consistently achieves state-of-the-art performance across representative agent benchmarks in success rate and grounding accuracy.The conclusion attributes these outcomes to more informative intermediate feedback for complex, long-horizon tasks.
- Conclusion: SPA provides a practical and scalable approach for training LLM agents on complex, long-horizon tasks.Its supported scope is training with more informative intermediate feedback.
A Theoretical Analyses of the Progress Estimation
The theoretical analysis trains a progress estimator to decompose terminal rewards into dense stepwise contributions. Under stated assumptions, this decomposition preserves the original policy-gradient objective while supplying intermediate feedback.
- Progress estimator: The progress estimator predicts each action’s contribution and aggregates these contributions to match the observed task-completion reward.The estimator uses state-action representations and is trained so cumulative predicted reward aligns with ground truth.
- Reward decomposition: SPA represents each predicted contribution as a difference between consecutive potential values, avoiding explicit potential computation.The estimator directly predicts the potential difference for efficient stepwise reward construction.
- Reward decomposition: Enforcing an initial potential of zero makes cumulative predicted reward equal the terminal potential and supports alignment with the ground-truth reward.The estimator is trained to ensure the trajectory-level prediction matches observed task completion.
- Assumptions: The analysis assumes sufficient model capacity, independent sampling, and convexity of the expected risk for the stated minimizer property.These assumptions support the theoretical alignment between predicted cumulative contributions and observed rewards.
- Policy optimization: Replacing sparse terminal rewards with dense predicted contributions yields intermediate rewards for policy-gradient optimization.The construction is applied under the REINFORCE objective to improve credit assignment signals.
- Policy optimization: Terminal-state potential terms telescope to zero in expectation, leaving a gradient identical to that under the original sparse reward.Thus, reward decomposition supplies dense feedback without changing the policy optimization objective.
B PPO’s Limitations in Addressing Delayed Rewards
PPO with GAE struggles to assign delayed terminal rewards to early actions in long-horizon tasks because the relevant gradient weights decay over time.
- Under sparse rewards, intermediate rewards are zero and only the terminal reward may be nonzero.
- The one-step TD-error is zero for intermediate steps but generally does not vanish at the penultimate step.
- GAE weights earlier actions by (γλ)^(n−1−t), which decays exponentially whenever γλ < 1.
- Consequently, a single terminal reward cannot effectively propagate to decisions made many steps earlier, making PPO with GAE ineffective for long-horizon sparse-reward tasks.
- The evaluated environments include WebShop, ALFWorld, and VirtualHome, which represent online shopping and household-task settings.
D Implementation Details
The implementation trains the base agent and progress estimator with separate optimization settings, then uses deterministic inference accelerated by vLLM on LLaMA-3.2-3B-Instruct.
- The base agent is trained for 3 epochs with AdamW and cosine annealing, while the progress estimator uses exploration data for 1 epoch with Adam.
- Inference uses greedy decoding with temperature 0, and vLLM accelerates LLM generation.
- Experiments run on eight NVIDIA A6000 GPUs with 48GB each, using the open-source LLaMA-3.2-3B-Instruct model.
E Limitations
The paper identifies task-specific tuning as a generalization constraint and leaves SPA’s behavior in extremely large real-world action spaces and very long trajectories unexplored.
- SPA’s progress estimator requires domain-specific tuning, which could limit generalization to unseen or highly diverse environments.
- SPA’s performance remains unexplored in real-world environments with extremely large action spaces or trajectories spanning thousands of steps.