Learning Reasoning Rewards from Expert Demonstrations with Inverse Reinforcement Learning

The evolution of large language models (LLMs) has revolutionized NLP, yet their reasoning capabilities lag behind. Prior work like Chain-of-Thought prompting and self-consistency methods improved performance but still rely on external verifiers or supervised signals. Reinforcement learning approaches, such as RLHF, depend on handcrafted reward functions or preference labels, which are costly and brittle. IRL offers a promising alternative by learning rewards directly from expert demonstrations, but its application in large models remains limited. This paper addresses these gaps by proposing a demonstration-based reward learning framework that does not require external verifiers, aiming to improve reasoning robustness, interpretability, and transferability.

Current methods face key challenges: designing explicit reward functions is difficult, especially for complex reasoning; off-policy imitation can lead to fragile performance when inference deviates from training trajectories; and existing reward signals lack granularity for diagnosing reasoning failures. These issues hinder the deployment of autonomous reasoning systems in real-world applications like medical diagnosis or scientific research, where robustness and interpretability are critical. The core problem is how to learn a stable, informative reward signal solely from expert demonstrations that can guide training, inference, and debugging without external supervision.

The paper introduces several innovations: 1) applying adversarial IRL to learn process-level rewards from demonstrations, avoiding external verifiers; 2) designing multi-granularity rewards—sparse, interval, and dense—to balance stability and feedback richness; 3) reusing a single reward model for training, reranking, and error localization, simplifying the pipeline; 4) integrating with GRPO for stable policy updates. These innovations enable the model to learn meaningful reasoning signals directly from expert traces, supporting multiple downstream tasks and improving robustness, interpretability, and transferability compared to prior imitation or preference-based methods.

• �� Construct positive and negative trajectory sets based on answer agreement and corruption, avoiding reliance on step labels.
• �� Train a token-level critic Dϕ to classify tokens as part of positive or negative reasoning traces, outputting probabilities that are converted into rewards.
• �� Implement multi-granularity reward masks (sparse, interval, dense) to control credit assignment and feedback detail.
• �� Formulate an adversarial IRL objective that maximizes the difference between expert and policy trajectory rewards, updating both the critic and policy iteratively.
• �� Use the critic logits to generate token rewards, which guide policy optimization via Group Relative Policy Optimization (GRPO), balancing exploration and stability.
• �� During inference, apply the learned reward for reranking sampled traces, selecting the highest scoring candidate.
• �� For error localization, analyze reward drops along traces to identify potential reasoning failures.

Experiments span three datasets: GSM8K for arithmetic reasoning, MMLU-Pro for scientific reasoning, and MedReason for medical reasoning. The models tested include instruction-tuned variants like Qwen2.5, Llama3, and Qwen3-4B. The evaluation metrics include pass@1, reranking improvements, and fault localization accuracy. Ablation studies compare different reward granularities and training settings, analyzing stability and transferability. The experiments demonstrate that R-AIRL-trained rewards outperform SFT in training, significantly boost reranking performance, and effectively localize reasoning errors, validating the framework’s versatility and robustness.

In GSM8K, R-AIRL achieved a maximum pass@1 of 85.8%, surpassing SFT by 15.7%. During inference, reranking with learned rewards improved pass@1 by up to 17.4 points across datasets. The reward models could transfer across tasks and backbones, with off-diagonal gains up to +13.0%. Dense rewards provided finer feedback but were less stable, while sparse rewards offered more consistent training. The reward-based reranking consistently outperformed random selection, confirming the utility of learned rewards for inference enhancement. Fault localization experiments showed over 86% accuracy in identifying reasoning failure points, demonstrating interpretability.

This approach is immediately applicable to domains requiring reliable reasoning, such as medical diagnosis, scientific research, and legal analysis. It enables models to learn reasoning signals solely from demonstrations, reducing reliance on costly external verifiers. The learned rewards facilitate better training, inference reranking, and debugging, making large models more autonomous and trustworthy. Long-term, integrating this framework with multimodal data and hybrid supervision could further expand its impact, enabling scalable, explainable AI systems capable of complex reasoning in real-world scenarios.

The method depends heavily on the quality of expert demonstrations; noisy or biased data can impair reward learning. Computational costs for critic training and policy updates are high, limiting scalability. The reward signals may still be insufficient for extremely complex or ambiguous tasks, requiring external validation. Additionally, the approach's effectiveness diminishes if demonstrations do not cover diverse reasoning paths, highlighting the need for better data curation and model robustness improvements.