PaperlayerView exact revision on arXiv

Source-linked AI summary

Two-Stage Reinforcement Learning for Sound and Adversarial Test Generation in Code LLMs

Jiacheng Xu, Wentao Zhang, Zhiyi Lyu, Fuxiang Zhang, Chaojie Wang, Yang Liu, Bo An

arXiv:2609.03955v1cs.CLcs.LG

TL;DR

High-quality tests for code LLMs must be both sound and discriminative, but informative tests depend on the solver’s changing failure modes. TCS addresses this with two-stage RL and improves code generation and inference-time selection across TACO and LiveCodeBench, including selection among external LLM outputs.

  • Problem

    Test cases for code LLMs are scarce because they must be sound for a reference solution and discriminative against plausible incorrect programs.

  • Method

    TCS jointly trains code and test generation with two-stage RL: Stage 1 learns ground-truth-verified tests, and Stage 2 learns candidate-conditioned counterexamples from a policy-aligned buffer.

  • Results

    Across TACO and LiveCodeBench, TCS improves both code generation and inference-time answer selection, with experiments comparing against joint SFT, code-only RL, and test-only RL.

  • Takeaways & Limitations

    Self-generated tests are most useful when reliability is controlled, complement curated public tests, and can transfer to strong external LLM outputs.

  • Takeaways & Limitations

    Inference currently generates only one test case per call, while multiple-test generation is hindered by reward hacking risks; training also uses a hard stage switch.

Abstract

from arXiv · show

Reinforcement learning (RL) has substantially advanced code generation with large language models (LLMs) through executable feedback. The feedback for coding problems mainly comes from specific test cases, where high-quality test cases are often scarce since they should be both sound and discriminative. We thus turn to study the auto-generation of test cases using the learned model. We find this is naturally an adversarial RL problem: the model is expected to generate effective test cases as counterexamples, depending on the solver's current failure modes. We propose Test Cases Scaling (TCS), a two-stage RL framework for effective test generation. Both stages train a test generator from a rolling policy-aligned buffer: Stage 1 generates tests consistent with the reference solution, and Stage 2 restricts the buffer to current failure modes and learns counterexample tests. Across TACO and LiveCodeBench, TCS improves both pass@1 and inference-time answer selection according to generated tests. We find the learned test generator also enables effective selection among other LLM outputs.

1 Introduction

The paper frames test generation for code LLMs as an adversarial, execution-verifiable RL problem requiring both soundness and discrimination. TCS addresses this with two-stage training and reports gains in code generation, inference-time selection, and transfer to external LLM outputs.

  • Generated tests must be reliable and targeted because unsound tests can mis-rank candidates while weak tests miss subtle bugs.Curated public tests can still outperform self-generated tests when used alone, so the goal is better rather than simply more tests.
  • Test generation is nonstationary because informative tests depend on the solver’s current failure modes.The objective is multi-solution and execution-defined, unlike exact target matching.
  • TCS formulates effective test generation as requiring soundness control and candidate-conditioned adversariality in an execution-verifiable setting.Generated tests can be checked against a ground-truth solution before serving as adversarial tests.
  • TCS uses two-stage RL: Stage 1 learns ground-truth-verified tests, while Stage 2 learns candidate-conditioned counterexamples through stage-specific rewards and a policy-aligned buffer.Stage 2 tests pass the reference solution but fail plausible incorrect programs.
  • TCS provides theoretical and empirical evidence for inference-time scaling, including an exponential reliability bound and improvements over joint SFT, code-only RL, and test-only RL.The experiments cover TACO and LiveCodeBench.
  • The learned verifier also improves selection among outputs from strong external LLMs.

2 Background

The background defines code generation as execution-based verification over test suites and motivates test quality as central to both training and evaluation. It also introduces RL’s group-relative optimization setting and the requirement that tests be sound and discriminative.

  • 2.1 Reinforcement Learning: GRPO fine-tunes a policy using the average reward of multiple outputs from the same prompt as a baseline, avoiding a separate value model.For each prompt, it samples a group of outputs and computes group-relative advantages.
  • 2.1 Reinforcement Learning: GRPO uses an importance-sampling ratio, clipped ratio, KL regularization strength, and KL divergence in its policy objective.The advantage is based on response rewards within the sampled group.
  • 2.2 Code and Test Case Generation: Code generation produces a program C from problem description P that satisfies the specification, with candidates verified against test suites containing input-output pairs.
  • 2.2 Code and Test Case Generation: The standard code-generation reward is 1 when a program matches every expected output in the ideal test suite and 0 otherwise.
  • 2.2 Code and Test Case Generation: A useful test suite must be sound for the reference solution and discriminative enough to expose failures of plausible incorrect implementations.Generating tests satisfying both properties is the challenge targeted by TCS.

3 Test Cases Scaling (TCS)

TCS jointly trains a solver and verifier with online, policy-aligned data, using a staged curriculum that first enforces soundness and then targets candidate-specific failures. The resulting self-generated tests support inference-time candidate selection through pooled pass-counts.

  • Shared solver–verifier setup: TCS assigns the same shared policy two roles: generating candidate code and producing tests that are sound for the reference solution and discriminative for incorrect candidates.The verifier conditions on the problem and candidate program, while the solver is trained to pass the test suite.
  • Online policy alignment: The online RL loop fills a rolling buffer with current solver outputs, allowing verifier training to track the solver’s evolving failure modes rather than a fixed offline distribution.The buffer retains recent policy-aligned items and supplies verifier prompts during joint training.
  • Stage 1: Soundness: Stage 1 rewards executable tests that match the ground-truth solution and are not exact copies of prompt examples, prioritizing soundness before harder adversarial targeting.The main guardrail is execution-based verification against C*, while the example-test filter blocks exact reuse.
  • Stage 2: Counterexamples: Stage 2 clears the buffer and admits executable incorrect candidates, then rewards tests that pass C* but fail a sampled incorrect candidate.This candidate-conditioned reward targets current model failure modes; direct optimization can be sparse early because a test must be both sound and adversarial.
  • Inference-time selection: At inference, TCS generates M tests for each of N candidate programs, pools K = N × M tests, and selects the candidate with the highest pass-count.The score sums indicator values for executions matching each test’s expected output.
  • Theoretical guarantee: Under the stated sampling assumptions, net-discriminative tests make pass-count selection exponentially reliable as the pooled test count increases.The expected correct-versus-incorrect pass-count gap grows linearly with K and concentrates by Hoeffding/Chernoff bounds.

4 Experiments

Experiments on TACO and LiveCodeBench evaluate TCS across training-time performance, inference-time selection, verifier quality, and staged rewards. TCS improves solver and verifier behavior, with the strongest gains when soundness-controlled self-generated tests support candidate selection.

  • Experimental Setup: TCS is evaluated on filtered, execution-verifiable TACO data and LiveCodeBench using 1.5B and 7B DeepSeek-R1-Distill-Qwen models.The TACO training and evaluation setup retains 6,318 curated instances, with evaluation on 1,000 validation problems and LiveCodeBench.
  • Main Results: TCS improves pass@1 and self-generated-test selection over base and joint-SFT baselines, indicating benefits for both solver and verifier capabilities.The joint-SFT baseline already uses offline candidate-conditioned, ground-truth-verified adversarial supervision, isolating the value of online RL.
  • Main Results: Self-generated tests become competitive with or stronger than reward-model ranking, especially without public tests, after TCS training.For weaker or untrained models, self-generated tests can remain unreliable, so soundness control is central to this improvement.
  • Main Results: 43.01 pass@1 is reached from 28.56 for the LiveCodeBench 7B base model using self-generated test selection at BoN-32.Curated public tests alone reach 45.99, while combining public and self-generated tests reaches 46.15.
  • Test Case Effectiveness: TCS-generated tests provide stronger filtering than CodeRM-8B tests and larger downstream selection gains than tests from strong external models.TCS-7B also achieves higher test-output prediction accuracy, used as a proxy for soundness and consistency.
  • Main Results: Code-only RL improves pass@1 but does not reproduce TCS’s test-based selection gains, while joint TCS achieves the strongest training-time and inference-time performance.The results support a distinct role for explicit verifier training and a synergistic interaction between solver and verifier improvement.
  • Test Case Effectiveness: Stage 1 establishes verifier soundness before Stage 2 learns counterexamples, because optimizing the Stage 2 reward from the outset produces sparse rewards and ineffective learning.Stage 1 alone yields modest scaling gains, whereas the full two-stage reward produces clearer gains, particularly when public and self-generated tests are combined.

5 Related Work

The paper situates TCS among work on code-LLM post-training, test generation, verifier learning, solver-verifier co-training, and inference-time scaling. Its narrower focus is matched-budget evidence for soundness control and candidate-conditioned adversariality in execution-verifiable settings.

  • Code-LLM Post-Training: Prior code-LLM post-training emphasizes SFT and RL for improving solver-side code generation.Recent work also studies unit-test generation, verifier learning, and solver-verifier co-training.
  • Paper Focus: TCS focuses on matched-backbone, matched-budget evidence for soundness control and candidate-conditioned adversariality.These properties are operationalized through a staged RL objective in an execution-verifiable setting.
  • Inference-Time Scaling: Inference-time scaling commonly samples multiple candidates and selects among them using reward models, interpreters, search, or self-generated tests.TCS studies test generation as the learned auxiliary signal for candidate selection.

6 Conclusion

The conclusion argues that useful test generation requires both soundness control and candidate-conditioned adversariality, which TCS addresses through RL-based self-verification. Across TACO and LiveCodeBench, TCS improves code generation and self-verification, while self-generated tests work best when reliability is controlled and they complement curated tests.

  • Conclusion: TCS improves both code generation and self-verification across TACO and LiveCodeBench.Theory and ablations support stronger inference-time scaling when soundness error decreases and counterexample rate increases.
  • Conclusion: Effective test generation combines soundness control with candidate-conditioned adversariality because test generation is execution-defined, multi-solution, and policy-dependent.The ground-truth solution supports soundness, while evolving solver failures determine useful counterexamples.
  • Conclusion: Self-generated tests are most useful when reliability is controlled and when they complement rather than replace curated public tests.The conclusion also reports gains beyond offline joint-SFT and transfer benefits for strong external models.

Limitations

The method is constrained by single-test generation per inference call, hard stage switching, and the need for ground-truth solutions during post-training. The evaluation also relies on selected execution-verifiable benchmark settings.

  • Method Limitations: TCS currently generates only one test case per inference call, limiting efficiency for simultaneous multi-test generation.Multiple-test generation is left for future work because simple counting or accuracy rewards may be vulnerable to reward hacking.
  • Method Limitations: The training procedure uses a hard switch between stages, while soft switching was deferred because reinforcement learning for LLMs is resource-intensive.Dynamically adjusting the proportions of the two reward types is proposed as a possible improvement.
  • Scope Boundary: TCS assumes access to a ground-truth solution during post-training to verify generated tests, limiting use where verified solutions are unavailable.The ground-truth solution is not required during inference.
  • Evaluation Scope: The study uses execution-verifiable TACO and a chronologically separated LiveCodeBench evaluation window.TACO is filtered for test coverage and passing Python solutions, while LiveCodeBench covers problems from August 2024 to February 2025.

A.2 Experimental Configuration

The experiments use RL and SFT configurations for code and test-case generation, with TCS implemented as a rolling-buffer procedure that transitions from code-oriented to candidate-conditioned test generation.

  • Configuration: TCS uses verl with batch size 128, PPO mini-batches of 64, GRPO groups of 16, temperature 0.8, and maximum response length 8192.
  • Configuration: The implementation omits KL loss and uses entropy loss to sustain model entropy.
  • Training schedule: For 7B models, Stage 1 and Stage 2 use 40 and 160 training steps, while the code-generation baseline uses 200 steps.
  • Training schedule: The Stage 1-to-Stage 2 transition is a hard switch triggered when training-batch test-generation accuracy reaches approximately its specified threshold.
  • Rolling buffer: Algorithm 1 generates response groups, computes code rewards, and appends samples meeting current buffer criteria while removing old data beyond the buffer window.
  • Test generation: Stage 2 rewards a test when it exposes its paired incorrect program, without requiring a universal counterexample across other incorrect programs.

B Additional Theoretical Analysis and Proofs

The analysis formalizes pooled, candidate-conditioned test selection and derives an exponential reliability guarantee when counterexample strength exceeds test unsoundness under independence and boundedness assumptions.

  • Assumptions: The analysis assumes at least one correct candidate, valid generated inputs, and execution outcomes where runtime errors count as failures.
  • Formal setup: At inference, M tests are sampled conditionally for each of N candidates, pooled, scored by pass count, and used to select the highest-scoring candidate.
  • Quantities and proof: The pooled-mixture view averages candidate-specific soundness errors and counterexample rates, while the proof requires independent bounded pooled tests rather than identical distributions.
  • Reliability bound: The expected score gap is at least K(δ − α), and Hoeffding concentration followed by a union bound yields exponential reliability across incorrect candidates.
  • Failure modes: The guarantee requires δ > α; otherwise unsound or inconsistent tests can overwhelm the counterexample signal and prevent meaningful exponential concentration.
  • Failure modes: If generated tests are dependent, the proof must use martingale or effective-sample-size concentration, while the qualitative dependence on K and δ − α remains.

C.1 Inference-Time Scaling Comparison

Additional experiments show that TCS improves inference-time selection as candidate and test counts grow, retains an advantage with public tests, and benefits from jointly training the solver and verifier.

  • Inference-time scaling: Without public tests, TCS maintains consistent gains as sampled candidates increase, while other methods can degrade from distractors, OOD candidates, or noisy tests.
  • Inference-time scaling: With strong public tests, all methods improve as candidate count increases, but TCS retains a clear advantage and can augment public cases with targeted tests.
  • Test-count scaling: Increasing independently generated tests per candidate from M = 1 to M = 4 improves selection performance monotonically in all four settings.
  • Test-count scaling: M = 1 already improves pass@1 by 6.84–11.76 points, motivating its default use in the main experiments.
  • Joint versus decoupled training: Joint TCS improves both solver and verifier behavior, whereas Code-RL mainly improves the solver and Test-RL mainly improves the verifier.
  • Reward ablations: Stage 1-only training has limited scaling gains, while the full two-stage reward becomes increasingly advantageous, especially when public and self-generated tests are combined.
  • Task specialization: RL applied only to code generation produces limited test-generation gains, indicating that test generation requires dedicated RL.
  • Difficulty and transfer: Across TACO difficulty levels, TCS improves code generation over the base and self-verification baselines, with larger benefits on harder tasks.

D Generated Adversarial Test Case Examples

The examples illustrate how Stage 2 generates valid tests that the reference solution passes but targeted incorrect programs fail, including tests that filter errors missed by public cases.

  • Stage 2 example: A Stage 2 adversarial test is valid, passes the ground-truth solution, and exposes an error in the targeted incorrect solution.
  • Stage 2 example: The football-score example motivates testing overlapping intervals that may reveal an incorrect draw-count calculation.
  • Inference example: On a difficult AtCoder problem, one generated test filtered two incorrect solutions that passed all public tests and selected the only correct answer among three public-test-passing responses.

E Prompt Template

The appendix provides templates for code generation, test-case generation, and test-case output prediction, with configurable prompts designed to produce valid, diverse, and flaw-revealing tests.

  • Code Generation Prompt: The code-generation template instructs the model to produce a Python program that reads stdin, solves the specified problem, and writes to stdout.The prompt encloses the requested implementation within designated Python delimiters.
  • Test Case Generation Prompt: The prompt uses four randomly sampled test-case types to improve generation diversity.The types include basic, edge, corner, and performance tests.
  • Test Case Generation Prompt: Generated tests must differ from example tests and satisfy the problem’s constraints.The template explicitly prohibits reproducing provided examples and requires validity under the stated constraints.
  • Test Case Generation Prompt: The test-case generation prompt asks for one valid test case targeting potential flaws in a supplied solution, without writing solution code.It requires the input, expected output, and a brief explanation, using a specified JSON format.
  • Test Case Output Prediction Prompt: The output-prediction prompt separately asks the model to infer the expected output for a given test input in JSON format.It supplies the problem description and test input as inputs to the prediction task.
  • Example Reasoning Process: The worked example illustrates reasoning from game rules, constraints, and candidate states when constructing a test case.Its analysis considers card configurations and simulates moves to predict the winner.
Loading 2609.03955v1…