Prompt Optimization for LLM Code Generation via Reinforcement Learning Ali Mohammadi Esfahani1(B) , Nafiseh Kahani1 , and Samuel A.Ajila1
arXiv:2605.19102v1 [cs.SE] 18 May 2026
Carleton University, Ottawa, ON, Canada [email protected]
Abstract. Large Language Models (LLMs) can generate code from natural language, but their performance is highly sensitive to prompt formulation. We propose a reinforcement-learning-based framework that models prompt refinement as a sequential decision-making problem. A Proximal Policy Optimization (PPO) agent iteratively improves prompts using a hybrid action space that combines direct generation, genetic lexical mutation and semantic rewriting, guided by shaped rewards derived from unit-test feedback. We evaluate the framework on MBPP+, HumanEval+, and APPS using CodeT5+, CodeLLaMA, and DeepSeekCoder as frozen code generators. On the 500-task MBPP+ test set, the PPO agent achieves strict Pass@1 scores of 57.58%, 64.80%, and 85.50%, respectively, outperforming EPiC, Reflexion, and Random-Hybrid. SoftPass@1 reaches 67.90%, 73.10%, and 88.20%, respectively. Similar improvements are observed on HumanEval+ and APPS across all backbone models. The results demonstrate that reinforcement learning with shaped test-driven rewards improves functional correctness in LLM-based code generation. Keywords: Reinforcement Learning · Prompt Engineering · Code Generation · Large Language Models · Genetic Algorithm.
1
Introduction
Large Language Models (LLMs) have improved many automated software engineering tasks, particularly code generation from natural language descriptions [1,2,9,6,11,13,14,36,17]. However, generated code still frequently contains functional and logical defects, which limits its practical use. Recent studies have highlighted that even advanced LLMs exhibit relatively low exact-match accuracy, with critical defects frequently introduced into the generated code, requiring manual debugging and correction [10,31]. In a recent study, defects in LLM-generated code were investigated and categorized, with prevalent issues such as functionality, algorithmic, and logic defects being identified, emphasizing the necessity for targeted improvements [10]. The authors also demonstrated that standard prompt engineering techniques, such as Chain-of-Thought and Structured Chain-of-Thought prompting, can mitigate common defects by providing clearer instructions and context to the model [10].
2
A. Mohammadi Esfahani et al.
However, these static prompting methods lack adaptability and often fall short when iterative refinement based on intermediate feedback is required. To address these limitations, several adaptive approaches have recently been proposed, leveraging evolutionary algorithms or reinforcement learning (RL) frameworks for prompt optimization. Techniques such as EPiC [30] employ lightweight evolutionary search algorithms to iteratively refine prompts based on the performance of generated code evaluated against test cases. Similarly, RL is used in Reflexion [29], where language agents iteratively reflect on task feedback, incorporating self-generated linguistic insights into future prompt formulations. In addition, frameworks such as the Large Language Model Debugger (LDB) [38] enhance LLM-generated code by integrating runtime execution insights to iteratively pinpoint and correct defects. Despite the effectiveness of these approaches, lexical or semantic prompt modifications are typically treated in isolation, and multiple dimensions of prompt transformations are not cohesively integrated. Furthermore, current strategies mainly rely on binary correctness signals, often overlooking the rich information contained in partial correctness or detailed feedback. The contribution lies in coordinating reinforcement learning, genetic prompt mutation, and semantic rewriting through a learned sequential policy for executable code generation. To fill these gaps, we propose an RL-driven prompt optimization approach specifically tailored for functional code generation tasks. This paper investigates whether RL can improve prompt optimization for code generation by adaptively combining lexical mutation and semantic rewriting. The proposed framework learns to select between lexical mutation and semantic rewriting based on intermediate test feedback. Prompt optimization is modeled as a sequential decisionmaking problem, where shaped rewards and SoftPass@1 is used as an auxiliary learning-sensitive metric to leverage partial correctness signals. The main contributions of this work are as follows.
– We formulate prompt optimization for code generation as a sequential decisionmaking problem and implement it as a PPO-based RL environment. – We define a hybrid action space that combines direct generation, lexical mutation, and semantic rewriting for adaptive multi-step prompt refinement. – We introduce a shaped reward function that uses partial test case correctness to provide denser learning signals than binary success alone. – We evaluate the framework on MBPP+, HumanEval+, and APPS using CodeT5+, CodeLLaMA, and DeepSeek-Coder, and compare it against direct generation, EPiC, Reflexion, and a Random-Hybrid baseline that uses the same action space without policy learning. The remainder of the paper is organized as follows. Section 2 reviews related work. Section 3 presents the proposed framework. Section 4 describes the experimental setup and results. Finally, Section 5 concludes the paper.
Prompt Optimization for LLM Code Generation via Reinforcement Learning
2
3
Related work
Recent advances in LLMs have stimulated extensive research on prompt engineering and optimization for tasks such as code generation, instruction following, and program synthesis. Early approaches relied on manual prompt design, where human-crafted instructions were iteratively refined to elicit accurate model responses [3], [6]. While effective in constrained settings, handcrafted prompts often lack generalization, motivating automated prompt optimization methods [37], [35], [12]. These approaches are broadly categorized into evolutionary search and RL frameworks. Evolutionary methods treat prompt optimization as population-based search. EPiC [30] applies lexical mutations (insertion, deletion, replacement) guided by test-case correctness as a fitness signal. EvoPrompt [15] incorporates LLMs into mutation and crossover operations to produce semantically coherent variants, while PromptBreeder [12] introduces self-referential co-evolution of prompts and mutation operators. Although effective for exploring lexical diversity, evolutionary approaches generally rely on predefined operators and lack adaptive learning for multi-step refinement. RL-based approaches formulate prompt optimization as sequential decisionmaking. RLPrompt [8] applies token-level edits using correctness-based rewards. PRewrite [19] extends the action space to include semantic rewriting via LLMbased paraphrasing. StablePrompt [20] improves training stability through Adaptive Proximal Policy Optimization, and PRL [4] explores RL-based prompt generation for few-shot tasks. However, most RL methods focus on narrow edit operations or fixed rewriting strategies and are typically evaluated on general language tasks rather than executable code generation with strict functional correctness constraints. Existing methods either rely on lexical search, fixed semantic rewriting, or sparse correctness signals. Few combine heterogeneous prompt transformations within a learned sequential policy, and even fewer evaluate them on executable code generation with unit-test-driven rewards. This work addresses those gaps by integrating lexical mutation and semantic rewriting in a unified PPO framework with shaped rewards based on partial correctness under unit-test evaluation.
3
Prompt Optimization via Reinforcement Learning
Prompt optimization is modeled as a sequential decision-making problem within an RL framework, where an agent iteratively refines natural language prompts in order to improve the functional correctness of code generated by an LLM. In this setting, each programming task corresponds to an episode of interaction with the environment. The state represents the current prompt, encoded as a semantic embedding. The action space consists of prompt transformation strategies that modify or reuse the current prompt. After a transformation is applied, the updated prompt is passed to a frozen code generation model to produce executable code. The generated program is evaluated against test cases,
4
A. Mohammadi Esfahani et al.
Fig. 1: Workflow of the RL–based prompt optimization framework
and the resulting correctness signal defines the reward. The objective of the agent is to learn a policy that selects prompt transformations maximizing functional correctness across refinement steps. The agent is trained using Proximal Policy Optimization [28], selected for its stability and effectiveness in LLM-interactive environments. PPO is widely used in alignment and instruction tuning due to its controlled policy updates, computational efficiency, and robustness to noisy rewards [25]. The action space includes three actions: (1) Direct Generation, which preserves the current prompt, (2) genetic algorithm–based lexical mutation inspired by EPiC [30], and (3) semantic rewriting inspired by Reflexion [29]. During each episode, the environment operates in an iterative loop: the agent observes the current prompt representation, samples a prompt transformation from its policy, generates code using the updated prompt, evaluates the output against test cases, receives a reward based on functional correctness, and updates its policy accordingly. Unit tests are required during training to compute the reward signal and during evaluation to assess correctness, but they are not needed when the trained policy is applied for prompt refinement at inference time. Figure 1 illustrates this test-driven prompt refinement process.
3.1
Environment Design
Prompt optimization is implemented within a custom Gymnasium-compatible RL environment. Each episode corresponds to a programming problem sampled from the training splits of MBPP+ [22], HumanEval+ [22], and APPS [16], covering a range of task difficulties from simple algorithms to more complex multistep reasoning tasks. An episode continues until either a fully correct program is generated or a predefined maximum number of refinement steps is reached. At time step t, the agent observes the current prompt pt . The prompt is embedded into a 384-dimensional vector st using the MiniLM sentence transformer [26]. This embedding serves as the state representation provided to the policy network and ensures a unified semantic representation across datasets. This prompt-only state isolates prompt-level optimization, but it does not encode generated code,
Prompt Optimization for LLM Code Generation via Reinforcement Learning
5
failed assertions, runtime exceptions, or tracebacks. Thus, the framework is best viewed as prompt-level refinement rather than traceback-guided debugging. At each step, the agent selects a prompt transformation, generating a modified prompt that is passed to the code generation model. The generated program at step t of episode i is denoted by fi,t . The resulting function is evaluated against i a set of test cases {(xi,k , yi,k )}K k=1 , where Ki denotes the number of unit tests associated with task i. For each test case, correctness is recorded as ci,t,k = Jfi,t (xi,k ) = yi,k K
(1)
where J·K denotes the indicator function. The pass ratio for step t in episode i is then computed as K
ρi,t =
i 1 X ci,t,k . Ki
(2)
k=1
This pass ratio measures the proportion of test cases passed by the generated program and provides an intermediate signal of functional correctness. It forms the basis of the shaped reward used to guide policy learning. 3.2
Action Space and Prompt Transformation Strategies
The agent operates in a discrete action space in which each action corresponds to a prompt transformation strategy. The proposed framework defines three actions. Direct generation is a no-edit control action, while lexical mutation and semantic rewriting are the two active refinement operators. The a priori motivation for combining the latter two is that they modify prompts along complementary dimensions: lexical mutation explores local surface-level variants, whereas semantic rewriting performs higher-level clarification and restructuring. Direct Generation generates code using the unchanged prompt. Although no modification is applied, this action serves as a no-edit option that allows the policy to exploit decoding stochasticity and preserve prompts that are already close to optimal. Because autoregressive language models can produce different outputs from the same prompt under stochastic decoding, this action still supports useful exploration [5]. Lexical Mutation applies a genetic algorithm inspired by EPiC [30]. The prompt is tokenized to form an initial population of candidate variants. Successive generations apply crossover, index shuffling, and mutation (with probability 0.2). Candidate prompts are evaluated based on the number of passed test cases, and tournament selection retains stronger variants for the next generation. Semantic Rewriting performs higher-level reformulation using a pretrained model [32], following the self-reflective paradigm of Reflexion [29]. A fixed metaprompt guides structured rewriting of the original prompt, while regex-based filtering enforces basic structural constraints and ensures that task-relevant keywords are preserved. At time step t, the agent samples an action ai,t from a stochastic policy πθ conditioned on the current state si,t :
6
A. Mohammadi Esfahani et al.
ai,t ∼ πθ (si,t )
(3)
Here, πθ outputs a probability distribution over the three available actions. For example, depending on the prompt embedding, the policy may assign higher probability to preserving the prompt, applying a lexical mutation, or invoking semantic rewriting. The selected action is applied via transformation operator T to obtain the updated prompt: pi,t+1 = T (pi,t , ai,t )
(4)
The operator T denotes the action-specific prompt update: it returns the unchanged prompt for direct generation, a locally mutated prompt for lexical mutation, or a semantically rewritten prompt for semantic rewriting. The transformed prompt is then used for code generation, and the policy learns to sequence actions based on functional correctness feedback. 3.3
Code Generation and Reward Structure
After transformation, the prompt is passed to a code generation model to proi duce a candidate function. Each task includes test cases {(xi,k , yi,k )}K k=1 , and correctness is determined by executing the generated code against these tests. Execution is performed in a sandboxed environment adapted from Reflexion [29], which safely captures syntax errors, runtime exceptions, and incorrect outputs. The outcome is used to compute a shaped reward based on the pass ratio ρi,t (Section 3.1): 1.0 ρ i,t Ri,t = −1 −2
if ρi,t = 1.0 if 0 < ρi,t < 1.0 if ρi,t = 0 if execution fails
(5)
This design rewards both full and partial correctness, penalizes complete functional failure, and strongly discourages execution errors. Partial pass ratios are used only as intermediate learning signals for policy optimization; they are not treated as final functional success. Final task success is evaluated separately using strict Pass@1, which requires all tests to pass. A stronger penalty is assigned to execution failure because syntactically invalid or non-executable outputs provide no usable functional signal for learning. In practice, execution failures receive a penalty of −2, which discourages the agent from producing syntactically invalid code while still allowing exploration of alternative prompt transformations. The cumulative episode reward is defined as Riepisode =
Ti X t=1
Ri,t .
(6)
Prompt Optimization for LLM Code Generation via Reinforcement Learning
7
This reward signal encourages the agent to discover sequences of prompt transformations that progressively improve functional correctness. The reward is therefore intended to guide search toward fully correct programs, not to accept partially correct programs as usable outputs. 3.4
Implementation
Algorithm 1 summarizes the training procedure of the RL-based prompt optimization framework. Each episode begins by selecting a programming task, which defines the environment instance. The initial prompt is embedded using MiniLM (Section 3.1) and fed to the policy network, providing a consistent semantic representation across tasks. At each step, the agent selects an action as defined in Section 3.2. The transformed prompt is passed to a code generation model to produce executable Python code. The generated code is executed in a sandbox and evaluated against predefined test cases. A shaped reward is computed based on the proportion of passed tests (Section 3.3), rewarding partial correctness and penalizing failures. Each transition is stored for mini-batch updates, and at the end of the episode the policy is updated using PPO [28], ensuring stable and sample-efficient learning through clipped objectives. PPO policies were trained independently for each dataset–backbone configuration, using the corresponding training split and frozen code generator for that setting.
4
Results and Analysis
This section presents the research questions, experimental setup, and results. The results are organized to answer the following research questions: RQ1: How do individual prompt transformation strategies perform in isolation when applied once without iterative feedback? RQ2: How does the use of shaped reward functions and step-wise feedback influence the learning dynamics and performance of the RL agent? RQ3: How effectively does multi-step RL improve prompt optimization compared to existing iterative strategies? 4.1
Experimental Setup
This section describes the experimental setup, including benchmarks, metrics, and baseline comparisons. Evaluation Metrics. To evaluate the trained agent, we report strict Pass@1 and SoftPass@1. Strict Pass@1 is treated as the primary measure of functional correctness because it counts a task as successful only when the generated solution passes all available tests. SoftPass@1 is used only as an auxiliary diagnostic metric for analyzing intermediate progress during multi-step refinement.
8
A. Mohammadi Esfahani et al.
Algorithm 1 Reinforcement Learning Training Loop
Algorithm 2 Evaluation Protocol for Prompt Optimization
Require: Dataset D, PPO agent πθ , environment E, episodes N 1: for i ← 1 to N do 2: Sample task (pi , Ti ) from D 3: si,1 ← Embed(pi ) 4: for t ← 1 to Ti do 5: ai,t ∼ πθ (si,t ) 6: pi,t+1 ← T (pi,t , ai,t ) ▷ transform prompt 7: fi,t ← LLM(pi,t+1 ) ▷ frozen code generator 8: Evaluate fi,t on Ti to compute ρi,t 9: Compute reward Ri,t and store transition in Bi 10: si,t+1 ← Embed(pi,t+1 ) 11: if ρi,t = 1.0 then 12: break 13: end if 14: end for 15: Update policy θ ← PPOUpdate(θ, Bi ) 16: end for
Require: Trained policy πθ , dataset Dtest , step cap T 1: Initialize Cstrict ← 0, Ssoft ← 0 2: for each task (pi , Ti ) in Dtest do 3: si,1 ← Embed(pi ), Pi ← 1, ζi ← 0 4: for t ← 1 to T do 5: ai,t ∼ πθ (si,t ) 6: pi,t+1 ← T (pi,t , ai,t ) ▷ transform prompt 7: fi,t ← LLM(pi,t+1 ) 8: Evaluate fi,t to compute ρi,t 9: Pi ← Pi · (1 − ρi,t ) ▷ accumulate soft success 10: if ρi,t = 1.0 then 11: ζi ← 1 ▷ mark strict success 12: break 13: end if 14: si,t+1 ← Embed(pi,t+1 ) 15: end for 16: Cstrict ← Cstrict + ζi 17: Ssoft ← Ssoft + (1 − Pi ) 18: end for 19: Compute metrics: P ass@1strict ← Ssoft Cstrict , P ass@1soft ← |D |Dtest | test |
Fig. 2: (Left) PPO-based training loop; (Right) Evaluation protocol. Both algorithms are model-agnostic; the code generator is denoted by LLM.
For task i, let ρi,t denote the fraction of test cases passed at refinement step t. Strict Pass@1 is computed as N
Pass@1strict =
1 X J∃t ∈ {1, . . . , Ti } : ρi,t = 1.0K. N i=1
(7)
This metric follows the standard interpretation of functional correctness: a partially correct program is not considered functionally correct or deployable. To capture partial progress during iterative refinement, SoftPass@1 is computed as Pass@1
soft
! Ti N Y 1 X = 1− (1 − ρi,t ) . N i=1 t=1
(8)
Prompt Optimization for LLM Code Generation via Reinforcement Learning
9
SoftPass@1 indicates whether the agent moves closer to full correctness across refinement steps. It should not be interpreted as evidence that partially correct code is practically usable. The evaluation protocol used to compute these metrics is detailed in Algorithm 2. Datasets. Experiments were conducted on MBPP+ [22], HumanEval+ [22], and APPS [16]. MBPP+ (974 tasks) was split into 374 training and 500 testing tasks. The remaining tasks were used for validation and excluded from training. For HumanEval+ and APPS, an 80%/20% train–test split was used. Model Selection. Following the survey of LLMs for software engineering by Hou et al. [18], model selection was guided by explicit criteria. Only open-source models were considered to ensure reproducibility, excluding proprietary systems such as GPT-4 [33], AlphaCode [3], PaLM-Coder [24], and Codex [36]. Models not explicitly pre-trained on large-scale code corpora, including T5 [7], GPTNeo [13], and GPT-J [21], were also excluded. Older models superseded by updated versions, such as CodeT5 [23], were not selected. Based on these criteria, CodeT5+ [34], CodeLLaMA [27], and DeepSeek-Coder [39] were chosen as code generation backbones. Baselines. In addition to direct generation, EPiC, and Reflexion, we include Random-Hybrid as a non-learning ablation baseline. Random-Hybrid uses the same three actions as the proposed PPO framework. However, instead of learning a policy, it samples one action uniformly at random at each refinement step. We run Random-Hybrid for 10 refinement steps per task, using the same evaluation protocol as the PPO setting. This baseline isolates whether the performance gains come merely from access to the hybrid action space or from learned action scheduling. 4.2
Experimental Results
The proposed RL-based prompt optimization framework was evaluated on MBPP+, HumanEval+, and APPS using held-out test sets to assess its effectiveness across benchmarks of varying difficulty. The following subsections analyze the results with respect to the research questions, comparing static, heuristic, and learningbased approaches across benchmarks of increasing difficulty. RQ1: Performance of Prompt Transformations We evaluate three single-step prompt transformation strategies without iterative refinement. Each benchmark task is processed once and evaluated using strict Pass@1. Quantitative results are provided in Table 1. Across most datasets and backbone models, semantic rewriting yields the strongest single-step improvements, followed by lexical mutation, while direct prompting performs the weakest. These gains suggest that reducing ambiguity and restructuring instructions increases the likelihood of generating functionally correct programs. Although lexical mutation is less effective than semantic rewriting in a single-step setting, it still improves over direct generation and provides useful
10
A. Mohammadi Esfahani et al.
Table 1: Comparison of Pass@1 and SoftPass@1 across datasets and prompt optimization actions. Semantic rewriting generally provides the strongest single-step improvement, while SoftPass@1 highlights partial functional progress beyond strict Pass@1. Dataset
Method
CodeT5+
CodeLLaMA
DeepSeek-Coder
Pass@1 SoftPass@1 Pass@1 SoftPass@1 Pass@1 SoftPass@1 Direct Generation 12.84% Genetic Mutation 20.23% Semantic Rewriting 23.73%
22.80% 33.40% 36.90%
41.91% 43.01% 46.20%
48.82% 51.70% 55.20%
76.10% 78.24% 78.49%
77.86% 78.79% 79.18%
Direct Generation 15.10% HumanEval+ Genetic Mutation 21.49% Semantic Rewriting 25.59%
19.61% 30.75% 37.12%
33.02% 34.63% 32.33%
35.61% 40.81% 38.91%
89.90% 89.24% 89.49%
90.01% 89.49% 89.94%
Direct Generation 8.20% Genetic Mutation 10.95% Semantic Rewriting 13.40%
14.60% 19.80% 23.70%
14.75% 17.90% 20.85%
22.10% 27.60% 31.20%
18.30% 21.75% 24.60%
27.40% 32.50% 36.80%
MBPP+
APPS
local exploration. This suggests that mutation and rewriting play complementary roles: rewriting introduces larger semantic clarification, whereas mutation explores smaller lexical variants around a promising prompt. Performance decreases on APPS due to its higher algorithmic complexity, yet the relative ordering of strategies remains largely consistent. This consistency indicates that structured prompt clarification provides a robust one-shot enhancement mechanism across difficulty levels. Overall, while single-step transformations improve performance, their impact remains limited without iterative feedback. These findings motivate the need for adaptive, multi-step refinement, as explored in the subsequent RL framework. The Random-Hybrid and PPO comparisons in RQ3 further test whether simply having both transformation types is sufficient, or whether learned action scheduling is needed to use them effectively. RQ2. Shaped Rewards and Feedback To evaluate the impact of reward design, we compare a sparse binary reward with the proposed shaped reward under identical environments and action spaces. In the binary setting, the agent receives a reward of 1.0 only if all test cases pass; otherwise, it receives zero. In contrast, the shaped reward assigns proportional credit based on the fraction of passed test cases at each step, with additional penalties for execution failures. This formulation provides continuous feedback rather than episodic success-only signals. Empirically, reward shaping leads to consistently stronger performance across datasets and backbone models. On MBPP+, the shaped-reward agent achieves substantial improvements under both strict Pass@1 and SoftPass@1 (see Table 2). In contrast, the binary-reward variant converges more slowly and exhibits higher instability during training, indicating that sparse success signals are insufficient for effective policy learning in multi-step prompt refinement. Step-wise feedback plays a central role in this improvement. Because the shaped reward provides proportional credit for partial correctness, the agent can detect
Prompt Optimization for LLM Code Generation via Reinforcement Learning
# Prompt # Write a function to find all tuples in a list # that match a given condition # Generated Code def find_tuples(l, k) : a = [] for i in range( len(l)): if l[i] == k: a.append( l[i]) return a # Pass Ratio: 0.00 # Reward: 0.00
# Prompt # Write a Python function named ’ find_tuples’ that # returns all tuples meeting a given condition # from the input list
11
# Prompt # Write a Python function that takes a list of tuples # and integer k, and returns tuples whose length is k
# Generated Code def find_tuples(tup, k): return [x for x in tup if len(x) == k]
# Generated Code def find_tuples(tup, k): return [x for x in tup if len(x) == k and isinstance(x, tuple)]
# Pass Ratio: 0.67 # Reward: 0.67
# Pass Ratio: 1.00 # Reward: 1.00
(a) Original Prompt (Am- (b) Partial Pass after Se- (c) Full Pass after Refinebiguous, No Pass) mantic Rewrite ment
Fig. 3: Illustrative progression from an ambiguous benchmark-style prompt to partial and full correctness, showing the corresponding generated code, pass ratios, and rewards.
when a semantic rewrite or lexical mutation increases the number of passing test cases. This enables refinement of promising transformations rather than discarding them prematurely. Under the binary condition, partially correct outputs are treated as complete failures, preventing incremental learning and discouraging multi-step exploration. The effect is further illustrated in Figure 3, which shows an illustrative benchmarkstyle refinement trajectory where intermediate prompt transformations progressively increase the pass ratio. The shaped reward reinforces these intermediate gains, allowing the policy to accumulate functional progress even when full correctness is not immediately achieved. Reward shaping stabilizes learning, supports incremental credit assignment, and is therefore used in all subsequent experiments. RQ3. Multi-Step RL Effectiveness We evaluate the proposed PPO-based RL framework against EPiC, Reflexion, and Random-Hybrid across MBPP+, HumanEval+, and APPS. All backbone models remain frozen during evaluation, and performance is measured using both
12
A. Mohammadi Esfahani et al.
Table 2: Comparison of PPO against EPiC, Reflexion, and Random-Hybrid across MBPP+, HumanEval+, and APPS. Bold values indicate the proposed PPO framework. Dataset
Method
CodeT5+
CodeLLaMA
DeepSeek-Coder
MBPP+
EPiC 41.89% Reflexion 41.63% Random-Hybrid 31.12% PPO (Ours) 57.58%
54.20% 55.10% 44.10% 67.90%
51.40% 52.70% 49.08% 64.80%
61.50% 63.60% 58.60% 73.10%
80.30% 81.60% 79.10% 85.50%
82.10% 83.40% 79.90% 88.20%
EPiC 33.80% Reflexion 35.40% HumanEval+ Random-Hybrid 31.10% PPO (Ours) 60.20%
46.90% 48.70% 42.60% 69.80%
41.20% 43.00% 38.10% 70.60%
52.60% 54.90% 45.80% 77.90%
90.10% 91.10% 89.60% 91.30%
91.80% 92.90% 90.10% 92.50%
EPiC 16.40% Reflexion 17.60% Random-Hybrid 13.90% PPO (Ours) 20.75%
28.75% 30.10% 24.10% 35.90%
24.80% 26.10% 21.50% 29.80%
38.60% 40.20% 31.90% 45.60%
29.90% 31.20% 24.70% 34.95%
45.10% 47.30% 37.50% 52.40%
Pass@1 SoftPass@1 Pass@1 SoftPass@1 Pass@1 SoftPass@1
APPS
strict Pass@1 and SoftPass@1. Detailed results are reported in Table 2. RandomHybrid uses the same three-action space as PPO but selects actions uniformly at random for 10 refinement steps, isolating the effect of learned action scheduling. Across all datasets and backbone models, PPO outperforms EPiC, Reflexion, and Random-Hybrid. Since Random-Hybrid has access to the same transformations as PPO, this result shows that the gains are not explained by the hybrid action space alone. On MBPP+, PPO improves strict Pass@1 over Random-Hybrid by 26.46 percentage points for CodeT5+, 15.72 points for CodeLLaMA, and 6.40 points for DeepSeek-Coder. Similar trends are observed on HumanEval+ and APPS, although absolute performance decreases on APPS due to its higher algorithmic complexity and stricter evaluation protocol. These results indicate that learned multi-step coordination is more effective than mutation-driven, rewritedriven, or randomly scheduled hybrid refinement. SoftPass@1 follows the same trend. Beyond increasing strict task completion rates, PPO accumulates stronger intermediate functional progress, suggesting that the learned policy can preserve and build upon partial improvements during iterative prompt refinement. This is particularly important because PPO and Random-Hybrid share the same available transformations; the performance gap therefore reflects learning when to apply each transformation rather than merely combining them. To assess statistical significance, paired comparisons were conducted across tasks in the evaluation sets. For strict Pass@1, McNemar’s test was applied on the MBPP+ test set (500 tasks), yielding statistically significant improvements over both EPiC and Reflexion (p < 0.001) with medium effect sizes (Cohen’s h approximately 0.27–0.35). For SoftPass@1, paired t-tests indicate highly significant gains (p < 0.0001) with medium-to-large effect sizes (Cohen’s d approximately 0.60–0.75). Overall, the results suggest that PPO improves prompt optimization
Prompt Optimization for LLM Code Generation via Reinforcement Learning
13
by adaptively sequencing lexical and semantic transformations based on feedback signals. 4.3
Threats to Validity
Several factors may affect the generality of the results. First, the evaluation relies on benchmark datasets and predefined test suites, which may not fully capture real-world programming complexity. Second, the code generation models remain frozen during training; therefore, improvements arise only from prompt optimization rather than model adaptation. Third, the quality of semantic rewriting depends on the rewriting language model and the structure of the meta-prompt. Finally, the RL state representation uses only prompt embeddings and does not explicitly incorporate previous actions or reward history. It also excludes generated code and execution tracebacks, limiting diagnostic refinement. In addition, semantic rewriting and PPO training increase computational cost because they require repeated rewriting, generation, execution, and reward computation.
5
Conclusion
This paper presented a reinforcement learning framework for multi-step prompt optimization in code generation. By modeling prompt refinement as a sequential decision-making process and combining lexical mutation with semantic rewriting, the proposed PPO-based agent consistently outperformed EPiC, Reflexion, and Random-Hybrid across MBPP+, HumanEval+, and APPS. The RandomHybrid comparison shows that the gains come not only from access to multiple transformation operators, but also from learned action scheduling. These results highlight the value of adaptive, feedback-driven prompt optimization, where the agent dynamically coordinates semantic rewriting and lexical refinement using shaped rewards and intermediate correctness signals. Consistent gains across multiple backbone models indicate robustness to architectural variation and task difficulty. The current framework remains limited by its prompt-only state representation, which does not explicitly encode generated code, failed tests, or execution tracebacks. Future work may extend the framework to larger benchmarks, transfer or curriculum learning, traceback-aware state representations, and practical development-tool integration.
Acknowledgements We thank the Natural Sciences and Engineering Research Council of Canada (NSERC) for supporting this work.
References 1. Ahmad, W.U., Chakraborty, S., Ray, B., Chang, K.W.: Unified pre-training for program understanding and generation. arXiv:2103.06333 (2021)
14
A. Mohammadi Esfahani et al.
2. Asare, O., Nagappan, M., Asokan, N.: Is github’s copilot as bad as humans at introducing vulnerabilities in code? Empirical Software Engineering 28(6), 129 (2023) 3. Austin, J., Odena, A., Nye, M., Bosma, et al.: Program synthesis with large language models. arXiv:2108.07732 (2021) 4. Batorski, P., Kosmala, A., Swoboda, P.: Prl: Prompts from reinforcement learning. arXiv:2505.14412 (2025) 5. Brown, T., Mann, B., Ryder, N., Subbiah, et al.: Language models are few-shot learners. Advances in neural information processing systems 33, 1877–1901 (2020) 6. Chen, M., et al.: Evaluating large language models trained on code. arXiv:2107.03374 (2021) 7. Chowdhery, A., et al.: Palm: Scaling language modeling with pathways. Journal of Machine Learning Research 24(240), 1–113 (2023) 8. Deng, M., et al.: Rlprompt: Optimizing discrete text prompts with reinforcement learning. arXiv:2205.12548 (2022) 9. Elnaggar, A., Ding, W., Jones, L., Gibbs, T., Feher, T., Angerer, C., Severini, S., Matthes, F., Rost, B.: Codetrans: Towards cracking the language of silicon’s code through self-supervised deep learning and high performance computing. arXiv:2104.02443 (2021) 10. Esfahani, A.M., Kahani, N., Ajila, S.A.: Understanding defects in generated codes by language models. In: 2024 34th International Conference on Collaborative Advances in Software and COmputiNg (CASCON). pp. 1–10 (2024). https://doi.org/10.1109/CASCON62161.2024.10837857 11. Feng, Z., Guo, D., Tang, D., Duan, N., Feng, X., Gong, M., Shou, L., Qin, B., Liu, T., Jiang, D., et al.: Codebert: A pre-trained model for programming and natural languages. arXiv:2002.08155 (2020) 12. Fernando, C., Banarse, D., Michalewski, H., Osindero, S., Rocktäschel, T.: Promptbreeder: Self-referential self-improvement via prompt evolution. arXiv:2309.16797 (2023), https://arxiv.org/abs/2309.16797 13. Fried, D., Aghajanyan, A., Lin, J., Wang, S., Wallace, E., Shi, F., Zhong, R., Yih, W.t., Zettlemoyer, L., Lewis, M.: Incoder: A generative model for code infilling and synthesis. arXiv:2204.05999 (2022) 14. Guo, D., Ren, S., Lu, S., Feng, Z., Tang, D., Liu, S., Zhou, L., Duan, N., Svyatkovskiy, A., Fu, S., et al.: Graphcodebert: Pre-training code representations with data flow. arXiv:2009.08366 (2020) 15. Guo, Q., Wang, R., Guo, J., Li, B., Song, K., Tan, X., Liu, G., Bian, J., Yang, Y.: Evoprompt: Connecting llms with evolutionary algorithms yields powerful prompt optimizers. In: International Conference on Learning Representations (ICLR) (2024), https://arxiv.org/abs/2309.08532 16. Hendrycks, D., Basart, S., Kadavath, S., Mazeika, M., Arora, A., Guo, E., Burns, C., Puranik, S., He, H., Song, D., et al.: Measuring coding challenge competence with apps. arXiv:2105.09938 (2021) 17. Holden, D., Kahani, N.: Code linting using language models. arXiv preprint arXiv:2406.19508 (2024) 18. Hou, X., Zhao, Y., Liu, Y., et al.: Large language models for software engineering: A systematic literature review (2024) 19. Kong, W., Hombaiah, S.A., Zhang, M., Mei, Q., Bendersky, M.: Prewrite: Prompt rewriting with reinforcement learning. In: Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL-HLT) (2024)
Prompt Optimization for LLM Code Generation via Reinforcement Learning
15
20. Kwon, M., Kim, G., Kim, J., Lee, H., Kim, J.: Stableprompt: Automatic prompt tuning using reinforcement learning for large language models. arXiv:2410.07652 (2024) 21. Li, C., Liang, J., Zeng, A., Chen, X., Hausman, et al.: Chain of code: Reasoning with a language model-augmented code emulator. arXiv:2312.04474 (2023) 22. Liu, J., Xia, C.S., Wang, Y., Zhang, L.: Is your code generated by chatgpt really correct? rigorous evaluation of large language models for code generation. Advances in Neural Information Processing Systems 36, 21558–21572 (2023) 23. McCabe, T.J.: A complexity measure. IEEE Transactions on software Engineering (4), 308–320 (1976) 24. Mohammadkhani, A.H., Tantithamthavorn, C., Hemmatif, H.: Explaining transformer-based code models: What do they learn? when they do not work? pp. 96–106 (2023) 25. Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, et al.: Training language models to follow instructions with human feedback. Advances in neural information processing systems 35, 27730–27744 (2022) 26. Reimers, N., Gurevych, I.: Sentence-bert: Sentence embeddings using siamese bertnetworks. In: EMNLP (2019) 27. Roziere, B., Gehring, J., Gloeckle, et al.: Code llama: Open foundation models for code. arXiv:2308.12950 (2023) 28. Schulman, J., Wolski, F., Dhariwal, P., et al.: Proximal policy optimization algorithms (2017) 29. Shinn, N., Cassano, F., Gopinath, A., Narasimhan, K., Yao, S.: Reflexion: Language agents with verbal reinforcement learning. Advances in Neural Information Processing Systems 36, 8634–8652 (2023) 30. Taherkhani, H., Sepidband, M., et al.: Epic: Cost-effective search-based prompt engineering of llms for code generation. arXiv:2408.11198 (2024) 31. Tambon, F., Dakhel, A.M., Nikanjam, A., Khomh, F., Desmarais, M.C., Antoniol, G.: Bugs in large language models generated code. arXiv:2403.08937 (2024) 32. Touvron, H.e.a.: Llama: Open and efficient foundation language models. arXiv:2302.13971 (2023) 33. Wang, J., Huang, Y., Chen, C., Liu, Z., Wang, S., Wang, Q.: Software testing with large language models: Survey, landscape, and vision (2024) 34. Wang, Y., Le, H., Gotmare, A.D., Bui, N.D.Q., Li, J., Hoi, S.C.H.: Codet5+: Open code large language models for code understanding and generation (2023) 35. Yang, C., Wang, X., Lu, Y., Liu, H., Le, Q.V., Zhou, D., Chen, X.: Large language models as optimizers. arXiv:2309.03409 (2023) 36. Zhang, Z., Chen, C., Liu, B., Liao, C., Gong, Z., Yu, H., Li, J., Wang, R.: Unifying the perspectives of nlp and software engineering: A survey on language models for code. arXiv:2311.07989 (2023) 37. Zhao, W.X., Zhou, K., Li, et al.: A survey of large language models. arXiv:2303.18223 (2023) 38. Zhong, L., Wang, Z., Shang, J.: Debug like a human: A large language model debugger via verifying runtime execution step-by-step. arXiv:2402.16906 (2024) 39. Zhu, Q., Guo, D., Shao, Z., Yang, D., Wang, P., Xu, R., Wu, Y., Li, Y., Gao, H., Ma, S., et al.: Deepseek-coder-v2: Breaking the barrier of closed-source models in code intelligence. arXiv:2406.11931 (2024)