SWE-Shepherd: Advancing PRMs for Reinforcing Code Agents

1. Problem Statement

Automating real-world software engineering (SE) tasks using Large Language Model (LLM) agents faces significant hurdles due to the long-horizon reasoning required to navigate large, evolving codebases. Existing approaches typically rely on:

• Static prompting strategies or handcrafted heuristics to select actions (e.g., file reading, code editing, test execution).
• Sparse feedback, where the agent only receives a signal upon final task completion (success/failure).

Limitations of Current Approaches:

• Inefficient Exploration: Without intermediate feedback, agents often explore irrelevant paths.
• Error Propagation: Early mistakes are not corrected until the end, leading to brittle solution trajectories.
• Lack of Fine-Grained Supervision: Agents struggle to evaluate the quality of intermediate decisions, making it difficult to distinguish between a promising step and a dead end.

2. Methodology: SWE-Shepherd Framework

The authors propose SWE-Shepherd, a framework that operationalizes Process Reward Models (PRMs) to provide dense, step-level supervision for repository-level code agents. The framework operates in six distinct stages (illustrated in Figure 1 of the paper):

1. Task Collection: Utilizing the SWE-Bench dataset, which contains real-world GitHub issues, repository snapshots, and test suites.
2. Trajectory Generation: A baseline LLM agent interacts with the codebase to generate solution trajectories consisting of alternating actions and observations.
3. Reward Computation: Intermediate steps are assigned scalar rewards based on heuristics reflecting progress toward resolution. Key factors include:
   • Successful test execution.
   • Relevant file access and target file modification.
   • Avoidance of repetitive actions.
   • Discounted cumulative rewards to capture long-term contributions.
4. Dataset Construction: A supervised dataset is created pairing input text (problem description, interaction history, candidate action) with the computed reward labels (normalized to $[0, 1]$). The dataset contains over 15,000 samples.
5. Process Reward Model (PRM) Training:
   • Architecture: A lightweight model built on a base LLM (specifically 1.5B parameters in the released model) with an MLP head projecting the final hidden token to a scalar reward.
   • Training: Fine-tuned using qLoRA (parameter-efficient) with Mean Squared Error (MSE) loss to predict step-level rewards.
6. Reward-Guided Inference: During deployment, the trained PRM scores candidate actions at each step. The agent selects the highest-scoring action, enabling reward-guided search without the computational overhead of full Reinforcement Learning (RL).

3. Key Contributions

• PRM for Code Agents: The first framework to apply Process Reward Models specifically to repository-level software engineering tasks, moving beyond web navigation (e.g., Web-Shepherd) to complex code editing.
• Action-Level Reward Dataset: Construction of a new dataset derived from SWE-Bench trajectories that annotates intermediate actions with scalar rewards, bridging the gap between sparse final outcomes and dense behavioral feedback.
• Lightweight Alternative to RL: Demonstrates that PRMs can serve as a practical middle ground between supervised imitation learning and RL-based optimization, offering dense feedback while remaining simple to train and deploy.
• Open Source Release: The authors released the codebase, the action-level reward dataset, and the trained 1.5B PRM model to the community.

4. Experimental Results

The framework was evaluated on 100 tasks from the SWE-Bench Verified subset, with a maximum of 30 interaction steps.

Comparative Performance:

Method	% Resolved	Avg. Cost ($)	Avg. Steps
SWE-Search (MCTS-based)	31%	0.274	–
mini-SWE-Agent (Baseline)	57%	0.029	15.2
SWE-Shepherd (Ours)	51%	0.053	12.2

Key Findings:

• Efficiency: SWE-Shepherd reduced the average number of interaction steps from 15.2 to 12.2, indicating more directed exploration and higher interaction efficiency compared to the baseline.
• Cost-Effectiveness: Both agent-based approaches (mini-SWE-Agent and SWE-Shepherd) significantly outperformed the search-based SWE-Search in resolution rate while requiring substantially lower computational cost.
• Resolution Trade-off: SWE-Shepherd showed a modest drop in resolution rate (57% $\to$ 51%) compared to the baseline. This suggests that while the PRM guides the agent toward "locally" high-reward actions, these do not always guarantee a "globally" correct patch.

Reward Analysis:

• The average reward for resolved tasks (0.4894) was only marginally higher than for unresolved tasks (0.4818).
• Conclusion: The current heuristic-based reward function has a weak correlation with final task success, highlighting a misalignment between intermediate behavioral signals and repository-level correctness.

5. Significance and Future Directions

Significance:
SWE-Shepherd validates that process-level supervision can make code agents more efficient and interpretable. It proves that dense, step-level feedback can reduce exploration costs without the complexity of full RL training. However, it also exposes a critical challenge: alignment. Current heuristics for intermediate rewards are insufficient to perfectly predict final task success, leading agents to optimize for locally consistent but ultimately incomplete solutions.

Future Work:

• Improved Reward Modeling: Developing reward functions that better correlate intermediate steps with final task success.
• Data Expansion: Training on more diverse and complex tasks to improve generalization.
• Hybrid Approaches: Combining reward-guided search with Reinforcement Learning to further boost resolution rates.
• Adaptive Strategies: Investigating dynamic step limits and candidate action generation to enhance overall efficiency.