🤖 machine learning

AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning

This paper introduces AIRL-S, a unified framework that combines Adversarial Inverse Reinforcement Learning with Group Relative Policy Optimization to infer dense, step-wise rewards from reference trajectories, thereby eliminating the need for labeled process data while enabling robust, cost-effective test-time scaling that achieves GPT-4o-level performance across mathematics, science, and code generation benchmarks.

Original authors: Can Jin, Yang Zhou, Qixin Zhang, Hongwu Peng, Di Zhang, Zihan Dong, Marco Pavone, Ligong Han, Zhang-Wei Hong, Tong Che, Dimitris N. Metaxas

Published 2026-08-27

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Can Jin, Yang Zhou, Qixin Zhang, Hongwu Peng, Di Zhang, Zihan Dong, Marco Pavone, Ligong Han, Zhang-Wei Hong, Tong Che, Dimitris N. Metaxas

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Large language models are the powerful engines behind many modern artificial intelligence systems, capable of generating text, solving problems, and writing code. For years, researchers have tried to make these systems smarter by teaching them to think through problems step-by-step, a method known as "chain of thought." However, teaching a model to reason well is difficult because the computer often only knows if the final answer is right or wrong, not whether the steps leading there were logical. To fix this, scientists have tried two main approaches. One involves training the model with rewards only at the very end of a task, which can be unstable and inefficient. The other uses a separate guide to check every single step of the reasoning process, but creating this guide usually requires expensive human experts to label thousands of examples, and the guide often fails when the model starts thinking in new ways.

A team of researchers has introduced a new system called AIRL-S that unifies these two approaches into a single, self-improving loop. Instead of relying on human labels for every step, the system learns to create its own internal guide by observing high-quality examples of reasoning. It then uses this guide to both train the model and help it search for the best answers during real-world use. In tests across eight different benchmarks involving mathematics, science, and coding, the researchers found that this method improved the model's performance by an average of nine percent over its starting version. The resulting system performed as well as the most advanced commercial models available, such as GPT-4o, without needing the costly human supervision that previous methods required.

The core innovation lies in how the system learns to judge its own work. Traditionally, to teach a model to reason step-by-step, researchers would need a "process reward model"—a separate program trained on data where humans had marked every correct and incorrect step. This is slow and expensive. The new method bypasses this by using a technique called adversarial learning. Imagine the system playing a game where one part tries to generate reasoning steps, and another part tries to distinguish between those generated steps and a set of high-quality reference examples. Through this competition, the system learns to recognize what a good reasoning step looks like without ever being told explicitly by a human. This learned guide, or reward model, becomes dense and detailed, offering feedback on every single step of the reasoning process rather than just the final result.

Once the system has learned this internal guide, it uses it in two ways simultaneously. First, it uses the guide to train the main model, encouraging it to take better steps during the learning phase. Second, it keeps the same guide to act as a verifier during the actual problem-solving phase. When the model is asked to solve a difficult problem, it can generate many different possible solutions. The learned guide then scores each step of these potential solutions, helping the system select the most logical path. This creates a unified pipeline where the tool used to teach the model is the same tool used to help it think during a test. The researchers demonstrated that this guide is robust; it works effectively even when applied to different models or different search strategies, suggesting that the knowledge it learned is general and transferable.

The results of this approach were measured across a wide range of challenging tasks. The researchers tested their model on eight standard benchmarks, including complex math competitions, scientific reasoning questions, and coding challenges. The model, which started as a standard open-source language model, improved its average accuracy by nine percent after training with this new method. In specific mathematical and scientific tasks, the improvement was even higher, reaching thirteen percent. When compared to other models that had been trained with expensive human-labeled data or other advanced reinforcement learning techniques, this new system consistently performed better. It even matched the performance of GPT-4o, a top-tier commercial model, despite having significantly fewer parameters. This suggests that the quality of the reasoning process, guided by this self-learned reward system, is more important than simply having a larger model.

A key finding of the study is how the two different types of learning signals work together. The system combines the detailed, step-by-step feedback from its self-learned guide with the final outcome rewards from traditional methods. The researchers found that relying on just one type of signal was less effective. The step-by-step guide helped the model explore better reasoning paths, while the final outcome signal ensured the model stayed focused on getting the correct answer. When balanced correctly, these two signals reinforced each other, leading to more stable training and better results. Furthermore, the system showed that it could use this learned guide to improve various search methods, such as those that try many different solutions at once or those that build a tree of possibilities. The guide was particularly effective in complex search scenarios where the model had to evaluate the value of intermediate steps to make the best choice.

The implications of this work extend beyond just getting higher scores on tests. By removing the dependency on expensive human annotations for step-by-step reasoning, this method offers a more scalable and cost-effective way to improve artificial intelligence. It suggests that models can learn to be their own teachers, refining their reasoning capabilities by observing high-quality examples and competing against themselves. The researchers noted that their approach is particularly well-suited for tasks where the final answer can be automatically verified, such as mathematics and coding. While the method is powerful, the authors acknowledge that it currently relies on these verifiable outcomes and may require further development to handle open-ended tasks where a single correct answer does not exist. Nevertheless, the ability to unify training and inference-time search into a single, efficient system marks a significant step forward in making large language models more reliable and capable reasoners.

Technical Summary: AIRL-S

Problem Statement

The paper addresses the limitations of current Test-Time Scaling (TTS) strategies for Large Language Models (LLMs) in complex reasoning tasks. Existing approaches generally fall into two categories, both of which suffer from specific drawbacks:

Reinforcement Learning (RL) with Outcome Rewards: Methods relying on sparse outcome rewards (e.g., correctness of the final answer) often experience training instability and sample inefficiency.
Search-Based Methods with Static Process Reward Models (PRMs): While search strategies like Monte Carlo Tree Search (MCTS) benefit from dense, step-wise guidance, training effective static PRMs requires expensive, granular human annotations or pseudo-labels. Furthermore, static PRMs are prone to "reward hacking" due to distributional shifts between the fixed verifier and the evolving policy, undermining their reliability as search verifiers.

The core challenge is to unify RL-based training and search-based inference without relying on costly step-level supervision, while ensuring the reward model remains robust and transferable.

Methodology: AIRL-S

The authors propose AIRL-S, a unified framework that integrates Adversarial Inverse Reinforcement Learning (AIRL) with Group Relative Policy Optimization (GRPO). The framework operates through a single pipeline where reward learning, policy optimization, and inference-time search are trained and evaluated together.

1. Learning a Step-Wise PRM via AIRL

Instead of using labeled process data, AIRL-S learns a dense, step-wise reward model directly from reference trajectories.

Data Generation: The system samples multiple Chains of Thought (CoTs) from the current policy. Trajectories yielding a correct final answer (binary outcome reward) are stored in a replay buffer as reference rollouts.
Adversarial Training: A discriminator $D_\phi$ is trained to distinguish between the reference rollouts (from the replay buffer) and the current policy's rollouts. The discriminator operates on state-action pairs (question + preceding steps, current step).
Reward Derivation: The intrinsic step-wise reward $r_\phi$ is derived from the discriminator's output using the log-odds ratio. This learned $r_\phi$ serves as the Process Reward Model (PRM).

2. Policy Optimization with Composite Objectives

The policy model $\pi_\theta$ is updated using a composite objective that balances dense process rewards and sparse outcome rewards:

AIRL Objective: The policy maximizes the rewards assigned by the learned discriminator, aligning its generation distribution with the reference distribution. This is formulated using a PPO-style objective based on advantages derived from $r_\phi$ .
GRPO Objective: To ensure final answer correctness, the framework incorporates the GRPO objective (used in DeepSeek-R1), which utilizes binary outcome rewards and group-relative advantages.
Combined Loss: The final objective is a weighted sum: $J(\theta) = \lambda J_{AIRL}(\theta) + (1-\lambda) J_{GRPO}(\theta)$ , where $\lambda$ balances the influence of process and outcome rewards.

3. PRM-Guided Test-Time Search

The learned PRM ( $r_\phi$ ) is reused as a verifier during inference to guide search algorithms.

Search Strategies: The framework supports Best-of-N sampling, Beam Search, and MCTS.
Scoring: During search, intermediate steps or complete solutions are scored using $r_\phi$ .
Aggregation: For Best-of-N and similar methods, the paper employs a PRM-Min-Sum aggregation strategy: the score of a chain is the minimum reward along its path, and the final answer is selected by summing the scores of all chains leading to that answer.

Key Contributions

Unified Framework: The introduction of AIRL-S, which combines AIRL-based process reward learning, GRPO-based outcome optimization, and PRM-guided test-time search without requiring costly step-level supervision.
Empirical Analysis of Objectives: An analysis demonstrating that a balanced objective (combining dense process rewards and sparse outcome rewards) improves both policy training and downstream search performance compared to using either signal alone.
Transferability and Performance: Extensive evaluations showing that the learned PRM consistently outperforms supervised baselines across different generators and search algorithms, establishing a robust methodology for scaling test-time computation.

Experimental Results

The authors evaluated the policy model and PRM on eight benchmarks covering mathematics (AIME2024, AMC, MATH500, GPQA), science, and code generation (HumanEval, LeetCode, LiveCodeBench, MBPP).

Policy Performance: The AIRL-S policy model (initialized from Qwen2.5-7B-Instruct) achieved an average performance improvement of 9% over the base model. It matched the performance of GPT-4o on these tasks and outperformed other local RL-trained baselines (e.g., Math-Shepherd, DeepSeek-Math-RL) and models trained via direct imitation learning (e.g., s1.1-7B).
PRM Effectiveness: The AIRL-derived PRM outperformed static PRMs trained on labeled data (e.g., Math-Shepherd-PRM, Llama3.1-PRM) when applied to Best-of-N search across multiple generative LLMs.
Search Compatibility: The learned PRM significantly improved performance across MCTS, Beam Search, and Best-of-N. The gains were most pronounced in MCTS, where step-value estimation is critical.
Ablation Studies: Experiments confirmed that $\lambda=0.5$ (balancing AIRL and GRPO) yields the best results. Pure GRPO lacked dense step-wise feedback, while pure AIRL lacked explicit outcome signals. The PRM also demonstrated transferability to weaker generators (e.g., Math-Shepherd-Mistral-7B-RL) and different search algorithms.

Significance and Claims

The paper claims that AIRL-S offers a robust and cost-effective methodology for scaling test-time computation in LLMs. By eliminating the dependency on labeled process data, the framework mitigates the high costs associated with training static PRMs. Furthermore, by unifying reward learning and policy optimization, the system avoids the distributional shifts that lead to reward hacking in static verifiers.

The authors emphasize that the primary contribution is not merely the reuse of existing components (AIRL, GRPO, search) but the system-level integration where the reward model learned during training is directly reused as a verifier during inference. This creates a cohesive loop where improvements in reward learning benefit both the training of the policy and the efficiency of the search process, enabling smaller models to achieve high-level reasoning performance comparable to much larger proprietary models.