Source-linked AI summary

Process Reward Models for LLM Agents: Practical Framework and Directions

Sanjiban Choudhury

arXiv:2502.10325v1cs.LGcs.AI

TL;DR

LLM agents must improve through interaction despite long-horizon decisions, sparse rewards, and limited supervision. The paper introduces AgentPRM and InversePRM, then evaluates them on ALFWorld, where small trained models outperform strong GPT-4o baselines and InversePRM reaches near-expert performance with fewer rollouts.

  • Problem

    LLM agents face long-horizon decision-making and sparse rewards, while process rewards remain underexplored for agentic settings and manual reward design is labor-intensive.

  • Method

    AgentPRM iteratively trains process reward models and policies using Monte Carlo rollouts, while InversePRM learns process rewards directly from expert demonstrations without explicit outcome rewards.

  • Results

    Small models trained with AgentPRM outperform strong GPT-4o baselines on ALFWorld, while InversePRM achieves near-expert performance with significantly fewer rollouts.

  • Takeaways & Limitations

    Process reward modeling offers a practical route for improving LLM agents through interaction, with demonstrations providing a more sample-efficient alternative in the reported setting.

  • Takeaways & Limitations

    Traditional stochastic exploration strategies do not scale well to high-dimensional LLM action spaces, and process rewards differ from one-step rewards because they represent cumulative rewards.

Abstract

from arXiv · show

We introduce Agent Process Reward Models (AgentPRM), a simple and scalable framework for training LLM agents to continually improve through interactions. AgentPRM follows a lightweight actor-critic paradigm, using Monte Carlo rollouts to compute reward targets and optimize policies. It requires minimal modifications to existing RLHF pipelines, making it easy to integrate at scale. Beyond AgentPRM, we propose InversePRM, which learns process rewards directly from demonstrations without explicit outcome supervision. We also explore key challenges and opportunities, including exploration, process reward shaping, and model-predictive reasoning. We evaluate on ALFWorld benchmark, show that small 3B models trained with AgentPRM and InversePRM outperform strong GPT-4o baselines, and analyze test-time scaling, reward hacking, and more. Our code is available at: https://github.com/sanjibanc/agent_prm.

1 Introduction

The paper presents process reward models as a scalable way for LLM agents to improve through interaction, addressing sparse feedback and limited self-correction. AgentPRM uses iterative Monte Carlo annotation and actor-critic-style training, while InversePRM learns process rewards from demonstrations and both are evaluated on ALFWorld.

  • LLM agents need alternatives to prompting and SFT because these approaches require manual effort or high-quality demonstrations and lack autonomous test-time self-correction.
  • Process reward models provide fine-grained supervision for intermediate actions, addressing sparse outcome rewards in long-horizon agent tasks.
  • AgentPRM computes PRM targets with asynchronous Monte Carlo rollouts and jointly refines PRMs and policies through iterative training.
  • AgentPRM follows a lightweight actor-critic paradigm and integrates with existing RLHF infrastructure by adding automatic reward annotation.
  • InversePRM learns process rewards directly from demonstrations without explicit outcome rewards and achieves higher sample efficiency than AgentPRM.
  • On ALFWorld, small 3B models trained with AgentPRM and InversePRM outperform strong GPT-4o baselines, alongside analyses of scaling, reward hacking, and related challenges.

2 Agent Process Reward Models: A Simple Framework

AgentPRM lifts RLHF-style training to turn-level agent interactions by learning process rewards from asynchronous rollouts and iteratively refining both the PRM and policy. On ALFWorld, iterative training improves success, while test-time scaling and reward optimization reveal clear limits and failure modes.

  • Formulation: AgentPRM models turn-level agent decisions with PRMs that provide intermediate state-action feedback instead of relying only on terminal outcome rewards.The PRM functions analogously to a Q-function in RL, while LLM actions contain both reasoning tokens and environment actions.
  • Training framework: The framework follows three stages: collect rollouts and compute targets, train the PRM, then update the policy against the PRM while regularizing toward the previous policy.The stages are designed to fit standard RLHF pipelines, with rollout-derived targets replacing preference labels.
  • Training framework: Asynchronous Monte Carlo rollouts automatically produce soft PRM targets, and iterative policy-PRM updates allow each component to refine the other.Rollouts are collected in parallel and aggregated into a dataset for supervised PRM training.
  • Results: 88.1% success rate and 91.0% in best-of-N mode are achieved by AgentPRM, with π2 improving from 73.9% to 85.8% and exceeding claude-3.5-sonnet at 76.1%.The reported comparison also gives 12.0 actions for π2 versus 19.0 for claude-3.5-sonnet.
  • Results: Success rises across policies from π0: 64.9% to π3: 88.1%, but later iterations saturate as policies over-optimize the PRM and approach model-capacity limits.The largest increase is from π1: 73.9% to π2: 85.8%, with gains appearing early in training.
  • Test-time scaling: Best-of-N scaling helps earlier policies substantially, whereas gains diminish for π2 and π3 because headroom is limited and reward over-optimization increases.The comparison varies N from 1 to 32.
  • Reward hacking: With 10k rollouts, success falls from 82% to 70% after 400 steps while validation process reward keeps increasing, indicating reward hacking.The process reward therefore diverges from the outcome reward during continued optimization.
  • Losses: Absolute and relative PRM losses yield similar performance, although the relative-loss dataset is smaller because states must be visited multiple times.The relative formulation ranks action pairs using a minimum value margin and a Bradley-Terry loss.

3 Inverse Process Reward Models

InversePRM learns process rewards from expert demonstrations and iteratively refines both the reward model and policy without explicit outcome rewards. On ALFWorld, it outperforms SFT and AgentPRM while benefiting from dense demonstrations and stronger policy quality at test time.

  • 3 Inverse Process Reward Models: InversePRM addresses settings without outcome rewards by learning process rewards from expert demonstrations and past learner rollouts.It frames reward inference as distinguishing expert behavior from learner behavior, then iteratively updates the reward model and policy.
  • 3 Inverse Process Reward Models: The method replaces one-step reward optimization with a PRM formulation based on Q-values and updates the policy to maximize the learned PRM while regularizing against the previous policy.The Q-value difference compares expert and learner transitions, while KL regularization limits policy drift and reward-hacking risks.
  • 3 Inverse Process Reward Models: InversePRM uses a three-stage loop: construct positive and negative transitions, train a PRM classifier, and train the policy with reinforcement learning.Expert transitions are positive examples, learner transitions are negative examples, and the classifier operates on differences in PRM values.
  • 3.3 Experiments: 86.6% vs. 63.4%: InversePRM outperforms SFT on ALFWorld using the same expert demonstrations, while reaching 82.8% vs. 73.9% against AgentPRM after one iteration.Its final policy approaches expert performance at 86.6% vs. 91.0%, despite AgentPRM using 70K rollouts.
  • 3.3 Experiments: 64.9% →82.8%: success rate improves sharply in InversePRM’s first iteration, whereas AgentPRM requires multiple iterations to reach similar performance.The passage attributes this difference to demonstrations that implicitly capture successful strategies and reduce exploration demands.
  • 3.3 Experiments: Policy quality matters more than Best-of-N scaling: BoN(π0, Q0) improves 64.9% →69.0%, whereas BoN(π1, Q0) reaches 88.0% and later performance saturates.The reported inference sweep varies N from 1 to 32, with saturation in the second iteration.

4 Challenges and Opportunities

The paper examines exploration, process reward shaping, and reducing costly environment interactions as central challenges for training LLM agents. It evaluates structured exploration and reference-policy shaping as practical responses.

  • Exploration: Traditional stochastic exploration methods do not scale well to high-dimensional, long-horizon tasks where reasoning quality matters.
  • Exploration: Reset-50-50 mixes states from the initial policy and an expert policy, exposing agents to good states early to bootstrap recovery from errors.The mixture uses 50% states from π0 and 50% from the expert policy.
  • Exploration: 82% > 73.9%: Reset-50-50 reaches a higher peak than the baseline while learning faster.The comparison uses OnlineDPO(π0, Q0), with the same starting policy and PRM; only the prompts used in Stage 3 change.
  • Exploration: 84% > 73.9%: Steered Exploration learns faster and reaches a higher peak than the baseline.It generates more diverse, higher-quality reason-actions than simply increasing sampling temperature, while training on the original prompt.
  • Process Reward Shaping: With 10k rollouts, non-shaped rewards are unstable at low sample regimes, whereas shaped rewards produce much more stable performance.
  • Process Reward Shaping: 64% → 32%: SFT initialization can still produce unstable RL because noisy PRM targets cause the policy to unlearn good behavior.The policy later climbs back up, but the reference policy is not used during RL after initialization.
  • Process Reward Shaping: 82.0% > 65.0%: process reward shaping lets the learned policy outperform the moderate-performance reference policy.For α = 0.5, the reference advantage counters noisy PRM targets and the learning curve rises steadily to 700 steps.
  • Model-Predictive Reasoning: Costly environment interactions motivate model-based approaches that simulate future trajectories using learned world models.

5 Related Work

Related work spans prompting and fine-tuning agents, privileged-feedback critics, hierarchical RL, and process reward models. The paper positions its contribution as an agent-focused, RLHF-compatible PRM framework for larger models and practical efficiency techniques.

  • Agent Training: Prompting alone is insufficient to correct test-time errors, while manually collecting successful reason-action demonstrations is difficult to scale.
  • Agent Training: LEAP uses privileged AI feedback to design critics, but privileged corrections can be unrealizable for the agent and lead to poor success rates.
  • Agent Training: ARCHER combines temporal-difference Q-value training with REINFORCE policy training, but its results are limited to small GPT2 models.
  • Process Reward Models: Prior PRM research mainly targets deterministic, known-transition math reasoning, using human annotations or automatically computed Q-value estimates.
  • Process Reward Models: In agent settings, PRMs show stronger gains because long horizons and limited environment access make outcome rewards infeasible, while reward hacking and value-estimation issues remain.

6 Conclusion

The paper introduces AgentPRM and InversePRM for training LLM agents with process rewards. On ALFWorld, small models outperform strong GPT-4o baselines, while InversePRM reaches near-expert performance with fewer rollouts.

  • AgentPRM is a simple, scalable framework for training LLM agents with process reward models, while InversePRM learns PRMs directly from demonstrations without explicit outcome rewards.
  • Small models trained with AgentPRM outperform strong GPT-4o baselines, and InversePRM achieves near-expert performance with significantly fewer rollouts.
  • The paper identifies exploration, process reward shaping, and model-predictive reasoning as key challenges and proposes methods combining RL techniques with LLM-specific capabilities.
Loading 2502.10325v1…