Learning to Refine Specifications for Code Generation

Abstract

The use of Large Language Models (LLMs) for code generation has advanced rapidly, yet current systems often struggle with complex algorithmic logic and edge cases. Standard methods for improving code generation, such as iterative self-refinement, often cause repetitive failure cycles, as models cannot independently escape their own logical blind spots. To address this gap, we propose a novel Reinforcement Learning from AI Feedback (RLAIF) framework. Our key insight is that specifications can be iteratively refined by leveraging feedback from the LLM's own outputs, creating a self-improving loop. This approach contrasts with prior methods that treat specifications as static inputs. Rather than updating the parameters of the base model, a proprietary LLM (the Worker) is accessed via API, while a smaller, trainable Student agent learns to guide generation by modifying natural language specifications. A Tutor agent supervises these refinements, ensuring they are precise and actionable, and provides a reward signal driven by failures identified by a synthetic test generator. Crucially, the Worker, Tutor, and Test Generator roles are all instantiated using the same proprietary model, isolating the learning capacity entirely within the Student agent. We evaluate our approach on N = 200 problems using a Qwen 2.5-14B-Instruct Student and a GPT-4o-mini model fulfilling the Worker, Tutor, and Test Generator roles on the LiveCodeBench and HumanEvalPlus benchmarks. Our method achieves relative pass@1 improvements of 3.2% on LiveCodeBench (64.0% vs. 62.0%) and 1.1% on HumanEvalPlus (96.0% vs. 95.0%).