rePIRL: Learn PRM with Inverse RL for LLM Reasoning

Official code for "rePIRL: Learn PRM with Inverse RL for LLM Reasoning" (ICML 2026).

Xian Wu*, Kaijie Zhu*, Ying Zhang, Lun Wang, Wenbo Guo (*Equal contribution. Meta AI; UC Santa Barbara; Google DeepMind.)

Paper (arXiv:2602.07832)

Overview

rePIRL is an inverse-RL-inspired framework for learning Process Reward Models (PRMs) for LLM reasoning. Given only expert trajectories — no token-level annotations, no preference labels, and no access to the expert policy — rePIRL jointly learns a PRM and a policy through a dual learning process: the PRM is updated via an MLE objective with importance-weighted policy rollouts, and the policy is updated by maximum-entropy RL (RLOO) under the current reward. The framework also unifies several existing PRM/offline-RL methods (PRIME, DPO, DQO, Math-Shepherd) as special cases under additional assumptions.

The trained PRM can be used standalone for test-time training, test-time scaling (verifier-based selection), and as an early signal on hard problems where outcome rewards are sparse.

Repository layout

recipe/irl/         # rePIRL training code (algorithm, trainer, workers, configs)
  run_irl_qwen.sh   # example training launcher (Qwen3-4B)
  main_irl.py       # entry point
  irl_core_algos.py # core IRL loss / importance weighting
  irl_ray_trainer.py
verl/               # underlying VeRL framework
eval/               # evaluation pipelines for math / coding benchmarks
scripts/, run_bash/ # auxiliary launch scripts

Setup

conda env create -f env.yaml
conda activate irl
pip install -e .

Training was run on 8× RTX PRO 6000 GPUs (see paper §4.1 / Appendix B).

Data

Training and validation parquet files (sampled from PRIME Eurus-2-RL-Data, 7k math + 7k coding, with four expert demonstrations per problem generated by Claude-3.7-Sonnet):

https://drive.google.com/drive/folders/119QT_hSkKEM330t4_kiPwCIOujGipxw-?usp=sharing

Place the files under data/ and update the paths at the top of recipe/irl/run_irl_qwen.sh if needed.

Training

Set your Weights & Biases API key (the launch scripts no longer ship one):

export WANDB_API_KEY=<your_key>   # or run: wandb login

Then launch training:

bash recipe/irl/run_irl_qwen.sh

The script trains both the policy and PRM. Key hyperparameters (policy lr 5e-7, PRM lr 3e-8, batch size 128, RLOO with 4 rollouts/prompt, entropy coef 0.001, outcome-to-PRM reward ratio 1:0.05 for Qwen2.5-3B-Instruct and 1:0.1 for Qwen3-4B-Base) follow the paper.

Evaluation

See eval/README.md. Math benchmarks (MATH500, AIME-2024, MinervaMath, AMC, OlympiadBench) and coding benchmarks (LeetCode, LiveCodeBench) each use a separate conda environment; pass@1 under greedy decoding is reported.

Citation

@inproceedings{wu2026repirl,
  title     = {re{PIRL}: Learn {PRM} with Inverse {RL} for {LLM} Reasoning},
  author    = {Wu, Xian and Zhu, Kaijie and Zhang, Ying and Wang, Lun and Guo, Wenbo},
  booktitle = {Proceedings of the 43rd International Conference on Machine Learning},
  year      = {2026}
}