Reinforcement learning method using code review model as reward model

Abstract

The invention discloses a reinforcement learning method using a code review model as a reward model, and the method comprises the steps: taking data containing a code and a corresponding manual review result label as a training sample, carrying out the supervised learning training to obtain a code review model, and enabling the code review model to receive the code and carry out the evaluation of the code, outputting an evaluation result; under a reinforcement learning framework, taking a code generation model as an intelligent agent, and inputting a code generated by the code generation model into the code review model; converting an evaluation result output by the code review model into a reward signal of a reinforcement learning agent; and updating strategy parameters of the code generation model by using the reward signal. According to the method, the code review process is formalized into a computable reward, more accurate and efficient optimization guidance is provided for the code generation task, and the method is mainly used for improving the code generation accuracy and reliability of an artificial intelligence programming assistant.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence technology, and in particular to a reinforcement learning method that uses a code review model as a reward model. Background Technology

[0002] In recent years, deep learning-based code generation models have made significant progress. However, effectively optimizing these models to generate higher-quality, more demanding code remains a challenge. Traditional supervised learning relies on large amounts of high-quality paired data (such as natural language description-code pairs), but acquiring such data is costly and it is difficult to cover all possible programming scenarios and best practices.

[0003] Reinforcement learning provides an alternative path for code generation by defining reward functions to guide model optimization. In existing technologies, reward functions are typically designed based on binary results such as whether the code compiles or passes predefined unit tests. This reward signal is too sparse and coarse, failing to provide fine-grained guidance on quality dimensions such as code style, readability, and efficiency. For example, even if a piece of code passes tests, it may still contain redundancy, potential vulnerabilities, or bad practices. Therefore, there is an urgent need for a reward mechanism that can deeply evaluate code quality like a human expert.

[0004] Therefore, it is necessary to improve one or more of the problems existing in the above-mentioned related technical solutions.

[0005] It should be noted that this section is intended to provide background or context for the technical solutions of this disclosure as set forth in the claims. The description herein does not constitute an admission that it is prior art simply because it is included in this section. Summary of the Invention

[0006] The purpose of this invention is to provide a reinforcement learning method that uses a code review model as a reward model, thereby overcoming, at least to some extent, one or more problems caused by the limitations and defects of related technologies.

[0007] This invention provides a reinforcement learning method that uses a code review model as a reward model, comprising: S1, using data containing code and corresponding manual review result labels as training samples, supervised learning is performed to train a code review model. The code review model can receive code, evaluate the code, and then output the evaluation result. S2, Under the reinforcement learning framework, the code generation model is used as an agent, and the code generated by the code generation model is input into the code review model; S3, the evaluation results output by the code review model are converted into reward signals for the reinforcement learning agent; S4, update the policy parameters of the code generation model using the reward signal.

[0008] In this invention, the manual review result label includes a code quality score, a Boolean indicator indicating whether a defect exists, and review comments text.

[0009] In this invention, the code review model uses the pre-trained model CodeBERT as the encoder and regression loss and classification loss as the loss functions.

[0010] In this invention, the loss function for:

[0011] in, The true label score for code quality. The code quality score predicted by the model. Mean square error, This is the loss balance coefficient. For binary cross-entropy, For the first i The true label of class defects The model predicts the first i The probability of the existence of a class of defects.

[0012] In this invention, the evaluation result is a scalar score or a review comment text.

[0013] In this invention, when the evaluation result is a scalar score, the scalar score is used as a reward, or the difference between the scalar score and the baseline score is used as a reward.

[0014] In this invention, when the evaluation result is a review opinion text, the matching degree between the review opinion text and the expected behavior template is used as a reward.

[0015] In this invention, the parameters of the code review model are fixed during the reinforcement learning training process.

[0016] The technical solution provided by this invention may include the following beneficial effects: This invention presents a reinforcement learning method that utilizes a code review model as a reward model. The rewards provided by the code review model are more refined, enabling optimization of quality dimensions such as code style, maintainability, and robustness that cannot be covered by compilers and simple tests, thus guiding the model to generate higher-quality code. Compared to sparse rewards that rely solely on the final test results, the reward signals provided by this invention are denser and more instructive, which can accelerate the convergence of the reinforcement learning model and improve training efficiency. The entire process requires no frequent manual intervention, has a high degree of automation, and achieves automated closed-loop optimization from code generation to quality assessment. Attached Figure Description

[0017] The accompanying drawings, which are incorporated in and form part of this specification, illustrate embodiments consistent with this disclosure and, together with the description, serve to explain the principles of this disclosure. It is obvious that the drawings described below are merely some embodiments of this disclosure, and those skilled in the art can obtain other drawings based on these drawings without any inventive effort.

[0018] Figure 1 A flowchart illustrating a reinforcement learning method that utilizes a code review model as a reward model in an exemplary embodiment of this disclosure is shown. Detailed Implementation

[0019] Exemplary embodiments will now be described more fully with reference to the accompanying drawings. However, these exemplary embodiments can be implemented in many forms and should not be construed as limited to the examples set forth herein; rather, they are provided so that this disclosure will be more comprehensive and complete, and will fully convey the concept of the exemplary embodiments to those skilled in the art. The described features, structures, or characteristics may be combined in any suitable manner in one or more embodiments.

[0020] Furthermore, the accompanying drawings are merely illustrative diagrams of embodiments of this disclosure and are not necessarily drawn to scale. The same reference numerals in the drawings denote the same or similar parts, and therefore repeated descriptions of them will be omitted. Some block diagrams shown in the drawings are functional entities and do not necessarily correspond to physically or logically independent entities.

[0021] This invention aims to address the problem in existing reinforcement learning technologies where reward signals are sparse and fail to accurately reflect the intrinsic quality of code when used for code generation. It provides a reinforcement learning method that utilizes the results of automated code review as a fine-grained reward signal. The core idea of ​​this invention is to integrate a trained code review model into the reinforcement learning loop as the reward model. Specifically, the method first trains a code review model offline, enabling it to understand code and provide evaluation results similar to those of a human reviewer (such as quality scores, defect labels, or textual comments). Subsequently, during reinforcement learning training, after each action (generating a piece of code) by the agent (i.e., the code generation model), the environment feeds the generated code into a fixed (or dynamically updated in some embodiments) code review model. The output of the code review model is converted into a scalar reward value that reflects the quality of the generated code. This reward signal is then used in reinforcement learning algorithms such as policy gradients to update the parameters of the code generation model, making it more likely to generate high-quality code that passes code review.

[0022] This example implementation provides a reinforcement learning method that uses a code review model as a reward model. Please refer to [link / reference]. Figure 1 This method may include: S1-S4, as follows: S1, using data containing code and corresponding manual review result labels as training samples, supervised learning is performed to train a code review model. The code review model can receive code, evaluate the code, and then output the evaluation result. S2, Under the reinforcement learning framework, the code generation model is used as an agent, and the code generated by the code generation model is input into the code review model; S3, the evaluation results output by the code review model are converted into reward signals for the reinforcement learning agent; S4, update the policy parameters of the code generation model using the reward signal.

[0023] In this embodiment, the rewards provided by the code review model are more refined, enabling optimization of quality dimensions such as code style, maintainability, and robustness that cannot be covered by compilers and simple tests, guiding the model to generate higher-quality code. Compared to sparse rewards that rely solely on the final test results, the reward signals provided by this invention are denser and more instructive, which can accelerate the convergence of reinforcement learning models and improve training efficiency. The entire process requires no frequent manual intervention, has a high degree of automation, and achieves automated closed-loop optimization from code generation to quality assessment.

[0024] This application utilizes a code review model as a reinforcement learning method for reward models, which can be applied to tasks such as code completion, code comment generation, code defect repair, or unit test generation.

[0025] The specific process of each step in the above embodiments will be described below.

[0026] S1, Train the Code Review Model (CRM).

[0027] Objective: To train a model that can evaluate code quality and output scalar scores or structured evaluation results.

[0028] Data preparation: - Dataset: Sample function code snippets from high-quality Python projects on GitHub.

[0029] - Tags: Manually labeled or rule-based labeled, including: - Code quality score y∈[0,1]; - Defect tags (e.g., whether there are null pointers, whether there are unused variables); - Optional: Review comments text (for text matching rewards).

[0030] Model structure: - Encoder: Uses the pre-trained model CodeBERT.

[0031] - Output head: - Regression Head: Outputs scalar quality score ∈[0,1]; - Classification Header: Outputs a defect probability vector P∈[0,1] k .

[0032] Loss function:

[0033] in, The true label score for code quality. The code quality score predicted by the model. Mean square error, This is the loss balance coefficient. For binary cross-entropy, For the first i The true label of class defects The model predicts the first i The probability of the existence of a class of defects.

[0034] Training parameters: - Optimizer: AdamW, learning rate 1e -5 ; Batch size: 32; - Training rounds: 5 epochs; - Early stopping: The validation set MSE no longer decreases.

[0035] S2, Constructing a reinforcement learning environment.

[0036] state space : -Currently generated code token sequence: s t =[x1,x2,……,x t ]; - Initial natural language description (prompt): d.

[0037] Action space : - Each token represents an action, and the vocabulary size is |V|≈50k; -action For autoregressive sampling: ,in This is the current state.

[0038] Environment components: - Code generation model (policy network): ; - Code review model (reward function): The parameters are fixed.

[0039] S3, Design the reward function.

[0040] enter: - Complete generated code snippet

[0041] Output: - Review the model output quality score: .

[0042] award design: Option 1: Use scores directly

[0043] Option 2: Use baseline bias (more stable)

[0044] in, Baseline score, which is the average score of the most recent M=10 samples; a discount factor can be introduced to obtain: Discount factor .

[0045] S4, Reinforcement Learning Training (PPO): During reinforcement learning training, the parameters of the code review model are fixed, and the code review model and the code generation model are trained adversarially and dynamically updated.

[0046] Policy network structure: -Using the GPT-2 architecture (CodeGPT-small) Input: Natural language description + Current code prefix Output: Probability distribution of the next token

[0047] Target loss function (Target loss function of PPO algorithm):

[0048] in, For time steps t The mathematical expectation, The parameterized policy of the current policy network, For the old policy network before the update, For the clipping function, For trimming parameters, This is the dominant function.

[0049] Advantage function : - Using Generalized Advantage Estimation (GAE):

[0050] in, As a discount factor, GAE attenuation coefficient, Represents the offset of the future time step. For time steps t + l The timing difference error, For time steps t Instant reward signal, For time steps t state The value function output.

[0051] Value function networks: - Shares the Transformer backbone with the policy network - Output Header: Linear layer output scalar

[0052] The hyperparameters used in training are shown in Table 1.

[0053] Table 1 Training Hyperparameters

[0054] Based on the above technical solutions, it is also necessary to evaluate the generation quality of the code generation model. The evaluation indicators are shown in Table 2.

[0055] Table 2 Evaluation Indicators

[0056] These include: consistency of generated code with reference standards (BLEU, CodeBLEU); code functional effectiveness (Pass@k, passing unit tests); and the quality level assessed by the code review model (CRM Score).

[0057] This application conducted experiments using the proposed reinforcement learning method.

[0058] The baseline method for the experiment is as follows: - Supervised fine-tuning (SFT); - RL + Unit Test Rewards (Pass@1); - RL + CRM Rewards (this application method).

[0059] The experimental results are shown in Table 3.

[0060] Table 3 Experimental Results

[0061] Experimental results show that the reinforcement learning method (RL+CRM) using code review model (CRM) as the reward model in this application significantly outperforms the comparison methods. This application offers superior code similarity, higher code quality, and guaranteed functional effectiveness. It addresses the problem of sparse reward signals in traditional methods, balancing training efficiency with generated code quality.

[0062] The training script in this application is based on the PyTorch and Transformers libraries to train a code generation model that combines "Code Review Model (CRM) + PPO Reinforcement Learning". The core process is as follows: (1) Construct a pseudo dataset (including natural language description, code snippets and quality scores) and configure hyperparameters (training epochs, learning rate, etc.); (2) Define a CRM model (with CodeBERT as the encoder) and train its code quality scoring ability using MSE loss; (3) Initialize the policy network (GPT-2 architecture) and value function header as the code generation agent; (4) Generate code based on prompt autoregressive sampling and obtain reward signals through the CRM model; (5) The PPO algorithm is used, combined with the GAE advantage estimation to calculate the loss and optimize the policy network parameters; (6) The trained code generation model is output by looping through the process of “training CRM → PPO sampling and updating”.

[0063] The training script code is as follows: #! / usr / bin / env python3 """ python train_crm_ppo.py\ --crm_epochs 3\ --ppo_epochs 5\ --rollouts 256\ --max_new_tokens 128\ --device cuda:0 """ import argparse, math, os, random, torch, torch.nn as nn from datasets import Dataset from transformers import (AutoTokenizer, AutoModel, GPT2LMHeadModel, GPT2Tokenizer, get_linear_schedule_with_warmup) from torch.distributions.categorical import Categorical from tqdm import tqdm import numpy as np # ---------- Hyperparameters ---------- parser = argparse.ArgumentParser() parser.add_argument("--crm_epochs", type=int, default=3) parser.add_argument("--ppo_epochs", type=int, default=5) parser.add_argument("--rollouts", type=int, default=256) parser.add_argument("--max_new_tokens", type=int, default=128) parser.add_argument("--device", type=str, default="cuda:0") parser.add_argument("--lr_crm", type=float, default=1e-5) parser.add_argument("--lr_ppo", type=float, default=3e-5) parser.add_argument("--batch_size", type=int, default=16) parser.add_argument("--eps_clip", type=float, default=0.2) parser.add_argument("--gamma", type=float, default=0.99) parser.add_argument("--lam", type=float, default=0.95) args = parser.parse_args() DEVICE = torch.device(args.device) # ---------- 1. Construct data ---------- def build_fake_data(n=5000): """Return HuggingFace Dataset: {'nl':str, 'code':str,'score':float}""" def _gen(): for i in range(n): nl = f"write a function to compute factorial of n (case {i})" code = f"def factorial(n):\nif n<=1:\nreturn 1\nreturn n*factorial(n-1)" # Fake score: randomly between 0.4 and 1.0 score = 0.4 + 0.6 * random.random() yield {"nl": nl, "code": code, "score": score} return Dataset.from_generator(_gen) fake_ds = build_fake_data() # ---------- 2. Code Review Model (CRM) ---------- class CRModel(nn.Module): def __init__(self, model_name="microsoft / codebert-base"): super().__init__() self.encoder = AutoModel.from_pretrained(model_name) self.reg = nn.Linear(self.encoder.config.hidden_size, 1) def forward(self, input_ids, attention_mask): h = self.encoder(input_ids=input_ids, attention_mask=attention_mask).pooler_output return torch.sigmoid(self.reg(h)).squeeze(-1)# [B] crm_tokenizer = AutoTokenizer.from_pretrained("microsoft / codebert-base") crm_model = CRModel().to(DEVICE) crm_optimizer = torch.optim.AdamW(crm_model.parameters(), lr=args.lr_crm) def train_crm(): crm_model.train() loader = torch.utils.data.DataLoader(fake_ds.shuffle(), batch_size=args.batch_size, drop_last=True) for epoch in range(args.crm_epochs): tot_loss = 0 for batch in tqdm(loader, desc=f"CRM epoch {epoch}"): code = [c.replace("\n", " ") for c in batch["code"]] score = torch.tensor(batch["score"]).float().to(DEVICE) tok = crm_tokenizer(code, padding=True, truncation=True, max_length=256, return_tensors="pt").to(DEVICE) pred = crm_model(tok.input_ids, tok.attention_mask) loss = nn.MSELoss()(pred, score) crm_optimizer.zero_grad(); loss.backward(); crm_optimizer.step() tot_loss += loss.item() print(f"CRM epoch {epoch} MSE={tot_loss / len(loader):.4f}") crm_model.eval()# Freeze the parameters after training # ---------- 3. Policy Network (GPT-2) ---------- policy_tok = GPT2Tokenizer.from_pretrained("gpt2") policy_tok.pad_token = policy_tok.eos_token policy_model = GPT2LMHeadModel.from_pretrained("gpt2").to(DEVICE) policy_optimizer = torch.optim.AdamW(policy_model.parameters(), lr=args.lr_ppo) value_head = nn.Linear(policy_model.config.n_embd, 1).to(DEVICE) # ---------- 4.rollout sampling ---------- @torch.no_grad() def rollout_one(prompt_text): """ Given the natural language prompt_text, autoregressive generation complete function Returns dict: tokens, log_probs, values, mask, reward """ prompt_tok = policy_tok(prompt_text, return_tensors="pt").to(DEVICE) input_ids = prompt_tok.input_ids max_len = input_ids.shape[1] + args.max_new_tokens generated = input_ids log_probs, values,

[0064] past = None while generated.shape[1]<max_len: out = policy_model(generated, past_key_values=past, output_hidden_states=True) logits = out.logits[:, -1, :]# [B, V] hidden = out.hidden_states[-1][:, -1, :]# [B, D] past = out.past_key_values val = value_head(hidden).squeeze(-1)# [B] dist = Categorical(logits=logits) action = dist.sample() generated = torch.cat([generated, action[:, None]], dim=1) log_probs.append(dist.log_prob(action)) values.append(val) mask.append(torch.ones_like(action)) if action.item() == policy_tok.eos_token_id: break # Calculate the reward code = policy_tok.decode(generated[0][prompt_tok.input_ids.shape[1]:], skip_special_tokens=True) reward = get_crm_reward(code) return { "tokens": generated, "log_probs": torch.stack(log_probs), "values": torch.stack(values), "mask": torch.stack(mask), "reward": reward } def get_crm_reward(code): """Score the code with CRM and return a float""" code = code.replace("\n", " ") ​​with torch.no_grad(): score = crm_model(tok.input_ids, tok.attention_mask).item() return score # ---------- 5. PPO Update ---------- def ppo_update(rollouts): """rollouts: list of dict""" # Concatenate into a batch all_logp = torch.cat([r["log_probs"] for r in rollouts]) all_vals = torch.cat([r["values"] for r in rollouts]) all_mask = torch.cat([r["mask"] for r in rollouts]) rewards = torch.tensor([r["reward"] for r in rollouts]).float().to(DEVICE) # Calculate GAE T = all_vals.shape[0] deltas = rewards - all_vals advantages =

[0065] adv = 0 for t in reversed(range(T)): adv = deltas[t] + args.gamma * args.lam * adv advantages.insert(0, adv) advantages = torch.tensor(advantages).float().to(DEVICE) returns = advantages + all_vals advantages = (advantages - advantages.mean()) / (advantages.std() +1e-8) # Multiple epochs, small batches idx = torch.arange(T) for _ in range(4): perm = idx[torch.randperm(T)] for start in range(0, T, args.batch_size): sel = perm[start:start+args.batch_size] logp_old = all_logp[sel] adv = advantages[sel] ret = returns[sel] # Rerun the model to get new logp and v # Simplified here: Use the old tokens and then forward (in reality, another forward is needed). # For simplicity, reuse the old tokens directly. tokens = torch.cat([rollouts[i]["tokens"]for i in sel.tolist()]) out = policy_model(tokens, output_hidden_states=True) logits = out.logits[:, :-1, :].reshape(-1, out.logits.size(-1)) hidden = out.hidden_states[-1][:, :-1, :].reshape(-1, out.hidden_states[-1].size(-1)) dist = Categorical(logits=logits) logp = dist.log_prob(tokens[:, 1:].reshape(-1)) ratio = torch.exp(logp - logp_old) surr1 = ratio * adv surr2 = torch.clamp(ratio, 1-args.eps_clip, 1+args.eps_clip) * adv policy_loss = -torch.min(surr1, surr2).mean() value_pred = value_head(hidden).squeeze(-1) value_loss = nn.MSELoss()(value_pred, ret) loss = policy_loss + 0.5*value_loss - 0.01*dist.entropy().mean() policy_optimizer.zero_grad() loss.backward() policy_optimizer.step() # ---------- 6. Main process ---------- if __name__ == "__main__": [[ID=Z3]]print("==== 1. Train CRM ====")[[ID=Z4]] train_crm() print("==== 2. PPO training ====") for epoch in range(args.ppo_epochs): rollouts = [] for _ in tqdm(range(args.rollouts), desc=f"PPO epoch {epoch}"): prompt = "write a python function to compute factorial of n" rollouts.append(rollout_one(prompt)) ppo_update(rollouts) # Print average reward avg_r = np.mean([r["reward"] for r in rollouts]) print(f"epoch {epoch}avg CRM reward = {avg_r:.3f}") print("Training complete! Model can be saved:") # policy_model.save_pretrained("crm_ppo_policy") It should be noted that although several modules of the system for executing actions are mentioned in the detailed description above, this division is not mandatory. In fact, according to embodiments of the present invention, the features and functions of two or more modules described above can be embodied in one module. Conversely, the features and functions of one module described above can be further divided into multiple modules for embodiment. Components shown as modules may or may not be physical units; that is, they may be located in one place or distributed across multiple network units. Some or all of the modules can be selected to achieve the purpose of the present invention according to actual needs. Those skilled in the art can understand and implement this without any inventive effort.

[0066] The embodiments of the present invention have been described above with reference to the accompanying drawings. However, the present invention is not limited to the specific embodiments described above. The specific embodiments described above are merely illustrative and not restrictive. Those skilled in the art can make many other forms under the guidance of the present invention without departing from the spirit and scope of the claims. All of these forms are within the protection scope of the present invention.

Claims

1. A reinforcement learning method that uses a code review model as a reward model, characterized in that, include: S1, The data containing code and corresponding manual review result labels are used as training samples to conduct supervised learning training to obtain a code review model. The code review model can receive code, evaluate the code, and then output the evaluation result. S2, Under the reinforcement learning framework, the code generation model is used as an agent, and the code generated by the code generation model is input into the code review model; S3, the evaluation results output by the code review model are converted into reward signals for the reinforcement learning agent; S4, update the policy parameters of the code generation model using the reward signal.

2. The reinforcement learning method using a code review model as a reward model according to claim 1, characterized in that, The manual review result labels include a code quality score, a Boolean indicator indicating whether a defect exists, and the review comments text.

3. The reinforcement learning method using a code review model as a reward model according to claim 1, characterized in that, The code review model uses the pre-trained model CodeBERT as the encoder and regression loss and classification loss as the loss functions.

4. The reinforcement learning method using a code review model as a reward model according to claim 3, characterized in that, The loss function for: in, The true label score for code quality. The code quality score predicted by the model. Mean square error, This is the loss balance coefficient. For binary cross-entropy, For the first i The true label of class defects The model predicts the first i The probability of the existence of a class of defects.

5. The reinforcement learning method using a code review model as a reward model according to claim 1, characterized in that, The evaluation results are either scalar scores or review comments.

6. The reinforcement learning method using a code review model as a reward model according to claim 5, characterized in that, When the evaluation result is a scalar score, the scalar score is used as a reward, or the difference between the scalar score and the baseline score is used as a reward.

7. The reinforcement learning method using a code review model as a reward model according to claim 5, characterized in that, When the evaluation result is a review comment text, the degree of matching between the review comment text and the expected behavior template will be used as a reward.

8. The reinforcement learning method using a code review model as a reward model according to claim 1, characterized in that, During reinforcement learning training, the parameters of the code review model are fixed.

Images