Multi-Turn Code Generation Through Single-Step Rewards

We thank the reviewer for the valuable feedback. We are glad the reviewer found the idea of one-step recoverability interesting. We respond to the questions below.

Results on POMDP setting.

This is an interesting ablation as for many prompts, some unit tests are present in the $x$. To evaluate the agents without this information in prompt, we removed the unit tests from the prompt $x$ in the test set of MBPP dataset and report the turn-wise BoN@3 score. We observe that the performance drops for all algorithms. We also observe that the performance gains for $\mu$Code is significantly higher than baselines across turns, demonstrating the efficacy of our proposed approach at leveraging execution feedback.

Other benchmarks beyond MBPP and HumanEval

We evaluate the methods on the CodeContests dataset and is discussed in the response to Reviewer QsKU. We show that $\mu$Code outperforms Multi-STaR by 1.6%.

The results in Table 2 where $\mu$-Code-R is worse than $\mu$Code.

To generate the dataset for training $\mu$Code-R, we concatenate the datasets used to train $\mu$-Code and MultiSTaR^2. We believe a better way of merging might lead to reduce this gap in performance and plan to add an additional analysis in the next version.

Whether $\mu$Code works for Iteration number >= 3.

In Table 1, we observe a performance drop with $\mu$Code on the MBPP dataset and observe improvements on HumanEval with training iterations. Given the simplicity of our approach, we believe more iterations would be needed for scaling to more challenging datasets and we leave this investigation for future research.

Theoretical claims

We thank the reviewer for pointing this out. Using the reward function defined in Section 2 (L100) to take values of either 0 or 1, we have updated the paper with the $-1 \leq A*(s,a) \leq 0$ in Definition 3.1, L218 and L733.

Whether StaR/Multi-StaR setting is closer to the original StaR or Rejection Sampling Finetuning.

The Multi-STaR baseline used in the paper is similar to the Rejection Sampling Finetuning defined in [2]. We plan to change the name of Multi-STaR to RFT for more clarity, and also change Multi-STaR^2 to RFT-R to denote RFT with relabeling.

Is the new policy for iteration N is always fine-tuned from the base model?

In this work, we followed a similar strategy as [1] for fine-tuning the base model at every iteration.

In Eq(3), how are the correct/incorrect pairs constructed?

The dataset for training the verifier was obtained through a sampling process. For each prompt, a pair of correct and incorrect response was obtained and we generated 16 pairs for each prompt.

In Eq(5), do you do the relabelling for all the problems?

Yes, $\mu$Code relabels all the problems even if there are no correct solutions to that problem. This is similar to how RL training would happen with a dense reward function, where the agent is updated to increase the likelihood of the response with higher reward values and climb uphill on the reward function with updates. $\mu$Code relabels the trajectories with solutions with higher reward values for fine-tuning.

What are the sizes of the datasets for each variants in Table 2 or other statistics of the dataset?

At each iteration, the size of the dataset for $\mu$Code is the total number of rollouts as we relabel every trajectory. For the Multi-STaR^2, we relabel ~70% of the rollouts denoting the fraction of prompts with at least one positive trajectory. The dataset for $\mu$Code-R is obtained by concatenating the two datasets. We ran hyperparameter tuning on the number of epochs to account for different sized datasets but finetuning with 2 epochs performed best.

Reproducibility:

We will be releasing the codebase and the model checkpoints with the paper. We provide details of important hyperparameters for training $\mu$Code below:

Essential References Not Discussed

[2] is a very interesting work that studies different CoT approaches and execution feedbacks. We have added [2] in references as the Multi-STaR baseline is similar to the Rejection Sampling Finetuning method described in [2].

Best Regards,

The Authors

References

[1] Qu et al.,Recursive Introspection: Teaching Language Model Agents How to Self-Improve, NeurIPS 2024

[2] Zheng et al., What Makes Large Language Models Reason in (Multi-Turn) Code Generation?, ICLR 2025

We thank the reviewer for their insightful feedback and are glad that you found the paper well organized. We are grateful for your questions and comments, and respond to your questions below.

While the paper claims that RL methods are inefficient, it fails to provide comparative experimental evidence.

Most approaches for finetuning LLM agents via RL rely on rule-based verifier like execution feedback in code or correctness in math problems. Such signals are sparse and credit assignment due to long-term rewards is challenging in multi-turn tasks. We propose $\mu$Code that leverages a learned verifier to provide dense signals and also uses the one-step recoverability property for better credit assignment.

Sparse Rewards: Let's consider a scenario where the agent is trained over the outcome of code execution. In such a setting, the reward can be either 0 or 1 denoting failure or success. When a correct solution is observed, an RL algorithm like PPO would increase the likelihood of this response. However, the agent is not updated for responses where it did not generate any correct solution for the prompt while sampling, and thereby the completions are discarded. This makes outcome based RL inefficient as the model updates relies on generating a correct solution, and exploration to a correct solution can be challenging. In many RL applications, this is mitigated with a dense reward signal, and $\mu$Code leverages a learned verifier for training. The generator is updated to predict solutions with the highest reward and with training should climb uphill this reward function. To demonstrate the benefits, we compare Multi-STaR^2 which is similar to $\mu$Code and performs relabeling with the sparse and ground truth reward (in Table 1 of the paper), and observed that $\mu$Code trained with dense rewards performed better.

Long-term Rewards: During RL training, the agent is trained to optimize the long-term rewards of the trajectories in multi-turn tasks. With the insight of one-step recoverability, we propose to relabel trajectories and improve credit assignment. To show this, we propose a baseline, Rejection Fine-Tuning, where we finetune the agent with the top rollouts ranked via the learned verifier (and call it RFT (LV)). We did a hyperparameter search over different number of top rollouts for this baseline and found picking top 6 solutions from 15 rollouts worked best. We observe that RFT (LV) does not perform well when compared with $\mu$Code.

MBPP

HumanEval

Verifier annotations may contain inherent biases.

We agree that the agent can exploit the inaccuracies within the learned verifier. This happens because the distribution of reponses from the generator shifts with training and the learned verifier can return inaccurate reward values. To account for the covariate shift, $\mu$Code updates the learned verifier after every iteration. Since the ground truth reward is already known (from the execution feedback), the verifier is updated using the oracle outcome. We would like to highlight that the learned verifier in $\mu$Code is trained with the concatenated dataset from prior iterations. In our experiments, we also observed that updating the verifier across iterations is beneficial for learning.

We additionally present emperical results with the generator of $\mu$Code and the learned verifier at the last iteration vs the learned verifier at iteration 0 (on rollouts from Instruct model). We observe upto 2% drops in performance on BoN@3 when using the verifier trained on Instruct rollouts vs the verifier updated with rollouts from current generator.

The conclusions, therefore, warrant further validation on more comprehensive datasets.

We conduct an experiment with the CodeContests dataset. We curated a 1K-sized dataset with randomly selected problems from the training set of CodeContests. We use Llama-3.1-8B-Instruct and compare with MultiSTaR and $\mu$Code which were trained for 1 iteration on 4 X A100 GPUs with 80GB memory. We report the Pass@3 and BoN@3 results on the test set in CodeContests with 165 problems. We observe that $\mu$Code outperforms Multi-STaR by 1.6% on BoN scores showing the applicability of $\mu$Code to more challenging benchmarks.

Best Regards,

The Authors

We thank the reviewer for the valuable feedback and are glad that you found our approach simple and theoretical justification insightful. We address your questions below.

I would like to see a more in-depth discussion regarding this claim, rather than treating it merely as a basic assumption.

For the task of code generation, there are two major sources of errors: 1) syntax errors and 2) algorithmic errors. For (1), we expect gradual improvement because fixing the errors should lead to a better solution at the next turn. However, the change might not be gradual for (2) when the approach is incorrect. In this scenario the agent might have to modify the algorithm and as pointed out by the reviewer can be difficult to recover from. Since, the task of code generation is one-step recoverable, this allows $\mu$Code to relabel with a better solution and thereby recover from harder scenarios. In the paper, we provided the theoretical justification to support this claim in Theorem 3.2 (Section 3.5) and our experiments demonstrate the efficacy of $\mu$Code. Below, we present an example where $\mu$Code successfully resolves an algorithmic error.

Prompt: Write a function to calculate the height of the given binary tree.

Turn 1

def max_height(root):
    if root is None:
        return 0
    
    queue = [(root, 1)]
    
    while queue:
        node, height = queue.pop(0)
        if node is None:
            return height
        queue.append((node.left, height + 1))
        queue.append((node.right, height + 1))
        
    return -1

Turn 2

def max_height(root):
    if root is None:
        return 0
    
    left_height = max_height(root.left)
    right_height = max_height(root.right)
    
    if left_height > right_height:
        return left_height + 1
    else:
        return right_height + 1

Note that the first solution takes an iterative approach and is incorrect because the function returns as soon as any leaf node is reached. However, the code solution at next iteration switches from an iterative to a simpler recursive approach and is correct.

Reproducibility

We will be releasing the codebase, model checkpoints and hyperparameters with the final version of the paper.

Can the method described in this paper be extended to scenarios beyond code generation to construct trajectories?

Yes, our method can be extended to any tasks where the correct solution can be generated in a single step (one-step recoverability). An application is theorem proving, where the the agent needs to write proof of theorems and the solution can be verified using external tools [1,2].

Could you consider using the additional tests introduced by HumanEval+ and MBPP+ as a private test set?

We thank the reviewer for bringing this up. In the submitted version, we used the HumanEval and MBPP datasets with its own set of public and private test split. Below, we provide comparison where private tests include the additional tests provided in MBPP+ and HumanEval+. We report results with the same model 1B-sized model checkpoints used to report results in Table 1 and present the results on Pass@3 and BoN@3.

MBPP+

HumanEval+

We observe that $\mu$Code outperforms baselines Multi-STaR by 5.1% and 5.7% on BoN@3 on MBPP+ and HumanEval+ datasets. Note that the numbers for MBPP are higher than what is reported in Table 1 because MBPP+ comprises of 378 problems as compared to 500 problems in the MBPP dataset.

Bigger benchmark

We conduct an experiment on CodeContests dataset and show that $\mu$Code outperforms Llama-3.1-8B-Instruct and MultiSTaR. The experiment is discussed in the response to Reviewer QsKU, where we show 1.6% performance gain over Multi-STaR at BoN@3.

Most of the reported results in Table 1 for HumanEval do not yield integers when multiplied by 164.

We are using a temperature of 0.7 for sampling to report our metrics. Because of the smaller dataset size, we observed variance in the results and resorted to generating 15 rollouts for each prompt. We report the mean scores in Table 1 and the values in Table 1 will be a multiple of .6 / 15 = 0.04. In the revised version, we plan to add the standard error with the mean results for clarity.

Best Regards,

The Authors

References

[1] STP: Self-play LLM Theorem Provers with Iterative Conjecturing and Proving, Dong el al. 2025.

[2] Learning Formal Mathematics From Intrinsic Motivation, Poesia et al., NeurIPS 2024.

Thank you for your constructive feedback. We are glad that you found the paper well structured and our experiments thorough. We address your questions below.

Investigate other LLMs that are not Llama, and include larger models.

We investigate the performance of the Qwen-2.5-1.5-Instruct model, Multi-STaR and $\mu$Code at Pass@3 and BoN@3 metrics on MBPP and HumanEval datasets. Below, we present the table and show that $\mu$Code outperforms the baselines.

MBPP

HumanEval

Our experiments are restricted to models that can be trained on 4 x A100 GPUs with 80 GB memory each and training of larger models greater than 8B requires more GPU memory. We expect $\mu$Code should scale to larger models with more compute.

Investigate the performance when the turn limit is above 3. > Same with the "training" time following algorithm 1.

In the paper, we ran BoN search for a turn limit of 3. We present the results with a turn limit of 6 below (with the 1B-sized model used in Table 1). Note that the model is trained with a turn limit of 3. $\mu$Code achieves the best performance on both datasets across most turn limits. This highlights the versatility of $\mu$Code at generalizing to larger turn limits at test-time. Additionally, our results show that $\mu$Code achieves the largest gains per turn, suggesting it is better at incorporating execution feedback and improving code solutions across turns. This experiment validates that $\mu$Code is performant with more turns during evaluation. We believe that training with a higher turn limit could further improve the performance.

MBPP

HumanEval

Fundamentally, the only baseline used is STaR.

We point the reviewer to our response to Reviewer QsKU where we have provided an additional baseline: Rejection Fine-tuning with a learned verifier, or RFT (LV). Here, we select the top trajectories ranked using the learned verifier for a prompt and finetune at each iteration. This is similar to the Multi-STaR baseline, where the learned verifier is used to filter trajectories instead of the ground truth reward. We observed that $\mu$Code performs better than RFT (LV) baseline.

Best regards,
The Authors