A code generation method based on process supervision reinforcement learning

Abstract

The application discloses a code generation method based on process supervision reinforcement learning, comprising the following steps: obtaining an original code dataset and preprocessing, generating a plurality of modified code segments based on the preprocessed original code dataset and automatically labeling to obtain a process supervision dataset; wherein the preprocessing comprises format standardization; training a process supervision reward model based on the process supervision dataset; generating a code segment through a strategy model and obtaining an overall reward signal of the code segment based on the trained process supervision reward model; and updating the strategy model based on the overall reward signal. The application realizes automatic construction of the process supervision dataset, greatly reduces the manual labeling cost, and significantly improves the quality and efficiency of the code generation model by introducing the process supervision reinforcement learning method.

Description

TECHNICAL FIELD

[0001] The present application belongs to the technical field of computer programming and artificial intelligence, and particularly relates to a code generation method based on process supervision reinforcement learning. BACKGROUND

[0002] With the advent of large models based on deep learning and natural language processing (such as GPT, LLaMA, etc.), the ability of code generation has been significantly improved. These large models can understand natural language descriptions and automatically generate corresponding code, and even solve complex programming problems, thereby significantly improving development efficiency and code quality.

[0003] However, the current code model development faces the following technical challenges:

[0004] High cost of data annotation: Building a high-quality code dataset requires professional programming knowledge and a lot of time for annotation to ensure the correctness, standardization and relevance of the code to the task. Not only does it require a detailed review of the code by humans, but also an evaluation of the code's functionality, performance, and other aspects, which makes the cost of obtaining large-scale high-quality annotated data extremely high.

[0005] Insufficient ability to handle complex tasks: When facing complex code generation tasks, such as involving multi-module interaction and complex algorithm implementation, the model is prone to logical errors or code structure confusion. Current methods rely heavily on simple heuristic methods or general optimization strategies, making it difficult to deeply analyze the logical relationships and structural requirements in complex tasks, resulting in code that is difficult to meet the actual application requirements.

[0006] Poor generality and adaptability of generated code: The generated code is often limited to specific scenarios or datasets, and its generality in different programming environments and frameworks is poor, making it difficult to adapt to diverse actual needs.

[0007] In order to make the model as close as possible to the complex needs of humans, reinforcement learning technology integrates human feedback and plays an important role in aligning large language models with human values. The advantage of reinforcement learning is that it can indirectly optimize non-differentiable reward signals such as CodeBLEU scores or human preferences through policy optimization and value function approximation. However, obtaining the required human feedback often requires a lot of manpower and resources. In the code generation task, reinforcement learning shows unique advantages: language models can automatically use compiler feedback in unit tests as reward signals, avoiding excessive reliance on human feedback. This approach not only optimizes output efficiently, but also significantly improves the performance of the model in code generation tasks. Although these methods have achieved some success, they mainly rely on compiler feedback signals for the entire code to train the reward model, making it difficult to deal with the reward sparsity problem. SUMMARY

[0008] To solve the above technical problems, the present application provides a code generation method based on process supervision reinforcement learning to solve the problems existing in the prior art.

[0009] To achieve the above purpose, the present application provides a code generation method based on process supervision reinforcement learning, comprising:

[0010] Obtain the original code dataset and preprocess it, generate several code variants based on the preprocessed original code dataset and label them, and obtain the process supervision dataset; wherein the preprocessing includes format standardization;

[0011] Train the process supervision reward model based on the process supervision dataset;

[0012] Generate a code snippet through a policy model, obtain the overall reward signal of the code snippet based on the trained process supervision reward model, and update the policy model based on the overall reward signal.

[0013] Preferably, the process of obtaining the process supervision dataset comprises:

[0014] Obtain several modified code snippets by modifying and restructuring the code in the preprocessed original code dataset line by line through the teacher model; label the code in the preprocessed original code dataset and the several modified code snippets to obtain the process supervision reward dataset.

[0015] Preferably, the labeling process comprises:

[0016] After the code in the preprocessed original code dataset is segmented line by line, mask processing is performed, and label marking is performed through the compiler to obtain positive samples; several modified code snippets are automatically labeled through the compiler and execution verification, and the positive samples are labeled as positive samples through all test cases, and the negative samples are labeled as negative samples; the positive samples and the negative samples constitute the process supervision reward dataset.

[0017] Preferably, based on the pre-trained language model, the policy model suitable for the code generation task is obtained by supervised fine-tuning through the code dataset containing natural language question description and corresponding code.

[0018] Preferably, the process of obtaining the overall reward signal of the code snippet comprises:

[0019] Divide the code snippet into multiple sub-snippets, independently assign a reward value to each sub-snippet, and aggregate the reward values assigned to each sub-snippet based on weighting to obtain the overall reward signal of the code snippet.

[0020] The process supervision reward model is:

[0021]

[0022] wherein w i represents the first i lines of the code sequence, the synchronous execution ends at time T i , (d, w i ) is the input segment, and φ represents the parameters of the process supervision reward model.

[0023] Preferably, the policy model is updated by combining the overall reward signal and the policy gradient using a proximal policy optimization algorithm.

[0024] Preferably, the method further comprises:

[0025] The benchmark test set is divided into different difficulty levels according to the code length, and the code generation accuracy of the policy model at different difficulty levels is evaluated by the pass@k index.

[0026] The present application provides a computer device, comprising a memory, a processor to store a computer program on the memory and executable on the processor, and the processor executes the computer program to implement the steps of the above method.

[0027] The present application provides a computer readable storage medium, which stores a computer program, and the computer program is executed by a processor to implement the steps of the above method.

[0028] The present application provides a computer program product, comprising a computer program, and the computer program is executed by a processor to implement the steps of the above method.

[0029] Compared with the prior art, the present application has the following advantages and technical effects:

[0030] The present application first acquires an original code data set and pre-processes it, generates a plurality of modified code segments based on the pre-processed original code data set and automatically labels them to obtain a process supervision data set; then, trains a process supervision reward model based on the process supervision data set; next, generates a code segment through a policy model and obtains an overall reward signal of the code segment based on the trained process supervision reward model; and finally, updates the policy model based on the overall reward signal.

[0031] The present application greatly reduces the workload of manual labeling by using the "statement mutation / reconstruction-compile and execution verification" strategy to automatically generate process supervision data, and provides high-quality benchmark data for subsequent training of the process supervision reward model. At the same time, through the process supervision reinforcement learning method, the policy model can be guided to generate higher-quality code, improve the performance of the code in terms of correctness, readability, maintainability, etc., and realize the intelligent and automated optimization of the code generation process. BRIEF DESCRIPTION OF DRAWINGS

[0032] The accompanying drawings, which are a part of this application, are used to provide further understanding of the application and are incorporated in and constitute a part of this application. The illustrative embodiments of the present application, and their description, are used to explain the present application. The present application is not limited to the illustrative embodiments described and depicted herein.

[0033] Figure 1 A process supervised reinforcement learning framework for embodiments of the present application;

[0034] Figure 2 A process supervised reward dataset generation method for embodiments of the present application;

[0035] Figure 3 Two structures of a result supervised reward model for embodiments of the present application, (a) a preference-based result supervised reward model, and (b) a compiler-based result supervised reward model. DETAILED DESCRIPTION

[0036] It should be noted that the embodiments and features of the embodiments in the present application can be combined with each other without conflict. The present application will be described in detail below with reference to the accompanying drawings and in combination with the embodiments.

[0037] It should be noted that the steps shown in the flowchart of the accompanying drawings can be executed in a computer system such as a set of computer executable instructions, and although a logical order is shown in the flowchart, in some cases, the steps shown or described herein can be executed in an order different from that shown herein.

[0038] Embodiment One

[0039] As shown in Figure 1 The present embodiment provides a code generation method based on process supervised reinforcement learning, which includes:

[0040] (I) Constructing a process supervised reward dataset

[0041] Using a published code dataset, the data is processed to construct a process supervised dataset for training.

[0042] The original data comes from the MBPP dataset, which consists of 974 crowdsourced Python programming problems. The problems in this dataset are designed to be solved by entry-level programmers and cover programming basics and standard library functions. Each problem includes a task description, a code solution, and 3 automated test cases.

[0043] After preprocessing the original data, the teacher model is used to perform line-by-line mutation and reconstruction of the code, and the compiler is used to automatically label each line of code based on the execution results. The labeled code data is used as the benchmark data for subsequent training of the process supervised reward model.

[0044] Specifically, the present embodiment re-divides the MBPP dataset to improve the generalization ability and robustness of the process supervision reward model. Specifically, ID: 601-974 is selected from the MBPP as the training set for the supervised fine-tuning stage and the seed set of the process supervision reward dataset, ID: 101-500 is selected as the training set for the reinforcement learning stage, and ID: 501-600 and 1-100 are selected as the validation set and the test set, respectively.

[0045] To establish a more standardized process supervision reward dataset, first, the code is standardized by using the standardization '\t' to ensure the uniformity of the code format. Subsequently, in order to make full use of the feedback signals provided by the compiler, test cases for the MBPP problem are supplemented. Specifically, the present embodiment uses a teacher model to generate test cases, aiming to achieve comprehensive path coverage. These test cases are executed using the compiler to verify their correctness and effectiveness. The prompt provided to the teacher model is designed as follows: "Given the following code and its existing test cases, supplement new test cases to achieve full path coverage. For another test set HumanEval, we do not make additional operations."

[0046] The present embodiment proposes to use a teacher model and compiler feedback to automatically build a process supervision reward dataset in the code field. The generation method of the dataset is as shown in Figure 2

[0047] Formally, let denote the code generation training dataset, where d i denotes the i-th problem description, and w i is the corresponding solution program code snippet. First, the reference codes are used to build positive samples. Specifically, the reference codes are segmented line by line to obtain where L i is the number of lines of code. Then, for each line of code, all the lines after it are masked, and the label corresponding to this line is directly marked as "Positive". In other words, the original reference code can be directly re-expressed as a positive sample for process supervision, and its format is as follows: But only positive samples are not enough to train the reward model; therefore, the present embodiment designs a novel strategy to build negative samples. Specifically, for each line of code, a teacher model is used to perform mutation and reconstruction operations according to a specific prompt example. Subsequently, the modified code line is verified together with the rest of the code by the compiler. According to the feedback of the compiler, if it passes all the test cases, it is marked as "Positive"; otherwise, it is marked as "Negative".

[0048] ​(ii) Reward model training

[0049] (1) Outcome supervised reward model

[0050] The outcome supervised reward model adopts a holistic reward approach, mapping the overall quality and reliability indicators corresponding to the problem description d i and the generated code w i to a scalar reward. Typically, this reward is only assigned to the last token in the generated sequence, which is defined as follows:

[0051]

[0052] where θ denotes the parameters of the outcome supervised reward model R O .

[0053] First, an original outcome supervised reward model is trained using the dataset constructed in the previous section. However, relying solely on this dataset to train the outcome supervised reward model has limitations: the active learning strategy exhibits strong bias towards incorrect answers in the dataset, thereby reducing the overall performance of the model. Therefore, this embodiment constructs a more robust outcome supervised reward model baseline.

[0054] Inspired by human feedback-based reinforcement learning, a preference-based outcome supervised reward model is designed in this embodiment. Specifically, for each question, multiple code snippets are uniformly sampled from the generator, and a teacher model is used to simulate human annotators to rank them according to code quality, thereby training the reward model. In addition, to more comprehensively evaluate the advantages and disadvantages of process supervision and outcome supervision in the field of coding, a compiler is introduced as a source of supervision signals, and four types of feedback signals generated by the compiler are used to optimize the generator model, thereby constructing a compiler-based outcome supervised reward model. Figure 3 The structures of the two outcome supervised reward models described above are shown in (a) and (b), respectively.

[0055] (2) Process supervised reward model

[0056] The process supervised reward model constructed in this embodiment rewards the quality of each code segment, thereby enabling more fine-grained adjustments and feedback at each step. The code sequence w is divided into k segments (w1, w2,..., w k ), where w i denotes the first i lines of the code sequence. The execution at time T i is denoted as Within this framework, the reward model rewards each input segment (d, w i) an award is assigned, with the highest award being given to the last fragment of w. Ultimately, the award is defined as follows:

[0057]

[0058] where φ denotes the parameters of the process supervision reward model.

[0059] This embodiment uses the process supervision dataset constructed by the “mutation / refactoring-verification” strategy to train the process supervision reward model. Through multiple iterations of training, the reward model can learn the advantages and disadvantages of each line of code, and thus accurately evaluate the quality of the code generation process. Under this setting, the training data of the process supervision reward model does not intersect with the training data of the result supervision reward model based on the compiler, so it is difficult to directly compare the results. Nevertheless, both the process supervision reward model and the result supervision reward model represent the optimal results under their respective training methods.

[0060] (Three) Integrate the reward model into the reinforcement learning framework

[0061] In this step, the process supervision reward model is integrated into the process supervision reinforcement learning framework. Specifically, first, the base model is supervised fine-tuned to adapt to the input and output formats of the MBPP dataset, which is used as the policy model in the reinforcement learning process. In the reinforcement learning phase, there are two reward models to score the code generated by the policy model. The result supervision reward model scores the entire code fragment generated by the policy model; the process supervision reward model scores each line of code in the code fragment generated by the policy model, and aggregates the weighted average of the scores of each line of code. Based on the above scoring results, the reward model provides feedback and reward signals to the policy model in real time. The policy model continuously updates based on these reward signals, aiming to generate code that can obtain higher scores from the reward model, and thus continuously evolves towards generating higher quality code.

[0062] In this process, the original result supervision reward model, the result supervision reward model based on preferences, and the result supervision reward model based on the compiler are used. These different types of result supervision reward models aim to explore the performance improvement of result supervision under different training methods, and through comparison, to demonstrate the superiority of the process supervision reward model proposed.

[0063] ​Specifically, the proximal policy optimization algorithm (PPO) is a policy gradient-based reinforcement learning algorithm. Its core idea is to limit the magnitude of changes between old and new policies to prevent the update speed from being too fast. This is particularly crucial in code generation, as the stability of the generated results directly affects the quality and consistency of the code.

[0064] In the code generation task, the PPO algorithm first uses the current policy to interact with the environment to obtain the state s t , select an action a t , and receive a reward r t and other data. Then, for each time step, the advantage function A t = ∑ t'>t γ t'-t (r t' + γV ψ (s t'+1 ) - V ψ (s t' )) is calculated, where the value function V ψ (s t ) represents the expected cumulative reward starting from state s t . In addition, a divergence penalty term is added for each token, which represents the ratio of the new and old policies. Thus, the reward function becomes:

[0065]

[0066] (Four) Effectiveness of the experimental verification method

[0067] In the last step of the whole process, first, the benchmark test set is divided into difficulty levels according to code length. Set clear code length interval standards to divide the benchmark test set into three difficulty levels: simple, medium, and difficult.

[0068] Then, for benchmark test sets of different difficulty levels, the integrated reward model is used to test the code generation framework pass@k index. The code generated by the code generation framework under each difficulty level is input into the test environment. For each test case, it is detected whether the generated code can successfully pass the test, and the proportion of at least one correct result in the first k generated results is calculated to obtain the pass@k index under different difficulty levels, so as to clearly evaluate the performance of the process supervision reinforcement learning method in different difficulty code generation tasks.

[0069] By dividing the benchmark set by code length and testing the pass@k index for each difficulty level, the performance of the code generation framework based on process supervision reinforcement learning can be systematically evaluated for tasks of different complexity from the perspective of code size, providing a key reference for accurately measuring model effectiveness.

[0070] The embodiment realizes the automatic construction of the code field process supervision reward data set. Specifically, the construction process includes data selection, data preprocessing, code mutation reconstruction and annotation modules. First, the data selection module selects appropriate data from the published code data set to ensure the diversity and representativeness of the data. Then, the data preprocessing module standardizes the code to ensure uniform code format. At the same time, in order to ensure the reliability of the reward data set, the selected data set is expanded with test cases.

[0071] In the key code mutation reconstruction and annotation link, the teacher model is used to perform line-by-line mutation and reconstruction operations on the code. Then, the compiler execution results are used to automatically annotate each line of code. The entire process is realized through automated scripts, without the need for a large amount of manual intervention, greatly improving the efficiency and accuracy of the data set construction, and providing high-quality benchmark data for subsequent process supervision reward model training.

[0072] The embodiment realizes the process supervision reinforcement learning code generation through specific algorithms and frameworks. The specific structure includes a process supervision data construction module, a reward model training module, a reinforcement learning framework integration module, and an experimental verification module. Through this process supervision reinforcement learning method, the policy model can generate higher quality code, improve the performance of the code in terms of correctness, readability, maintainability, and other aspects, and realize the intelligent and automated optimization of the code generation process.

[0073] The above is only the preferred specific embodiment of the present application, but the protection scope of the present application is not limited thereto, any person skilled in the art can easily think of changes or replacements within the technical scope disclosed in the present application, which should be covered within the protection scope of the present application. Therefore, the protection scope of the present application should be subject to the protection scope of the claims.

Claims

1. A code generation method based on process-supervised reinforcement learning, characterized in that, Includes the following steps: The process involves acquiring the original code dataset and preprocessing it, generating several modified code snippets based on the preprocessed original code dataset and automatically annotating them to obtain a process supervision dataset; the preprocessing includes format standardization. The process of obtaining a process-supervised dataset includes: The teacher model is used to mutate and reconstruct the code in the preprocessed original code dataset line by line to obtain several modified code snippets; the code in the preprocessed original code dataset and several modified code snippets are labeled to obtain the process supervision reward dataset. The annotation process includes: The code in the preprocessed original code dataset is segmented line by line and masked. The compiler then labels the code to obtain positive samples. Several modified code snippets are automatically labeled through compiler and execution verification. If all test cases are passed, the code is labeled as a positive sample; if not, it is labeled as a negative sample. The positive and negative samples together form the process supervision reward dataset. The process supervision reward model is trained based on the process supervision dataset; Code snippets are generated using a policy model, and the overall reward signal of the code snippets is obtained based on a trained process-supervised reward model; the policy model is then updated based on the overall reward signal. The process of obtaining the overall reward signal for a code snippet includes: The code snippet is divided into multiple sub-snippets, and a reward value is independently assigned to each sub-snippet. The reward values ​​assigned to each sub-snippet are then weighted and aggregated to obtain the overall reward signal for the code snippet. The process supervision reward model is as follows: ； in, Indicates the beginning of the code sequence Okay, execute synchronously in time Finish, For the input fragment, This represents the parameters of the process supervision reward model.

2. The code generation method based on process-supervised reinforcement learning according to claim 1, characterized in that, Based on a pre-trained language model, a policy model adapted to the code generation task is obtained through supervised fine-tuning using a code dataset containing natural language problem descriptions and corresponding codes.

3. The code generation method based on process-supervised reinforcement learning according to claim 1, characterized in that, The policy model is updated by using a proximal policy optimization algorithm, which combines the overall reward signal and the policy gradient.

4. The code generation method based on process-supervised reinforcement learning according to claim 1, characterized in that, Also includes: The benchmark test suite is divided into different difficulty levels based on code length; The pass@k metric was used to evaluate the code generation accuracy of the strategy model at different difficulty levels.

5. A computer device, comprising: A memory, a processor, and a computer program stored in the memory and executable on the processor, characterized in that the processor executes the computer program to implement the steps of the method according to any one of claims 1-4.

6. A computer-readable storage medium having a computer program stored thereon, characterized in that, When executed by a processor, the computer program implements the steps of the method described in any one of claims 1-4.

7. A computer program product, comprising a computer program, characterized in that, When executed by a processor, the computer program implements the steps of the method described in any one of claims 1-4.

Images