Technical Summary: rePIRL: Learn PRM with Inverse RL for LLM Reasoning

Problem Statement

In Large Language Model (LLM) reasoning, Process Reward Models (PRMs) are utilized to provide token-level intermediate rewards, aiming to improve training efficiency, reduce variance, and prevent reward hacking. However, existing methods for learning effective PRMs face significant limitations:

• Offline Methods: Supervised approaches (e.g., Uesato et al., 2022; Lightman et al., 2023) and DPO-type works require access to token-level annotations or preference labels from expert trajectories, which are costly to obtain. Monte Carlo Tree Search (MCTS) based methods (e.g., Wang et al., 2024) require running an expert policy to generate trajectories, which is computationally expensive and infeasible if the expert policy is unavailable.
• Online Methods: Approaches like PRIME (Cui et al., 2025a) rely on strong assumptions regarding the reward functions and the policy model under training. Intrinsic reward methods (e.g., Zhao et al., 2025b) use model entropy or confidence as rewards but suffer from "entropy collapse," where performance degrades in later training stages due to excessive mode-seeking behavior.

The core challenge is to learn an effective PRM with minimal assumptions: specifically, without requiring token-level reward annotations, preference labels, access to the expert policy, or strong assumptions about the reward function's form.

Methodology: rePIRL Framework

The authors propose rePIRL, a framework inspired by Inverse Reinforcement Learning (IRL) that learns PRMs from expert trajectories alone. The method models multi-step LLM reasoning as a token-level Markov Decision Process (MDP) and employs a dual learning process that updates the policy and the PRM interchangeably.

1. Theoretical Foundation

• Objective Construction: Following the classical IRL framework, the authors define a hidden variable $o_t \in \{0, 1\}$ indicating whether an action is sampled from the expert policy. They model the probability of an action being expert-sampled using an energy-based function: $p(o_t=1|s_t, a_t) \propto \exp(r_\phi(s_t, a_t))$.
• Likelihood Maximization: The reward function $r_\phi$ is learned by maximizing the likelihood of expert trajectories. This leads to an objective involving a partition function $z(\phi)$, which is intractable for LLMs due to the vast state and action spaces.
• Tractability via Sampling: To resolve the intractability of the partition function, the authors apply a gradient trick and importance sampling. Instead of sampling from the optimal policy (which is unknown), they sample trajectories $\bar{\tau}$ from the current policy $\pi_\theta$. The gradient is approximated using importance weights $w = \exp(r_\phi(\bar{\tau})) / \pi_\theta(\bar{\tau})$.
• Policy Optimization: To prevent model collapse (deterministic policies), the policy is updated using a Maximum Entropy RL objective (Haarnoja et al., 2018), which balances reward maximization with entropy regularization. This is solved using algorithms like RLOO, PPO, or GRPO.

2. Learning Algorithm

The training loop alternates between updating the PRM and the policy:

1. Rollout: Generate policy rollouts and fetch expert demonstrations.
2. Reward Calculation: Compute outcome rewards and average process rewards.
3. PRM Update: Update the reward model using a loss function that contrasts expert trajectories against policy rollouts, weighted by importance sampling scores. Crucially, the authors use average process rewards rather than cumulative rewards to mitigate reward hacking.
4. Policy Update: Update the policy model using the learned process rewards combined with outcome rewards (if available), employing maximum entropy objectives.

3. Unification of Existing Methods

A key theoretical contribution is demonstrating that rePIRL can unify various existing PRM learning methods (DPO, DQO, Math-Shepherd, PRIME) under a single framework by introducing specific additional assumptions. This positions rePIRL as the most general method, requiring the fewest assumptions.

Key Contributions

1. Novel Framework: Introduction of rePIRL, an IRL-inspired framework for learning PRMs from expert trajectories without token-level annotations, preference labels, or access to the expert policy.
2. Unified Theory: A theoretical analysis showing that rePIRL unifies state-of-the-art offline and online PRM methods (including DPO, DQO, PRIME, and MCTS-based approaches) under a single learning framework, clarifying their underlying assumptions.
3. Scalable Algorithm: Development of a scalable learning algorithm tailored for LLMs, utilizing importance sampling to handle intractable partition functions and maximum entropy RL to ensure training stability.
4. Empirical Validation: Comprehensive evaluation on math and coding benchmarks demonstrating superior performance over baselines.

Experimental Results

The authors evaluated rePIRL using Qwen2.5-3B-Instruct and Qwen3-4B-Base on seven standard benchmarks (AIME-2024, AMC, Math-500, MinervaMath, OlympiadBench, Leetcode, LiveCodeBench).

• Performance: rePIRL achieved state-of-the-art results among 3B- and 4B-scale models. For instance, on Qwen3-4B-Base, it reached 41.6% average accuracy on math benchmarks and 39.8% on coding benchmarks, outperforming strong baselines like PRIME, MCTS, and vanilla RLOO.
• Comparison: It showed absolute improvements of 3.2% (math) and 4.3% (coding) over vanilla RLOO with Qwen3-4B-Base.
• Applications:
  • Test-time Training: When trained solely on PRM signals without outcome rewards, rePIRL matched the performance of models trained with verifiable outcome rewards.
  • Test-time Scaling: Using the PRM as a verifier for selecting better rollouts yielded monotonic performance improvements, outperforming majority voting.
  • Hard Problems: PRM provided informative early signals for difficult problems where outcome rewards were sparse, leading to faster convergence compared to outcome-only training.
• Efficiency: rePIRL introduced only a ~1.4x training time overhead compared to outcome-based RL (e.g., RLOO) while avoiding the need for separate value networks.

Significance and Claims

The paper claims to be the first work to generalize Inverse Reinforcement Learning for LLM reasoning and the first to unify existing PRM learning methods into a single framework.

• Minimal Assumptions: The authors emphasize that rePIRL operates under the minimal necessary assumptions, making it more generalizable than methods requiring token-level annotations or specific expert policy access.
• Stability: The framework addresses the entropy collapse issue common in intrinsic reward methods by leveraging maximum entropy RL, ensuring stable training dynamics.
• Utility: The trained PRM is presented as a reusable, standalone module that benefits downstream applications such as test-time training and scaling, particularly in scenarios where outcome rewards are unavailable or sparse.

The authors conclude that while rePIRL adds a unified theoretical perspective and practical utility, future work could explore combining it with intrinsic reward methods, systematic comparisons of online RL heuristics, and scaling to agentic tasks and larger models.