Enhancing LLM Agents for Code Generation with Possibility and Pass-rate Prioritized Experience Replay
Yuyang Chen, Kaiyan Zhao, Yiming Wang, Ming Yang, Jian Zhang, Xiaoguang Niu
TL;DR
This paper tackles inefficiency in code-generation by addressing sparse rewards in LLMs and introduces the BTP pipeline, which couples beam search sampling, a testing phase, and a P2Value-guided prioritized experience replay (PPER) to fine-tune models. A replay buffer stores tuples $(X,S,t_i,P(t_i),pass\_rate_i)$ and sampling is driven by $\text{P2Value} = \alpha \cdot P(t_i) + (1 - \alpha) \cdot pass\_rate_i$, balancing output likelihood with test success. Empirically, BTP improves code-generation performance across APPS, CodeContests, and HumanEval for baseline models and also enables self-improvement through self-sampling, with mixing datasets further boosting generalization. The approach offers a practical path to more efficient and robust code-generation LLMs, especially in scenarios with sparse rewards or limited high-quality fine-tuning data. Limitations include dependence on available test suites, prompting future work to expand test resources and integrate with complementary learning signals.
Abstract
Nowadays transformer-based Large Language Models (LLM) for code generation tasks usually apply sampling and filtering pipelines. Due to the sparse reward problem in code generation tasks caused by one-token incorrectness, transformer-based models will sample redundant programs till they find a correct one, leading to low efficiency. To overcome the challenge, we incorporate Experience Replay (ER) in the fine-tuning phase, where codes and programs produced are stored and will be replayed to give the LLM agent a chance to learn from past experiences. Based on the spirit of ER, we introduce a novel approach called BTP pipeline which consists of three phases: beam search sampling, testing phase, and prioritized experience replay phase. The approach makes use of failed programs collected by code models and replays programs with high Possibility and Pass-rate Prioritized value (P2Value) from the replay buffer to improve efficiency. P2Value comprehensively considers the possibility of transformers' output and pass rate and can make use of the redundant resources caused by the problem that most programs collected by LLMs fail to pass any tests. We empirically apply our approach in several LLMs, demonstrating that it enhances their performance in code generation tasks and surpasses existing baselines.