S tep C oder: Improving Code Generation with Reinforcement Learning from Compiler Feedback (2024)

Abstract: The advancement of large language models (LLMs) has significantly propelled the field of code generation. Previous work integrated reinforcement learning (RL) with compiler feedback for exploring the output space of LLMs to enhance code generation quality. However, the lengthy code generated by LLMs in response to complex human requirements makes RL exploration a challenge. Also, since the unit tests may not cover the complicated code, optimizing LLMs by using these unexecuted code snippets is ineffective. To tackle these challenges, we introduce StepCoder , a novel RL framework for code generation, consisting of two main components: CCCS addresses the exploration challenge by breaking the long sequences code generation task into a Curriculum of Code Completion Subtasks, while FGO only optimizes the model by masking the unexecuted code segments to provide Fine-Grained Optimization. In addition, we furthermore construct the APPS+ dataset for RL training, which is manually verified to ensure the correctness of unit tests. Experimental results show that our method improves the ability to explore the output space and outperforms state-of-the-art approaches in corresponding benchmarks. The code and dataset will be made available upon publication.

## introduction code generation or program synthesis aims to automatically generate source code that adheres to a specified programming requirement, which is typically described in natural language @xcite . recently, with the development of large language models (llms), techniques based on llm @xcite have demonstrated impressive ability in code generation. however, challenges persist in aligning these models with complex human requirements @xcite , indicating a gap that still exists in fully meeting user expectations. in this context, learning from compiler feedback exhibits impressive potential to improve the comprehension of complicated human requirements and the quality of generated codes @xcite . this feedback from compilation and execution results is instrumental in directly ascertaining the functional correctness of programs @xcite . researchers @xcite introduce reinforcement learning (rl) and leverage compiler feedback from unit tests as a reward metric to guide the exploration of the output space of llms. the intention is for the policy model to favor actions that yield higher rewards increasingly. nevertheless, the optimization of llms for code generation via rl presents several hurdles. first, the increasing complexity of human requirements often results in the generation of longer code sequences, which makes exploration struggle @xcite . second, in cases where a single unit test fails to cover the complex code, unexecuted code snippets may emerge that are not relevant to the reward. rendering optimization based on the entire code sequence is potentially imprecise. additionally, our analysis reveals quality limitations in existing datasets like apps @xcite for rl training, which impedes accurate learning from compiler feedback through rl. to tackle these challenges, we first introduce stepcoder, an innovative framework developed for enhancing code generation through reinforcement learning. stepcoder integrates two key components: curriculum of code completion subtasks (cccs) and fine-grained optimization (fgo). cccs is designed to alleviate the complexities associated with exploration in code generation, while fgo is designed to provide more precise and effective optimization strategies. specifically, cccs employs a step-by-step strategy to break down complex exploration problems (i.e., code generation) into a curriculum of easier sub-tasks (i.e., code completion). as the training progresses, the difficulty of code completion tasks rises by increasing the portion of code that needs to be completed. eventually, the aim is for the model to evolve to a stage where it can effectively generate code solely from human requirements, thus fulfilling the original training goal of code generation. on the other hand, the key insight of fgo is that code snippets that are not executed in a unit test do not contribute to the final reward calculation. therefore, fgo uses a dynamic masking technique to mask unexecuted snippets from unit test evaluations, ensuring that the model is optimized utilizing only the relevant code segments. subsequently, our endeavor involves the development of apps+, a dataset of superior quality specifically curated for code generation. apps+ is meticulously designed to exclude code segments that exhibit syntax errors, are irrelevant to the stipulated problem, or fail to produce any output. additionally, we have taken measures to standardize the format of inputs and outputs in unit tests to guarantee deterministic output comparisons. we evaluate the effectiveness of popular llms on apps+. the results reveal that although llms show progressive improvements, they face difficulties with complex human requirements. we further evaluate our method on several extensively used benchmarks including mbpp @xcite and humaneval @xcite . the experimental results show that stepcoder effectively eases the exploration difficulty in code generation, outperforming other reinforcement learning-based methods in effectiveness. the main contributions of our paper are as follows: ## motivation in this section, we clearly illustrate the challenges faced by reinforcement learning in code generation using a simplified example from apps @xcite , which was widely used for rl training in code generation. exploration problems of rl in code generation. in tackling complicated sequence but sparse reward challenges, exploration methods play a crucial role @xcite . when a policy model explores a trajectory with high returns, it undergoes optimization, making it inclined to take similar actions in the future @xcite . consider the code shown in figure 1 , aimed at fulfilling a given human requirement. we first collect the conditional statements (cs) that are indicated by the dashed box by analyzing its abstract syntax tree. conditional statement introduces new independent paths, increasing the complexity of the program @xcite . suppose p θ (cs i ) denotes the probability that the policy model with parameter θ completes the i-th conditional statement. the probability that the policy model correctly generates this code according to human requirements can be expressed as follows: ## method in this section, we elaborate on the methodological details of stepcoder, which provide an easier exploration and fine-grained optimization for rl in code generation, respectively, as shown in figure 2 . ## experiments in this section, we first introduce apps+, a highquality dataset for code generation by manually verifying based on the apps dataset. then, we elaborate on the experiment details and the experimental results. ## related work large language models for code generation. recently, pre-trained language models have shown remarkable ability in understanding natural language and code generation by training on large text corpora containing code data @xcite . in addition, sft models achieve more competitive performance such as starcoder @xcite , wizardcoder @xcite @xcite . reinforcement learning is a method of learning the optimal policy by exploring the environment and obtaining rewards @xcite . recently, some researchers have introduced rl to llms and improved the quality of the generated code by utilizing the unit test feedback to explore the output space of the pol-icy model. for instance, coderl @xcite leverages unit test signals for rewards and employs actor-critic methods @xcite to enhance models on code generation. ppocoder @xcite refines coderl by employing the ppo algorithm @xcite and rltf @xcite provides fine-grained rewards through the error locations, but the reward space is still sparse. however, the exploration of complex tasks in an environment characterized by a sparse reward is challenging, limiting the effectiveness of rl in boosting code generation model performance exploration in reinforcement learning. exploration is crucial in addressing long sequences and sparse reward problems @xcite . in the sequence generation task, researchers improved exploration by initializing the policy model using the sft model @xcite . our proposed approach incorporates similar methods, but additional methods are necessary to ensure effective exploration, especially when tackling complex human-driven requirements, where the limited quality of generated code makes exploration still challenging. other notable methods introduce the process-supervised reward model to provide step-by-step rewards for complex sequence generation tasks such as mathematical reasoning and code generation @xcite @xcite . however, these methods require labelling a large preference dataset to train the reward model. similar to our approach, some methods construct a learning curriculum by initiating each episode from a sequence of progressively more challenging starting states (salimans and @xcite . in contrast to our approach, these methods are designed to address the problem of exploration in other fields, such as gaming and robotic manipulation. meanwhile, our approach combines software engineering features to dynamically determine the starting states through conditional statements and introduces fgo to enhance fine-grained optimization with the coverage information. ## conclusion in this paper, we introduce stepcoder, a novel training framework via reinforcement learning (rl). stepcoder breaks down complicated exploration problems to reduce the difficulty of exploring en-vironments with sparse rewards while providing fine-grained optimization. in addition, we also construct a high-quality dataset apps+, specifically for code generation. experiments indicate that our method can effectively improve the quality of generated code via rl compared to other approaches.