PaperlayerView exact revision on arXiv

Source-linked AI summary

Multi-Turn Code Generation Through Single-Step Rewards

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

arXiv:2502.20380v2cs.LGcs.AIcs.CL

TL;DR

Multi-turn code generation must learn from execution feedback, but prompting alone is insufficient and existing multi-turn approaches rely on complex reinforcement learning. µCODE uses one-step recoverability with an iteratively trained generator and verifier, and outperforms leading multi-turn approaches on several code benchmarks.

  • Problem

    Prompting alone cannot teach recovery from all execution errors within limited context, and existing multi-turn methods rely on complex reinforcement learning.

    Source: §1 · p. 1

  • Method

    µCODE treats code generation as one-step recoverable and iteratively trains a generator and verifier using execution-feedback rollouts, verifier-guided local search, and expert labels.

    Source: §1 · p. 1

  • Results

    µCODE outperforms leading multi-turn approaches on MBPP, HumanEval, and CodeContests.

    Source: §1 · p. 2

  • Takeaways & Limitations

    One-step recoverability permits greedy single-step reward optimization instead of multi-step reinforcement learning, reframing training as imitation learning.

    Source: §1 · p. 1

    Models were trained only up to eight billion parameters and on Python datasets, so the findings might not generalize to larger models or other programming languages.

    Source: §6 · p. 9

Abstract

from arXiv · show

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 the state-of-the-art baselines. 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. Our code is available at https://github.com/portal-cornell/muCode.

1. Introduction

Prompting alone cannot teach recovery from all execution errors within limited context, while multi-turn reinforcement-learning approaches face sparse learning signals. µCODE uses one-step recoverability to train code generation from execution feedback with a generator and learned verifier.

1. Introduction

  • Prompting alone is insufficient for learning recovery from all execution errors within limited context, while multi-turn reinforcement learning uses sparse long-term rewards.The paper identifies a need to train models to learn from execution feedback during training.Source: §1 · p. 1
  • One-step recoverability means correct code can be recovered from any intermediate state in a single step, allowing greedy single-step reward optimization instead of multi-step reinforcement learning.The paper says this reframes the problem as imitation learning, producing a more stable and efficient training process.Source: §1 · p. 1
  • µCODE iteratively collects execution-feedback rollouts, trains a single-step verifier, uses it to guide a local-search expert, and fine-tunes the generator on the resulting labels.The verifier also ranks sampled trajectories for multi-turn Best-of-N selection at inference time.Source: §1 · p. 1, Figure 1 · p. 2
  • The contributions include theoretical performance-bound analysis using one-step recoverability and multi-turn Best-of-N inference-time scaling with learned-verifier selection.Source: §1 · p. 1
  • µCODE outperforms leading multi-turn approaches on MBPP, HumanEval, and CodeContests.The paper also reports that ablations show learned verifiers aid generator learning and that higher inference budgets show promising scaling trends.Source: §1 · p. 2

2. Background

The paper models multi-turn code generation as an MDP in which the interaction history is the state and a complete code snippet is the action. It optimizes discounted correctness reward and relies on the ability to recover the correct code in one step.

2. Background

  • At each turn, the agent generates a complete code snippet from the prompt and prior execution outcomes; public-test outcomes guide refinement until success or the turn limit.The first code passing public tests is evaluated on private tests for a binary correctness score.Source: §2 · p. 2
  • The MDP state is the interaction history and its action is the current code snippet; reward is correctness on public and private tests, otherwise zero.Source: §2 · p. 2
  • The training objective maximizes cumulative discounted reward, favoring a correct solution in fewer turns.The stated discount factor γ is in [0, 1).Source: §2 · p. 2
  • Because the correct code can be generated at any step, the process is one-step recoverable; the algorithm iteratively collects rollouts, trains a verifier, constructs a local-search expert, and relabels data.Source: §2 · pp. 2–3

3. µCODE: Multi-turn Code Generation

µCODE alternates between training a verifier to score code and a generator to imitate verifier-guided local search, then uses both for multi-turn inference. Its analysis shows how one-step recoverability supports imitation learning with a tighter performance bound than for unrecoverable MDPs.

3. µCODE: Multi-turn Code Generation

  • µCODE iteratively trains a verifier and generator through expert iteration, then uses both for inference-time Best-of-N search.Rollouts supply interaction data; a local search expert relabels code solutions for generator fine-tuning.Source: §3 · p. 3

3.1. The µCODE Algorithm

  • At each iteration, the current generator produces interaction data, the verifier and expert are updated, and the generator is fine-tuned on relabeled data.The best generator-verifier pair on validation data is returned after M iterations.Source: §3.1 · p. 3

3.2. Training Verifier

  • The verifier scores code from the prompt and current solution, using supervised oracle-reward labels rather than prior interaction states.The paper compares BCE reward prediction with Bradley-Terry training on correct and incorrect solution pairs, hypothesizing relative ranking is easier to optimize.Source: §3.2 · p. 3

3.3. Training Generator

  • The expert selects each prompt’s best collected solution using a weighted combination of oracle and learned verifier scores, then relabels data for generator fine-tuning.The weights are βO = 1.0 and βL = 0.1; combining both verifiers yielded better experimental performance.Source: §3.3 · p. 3

3.4. Inference: Multi-turn Best-of-N

  • At each inference turn, the generator proposes N solutions and the verifier selects one for execution; feedback conditions the next turn until public tests pass or the turn limit is reached.Even a high verifier score does not end the search before public tests pass.Source: §3.4 · p. 4

3.5. Analysis

  • One-step recoverability gives µCODE a performance bound of O(ϵT), compared with the O(ϵT^2) worst case for unrecoverable MDPs.The analysis attributes this setting to history-independent optimal actions and recovery from the next state, supporting imitation learning without hierarchical credit assignment.Source: §3.5 · p. 4

4. Experiments

Across MBPP, HumanEval, and CodeContests, µCODE outperforms single- and multi-turn baselines, with larger gains under multi-turn Best-of-N search. The analyses examine verifier-guided training and selection, execution feedback, inference-time scaling, and verifier loss choices.

4.1. Setup

  • Experiments train on MBPP and sampled CodeContests problems, then evaluate on MBPP test, HumanEval, and CodeContests test sets.The MBPP training set contains 374 problems; the CodeContests training sample contains 1,000 problems.Source: §4.1 · p. 5

4.2. Results

  • 2.2% on MBPP and 1.3% on HumanEval: µCODE beats RFT at 1B BoN N=1 and also outperforms single- and multi-turn baselines across benchmarks.At 1B multi-turn BoN, µCODE exceeds baselines by 4.4% on MBPP and 5.8% on HumanEval; on CodeContests it beats RFT and the base model by 1.4%.Source: Table 1 · p. 5, §4.2 · pp. 5–6

4.3.1. WHAT MAKES A GOOD GENERATOR IN µCODE?

  • RFTLV using learned-verifier scores beats oracle-filtered RFT, while combining oracle and learned scores for relabeling outperforms either variant by around 2%.The learned verifier can select top-K rollouts even for prompts with no correct solutions, supplying dense rewards for generator training.Source: §4.3.1 · p. 6

4.3.2. DOES LEARNED VERIFIER AID BON SEARCH?

  • LV beats random selection, PT beats LV except similarly on 8B HumanEval, and PT+LV performs best across methods and datasets.The authors hypothesize that learned scores help distinguish candidates that pass public tests and select the most promising candidate when all are incorrect.Source: §4.3.2 · pp. 6–7, Table 3 · p. 6

4.3.3. CAN µCODE UTILIZE EXECUTION FEEDBACK?

  • 15.9%: µCODE's BoN accuracy rises from turn 1 to turn 6 on MBPP (POMDP), matching MBPP performance, while base and RFT remain about 6% lower.All methods initially lose up to 13% when public tests are removed from the prompt; µCODE closes the gap as turns accumulate.Source: §4.3.3 · p. 7, Figure 2 · p. 7

4.3.4. DOES µCODE SCALE WITH INFERENCE BUDGET?

  • BoN performance improves with larger N on both datasets, plateauing at N ≥ 5 on HumanEval while continuing to gain on MBPP.The verifier selects among N candidates at each turn to avoid the exponentially growing candidate count of multi-turn search.Source: §4.3.4 · pp. 7–8, Figure 3 · p. 7

4.3.5. LOSS FUNCTION FOR VERIFIER

  • BT-trained verifiers consistently outperform BCE-trained verifiers on MBPP and HumanEval as candidate count increases.The comparison uses base-model-generated samples and reports BoN performance after the verifier ranks candidates.Source: §4.3.5 · p. 8, Figure 4 · p. 8

4.3.6. QUALITATIVE RESULT

  • In the qualitative example, execution feedback yields two correct responses at turn three, and the verifier selects one over an incorrect candidate.The example uses five candidates per turn and displays the top three ranked solutions.Source: §4.3.6 · p. 8

5. Related Work

Prior work tackles multi-step coding through prompting agents, trained critics or reward models, and methods without learned critics. The paper contrasts these approaches with µCODE’s use of verifier-guided search and execution feedback.

5. Related Work

  • Prompting-based agents use different strategies for iterative coding, but the cited work reports that they yield limited improvements.Examples include Self-Debugging, CodeT, AlphaCodium, and MapCoder.Source: §5 · p. 8
  • CodeRL trains a repairer for one erroneous completion, whereas µCODE’s generator incorporates multiple turns; ARCHER instead uses a two-level hierarchical MDP.The passage also describes MCTS-based training and inference-time verifier approaches.Source: §5 · p. 8
  • In the Figure 5 example, the verifier selects a better solution at each turn, and execution feedback accompanies successive improvements until a solution passes all tests.Two correct solutions appear at the final turn.Source: Figure 5 · p. 9
  • RLEF and PPOCoder omit learned critics or reward models, while µCODE’s comparison with RFT is presented as empirical evidence favoring a reward model.The passage also contrasts µCODE’s training efficiency with SCoRe and RISE.Source: §5 · p. 9

6. Conclusion

The conclusion presents µCODE as a single-step-reward approach to multi-turn code generation, using a learned verifier to guide iterative code improvement. The authors report strong comparisons while noting limits on model scale, training data, and language coverage.

6. Conclusion

  • µCODE models code generation as a one-step recoverable MDP and iteratively improves code using a learned verifier to guide search.The conclusion describes the method as simple and scalable.Source: §6 · p. 9
  • The authors report that µCODE outperforms methods using oracle verifiers by a large margin.Source: §6 · p. 9
  • The evaluation trains models only up to eight billion parameters, uses MBPP training data with 374 examples, and covers Python alone.The authors state that findings may not generalize to larger models or other programming languages.Source: §6 · p. 9

Impact Statement

The impact statement describes potential productivity gains from automating routine coding tasks, alongside a risk that bugs could spread without quality control.

Impact Statement

  • Automating routine coding tasks could reduce human labor and accelerate production timelines.Source: Impact Statement · p. 9
  • Without proper quality control, bugs introduced by these advances could propagate at scale.Source: Impact Statement · p. 9

A.1. Proof of Theorem 3.2

The proof invokes the Performance Difference Lemma and DAGGER’s connection between interactive imitation learning and no-regret online learning. It also uses the one-step recoverability advantage bound in its performance analysis.

A.1. Proof of Theorem 3.2

  • The proof uses the Performance Difference Lemma, which expresses the performance difference between policies as a sum of advantages.Source: §A.1 · p. 13
  • The appendix defines the advantage with respect to π′ as Qπ′(st, at) − Vπ′(st), and identifies the induced state distribution under π.Source: §A.1 · p. 13
  • For a one-step recoverable MDP, the optimal-policy advantage is bounded between -1 and 0 for every state-action pair.The passage says this bound is then used to bound the performance difference.Source: §A.1 · p. 13
  • The proof invokes DAGGER to reduce interactive imitation learning to no-regret online learning and cites its guarantee after N iterations.Source: §A.1 · p. 13

B.1. Qwen Model

This section presents Qwen-model comparisons, evaluations with additional private tests, and hyperparameters for SFT and reward-model training.

B.1. Qwen Model

  • Table 4 compares µCODE with baselines on MBPP and HumanEval using Qwen-2.5-1.5B-Instruct.The caption does not report the table's cell values or identify a winning method.Source: Table 4 · p. 14

B.2. Additional private tests

  • MBPP+ and HumanEval+ use 35x and 80x more tests than their original counterparts, respectively.The caption also notes that MBPP+ scores higher than MBPP because it contains 30% fewer problems.Source: Table 5 · p. 14

C. Hyperparameters

  • Table 6 lists hyperparameters for SFT and reward-model training.The caption does not provide the hyperparameter values.Source: Table 6 · p. 14

C.0.1. TRAINING PARAMETERS

The appendix describes training and inference settings, prompt templates, and the public-private test split used in the experiments.

C.0.1. TRAINING PARAMETERS

  • µCODE training uses two iterations, each starting from the base model, with runs on four RTX 6000 Ada GPUs for 1B models or four H100 GPUs for 8B models.The reported training covers Llama-3.1-8B-Instruct and Llama-3.2-1B-Instruct on MBPP and HumanEval.Source: §C.0.1 · p. 14

C.0.2. INFERENCE PARAMETERS

  • Inference uses SGLang, with temperature 0 for greedy experiments and 0.7 for BoN search; completions are capped at 1000 tokens per turn.Greedy settings are intended to ensure deterministic results.Source: §C.0.2 · p. 14

C.1.1. SINGLE STEP PROMPT

  • The single-step prompt section provides code-completion prompts and problem and test examples, with instructions to return implementation code in a single code block.Examples include HumanEval, MBPP, and CodeContests; the instructions also say not to include tests in the code.Source: §C.1.1 · pp. 15–16

C.1.2. FEEDBACK PROMPT

  • Multi-step feedback consists only of executor error traces, illustrated with a HumanEval example.The example includes an assertion error.Source: §C.1.2 · pp. 16–17

C.2. Public Private Tests

  • For HumanEval and MBPP, one public test is used, while remaining tests and specified official-suite tests are reserved as private tests.HumanEval uses a docstring test publicly and the remaining tests plus the official suite privately; MBPP uses one official test publicly and the rest plus challenge-list tests privately.Source: §C.2 · p. 17

Related papers

Loading 2502.20380v2…