Multi-Turn Code Generation Through Single-Step Rewards

Arnav Kumar Jain, Gonzalo Gonzalez-Pumariega, Wayne Chen, Alexander M Rush, Wenting Zhao, Sanjiban Choudhury

Keywords: Code Generation, Language Models, Reinforcement Learning

Abstract:

We address the problem of code generation from multi-turn execution feedback.
Existing methods either generate code without feedback or use complex, hierarchical reinforcement learning to optimize multi-turn rewards.
We propose a simple yet scalable approach, $\mu$Code, that solves multi-turn code generation using only single-step rewards.
Our key insight is that code generation is a one-step recoverable MDP, where the correct code can be recovered from any intermediate code state in a single turn.
$\mu$Code iteratively trains both a generator to provide code solutions conditioned on multi-turn execution feedback and a verifier to score the newly generated code.
Experimental evaluations show that our approach achieves significant improvements over state-of-the-art baselines such as STaR.
We provide analysis of the design choices of the reward models and policy, and show the efficacy of $\mu$Code at utilizing the execution feedback.

Submission Number: 29