Source-linked AI summary
AgentPRM: Process Reward Models for LLM Agents via Step-Wise Promise and Progress
Zhiheng Xi, Chenyang Liao, Guanyu Li, Yajie Yang, Wenxiang Chen, Zhihao Zhang, Binghai Wang, Senjie Jin, Yuhao Zhou, Jian Guan, Wei Wu, Tao Ji, Tao Gui, Qi Zhang, Xuanjing Huang
TL;DR
LLM agents need to make sequential, context-sensitive decisions, but existing process supervision does not adequately capture action quality, dependencies, or scalable training costs. AgentPRM evaluates both each step’s promise and progress using TD-based estimation with GAE, then guides search and can support reinforcement learning. Across agent tasks, it outperforms baselines, is over 8× more compute-efficient, and improves more robustly as inference compute increases.
Problem
Agent-task actions lack clear-cut correctness, existing PRMs overlook dependencies between decisions, and prior training methods rely on costly annotations or Monte Carlo sampling.
Method
AgentPRM models each decision’s goal-achievement promise and sequential progress, using TD-based estimation with GAE for scalable training-data acquisition.
Results
AgentPRM consistently outperforms baselines across models, sampling strategies, and tasks, achieving over 8× greater compute efficiency and more stable improvement with increasing inference compute.
Takeaways & Limitations
AgentPRM provides a process-supervision approach for guiding multi-step LLM agents and shows potential for reinforcement learning and generalization to mathematical tasks.
Takeaways & Limitations
Monte Carlo-based estimation remains resource-intensive and requires many rollouts, motivating the paper’s TD-based alternative.
Abstract
from arXiv · showhide
Despite rapid development, large language models (LLMs) still encounter challenges in multi-turn decision-making tasks (i.e., agent tasks) like web shopping and browser navigation, which require making a sequence of intelligent decisions based on environmental feedback. Previous work for LLM agents typically relies on elaborate prompt engineering or fine-tuning with expert trajectories to improve performance. In this work, we take a different perspective: we explore constructing process reward models (PRMs) to evaluate each decision and guide the agent's decision-making process. Unlike LLM reasoning, where each step is scored based on correctness, actions in agent tasks do not have a clear-cut correctness. Instead, they should be evaluated based on their proximity to the goal and the progress they have made. Building on this insight, we propose a re-defined PRM for agent tasks, named AgentPRM, to capture both the interdependence between sequential decisions and their contribution to the final goal. This enables better progress tracking and exploration-exploitation balance. To scalably obtain labeled data for training AgentPRM, we employ a Temporal Difference-based (TD-based) estimation method combined with Generalized Advantage Estimation (GAE), which proves more sample-efficient than prior methods. Extensive experiments across different agentic tasks show that AgentPRM is over $8\times$ more compute-efficient than baselines, and it demonstrates robust improvement when scaling up test-time compute. Moreover, we perform detailed analyses to show how our method works and offer more insights, e.g., applying AgentPRM to the reinforcement learning of LLM agents.
1. Introduction
AgentPRM addresses the difficulty of evaluating and guiding sequential agent decisions by modeling both each step’s promise and its progress toward the goal. Its TD-based estimation with GAE supports scalable training, while experiments report improved performance, compute efficiency, and inference-scaling behavior.
- Motivation: Agent tasks require context-sensitive, forward-looking decisions based on environmental feedback rather than isolated static predictions.Effective performance depends on interpreting environmental cues and anticipating the consequences of decisions.
- Motivation: Prior approaches rely on scarce expert-labeled trajectories or costly, inflexible prompt engineering with commercial models.Supervised fine-tuning is difficult to scale, while prompt engineering is constrained by APIs and customization costs.
- Motivation: Existing PRMs struggle because agent actions lack clear-cut correctness, sequential dependencies are overlooked, and training often requires costly annotations or Monte Carlo sampling.These challenges make step evaluation and scalable reward-model training non-trivial.
- AgentPRM: AgentPRM predicts each decision’s contribution to the final goal while capturing interdependencies between sequential decisions to track progress and balance exploration with exploitation.The model evaluates both immediate progress and long-term promise rather than promise alone.
- Training: TD-based estimation with GAE provides an automated, scalable training-data method that is more efficient than prior Monte Carlo-based methods.The approach is designed to provide a better variance–bias trade-off in estimation.
- Results: AgentPRM achieves over 8× greater compute efficiency than baselines across three agent tasks and multiple sampling strategies, with stable improvement as inference compute scales.The experiments cover various models and tasks and report consistent gains in performance and compute efficiency.
2. Preliminary and Background
The paper formalizes agent interaction as a partially observable decision process and distinguishes trajectory-level outcome models from step-level process supervision. It then uses reward models to evaluate intermediate decisions and guide search over candidate actions.
- Agent-task formulation: An agent task is formalized as a POMDP with instruction, state, action, observation, transition, and reward spaces.The agent repeatedly selects actions from interaction history, receives observations, and forms a trajectory until completion or a step limit.
- Reward models: An outcome reward model takes an instruction–trajectory pair and predicts whether the completed trajectory satisfies the task.ORM training labels sampled trajectories with their corresponding outcome rewards.
- Reward models: Process reward models evaluate actions or intermediate states, but correctness-based step scoring from reasoning tasks is unsuitable for agent tasks.PRMs are trained from step-level annotations and optimized to predict those labels.
- Search with reward models: Best-of-N samples N trajectories, evaluates them with a reward model, and selects the highest-scoring trajectory as the final output.In this setting, the score of the final step represents trajectory quality.
- Search with reward models: Step-level beam search expands M candidate actions per node, scores them with a PRM, retains the top N candidates, and returns the highest-scoring terminal trajectory.The method is presented as a search procedure for guiding agent decisions with process rewards.
3. Methodology
AgentPRM redefines process rewards for agent tasks by evaluating both each decision’s contribution toward the goal and its progress relative to sequential decisions. It is trained with TD-based estimation and GAE, then used for inference-time search or Best-of-N evaluation.
- AgentPRM motivation: Agent tasks require rewards that capture both a decision’s promise of reaching the goal and its progress through sequential dependencies.This addresses unclear action correctness and the limitations of treating states independently.
- Value-based rewards: AgentPRM models action values as expected future success probabilities for state-action pairs.The PRM is trained to predict action values from annotated state-action labels using mean squared error.
- Inference-time use: AgentPRM predictions support inference-time search and Best-of-N selection for agent tasks.Beam search expands candidate actions, scores them with a PRM, retains top candidates, and returns the highest-scoring terminal trajectory.
- Capturing progress: Advantage-based rewards measure local progress by comparing an action’s contribution with the state’s expected value.Positive advantage indicates progress, while negative advantage suggests the action moves away from progress; the resulting loss combines value and advantage fitting.
- TD-based estimation with GAE: TD-based estimation with GAE replaces costly MC-based rollouts for scalable AgentPRM training.GAE reduces variance and improves stability, while avoiding additional rollouts from each state.
4. Experiments
Experiments evaluate AgentPRM across WebShop, BabyAI, and TextCraft using Best-of-N and search, showing consistent gains over reward-model baselines and stable scaling with inference compute.
- Experimental Setup: Experiments cover WebShop, BabyAI, and TextCraft with Best-of-N and beam-search evaluation across multiple models and sampling budgets.WebShop contains 1.18 million products; BabyAI tests grid-world instruction following; TextCraft tests multi-step crafting and fetching.
- Main Results: Reward models improve agent-task performance over greedy decoding, especially as inference compute increases for additional sampling.This pattern is reported for Best-of-N and search strategies compared with greedy SFT and RFT decoding.
- Main Results: AgentPRM consistently outperforms ORMs and PVMs in both Best-of-N and test-time search across tasks and sampling budgets.With Qwen2.5-3B on WebShop under 8 × 8 sampling search, AgentPRM surpasses PVM by more than 20.0 points.
- Main Results: 8× greater compute efficiency is achieved by AgentPRM compared with PVMs and ORMs on average.The comparison is reported across the evaluated agent tasks and sampling strategies.
- Scaling Behavior: AgentPRM maintains a more robust and stable improvement trend as inference compute scales, while PVMs and ORMs can plateau or degrade.The authors associate baseline degradation with issues such as false positives or reward hacking.
5. Discussion and Analysis
Analyses examine AgentPRM’s progress term, reinforcement-learning behavior, mathematical-reasoning transfer, labeling efficiency, credit assignment, and generalization across model sizes.
- Ablation Analysis: Removing L_A(ϕ) reduces agent-task performance across sampling strategies, supporting the importance of progress modeling.The ablation evaluates the advantage term used to capture dependencies between steps.
- Reinforcement Learning: AgentPRM produces more stable and effective reinforcement-learning optimization than baseline reward models on BabyAI and TextCraft.The experiments use Qwen2.5-3B and PPO; baseline methods show optimization instability or slower improvement.
- Mathematical Reasoning: AgentPRM surpasses other baselines on GSM8K after mathematical reasoning is formulated as multi-turn decision making.Partial solutions serve as states and successive reasoning steps serve as actions.
- Sampling Efficiency: TD-based estimation with GAE uses fewer labeling tokens than MC-based estimation while achieving better Best-of-N and beam-search performance.The comparison evaluates both sampling cost and downstream performance.
- Credit Assignment: Action-value estimates assign higher scores to actions leading to positive goals and lower scores to actions leading to negative goals.Visualizations on WebShop and BabyAI indicate effective credit assignment across successful and unsuccessful trajectories.
- Generalization: AgentPRM consistently outperforms baselines on Qwen-2.5-7B-Instruct and Llama3.1-8B-Instruct across WebShop, BabyAI, and TextCraft.The result supports generalization across model sizes and architectures.
6. Related Work
Related work includes training-based agents, prompt-engineering approaches, and process reward models, while highlighting challenges in scalability, customization, and step-level supervision.
- Training-Based Agents: Training-based agent methods imitate expert-labeled trajectories but are difficult to scale and provide limited environmental exploration.These methods collect expert trajectories and train learners to imitate them step by step.
- Prompt Engineering: Prompt-engineering approaches rely on commercial models such as GPT-4o, creating API constraints that hinder customization.The cited limitation concerns dependence on external commercial-model access.
- Process Reward Models: Process reward models score intermediate actions or states and can provide dense rewards for reinforcement learning, test-time search, or reranking.In LLM reasoning, their scoring criterion typically concerns the correctness of individual steps.
- Process Reward Models: PRM labeling is expensive and difficult to scale, motivating automated annotation methods based on Monte Carlo sampling.The related-work discussion identifies sampling-based annotation as a response to labeling cost.
7. Conclusion
AgentPRM is introduced as process supervision for multi-step LLM-agent decision-making, modeling both each step’s goal-achievement promise and sequential progress. Experiments report improvements across varied settings, compute efficiency, scaling with inference compute, mathematical-task generalization, and analyses of data efficiency and architecture adaptability.
- AgentPRM captures each step’s probability of achieving the goal and the interdependence between sequential steps.These dimensions are termed promise and progress.
- AgentPRM outperforms other baselines across various sampling strategies, models, and tasks.
- AgentPRM is more compute-efficient, with performance improving robustly as inference compute increases.
- AgentPRM generalizes to mathematical tasks, demonstrating versatility beyond LLM-agent settings.
- Additional analyses and ablations examine AgentPRM’s operation, data efficiency, and adaptability across model architectures and sizes.
A. More Detailed Discussion of Related Work
Related work spans outcome and process supervision for LLM agents and reasoning, using trajectory reranking, Bellman estimation, planning trees, Monte Carlo estimation, and advantage-based rewards. Table 4 compares process-supervision paradigms by supervision granularity and progress.
- ARMAP labels data to construct outcome reward models that rerank trajectories and improve Best-of-N performance.
- Q* AGENT estimates each step’s Q-value with the Bellman equation to train process reward models.
- DPO-Q builds a planning tree with MCTS and uses DPO to estimate the value of each step.
- PQM models relationships between reasoning steps using Monte Carlo-based estimation and a ranking loss.
- PAV estimates whole-trajectory rewards through an outcome reward model and uses individual-step advantages for reinforcement learning and search.
- Table 4 compares process-supervision paradigms using supervision granularity and progress as dimensions.
B. Algorithm
The algorithm section presents AgentPRM training and beam search procedures, including trajectory collection and model training. Its stated inputs include an initialized model, a reward function, and a per-query temporal-difference sample count.
- AgentPRM training and beam search are presented in Algorithm 1 and Algorithm 2, respectively.
- Algorithm 1 is titled “Training of AgentPRM.”
- The algorithm takes an initialized AgentPRM model, a reward function, and sample number per query N_TD as inputs.
- The training procedure includes trajectory collection and AgentPRM model training.
C. More Implementation Details for RL
For reinforcement learning, the authors train LLM agents with PPO under specified optimization and interaction settings, using the PRM’s final-step score as the trajectory reward.
- LLM agents are trained with PPO using batch size 16, learning rate 1 × 10^-6, KL coefficient 1 × 10^-3, and sampling temperature 1.0.
- BabyAI and TextCraft use a maximum interaction horizon of 20.
- In PRM-based reinforcement learning, the PRM’s predicted score at the final step serves as the trajectory reward.
D. Qualitative Analysis
AgentPRM guides beam search by assigning different values to candidate actions, helping the agent explore alternatives and select successful steps. In the illustrated task, the policy completes the goal in 10 steps under AgentPRM guidance.
- D. Qualitative Analysis: AgentPRM-guided beam search distinguishes good actions from bad actions through high and low action scores.The qualitative analysis describes this scoring as the mechanism by which AgentPRM guides search.
- D. Qualitative Analysis: The beam-search procedure expands candidate actions at each node using a trained PRM and policy.The algorithm is presented as step-level beam search with PRM, with candidate steps selected or discarded during search.
- D. Qualitative Analysis: The search begins with an exploratory move toward blue ball 1 while the goal is to pick up a grey ball.The initial state shows the agent facing a wall near blue ball 1.
- D. Qualitative Analysis: A later action picks up grey box 1 to create more options while continuing exploration toward grey ball 2.This candidate receives value 0.1931 in the illustrated trajectory.
- D. Qualitative Analysis: The successful terminal action picks up grey ball 1 and receives value 0.8913 after the task is completed.This is the highest explicitly shown value among the cited qualitative steps.
- D. Qualitative Analysis: The figure reports that the policy solves the task in 10 steps under AgentPRM-guided beam search.Figure 9 presents this as a successful qualitative example.