rePIRL: Learn PRM with Inverse RL for LLM Reasoning

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.

When training artificial intelligence systems to solve complex math and coding problems, researchers often use step-by-step graders, known as Process Reward Models (PRMs), to reward strong intermediate reasoning rather than just final answers. However, existing approaches for building these graders have important limitations: they either depend on unrealistic assumptions about expert policies or struggle to generalize to new problems. To address these challenges, we introduce rePIRL, a new framework inspired by the way humans learn through observing experts. With minimal assumptions about expert behavior, rePIRL employs a dual learning process in which both the problem-solving model and the step-grading model learn jointly from expert demonstrations, continuously improving one another over time. We further adapt this framework to operate effectively at the scale of modern large language models. This matters because our framework unifies previous approaches and proves significantly more effective on standard math and coding tests. Crucially, our highly accurate graders allow AI models to perform better during actual testing, and they provide valuable early feedback when training models to tackle exceptionally hard problems.