Technical Summary: Prompt Optimization for LLM Code Generation via Reinforcement Learning

Problem Statement

While Large Language Models (LLMs) have advanced automated software engineering, particularly in generating code from natural language, their output frequently contains functional and logical defects. These errors limit practical deployment and necessitate manual debugging. Although static prompt engineering techniques (e.g., Chain-of-Thought) can mitigate common issues, they lack adaptability when iterative refinement based on intermediate feedback is required. Existing adaptive approaches, such as evolutionary algorithms (e.g., EPiC) or reinforcement learning (RL) frameworks (e.g., Reflexion), often treat lexical and semantic prompt modifications in isolation or rely on sparse binary correctness signals, failing to leverage rich information from partial correctness or integrate multiple transformation dimensions cohesively.

Methodology

The authors propose a reinforcement learning-based framework that models prompt refinement as a sequential decision-making problem. The core components of the methodology include:

1. Framework Architecture

The system utilizes a Proximal Policy Optimization (PPO) agent to iteratively improve prompts. The process operates within a custom Gymnasium-compatible environment where:

• State: The current prompt is encoded as a 384-dimensional semantic embedding using the MiniLM sentence transformer. Notably, the state representation is prompt-only; it does not encode generated code, runtime exceptions, or execution tracebacks.
• Action Space: The agent selects from a hybrid action space comprising three strategies:
  1. Direct Generation: Preserving the current prompt (exploiting decoding stochasticity).
  2. Lexical Mutation: A genetic algorithm-inspired approach involving token-level insertion, deletion, replacement, and crossover, guided by test-case performance.
  3. Semantic Rewriting: Higher-level restructuring and clarification of the prompt using a pretrained model, following a self-reflective paradigm.
• Environment Loop: The agent observes the prompt, samples an action, applies the transformation, passes the new prompt to a frozen code generation model, and evaluates the output against unit tests.

2. Reward Structure

A key innovation is the use of a shaped reward function derived from unit-test feedback, rather than relying solely on binary success/failure signals. The reward $R_{i,t}$ is defined as:

• $1.0$ if the pass ratio $\rho_{i,t} = 1.0$ (all tests passed).
• $\rho_{i,t}$ if $0 < \rho_{i,t} < 1.0$ (partial correctness).
• $-1$ if $\rho_{i,t} = 0$ (no tests passed).
• $-2$ if execution fails (syntax errors or runtime exceptions).

This design provides denser learning signals, allowing the agent to receive credit for incremental improvements and penalizing execution failures more heavily to discourage syntactically invalid code.

3. Evaluation Metrics

The framework is evaluated using two metrics:

• Strict Pass@1: The primary metric, counting a task as successful only if a solution passes all tests at any step.
• SoftPass@1: An auxiliary metric measuring the cumulative probability of passing tests across refinement steps, used to analyze intermediate progress.

Key Contributions

The paper outlines four primary contributions:

1. Sequential Formulation: Prompt optimization for code generation is formulated as a sequential decision-making problem and implemented as a PPO-based RL environment.
2. Hybrid Action Space: The definition of a combined action space integrating direct generation, genetic lexical mutation, and semantic rewriting for adaptive multi-step refinement.
3. Shaped Rewards: The introduction of a reward function utilizing partial test case correctness to provide denser learning signals than binary success alone.
4. Comprehensive Evaluation: Extensive evaluation on MBPP+, HumanEval+, and APPS benchmarks using CodeT5+, CodeLLaMA, and DeepSeek-Coder as frozen backbones, comparing against direct generation, EPiC, Reflexion, and a Random-Hybrid baseline.

Experimental Results

The framework was tested on 500 tasks from the MBPP+ test set, as well as HumanEval+ and APPS.

• Performance Gains: On MBPP+, the PPO agent achieved strict Pass@1 scores of 57.58% (CodeT5+), 64.80% (CodeLLaMA), and 85.50% (DeepSeek-Coder). These results outperform EPiC, Reflexion, and a Random-Hybrid baseline across all backbone models.
• Comparison with Baselines:
  • The PPO agent significantly outperformed the Random-Hybrid baseline (which uses the same action space but selects actions uniformly at random). For instance, on MBPP+ with CodeT5+, PPO improved strict Pass@1 by 26.46 percentage points over Random-Hybrid. This indicates that the performance gains stem from learned action scheduling rather than merely the availability of multiple transformation types.
  • PPO also surpassed EPiC and Reflexion, demonstrating that coordinating lexical and semantic transformations via a learned policy is more effective than isolated evolutionary search or reflection-based methods.
• Reward Shaping Impact: Experiments comparing sparse binary rewards against shaped rewards showed that the latter leads to faster convergence and higher stability. Binary rewards resulted in slower learning and higher instability, confirming that partial correctness signals are essential for effective multi-step refinement.
• Statistical Significance: McNemar's tests on the MBPP+ test set confirmed statistically significant improvements ($p < 0.001$) over EPiC and Reflexion for strict Pass@1, with medium effect sizes.

Significance and Claims

The paper claims that its primary significance lies in demonstrating that reinforcement learning with shaped, test-driven rewards can improve the functional correctness of LLM-based code generation.

Key takeaways emphasized by the authors include:

• Adaptive Coordination: The ability of the agent to dynamically coordinate between lexical mutation (local exploration) and semantic rewriting (global restructuring) based on intermediate feedback is crucial for success.
• Value of Partial Signals: Utilizing partial correctness as a learning signal is superior to binary success/failure for training agents in multi-step refinement tasks.
• Robustness: The framework shows consistent improvements across different backbone models and varying task difficulties (from MBPP+ to the more complex APPS), suggesting robustness to architectural variations.

The authors maintain a modest stance, acknowledging limitations such as the prompt-only state representation (excluding code and tracebacks), the reliance on frozen models, and the computational cost of repeated generation and execution. They suggest future work could involve traceback-aware state representations and integration into practical development tools.