PAIR: Prefix-Aware Internal Reward Model for Multi-Turn Agent Optimization

• URL: http://arxiv.org/abs/2605.17877v1
• Date: Mon, 18 May 2026 05:39:30 GMT
• Title: PAIR: Prefix-Aware Internal Reward Model for Multi-Turn Agent Optimization
• Abstract summary: Group Relative Policy Optimization (GRPO) has been emerging as a leading choice, but its reliance on sparse outcome rewards severely limits credit assignment across intermediate steps.<n>We propose a two-stage model with a frozen hidden-state probe estimating belief-consistency and a lightweight attention-based head correcting it toward grounded correctness.<n> Experimental results show that PAIR achieves the highest AUROC on contaminated trajectories while operating at negligible inference cost.
• Score: 20.910779784518798
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: A significant hurdle for current LLMs is the execution of complex, multi-stage tasks. Group Relative Policy Optimization (GRPO) has been emerging as a leading choice, but its reliance on sparse outcome rewards severely limits credit assignment across intermediate steps. Existing remedies such as running full rollouts to assign step-level advantages, calling external LLM judges at each step, or computing intrinsic rewards that require ground-truth answers at every evaluation introduce significant costs or practical constraints. We hypothesize that internal correctness probing over LLM hidden states can be repurposed as a step-level reward signal, potentially addressing all of these limitations at once. However, existing probing research assumes clean inputs, and we first show that this assumption breaks down in multi-step settings: hidden-state probes degrade severely under prefix contamination tracking coherence with the (possibly corrupted) prefix rather than grounded correctness, while attention-based features remain robust to contamination but underperform on clean prefixes. Building on this complementary relationship, we propose the Prefix-Aware Internal Reward (PAIR), a two-stage model with a frozen hidden-state probe estimating belief-consistency and a lightweight attention-based head correcting it toward grounded correctness. Experimental results show that PAIR achieves the highest AUROC on contaminated trajectories while operating at negligible inference cost, enabling dense step-level reward signals for GRPO training without external model calls, ground-truth dependencies, or full-trajectory rollouts.

Related papers

• ARISE-RL: Agentic Rubric-Grounded Iterative Self-Evolution with Reinforcement Learning [74.19730275669227]
  ARISE-RL is a novel full-cycle self-evolution framework for training open-ended agents.<n>We present Reward-Gated Self-Evolution Distillation (RG-SED), which distills a memory-augmented variant of the same policy back into itself.<n>We also present ECR-Bench, an expert-calibrated rubric benchmark suite covering single-tool deep research and multi-tool travel planning.
  arXiv  Detail & Related papers  (2026-09-01T10:54:13Z)
• Uncertainty-Aware LLM-Guided Policy Shaping for Sparse-Reward Reinforcement Learning [2.384534878752428]
  We propose Uncertainty-Aware LLM-Guided Policy Shaping (ULPS), a framework that integrates a calibrated Large Language Model into the Reinforcement Learning training loop.<n>ULPS employs an A*-based oracle to synthesize optimal symbolic trajectories, which are used to fine-tune a BERT-based language model.<n>We observe consistent improvements in success rate, reward efficiency, and sample complexity over unguided, uncalibrated, and standard RL baselines.
  arXiv  Detail & Related papers  (2026-06-04T19:46:45Z)
• PAC-Bayesian Reward-Certified Outcome Weighted Learning [0.0]
  Estimating optimal individualized treatment rules (ITRs) via outcome weighted learning (OWL) often relies on observed rewards that are noisy or optimistic proxies for the true latent utility.<n>We propose PAC-Bayesian Reward-Certified Outcome Weighted Learning (PROWL)<n>Given a one-sided uncertainty certificate, PROWL constructs a conservative reward and a strictly policy-dependent lower bound on the true expected value.
  arXiv  Detail & Related papers  (2026-04-02T12:08:56Z)
• Improving Search Agent with One Line of Code [68.58667107354253]
  Tool-based Agentic Reinforcement Learning (TARL) has emerged as a promising paradigm for training search agents to interact with external tools.<n>We propose textbfSearch textbfAgent textbfPolicy textbfOptimization (textbfSAPO), which stabilizes training via a conditional token-level KL constraint.
  arXiv  Detail & Related papers  (2026-03-10T04:07:39Z)
• Truncated Step-Level Sampling with Process Rewards for Retrieval-Augmented Reasoning [32.295907409325615]
  Training large language models to reason with search engines via reinforcement learning is hindered by a credit assignment problem.<n>We propose SLATE, a framework built on two complementary ideas.<n> Experiments on seven QA benchmarks confirm that SLATE consistently outperforms both sparse-reward and process-reward baselines.
  arXiv  Detail & Related papers  (2026-02-26T19:05:40Z)
• Save the Good Prefix: Precise Error Penalization via Process-Supervised RL to Enhance LLM Reasoning [59.76691952347156]
  Reinforcement learning (RL) has emerged as a powerful framework for improving the reasoning capabilities of large language models (LLMs)<n>Most existing RL approaches rely on sparse outcome rewards, which fail to credit correct intermediate steps in partially successful solutions.<n>We propose Verifiable Prefix Policy Optimization (VPPO), which uses PRMs only to localize the first error during RL.
  arXiv  Detail & Related papers  (2026-01-26T21:38:20Z)
• BAPO: Boundary-Aware Policy Optimization for Reliable Agentic Search [72.87861928940929]
  Boundary-Aware Policy Optimization (BAPO) is a novel RL framework designed to cultivate reliable boundary awareness without compromising accuracy.<n>BAPO introduces two key components: (i) a group-based boundary-aware reward that encourages an IDK response only when the reasoning reaches its limit, and (ii) an adaptive reward modulator that strategically suspends this reward during early exploration, preventing the model from exploiting IDK as a shortcut.
  arXiv  Detail & Related papers  (2026-01-16T07:06:58Z)
• ArenaRL: Scaling RL for Open-Ended Agents via Tournament-based Relative Ranking [84.07076200941474]
  ArenaRL is a reinforcement learning paradigm that shifts from pointwise scalar scoring to intra-group relative ranking.<n>We construct an intra-group adversarial arena and devise a tournament-based ranking scheme to obtain stable advantage signals.<n>Experiments show that ArenaRL substantially outperforms standard RL baselines.
  arXiv  Detail & Related papers  (2026-01-10T08:43:07Z)
• Prune-Then-Plan: Step-Level Calibration for Stable Frontier Exploration in Embodied Question Answering [52.69447404069251]
  Large vision-language models (VLMs) have improved embodied question answering (EQA) agents by providing strong semantic priors for open-vocabulary reasoning.<n>We propose Prune-Then-Plan, a framework that stabilizes exploration through step-level calibration.
  arXiv  Detail & Related papers  (2025-11-24T22:50:50Z)
• Repurposing Synthetic Data for Fine-grained Search Agent Supervision [81.95597592711688]
  LLM-based search agents are increasingly trained on entity-centric synthetic data.<n> prevailing training methods discard this rich entity information, relying instead on sparse, outcome-based rewards.<n>We introduce Entity-aware Group Relative Policy Optimization (E-GRPO), a novel framework that formulates a dense entity-aware reward function.
  arXiv  Detail & Related papers  (2025-10-28T17:50:40Z)
• FAPO: Flawed-Aware Policy Optimization for Efficient and Reliable Reasoning [42.18368547352248]
  Reinforcement learning with verifiable rewards (RLVR) has emerged as a promising paradigm for enhancing the reasoning capabilities of large language models (LLMs)<n>We propose Flawed-Aware Policy Optimization (FAPO), which presents a parameter-free reward penalty for flawed-positive rollouts.<n>Experiments show that FAPO is effective in broad domains, improving outcome correctness, process reliability, and training stability without increasing the token budget.
  arXiv  Detail & Related papers  (2025-10-26T05:49:38Z)
• PACR: Progressively Ascending Confidence Reward for LLM Reasoning [55.06373646059141]
  We propose Progressively Ascending Confidence Reward (PACR)<n>PACR is a dense, model-intrinsic reward computed directly from the model's evolving belief in the correct answer.<n>Our results suggest that dense, model-intrinsic shaping signals can make RLVR training more effective and reliable.
  arXiv  Detail & Related papers  (2025-10-25T11:25:35Z)
• Hybrid Reward Normalization for Process-supervised Non-verifiable Agentic Tasks [12.31210445905605]
  We introduce Principle Process Reward (PPR), an RL approach that unifies step-level assessment and outcome verification.<n>PPR achieves state-of-the-art performance across a wide range of benchmarks, demonstrating its impressive robustness and generalization.
  arXiv  Detail & Related papers  (2025-09-29T23:44:55Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.