rePIRL: Learn PRM with Inverse RL for LLM Reasoning

TL;DR

rePIRL introduces a novel inverse reinforcement learning framework to learn process reward models for large language model reasoning, addressing limitations of existing methods and enhancing reasoning capabilities.

Contribution

It proposes a dual learning algorithm for PRM and policy that requires minimal assumptions about expert policies, unifying online and offline learning methods.

Findings

rePIRL outperforms existing methods on math and coding reasoning datasets.

The learned PRM improves test-time training and scaling.

Ablation studies validate the effectiveness of key design choices.

Abstract

Process rewards have been widely used in deep reinforcement learning to improve training efficiency, reduce variance, and prevent reward hacking. In LLM reasoning, existing works also explore various solutions for learning effective process reward models (PRM) with or without the help of an expert policy. However, existing methods either rely on strong assumptions about the expert policies (e.g., requiring their reward functions) or suffer intrinsic limitations (e.g., entropy collapse), resulting in weak PRMs or limited generalizability. In this paper, we introduce rePIRL, an inverse RL-inspired framework that learns effective PRMs with minimal assumptions about expert policies. Specifically, we design a dual learning process that updates the policy and the PRM interchangeably. Our learning algorithm has customized techniques to address the challenges of scaling traditional inverse RL to LLMs. We theoretically show that our proposed learning framework can unify both online and offline PRM learning methods, justifying that rePIRL can learn PRMs with minimal assumptions. Empirical evaluations on standardized math and coding reasoning datasets demonstrate the effectiveness of rePIRL over existing methods. We further show the application of our trained PRM in test-time training, test-time scaling, and providing an early signal for training hard problems. Finally, we validate our training recipe and key design choices via a detailed ablation study.