Process Supervision-Guided Policy Optimization for Code Generation

## Abstract

Abstract Reinforcement learning (RL) with unit test feedback has enhanced large language models’ (LLMs) code generation, but relies on sparse rewards provided only after complete code evaluation, limiting learning efficiency and incremental improvements. When generated code fails all unit tests, no learning signal is received, hindering progress on complex tasks. To address this, we propose a Process Reward Model (PRM) that delivers dense, line-level feedback on code correctness during generation, mimicking human code refinement and providing immediate guidance. We explore various strategies for training PRMs and integrating them into the RL framework, finding that using PRMs both as dense rewards and for value function initialization significantly boosts performance. Our approach increases our in-house LLM’s pass rate from 28.2% to 29.8% on LiveCodeBench and from 31.8% to 35.8% on our internal benchmark. Our experimental results highlight the effectiveness of PRMs in enhancing RL-driven code generation, especially for long-horizon scenarios.

## 1 Introduction

The rapid advancement of large language models (LLMs) has revolutionized code generation, enabling models to achieve near-human performance on programming tasks (Chen et al., ; Li et al., ; OpenAI, ). These models have demonstrated remarkable abilities to generate syntactically correct and functionally viable code snippets, significantly aiding software development processes.
Building upon these successes, recent research has explored the use of reinforcement learning (RL) from unit test feedback to further enhance the code generation capabilities of LLMs (Le et al., ; Shojaee et al., ; Liu et al., ; Dou et al., ). By incorporating unit tests as a reward mechanism, these methods aim to guide LLMs toward generating code that not only compiles but also passes specified test cases, thereby improving overall code reliability and quality.

However, a significant challenge arises from the nature of the reward signals derived from unit tests. These signals are inherently sparse, as they are only received at the end of an episode after the entire code snippet has been generated and evaluated. This delay in feedback impedes learning efficiency and limits the model’s ability to make incremental improvements during code generation. When an LLM fails to generate code that passes any unit tests, it receives no meaningful learning signal, making it difficult to learn to solve more complex coding problems.
In contrast, human programmers typically do not rewrite code from scratch when their programs fail unit tests. Instead, they analyze the code to pinpoint and fix errors, leveraging their understanding of programming logic and structure to iteratively improve upon the current version. This process of step-by-step refinement, which involves receiving and acting upon fine-grained feedback, is missing in the current RL training loop for code generation from unit test feedback.

To address this limitation, we propose integrating a Process Reward Model (PRM) (Lightman et al., ; Wang et al., ) into the RL training framework for code generation. A PRM provides dense signals by offering line-level feedback that indicates the correctness of each generated line of code. This fine-grained feedback mechanism mimics the human approach to code refinement and has the potential to enhance learning efficiency by providing immediate guidance during code generation.
While the concept of using PRMs is intuitive, training an effective PRM and integrating it into RL training is non-trivial. Challenges include accurately modeling the correctness of partial code snippets and ensuring stable and effective training dynamics when combining PRM-generated signals with traditional RL methods. Although previous research has attempted to incorporate PRMs into LLM RL training (Wang et al., ), these efforts have been limited to the mathematical domain and have not fully explored the complexities involved.

In this work, we conduct a comprehensive analysis of how PRMs can be integrated into RL training for code generation. We explore various strategies for training a robust code PRM and investigate different methods of utilizing PRMs to improve code generation performance. Based on our experiments, we provide a practical recipe for successfully using PRMs and integrating them into RL training in the context of code generation problems. Notably, one of our key findings is that using PRMs concurrently as both dense rewards and value function initialization in RL training leads to a significant performance improvement.
Our contributions can be summarized as follows:

- •
We propose an effective approach that automatically generates process-level supervision data by identifying the first error line in generated code using binary search. We then train a PRM on this data to generate dense signals during RL training. To the best of our knowledge, we are the first to demonstrate that PRMs can benefit RL from unit test feedback in code generation.
- •
We conduct systematic experiments to determine how to properly and effectively integrate PRMs into RL. Our analysis explores various strategies for training a high-quality code PRM and utilizing PRMs to improve code generation. We summarize our findings into a practical recipe for successfully using PRMs in the context of code generation.
- •
By following the recipe, we significantly improve our in-house proprietary LLM’s pass rate from 28.2% to 29.8% on the LiveCodeBench dataset and from 31.8% to 35.8% on our in-house benchmark. Besides, we find that integrating PRMs into RL training benefits code generation in long-horizon scenarios.

## 2 Problem Formalization

In code generation tasks, we define a code generation problem as a sequence of tokens $\mathbf{x}=(x_{1},x_{2},\ldots,x_{m})$, where each $x_{i}$ denotes the $i$-th element or token of the input prompt, which may include problem descriptions.The primary objective for the model in this context is to process the given input $\mathbf{x}$ and generate a coherent and syntactically correct sequence of code tokens. This sequence is denoted as $\mathbf{y}=(\mathbf{y}^{(1)},\mathbf{y}^{(2)},\ldots,\mathbf{y}^{(T)})$, where $T$ represents the total number of code generation steps.
Each individual code generation step, $\mathbf{y}^{(t)},\ t=1,2,\ldots,T$, is composed of a series of tokens $y^{(t)}_{1},y^{(t)}_{2},\ldots,y^{(t)}_{n_{t}}$, where $y^{(t)}_{i}$ corresponds to the $i$-th token within the $t$-th step, and $n_{t}$ denotes the number of tokens in this step.

Typically, a pre-trained language model (LM), denoted as $p_{\theta}$, is employed to model the conditional probability distribution of the code generation steps $\mathbf{y}$, given the code generation problem $\mathbf{x}$, which is mathematically represented as $p_{\theta}(\mathbf{y}\mid\mathbf{x})$, parameterized by $\theta$. The model is optimized through training on a dataset $\mathcal{D}_{\mathbf{x}\mathbf{y}}$ containing pairs of prompts and their corresponding code solutions. This training process, often referred to as Supervised Fine-Tuning (SFT), involves maximizing the log-likelihood of the dataset.

### 2.1 Baseline Method: Reinforcement Learning from Unit Test Feedback

Code generation tasks can be formulated within a Reinforcement Learning (RL) framework, where code generation is treated as a sequence of decision-making steps.
Once the model has undergone SFT,
the RL phase is employed to refine the model’s ability to generate functionally correct code using feedback from unit tests (Liu et al., ). Unit test feedback is derived by executing the generated program on predefined test cases. The feedback serves as a signal for learning and can be transformed into a reward. A simple reward function based on the outcome of the unit tests could be defined as follows:

$$ $R_{\texttt{UT}}(\mathbf{x},\mathbf{y})=\begin{cases}1,&\text{if the program $ \mathbf{y}$ passes all unit test cases}\\ 0,&\text{otherwise}\end{cases}$ $$

This binary reward formulation encourages the model to generate programs that can successfully pass all unit test cases.
Given a collection of unlabeled code generation prompts $\mathcal{D}_{\mathbf{x}}$, the model $p_{\theta}$ is optimized to maximize the expected reward over all possible code generation trajectories.

## 3 Process Supervision-Guided Policy Optimization

Figure: Figure 1: Overview of our method. The approach consists of two main components: (1) A binary search-based method to automate PRM training data labeling, which is used to train a code PRM; and (2) Integration of the PRM into RL training as both dense rewards and value function initialization.
Refer to caption: x1.png

While the Reinforcement Learning from Unit Test Feedback (RLTF) offers a framework for improving code generation models, it suffers from significant limitations due to the sparsity of its reward signal. The binary nature of unit test feedback—indicating only whether the entire program passes or fails—provides no guidance on which specific parts of the code contributed to the outcome. This lack of intermediate feedback makes it challenging for the model to identify and correct errors during training, leading to slow convergence and suboptimal performance.
In contrast, human programmers iteratively develop and refine their code. When a program fails to pass unit tests, they do not typically rewrite it from scratch. Instead, they analyze the code to pinpoint and fix errors, leveraging their understanding of programming logic and structure. This process of step-by-step refinement is crucial for efficient problem-solving.

Motivated by this observation, we propose Process Supervision-Guided Policy Optimization, a method that integrates fine-grained feedback into the RL framework.
Figure 1 illustrates the overview of our approach.
By providing intermediate rewards that assess the correctness of partial code sequences, our approach guides the model more effectively toward generating functionally correct programs. This is achieved through a Process Reward Model (PRM) (Lightman et al., ) that evaluates each code generation step, offering dense reward signals that addresses the limitations of sparse end-of-trajectory rewards.

### 3.1 Process Supervision via Process Reward Models

Our method introduces a PRM to assess the correctness of each line of the code during the generation process. The PRM serves as an oracle that provides intermediate rewards based on the potential of the current code prefix to be extended into a correct program. By offering intermediate feedback, the PRM helps the model identify and reinforce beneficial code generation patterns while discouraging those that introduce errors. This fine-grained feedback mirrors the human approach to coding, where programmers continuously evaluate and adjust their code.

#### 3.1.1 Data Collection

Figure: Figure 2: Binary search over code steps at line level to label prefixes. The first midpoint at $m=3$ is accepted, so the search interval moves to $[4,5]$. The next midpoint at $m=4$ is rejected, indicating unrecoverable errors occur after step 3.

To effectively train the PRM, we require a dataset that provides fine-grained annotations indicating the correctness of partial code sequences. However, manually annotating the correctness of each line of code generated by the model is costly and not scalable. Instead, we employ an automated approach inspired by techniques used in recent works (Wang et al., ; ; Luo et al., ). Our method leverages the model’s own capabilities to generate completions for partial code prefixes and uses automated testing to assess their correctness.
The key idea is to determine whether a partial code prefix can be extended—by any means—into a complete program that passes all unit tests. If so, we consider the prefix as potentially correct; otherwise, it is labeled as incorrect.

Given a prompt $\mathbf{x}$, we generate a complete code response $\mathbf{y}=(\mathbf{y}^{(1)},\mathbf{y}^{(2)},\ldots,\mathbf{y}^{(T)})$ using the current policy $p_{\theta}$. Our goal is to determine the correctness of each partial code prefix $\mathbf{y}^{\leq t}$ for $t=1,2,\ldots,T$. To achieve this, we employ a best-of-$K$ sampling strategy to approximate an oracle capable of completing partial code prefixes.
For each partial code prefix $\mathbf{y}^{\leq t}$, we generate $K$ potential completions $\{\mathbf{c}_{k}\}_{k=1}^{K}$ using the current policy. We then form full programs $\mathcal{P}_{k}=(\mathbf{y}^{\leq t},\mathbf{c}_{k})$ and execute them against the unit tests $\mathcal{U}$. If any of these programs pass all unit tests, we label the partial code prefix as correct; otherwise, it is labeled as incorrect.
To efficiently identify the transition point where errors occur, we employ a binary search over the code generation steps (Luo et al., ), which is formalized in Algorithm 1.
For example, consider a code response divided into five steps ($T=5$), as shown in Figure 2. The partial prefix up to $\mathbf{y}^{(3)}$ can be completed to pass all unit tests, so it is labeled as correct. The prefix up to $\mathbf{y}^{(4)}$ cannot, meaning steps beyond $\mathbf{y}^{(3)}$ are labeled as incorrect.
For each partial code prefix $\mathbf{y}^{\leq m}$, the label $l_{m}$ is assigned based on the outcome of the completion attempts:

$$ $l_{m}=\begin{cases}+1,&\text{if any }\mathcal{P}_{k}\text{ passes all unit tests}\\ -1,&\text{otherwise}\end{cases}$ (1) $$

which indicate whether the prefix is potentially correct (can be completed to a correct program) or incorrect (contains unrecoverable errors).

#### 3.1.2 Training

Using the collected data $\{(\mathbf{x},\mathbf{y}^{\leq m},l_{m})\}$, we train the PRM $R_{\phi}$ to predict the correctness of partial code prefixes. The PRM learns to assign higher rewards to prefixes labeled as correct and lower rewards to those labeled as incorrect. The training objective, i.e., Mean Squared Error (MSE), minimizes the discrepancy between the PRM’s predictions and the assigned labels:

$$ $\min_{\phi}\sum_{(\mathbf{x},\mathbf{y}^{\leq m})}\left(R_{\phi}(\mathbf{x}, \mathbf{y}^{\leq m})-l_{m}\right)^{2}$ (2) $$

This regression formulation allows the PRM to estimate the likelihood that a given prefix can be successfully completed.
Notably, aside from employing Mean Squared Error (MSE) loss, we also experimented with Cross-Entropy Loss and empirically found that MSE loss yielded better performance in our case.

### 3.2 Integrating PRM into RL Training

Given a learned PRM, we aim to identify best practices for enhancing code generation during RL training. While prior work Lightman et al. (); Wang et al. (); Jiao et al. (); Wang et al. (); Luo et al. () uses PRMs to verify intermediate steps in mathematical tasks, this approach is less relevant for code generation. In mathematical domains, LLMs may generate correct answers with faulty reasoning Lightman et al. (), making intermediate verification essential. However, in code generation, problems are typically accompanied by multiple unit tests, making it improbable for incorrect code to pass all tests. As a result, the emphasis on intermediate verification is less applicable.
Instead, we propose leveraging PRMs as auxiliary sources of dense signals to facilitate better exploration during RL training. While preliminary attempts have been made to incorporate PRMs into RL trainingWang et al. (), these efforts are limited and warrant a more thorough investigation. We explore the following methods to integrate PRMs effectively:

##### PRM as Dense Rewards.

Similar to Wang et al. (), we use PRMs to provide step-level reward signals that guide more efficient policy exploration during RL training. By rating the correctness of each line in the code response, the PRM supplies “dense” rewards that encourage the policy to explore more promising code paths, leading to improved performance.

##### PRM as Value Initialization.

The PRM’s method of annotating code, by fixing a prefix $\mathbf{y}^{\leq t}$ and rolling out the policy to sample correct responses, can be viewed as a “hard” value estimation of $\mathbf{y}^{\leq t}$ . We hypothesize that the PRM’s capability to provide line-level feedback could serve as a useful inductive bias for initializing the value function in RL algorithms, which can ease the credit assignment problem by offering a more informed starting point.

##### PRM as Both Dense Rewards and Value Initialization.

To fully capitalize on the advantages of PRMs, we combine both approaches. By using PRMs for dense rewards and as an initialization for the value model, we aim to enhance RL training through improved exploration and more effective credit assignment.

## 4 Experimental Results

### 4.1 Experimental Setup

##### Datasets and Evaluation.

We utilize in-house datasets to train our model for code generation.
Specifically, the training set, $\mathcal{D}_{\text{train}}$, is a comprehensive Reinforcement Learning with Human Feedback (RLHF) dataset that includes, as a subset, approximately $30,000$ diverse coding problems. Each of these problems is paired with unit tests designed to validate the functional correctness of the generated code.
For evaluation, we employ two benchmarks: the publicly available LiveCodeBench( Jain et al. ()) and our proprietary code generation benchmark (InHouseBench).
LiveCodeBench is a rigorous benchmark consisting of 612 coding tasks, collected between May 1, 2023, and September 1, 2024, designed to assess the code generation capabilities of LLMs.
InHouseBench comprises 245 challenging coding problems in Chinese, spanning both Python and C++.
All test problems are novel and unseen in public datasets, ensuring there is no risk of data contamination.
The benchmark contains two problem categories: (1) Contest, which includes 169 coding competition-style problems, and (2) NL2Alg, which comprises 76 problems focused on translating natural language algorithm descriptions in Chinese into executable code.
For evaluation, we let each model generate 10 candidate responses for each problem, using a temperature of 0.2, nucleus sampling with top-$p\text{=}0.95$, and top-k sampling with $k\text{=}128$, following common practice. We adopt Pass@1 as the evaluation metric, in line with previous work (Kulal et al., ; Chen et al., ; Jain et al., ).

##### Base Models and RL Baseline.

The base model used in our experiments is a small in-house model, referred to as Doubao-Lite throughout the remainder of this paper. Initially, Doubao-Lite was fine-tuned on our Supervised Fine-Tuning (SFT) dataset, resulting in Doubao-Lite-SFT, which then served as the initialization for the subsequent RLHF training phase.
We fine-tune Doubao-Lite-SFT ($\pi_{\mathrm{ref}}$) on the RLHF dataset $\mathcal{D}_{\text{train}}$ using Proximal Policy Optimization(PPO) (Schulman et al., ) to obtain Doubao-Lite-RL ($\pi_{\theta}$). In our setup, two types of Outcome Reward Models (ORMs) are employed as the objective functions for RL training. For non-coding prompts, we use a general reward model, $R_{\texttt{general}}(\mathbf{x},\mathbf{y})$, derived from preference learning on a human-annotated dataset (Ouyang et al., ). For coding prompts, the ORM is defined as a binary indicator of whether the response passes all unit tests, $R_{\texttt{UT}}$ ($\mathbf{x}$, $\mathbf{y}$). Following  Ouyang et al. (), RLHF optimization objective is defined as:

$$ $\max_{\theta}\sum_{\mathbf{x}\in\mathcal{D}_{train}}\mathbb{E}_{\mathbf{y}\sim \pi_{\theta}(\mathbf{y}\mid\mathbf{x})}\left[R(\mathbf{x},\mathbf{y})-\beta \mathrm{KL}(\pi_{\theta}\parallel\pi_{\mathrm{ref}})\right],$ $$

with $R(\mathbf{x},\mathbf{y})=R_{\texttt{general}}(\mathbf{x},\mathbf{y})$ for non-coding prompts and $R(\mathbf{x},\mathbf{y})=R_{\texttt{UT}}(\mathbf{x},\mathbf{y})$ for coding prompts.

##### PRM Training.

To ensure that the PRM training data effectively covers the state space the language model may encounter during the next RL training phase, we sample policy models from various stages of the RL baseline training. Specifically, we select 4 checkpoints evenly spaced throughout the RL baseline model’s training process. For each checkpoint, we sample $n$ responses for each coding prompt in the training dataset $\mathcal{D}_{\text{train}}$. For each sampled response, we apply the binary search labeling procedure described in Algorithm 1, using $K=20$ completions for each partial code prefix.
The data collected from all checkpoints is then aggregated into a PRM training set, denoted as $\mathcal{D}_{\text{PRM}}$. We initialize the PRM with the value model from the RL baseline and fine-tune it on the aggregated dataset, $\mathcal{D}_{\text{PRM}}$, using the objective function defined in Eq. (2).

##### Integrating PRM into RL.

As described in Section 3.2, we explore two methods for integrating the Process Reward Model (PRM) into RL training: (1) using PRM as a source of dense reward signals (DenseReward) and (2) initializing the value function in PPO with PRM (ValueInit). In the DenseReward approach, PRM assigns additional reward signals at each end-of-line token (\n) in the code response for coding prompts. Thus, the RL optimization objective for coding prompts is modified to the weighted sum of $R_{\texttt{UT}}$ and $R_{\texttt{PRM}}$, as defined below:

$$ $\max_{\theta}\sum_{\mathbf{x}\in\mathcal{D}_{train}}\mathbb{E}_{\mathbf{y}\sim \pi_{\theta}(\mathbf{y}\mid\mathbf{x})}\left[R_{\texttt{UT}}(\mathbf{x}, \mathbf{y})+\lambda R_{\texttt{PRM}}(\mathbf{x},\mathbf{y})-\beta\mathrm{KL}( \pi_{\theta}\parallel\pi_{\mathrm{ref}})\right],$ (3) $$

where $\lambda$ controls the relative importance of PRM in shaping the reward. Specifically, we set $\lambda=0.25$ when the code response does not pass all unit tests, i.e., $R_{\texttt{UT}}(\mathbf{x},\mathbf{y})=0$, and $\lambda=0.025$ when the response passes all unit tests, i.e., $R_{\texttt{UT}}(\mathbf{x},\mathbf{y})=1$. The intuition behind this reward shaping is to leverage PRM to provide informative signals when the RL policy fails to generate a valid solution, while minimizing the risk of PRM over-optimization (Rafailov et al., ; Skalse et al., ) once a correct solution is found. Our empirical results indicate that this reward shaping strategy performs effectively in our experimental setting.
In the ValueInit setting, PRM is simply used to initialize of the value function in PPO.
Notably, these two approaches-DenseReward and ValueInit—are complementary and can be applied concurrently.

### 4.2 Key Considerations for Integrating PRM into RL Training

While integrating PRM into RL training might seem straightforward, we found that achieving effective results requires careful attention to several critical factors. In this section, we highlight key implementation details essential for the successful application of PRM in RL training.

#### 4.2.1 PRM Training: More Data or Better Data?

Recent research on LLMs highlights that data quality often outweighs quantity (Gunasekar et al., ; Li et al., ). We found the same holds true for PRM training data selection. Although automated data annotation allows for the generation of large volumes of PRM training data through model sampling, our experiments showed that increasing data volume can sometimes degrade PRM performance when integrated into RL. In contrast, a smaller, carefully curated subset of the full dataset led to better supervision and improved outcomes.
For example, when all sampled responses to a given prompt either consistently pass (or fail) unit tests, PRM gains little useful information. In such cases, the model can only learn to predict the correct (or incorrect) label when it encounters the same prompt again, limiting its ability to generalize. We explored various strategies for selecting and filtering data, as detailed in Section 4.3.

#### 4.2.2 RL Training: Alleviating PRM Hacking

Reward model hacking (Skalse et al., ) is a well-known issue in RLHF training, where the policy learns to exploit flaws in the reward model to achieve high rewards without genuinely improving the quality of response.
Similarly, we observed that PRM is also susceptible to such exploitation.
Here we discuss two key practical strategies to mitigate the risk of PRM hacking and ensure the reward signals remain aligned with the intended task objectives.

##### PRM reward length normalization.

As described in Section 4.1, when used to provide dense rewards, PRM assigns line-level reward signals at the end-of-line tokens in the LLM-generated response.
However, if we directly use the predictions of the learned PRM, denoted as $R_{\phi}$, as the reward signal $R_{\mathrm{PRM}}$ in 3, we observed that this can be exploited.
Specifically, the policy may generate numerous lines for which PRM predicts positive rewards, thus inflating the overall reward.
This occurs because writing more lines allows the model to accumulate excessive intermediate rewards, effectively hacking the optimization objective.
To mitigate this issue, we apply length normalization to the PRM predictions. Given a prompt $\mathbf{x}$ and a response $\mathbf{y}$ with $T$ lines, $\mathbf{y}$ = ($\mathbf{y}^{(1)}$, $\mathbf{y}^{(2)}$, …, $\mathbf{y}^{(T)}$), we define the PRM dense reward signal at the $m$-th line as:

$$ $R_{\mathrm{PRM}}(\mathbf{y}^{(m)})=\frac{1}{T}\cdot R_{\phi}(\mathbf{x}, \mathbf{y}^{\leq m}).$ $$

This normalization ensures that the policy does not gain higher cumulative rewards by generating trivial or unnecessarily long responses, as the accumulated reward is bounded within the range of $[-1,1]$ regardless of the response length.

##### Assigning additional neutral label into PRM training.

While length normalization helps reduce PRM exploitation, it is insufficient to fully prevent PRM hacking. Even with normalization, we empirically observed that the model can still exploit PRM by generating excessive comment lines within the code. The underlying issue is that writing a correct comment is often far easier than producing correct code. As a result, the model can artificially inflate the PRM reward by including unnecessary comment lines.
To address this issue, we introduce an additional neutral label in the PRM annotation, as defined in Equation 1:

$$ $l_{m}=\begin{cases}+1,&\text{if any }\mathcal{P}_{k}\text{ passes all unit tests}\\ 0,&\text{if the line is a comment}\\ -1,&\text{otherwise}\end{cases}$ $$

By assigning a neutral label (0) to comment lines, we remove the reward bias that encourages the model to generate unnecessary comments. This adjustment ensures that only meaningful contributions to the code, rather than extraneous comments, are rewarded by the PRM.

**Table 1: Comparison of model performance (Pass@1) across LiveCodeBench and InHouseBench with PRM as DenseReward and ValueInit settings. The performance of our models (InHouse-Lite series) on both LiveCodeBench and InHouseBench are averaged over 10 independent runs. We also report the performance of other public models for comparative purpose. The performance of Gemini-Flash-1.5 on InHouseBench is omitted due to legal considerations.**
| Model | Setting | Dataset |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Dense | Value | LiveCodeBench | InHouseBench |  |  |  |  |  |  |
| Reward | Init. | Easy | Medium | Hard | Overall | Contest | NL2Alg | Overall |  |
| GPT-4o-mini | - | - | 81.9 | 27.2 | 3.6 | 40.7 | 43.8 | 68.4 | 51.4 |
| Qwen2-72B | - | - | 65.0 | 21.3 | 2.8 | 32.2 | 14.8 | 51.3 | 26.1 |
| Gemini-Flash-1.5 | - | - | 67.7 | 13.1 | 1.9 | 29.6 | - | - | - |
| DeepseekCoder-33B | - | - | 60.8 | 14.8 | 1.2 | 27.7 | 10.3 | 50.3 | 22.7 |
| Ours-SFT | - | - | 55.3 | 9.3 | 0.3 | 23.5 | 10.4 | 41.4 | 20.0 |
| Ours-RL | × | × | 70.0 | 7.2 | 1.7 | 28.2 | 24.4 | 48.7 | 31.8 |
| × | ✓ | 67.9 | 8.9 | 1.9 | 28.2 | 25.0 | 45.4 | 31.4 |  |
| ✓ | × | 68.5 | 9.9 | 2.5 | 28.9 | 25.2 | 48.1 | 32.3 |  |
| ✓ | ✓ | 69.3 | 12.0 | 1.6 | 29.8 | 27.9 | 53.5 | 35.8 |  |

### 4.3 Main Results and Analysis

Figure: Figure 3: Pass@1 difference between policies trained with and without PRM across varying response lengths. Policies trained with PRM exhibit consistent improvements over those without PRM for longer-horizon responses (greater than 100 tokens). This demonstrates PRM’s effectiveness in providing intermediate feedback, thereby enabling RL to do more explorations.
Refer to caption: x2.png

##### Comparing Different Strategies of Using PRM in RL Training.

We explore three strategies for integrating the PRM into RL training, as described in Section 4.1: DenseReward, ValueInit, and a combined approach of DenseReward & ValueInit. The performance of RL models trained using these strategies on LiveCodeBench and InHouseBench, in comparison to SFT and RL baselines, is summarized in Table 1. For further comparison, we also include results from several publicly available models OpenAI (); Bai et al. (); Reid et al. (); Guo et al. ().
Our experimental results reveal that using PRM solely as dense rewards significantly outperforms the RL baseline, consistent with findings from Wang et al. (). This suggests that the granular feedback provided by PRM helps the policy explore more promising solutions by offering continuous corrections at intermediate steps.
Moreover, we observe that combining PRM as both dense rewards and value function initialization results in substantial performance gains, with a $5.7\%$ improvement on LiveCodeBench and a $12.6\%$ improvement on InHouseBench, relative to the RL baseline.

Interestingly, using PRM solely for value function initialization does not provide notable benefits. We hypothesize that while value function initialization can enhance credit assignment and stabilize the learning process, it does not address the underlying issue of sparse reward signals. Without dense feedback, the policy may fail to explore the solution space effectively, resulting in limited improvement.
In contrast, the combination of DenseReward and ValueInit offers a synergistic effect: dense rewards enhance exploration by providing rich intermediate feedback, while value function initialization improves stability and credit assignment.
Together, these mechanisms enable the policy to converge more efficiently toward optimal solutions, which explains the significant performance improvements we observed.

##### PRM enhances code generation in long-horizon scenarios.

To better understand when PRM benefits code generation most, we analyze its effect based on the length of the generated responses. Intuitively, the dense nature of the reward signals provided by PRM is particularly advantageous for long-horizon tasks, where intermediate feedback can guide policy exploration more effectively.
To validate this, we quantitatively compared the pass rate (Pass@1) of models trained with and without PRM across different response lengths. The results are visualized in Figure 3.
Overall, the model trained with PRM demonstrates a $9\%$ improvement in Pass@1 compared to the baseline model trained without PRM.
Notably, PRM consistently improves Pass@1 for responses longer than 100 tokens. However, for responses shorter than 100 tokens, PRM yields comparable or slightly inferior results. We hypothesize that in short-horizon scenarios, PRM may function similarly to a biased ORM, which limits its ability to provide meaningful improvements. Shorter responses inherently benefit less from the dense feedback signals PRM offers, as these tasks may already be well-explored by the policy without needing extensive guidance.
In contrast, for more complex, long-horizon tasks, PRM offers valuable intermediate signals that effectively guide the policy to explore better solutions. These signals provide a more nuanced understanding of the correctness of individual code lines, helping the policy navigate the larger solution space with the same amount of optimization compute.

##### The Importance of PRM Training Data Selection

PRM training data can be categorized at two levels: At the response level, responses are classified as Correct (passes unit tests immediately), Revised (initially fails but can find a correct prefix), and Wrong (cannot find any correct prefix by binary search within the given budget). At the prompt level, prompts are categorized as Easy (all responses are Correct), Medium (mixed response types), and Hard (all responses are Wrong). We tested the following data selection strategies: Full (use all collected data), Remove Hard (exclude Hard prompts and their responses), Medium Only (include only prompts with mixed response types), and Revised Only (use only Revised responses).
We empirically found that Revised Only, which includes the richest process-level correction signals, performs best in our setting.

| Strategy | LCB | IHB |
| --- | --- | --- |
| Full | 26.9 | 34.6 |
| Remove Hard | 27.8 | 33.8 |
| Medium Only | 26.9 | 32.5 |
| Revised Only | 29.8 | 35.8 |

Table: Table 2: Comparison of different PRM data selection strategies on two datasets: LCB (LiveCodeBench) and IHB (InHouseBench).

Figure: Figure 4: Pass@1 on LiveCodeBench as the average number of responses per prompt for PRM data collection increases (logarithmic scale).

##### How much data needed to train a PRM that benefits RL training?

Given that automatic PRM data collection is computationally expensive, we examine how the performance of policies trained with PRM scales with the number of training samples. Figure 4 shows how the pass rate of models trained with varying amounts of PRM data changes along the average number of responses collected per prompt for PRM data collection, as described in Section 3.1.1. The key finding is that the performance of models trained with PRM improves consistently as the number of PRM training samples increases, highlighting the effectiveness and scalability of our approach.

## 5 Related Works

### 5.1 LLMs for Code Generation

Recently, large language models (LLMs) have demonstrated impressive capabilities in code generation by pre-training on vast text datasets that include code (Lu et al., ; Christopoulou et al., ; Allal et al., ; Zheng et al., ; Li et al., ). Additionally, models fine-tuned through supervised fine-tuning (SFT) have achieved competitive results in code generation tasks (Chen et al., ; Li et al., ; Luo et al., ; Rozière et al., ; Guo et al., ).
Reinforcement Learning (RL) optimizes policies by interacting with an environment and receiving rewards (Williams, ). Recently, RL has been incorporated into LLMs to enhance code generation using unit test feedback (Shojaee et al., ; Liu et al., ; Le et al., ). CodeRL (Le et al., ) applies unit test signals as rewards with an actor-critic method, while PPOCoder (Shojaee et al., ) builds on this by using the PPO algorithm. RLTF (Liu et al., ) improves precision by locating errors, though the reward space remains sparse. Despite progress, RL’s potential to significantly boost code generation in sparse reward environments remains underexplored.

### 5.2 Process Reward Models

Process reward models (PRMs) have garnered significant attention in recent LLM developments, particularly in the mathematical reasoning domain, where they provide verification for intermediate reasoning steps (Lightman et al., ; Wang et al., ; Jiao et al., ; Wang et al., ; Luo et al., ). While some approaches rely on costly and resource-intensive human-annotated process data (Lightman et al., ), recent research has focused on automating the collection of process supervision data Wang et al. (); Jiao et al. (); Wang et al. (); Luo et al. (). Building on these efforts, we similarly automate process supervision but differ in our primary objective. Rather than using PRMs solely as enhanced verifiers compared to Outcome Reward Models (ORMs), we focus on their integration into RL training for code generation. While Wang et al. () provides preliminary results on PRMs improving RL training in the mathematical domain, their findings are limited. Our work offers a more thorough and systematic investigation of how PRMs can be leveraged in RL for code generation tasks.

## 6 Conclusions and Limitations

In this work, we addressed the challenge of sparse reward signals in reinforcement learning (RL) for code generation by introducing a Process Reward Model (PRM) that provides dense, line-level feedback. This approach mirrors human-like code refinement and improves learning efficiency. Our experiments demonstrate that integrating PRMs significantly enhances the pass rates of code generation models on both the LiveCodeBench dataset and proprietary benchmarks. This method has the potential to improve long-horizon code generation scenarios, advancing the state-of-the-art in LLM-based code generation.

Despite the promising results, our approach has several limitations that warrant further investigation. First, the effectiveness of the PRM relies heavily on the quality of the collected data. While we automated data collection using binary search and unit tests, this method may not capture all nuances of code correctness and may introduce noise, especially in more complex or ambiguous programming tasks. Second, the computational cost of collecting PRM training data is still substantial, even though we employed binary search to mitigate it. Third, our automated PRM data collection method requires external verification (unit tests in our case), which is not applicable to many other domains such as creative writing or open-ended generation problems. This could limit the applicability of our approach in those areas.

## Appendix A Appendix

Figure: Figure 5: RL training curve of our method (DenseReward&ValueInit) compared to other settings. Using PRM as both DenseReward and ValueInit yields the best result.
Refer to caption: extracted/5947965/figs/rl_training.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: - > To ensure that the PRM training data effectively covers the state space the language model may encounter during the next RL training phase, we sample policy models from various stages of the RL baseline training. Specifically, we select 4 checkpoints evenly spaced throughout the RL baseline model’s training process. For each checkpoint, we sample n responses for each coding prompt in the training dataset Dtrain. For each sampled response, we apply the binary search labeling procedure described in Algorithm 1

  For the oracle, is this also using the policy model checkpoints? I see "Our method leverages the model’s own capabilities to generate completions for partial code prefixes and uses automated testing to assess their correctness" which suggests that the policy model itself is used to generate the best-of-K samples. I can appreciate that not requiring a separate oracle model is nice because it is self-contained, but I think this will result in a suboptimal dataset compared to e.g. using a fully-trained RLTF code model as the oracle.

- Reading Section 4.2.2 on RL Training makes me think: Wouldn't a different scheme of data labelling, which labels each line with the "marginal contribution of the line toward success" be more effective than simply rating [0: infeasible, 1: feasible]? Specifically, my intuition is that each added line should improve the success rate of the oracle given K attempts, so my proposed reward is something like "the reward for step M should be the `(success_rate_of_oracle_at_step_M - success_rate_of_oracle_at_step_Mminus1)`". This captures the idea that each line should increase the likelihood of the program succeeding, and naturally avoids reward hacking by simply adding more lines. Of course, this is easier said than done, I expect the process to be noisy, but this formulation for the dataset seems to be better aligned to the true objective.

- (Not a weakness, just a typo) In Section 3.2, I think the authors meant to do citep instead of just cite: "In mathematical domains, LLMs may generate correct answers with faulty reasoning Lightman et al. (2023), making intermediate verification essential" and "While preliminary attempts have been made to incorporate PRMs into RL trainingWang et al. (2024a)"
soundness: 4
strengths: Nice work overall! I enjoyed reading this and have become more bullish on PRMs as a general research direction after this.
- Very readable and easy to understand.
- Clearly a lot of work was put into it: Devises a novel practical approach for training models with PRMs for code generation, with careful work put into ablations and studying each component.
- Competitive with the SOTA RLTF approach
- Clever method to generate process-level supervision data to train PRMs! The idea of using a best-of-K oracle to label whether each code prefix is feasible or contains unrecoverable errors is a non-obvious but effective way to generate data.
- Section 4.2 and 4.3 were great! Well-reasoned experiments and execution, with great experiment-backed insights on how best to use PRMs in this domain.
  - Particularly appreciated the detail-to-attention uncovering PRM Hacking and implementing mitigations.
  - I like that the main results (Table 1) independently shows the effect of introducing [Value Init], [Dense Reward], and [Value Init + Dense Reward]! Very clear.
summary: The paper introduces a practical approach for training models with process-oriented feedback for code generation. The approach uses a novel automatic (LLM + unit test -based) data generation process to create a dataset of code with quality labels on every line. The line-by-line labels can be used to trained a Process Reward Model (PRM), which can in turn be used as a dense reward signal for training code models with RL. The approach is shown to outperform the baseline RL-trained code models which are trained on sparse rewards from unit tests. The authors perform careful experiments and ablations to motivate each part of their design.
weaknesses: No clear weaknesses come to mind. There were some choices made which I had different ideas about, but this is not a critique of the work or the claims, so I have put those into the Questions section.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. On the sensitivity to the $\lambda$: L269 mentions that the authors use different $\lambda$ for correct/incorrect code. I'm curious: Did the author try using the same $\lambda$ or other $\lambda$ during hyperparam sweep? Whether the exp. results deviate much from the ones reported in Table 1?

2. Why the x-axis in Figure 4 could have <1 value as the avg. number of collected response per prompt?

3. L366 says the gain is 5.7% on LiveCodeBench and 12.6% on in-house benchmark relative to the RL baseline. What lines do these numbers correspond to in Table 1? I'm guessing I should be comparing the row of Ours-RL (Dense Reward x, Value Int. x) and the last row. But from Table 1 it's 29.8% - 28.2% = 1.6% and 35.8% - 31.8% = 4%.
soundness: 3
strengths: The contributions of the paper are the following:
1. Adapting the PRM to the PPO pipeline for code generation
2. Details of how the authors design the binary search, collect, and choose PRM training data
3. Empirical results on LiveCodeBench and in-house benchmark show the effectiveness of PRM and analysis to investigate the source of gain

The paper is well-written and easy to follow. It also highlights challenges when incorporating the PRM into the PPO framework (such as reward hacking in Sec 4.3.3). To mitigate reward hacking, it proposes length normalization and neutral labels for comment lines.
summary: This paper proposes integrating the Process-Reward Model (PRM) that provides line-level rewards into PPO training for single-turn code generation. The authors design a binary search procedure to gather line-level rewards, which is defined as being able to lead to correct code (the authors use best-of-K sampling to approximate it) for PRM training. The authors propose using the PRM as a source of dense reward and/or value function initialization and evaluating the methods' effect on LiveCodeBench and the in-house benchmark using the in-house model.
weaknesses: 1. My primary concerns are reproducibility, the lack of experimental details, and data transparency, which hinder the community from reproducing the results presented in the work and accessing the needed efforts. This paper uses an in-house model, and part of the results are reported on the in-house benchmark. This is not inherently an issue but means that there needs to be more details in the text/contributions in other components. Therefore, I encourage the authors to include the following:
-----
#### Description of the In-house model
1. What is the parameter size and architecture? 
2. What is the size of the SFT set, if available (number of tokens/data points) mentioned in L241 and L337? Could the authors show an example of the SFT data? Could the authors provide some basic configuration of the SFT (e.g., lr, number of gradient steps)?

#### PRM data and training
1. Could the authors give some statistics on the different strategies of PRM training data, if available (number of tokens/data points, avg. lines of code, the distribution of the % or the line number where the line-level label turns -1 from +1)? These statistics could be a valuable add-on to Table 2 and the plain strategy. Could the authors show one example of the reward produced by the binary search?

#### PPO training
1. what is the KL penalty (the value of $\beta$ in Eq(3)) used in the exp.? Some basic hyperparams (lr, steps) would be appreciated.
------
2. The proposed PRM data collection requires model checkpoints from RL baseline training (L252). Therefore, the whole pipeline is: RL baseline training -> PRM data generation collection -> PRM training -> RL training w/ PRM. It makes the first RL baseline training mandatory and increases the computation cost. Did the author explore alternatives, such as using a base policy (e.g., in-house-SFT) to conduct PRM data generation and collection?

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Better articulate the contributions:
1) What unique challenges arise when applying PRMs to code generation compared to other domains? 
2) How does the approach specifically address these challenges? What are the specific differences between the proposed methods and previous methods that inspire this method?
3) The unique challenges in collecting process supervision for code generation versus mathematical reasoning tasks. A point-by-point comparison of their method with OmegaPRM and other relevant approaches, highlighting any novel aspects specific to code generation. 
4) Analysis of code LLMs trained with the PRM. Any comparisons on how training with/without PRM improve generated code responses. How does the learned PRM generalize? How does PRM help the value estimate during training?
5) Include results on these specific benchmarks (APPS, MBPP, HumanEval) in their evaluation section. If these benchmarks were not used, provide a clear explanation for why they were omitted and how their chosen benchmarks compare in terms of difficulty and relevance.
Discuss how their results on the chosen benchmarks might translate to performance on these more standard benchmarks.
soundness: 2
strengths: First attempt to use PRM for code LLMs.
summary: The paper studied using PRM for code generation.
weaknesses: The current paper delivery is poor, many critical questions is not clearly explained.

Weakness 1: 
The paper discusses using PRM for code generation. But I don't see what specific challenges the proposed method addresses regarding code generation.

The motivation for using PRM in code LLMs (Sec.1 para.2 line 035-043) is that the reward signal is sparse. Sparse reward (or more accurately, the use of bandit feedback) is a common challenge in multi-step reasoning tasks such as math reasoning and code, many research have been conducted to address this challenge, e.g., introducing PRM. 
- If the paper limits its scope to code generation, the paper should clearly explain what are the specific challenges it address regarding the code generation task. I notice that unit test feedback is mentioned as a cause for the sparse reward, but I don't consider it as a particular challenge specific in the code generation task. In its essence, test cases serve as a means of verification of the holistic response -- same in its role as evaluating the correctness of a holistic response using the gold-answer in math reasoning. The natural question is: Is previous methods to deal with sparse rewards in math reasoning directly applicable to code generation? If not, the current paper does not clearly explain this. 
- If the paper targets at proposing new methods for PRM, the paper should clearly state what is new about the methods. 

Based on my current understanding of the paper, the paper answers neither of these.


Weakness 2:
For the current version of this paper, I have doubts on the claimed contributions (line 061)
- From my understanding based on the current paper delivery, the paper only applies previous methods to the code generation task. Reasons detailed as follows:
	- Most importantly, the paper do not clearly explain what additional challenges arise when collecting process supervision for code generation, in contrast to existing works that collect process supervision in reasoning tasks (e.g., OmegaPRM).
	- The paper do not clearly explain how the proposed method differ with previous methods (OmegaPRM), except for the application domain (math versus code). 

         OmegaPRM (https://arxiv.org/pdf/2406.06592) introduces the method of automated process supervision using binary search. The paper mentions in line 147: "Instead, we employ an automated approach inspired by techniques used in recent works", but they do not mention how the proposed method differ with OmegaPRM.
- The paper claims to conduct empirical study of "how to properly and effectively integrate PRMs into RL". I see from section 4.3 that, the paper experiment on using PRM as dense rewards, or/and as Value initialization. Using PRM as dense rewards is studied in MathShepherd (https://arxiv.org/abs/2312.08935). The paper draws a conclusion that using PRM to initialize the value function of PPO does not work, but this paper (https://arxiv.org/abs/2406.03816) explores an effective way to employ PRM for value initialization of PPO. So I don't see valuable contributions in this empirical study. 
- I question the evaluation of the proposed method. There are no results on common benchmarks for code generations, such as APPS, MBPP, HumanEval

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 3
questions: In Table 1, you compare "Ours-SFT" against "Ours-RL", but the text in L366/367 refers to it as "RL baseline"? Relatedly, since you train the PRM on data produced during the RLHF phase, does that mean "Ours-RL" numbers are the product of two RL training stages? Where is the model after RLHF then (i.e., the "RL baseline")?
soundness: 2
strengths: With direct experience in the paper's domain I found it, for the most part, easy to read and understand. I liked how the paper ablates different parameters for PRM training and usage. The subject of the paper is also of interest within the community as well as for the deployment of code LLMs in products.
summary: The paper at hand concerns itself with RL fine-tuning of LLMs from unit test feedback, i.e., obtained in code generation tasks by testing the LLM output against a set of unit tests. The main feature here is the use of a process reward model (PRM) to supply dense rewards at multiple points within a generation (token sequence), rather than just at the end (unit test feedback). While PRMs have been used in several domains before, the key contributions are a recipe to obtain a PRM for code as well as ablations on how to best utilize it during RL fine-tuning.
weaknesses: My main grievance with this paper is that it has very low reproducibility. Models, training data, and evaluation sets are all proprietary; we don't even learn about the size of the models. It does not seem to be a large model, though, indicated by the "lite" name and by the fact that Table 1 only compares to "mini" models from OpenAI and Google. To add to that, a number of crucial details are missing, e.g., hyper-parameters for SFT and RL.

Throughout the paper, the authors claim that their PRM improves exploration. I don't see this claim verified experimentally. Instead, the experiments show better *generalization* to LiveCodeBench and their in-house benchmark, which is not the same. It would be interesting to see if indeed a larger fraction of training problems is solved during training, or whether more diverse solutions are found.

Meta-review: > The paper at hand concerns itself with RL fine-tuning of LLMs from unit test feedback, i.e., obtained in code generation tasks by testing the LLM output against a set of unit tests. The main feature here is the use of a process reward model (PRM) to supply dense rewards at multiple points within a generation (token sequence), rather than just at the end (unit test feedback). While PRMs have been used in several domains before, the key contributions are a recipe to obtain a PRM for code as well as ablations on how to best utilize it during RL fine-tuning.

This looks like it's good work, but the paper is problematic. It's not clear from the paper alone what are the differences with OmegaPRM, and this a notable part of the contribution. Moreover, the main problem is that the paper is not easy to compare to (let alone reproduce). During rebuttal, the authors worked on applying their RL+PRM method on Qwen 2.5 7B and LiveCodeBench, but before that it was just a proprietary (in house) model on a proprietary (in house) dataset. The two scores of "3" were made by honest reviews, and while those reviewers didn't change their score in the light of the rebuttal, I do not think it would have been likely they both change them to 6+, even with the changes provided by the authors.

I suggest to rework the paper to put the focus on the algorithmic contribution with at least a comparable setup, or better yet a reproducible one, so that the paper gets read from the get go on its merits, not its (major) flaws.

Decision: Reject