Prompt Optimization for LLM Code Generation via Reinforcement Learning

2026-05-18

AI Summary

This work addresses the strong dependence of large language models’ code generation performance on prompt quality and the lack of efficient automated prompt optimization methods. The authors formulate prompt optimization as a sequential decision-making problem and propose a reinforcement learning framework that iteratively refines prompts through a hybrid action space comprising direct generation, genetic-style mutation, and semantic rewriting. A shaping reward mechanism based on unit test feedback guides the optimization process. To the best of our knowledge, this is the first approach to combine reinforcement learning with a hybrid action space for prompt optimization. Using the PPO algorithm, the method is trained on frozen-weight models—including CodeT5+, CodeLLaMA, and DeepSeek-Coder—and achieves substantial improvements over baselines such as EPiC and Reflexion on MBPP+, HumanEval+, and APPS benchmarks, attaining Strict Pass@1 scores of 57.58%–85.50% and Soft-Pass@1 scores of 67.90%–88.20% on MBPP+.

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 DeepSeek-Coder 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. Soft-Pass@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.

Problem

• Prompt Optimization
• LLM Code Generation
• Functional Correctness
• Reinforcement Learning
• Test-Driven Rewards

Innovation

• Reinforcement Learning
• Prompt Optimization
• Code Generation
• PPO
• Test-Driven Reward