Automatic Unit Test Data Generation and Actor-Critic Reinforcement Learning for Code Synthesis (2023)

Abstract: The advent of large pre-trained language models in the domain of Code Synthesis has shown remarkable performance on various benchmarks, treating the problem of Code Generation in a fashion similar to Natural Language Generation, trained with a Language Modelling (LM) objective. In addition, the property of programming language code being precisely evaluable with respect to its semantics – through the use of Unit Tests to check its functional correctness – lends itself to using Reinforcement Learning (RL) as a further training paradigm. Previous work has shown that RL can be applied as such to improve models’ coding capabilities; however, such RL-based methods rely on a reward signal based on defined Unit Tests, which are much harder to obtain compared to the huge crawled code datasets used in LM objectives. In this work, we present a novel approach to automatically obtain data consisting of function signatures and associated Unit Tests, suitable for RL training of Code Synthesis models. We also introduce a straightforward, simple yet effective Actor-Critic RL training scheme and show that it, in conjunction with automatically generated training data, leads to improvement of a pre-trained code language model’s performance by up to 9.9% improvement over the original underlying code synthesis LM, and up to 4.3% over RL-based models trained with standard PPO or CodeRL.

## introduction large language models (llms) have been dominating the field of natural language processing (nlp) since the introduction of the transformer architecture @xcite and the first large encoder-based @xcite , decoder-based @xcite , and full transformer @xcite pretrained models. more recently, large transformerbased models began to expand beyond natural lan-* equal contribution guages, most notably to the area of programming languages, where they are being applied to tasks such as code understanding and code synthesis (also referred to as text-to-code generation), or code translation @xcite @xcite . often, models addressing specific tasks build upon large pre-trained (code) language models, i.e. "foundation models" @xcite , and apply further fine-tuning on task-specific data under a fully supervised lm or imitation learning paradigm. for code synthesis, that is, the generation of programming language code conditioned on a natural language prompt, reinforcement learning (rl) has recently gained traction as an alternative or complementary training method @xcite . these approaches make use of the fact that, as opposed to natural language, correctness of code is relatively straight forward to evaluate through compilation (syntax) and unit tests/functional testing (semantics). however, rl approaches that train on code functionality rely on data with corresponding unit tests, that are often only available in very limited quantities or not at all. this is in stark contrast to approaches that employ language modelling objectives, trainable on large quantities of crawled data. overcoming this lack of unit test paired data has the potential to greatly improve the results achievable with such methods. in this work, we also employ reinforcement learning to improve the performance of a pretrained code language model on the code synthesis task. modelling the reward based on unit test pass rates, we introduce a simple yet effective rl approach to fine-tuning code synthesis models, based on policy gradients and a simple feed-forward critic model. in an effort to overcome the data sparseness issues of previous approaches, we introduce a heuristic approach to generate large data of natural language problems, function signatures, and unit tests, and show that the automatically generated data can help further improve code synthesis performance. finally, we release our code, models, new dataset of problems, signatures, and unit tests. foot_0 ## related work the advent of large language models for programming languages was initially driven by repurposing general natural language understanding (nlu) and generation (nlg) approaches to the new domain. models such as cubert @xcite , codegpt @xcite , and codet5 @xcite greatly reused the natural language models, but trained on source code to provide them skills in the new context. ## experiments setting with our overall procedures to rl for code synthesis and unit test generation in place, we run a battery of experiments to evaluate the efficacy of our approach. as a general dataset for the task, we use the mostly basic python problems dataset (mbpp, @xcite . mbpp provides a total of 964 instances of natural language problem description, an associated gold solution, and unit tests. we follow the official train and validation splits of 374 and 90 instances respectively, and combine the original few-shot prompting and test splits into a test set totalling 500 examples. for all experiments, we use the pre-trained pangu-coder @xcite code lm with 300m parameters to initialise the policy. we initialise the critic network on top of a copy of the same pre-trained model, as a two-layer mlp, taking as input the final representation of the end-of-sequence token of the underlying lm, downsampling to half the number of dimensions and again to a single output unit, with the relu activation function between layers. during critic training, we freeze the critic's copy of the code lm, and only update mlp. we set the policy learning rate α = 5e -7 and critic lr β = 1e -6 . the maximum number of training epochs is 100, and we generate n @xmath0. the reward hyper parameters are set to λ = 50 and η = 0.5. the target kl divergence is set to ρ = 0.07. when sampling solutions from the policy, we use nucleus sampling @xcite , with @xmath1.8 and @xmath2.95. each training minibatch contains 32 solutions with a max length @xmath3. for each model under consideration, we evaluate performance on the mbpp validation set after each epoch, and select the checkpoint with best validation score (based on greedy decoding) for test set evaluation. we report more details on the experiments and the baselines in appendix a. finally, ablation studies are carried out in section 5. ## ablation studies in this section, we verify if several components of our method are necessary. in those experiments, to reduce the needed computations, we only train on mbpp train without the augmented unit tests. the validation process is kept the same. ## conclusions and future work we have presented two novel ways of improving code synthesis models. on the training side, we introduced a practical reinforcement learning approach, based on an simple actor-critic rein-force algorithm. this approach shown to be simple yet effective, improving over clm as well as ppo-trained models. to provide additional learning signals, we furthermore introduced a method of extracting function-level strongly-typed code from large crawled code datasets, automatically generating unit tests for extracted functions, and conversion for use with a weakly typed target language. the generated data was able to further improve model performance, getting the overall best results for greedy code generation. in the future, we would like to explore our methods in settings with wider scope, i.e., go from function-level code to class-or even project-level. this poses significant challenges, as rewards will become even sparser, and data augmentation will have to take much larger scope into account.