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

AI summary

Overview

• Research area: Reinforcement learning for code large language models (LLMs), specifically automatic test-case generation, self-verification, and inference-time scaling.
• Technical level: Advanced. The paper assumes familiarity with policy-gradient RL (GRPO/PPO-style objectives, importance ratios, KL regularization) and with execution-based code evaluation.
• Scope: The paper proposes and evaluates Test Cases Scaling (TCS), a two-stage RL framework that trains a single LLM to act both as a code solver and as a candidate-conditioned test generator, and analyzes when the resulting self-generated tests make pass-count selection reliable.

What This Paper Is About

Code LLMs can be trained with execution feedback, but that feedback depends on test cases that are often scarce, and useful tests must be both sound (agreeing with a ground-truth solution) and discriminative (exposing plausible incorrect code). The authors recast test generation as an adversarial RL problem, where tests should target the solver's current failure modes rather than a fixed offline distribution. TCS is their answer: a two-stage RL framework that first learns ground-truth-verified (sound) tests and then learns candidate-conditioned counterexample tests, with the learned tests also used to select among candidate programs at inference time.

Key Contributions

1. Problem framing. The authors argue that effective test generation for code LLMs requires both soundness control and candidate-conditioned adversariality, and formulate this in an execution-verifiable post-training setting where generated tests are checked against a ground-truth solution $C^{*}$.
2. TCS framework. They propose Test Cases Scaling, a two-stage RL framework with stage-specific rewards — Stage 1 for ground-truth-verified soundness (Eq. 3) and Stage 2 for candidate-conditioned counterexamples (Eq. 4) — trained from a rolling, policy-aligned buffer that retains only items from the most recent $T_b$ training steps.
3. Theory for test-based inference-time scaling. They derive an exponential reliability bound for pass-count selection under self-generated tests, showing mis-selection probability is at most $(N-1)\exp(-K(\delta-\alpha)^2/2)$ under an independent stratified sampling and net-discriminativeness assumption ($\delta > \alpha$).
4. Empirical validation. On TACO and LiveCodeBench, TCS improves pass@1 and inference-time answer selection over joint SFT, code-only RL, and test-only RL, and the learned verifier also improves selection among outputs from other (including strong external) LLMs.

Main Findings

• Training-time and inference-time gains. TCS improves pass@1 for both DeepSeek-R1-Distill-Qwen-1.5B and 7B relative to base and joint-SFT baselines, and yields the strongest gains when self-generated tests are used for selection. For example, the 7B model's average across TACO/LiveCodeBench settings with public tests rises from 37.32 (SFT + reward model) to 41.30 (TCS + reward model), and from 37.40 (SFT + self-generated tests) to 44.57 (TCS + self-generated tests).
• Self-generated tests become competitive after TCS. For the base model, self-generated tests are often weaker than reward-model ranking (e.g., 1.5B average 17.95 vs. 20.18). TCS reverses this: 1.5B average 28.02 for self-generated tests vs. 25.21 for reward-model ranking; 7B average 44.57 vs. 41.30.
• Public tests and self-generated tests are complementary, not substitutes. On LiveCodeBench, the 7B base model improves from pass@1 28.56 to 43.01 with self-generated test selection at BoN-32; curated public tests alone reach 45.99, and combining public with self-generated tests works best at 46.15.
• A small TCS model can beat a much larger baseline at inference time. Using DeepSeek-R1-Distill-Qwen-14B as a strong baseline, the TCS-fine-tuned 1.5B model can outperform the 14B baseline under the same selection rule, and TCS gives a more robust test-based scaling signal as $N$ grows, while reward-model selection becomes less stable as $N$ grows.
• Online RL beats offline adversarial supervision. The joint-SFT baseline already receives candidate-conditioned, ground-truth-verified adversarial supervision from DeepSeek-R1-Distill-Qwen-32B samples, but TCS's policy-aligned online updates produce a consistent gap over it.
• Joint training beats decoupled training. On R1-Distill-Qwen-1.5B, code-only RL improves pass@1 but adds little under "+TC"; test-only RL yields a stronger verifier (larger gains with "+TC" but modest direct pass@1); joint TCS achieves both the best training-time and the largest inference-time gains.
• Both stages are necessary. Stage 1-only training yields only modest test-based scaling gains under the same budget. Directly optimizing the Stage 2 reward $R_2^t$ from the outset yields sparse rewards and ineffective counterexample learning for many steps (e.g., for R1-Distill-Qwen-1.5B), motivating the curriculum.
• Not just any external generator helps. Under the same inference-time budget, tests generated by CodeRM-8B yield weaker filtering performance than both reward-model ranking and TCS's self-generated tests.
• Evidence on test quality. Tests generated by TCS-7B produce larger downstream selection gains on LiveCodeBench than those from strong external models, and TCS models achieve higher accuracy on the official LiveCodeBench Test Output Prediction task, which the authors use as a proxy for soundness and consistency.
• Harder problems benefit more. Gains persist across TACO difficulty levels and become more pronounced on harder subsets (Appendix Table 7).
• Cost asymmetry. Using decoded tokens as a hardware-independent proxy, one generated test uses 22–34% as many tokens as one code candidate in the Table 1 evaluation responses (22–24% for 1.5B and 32–34% for 7B), and independent test-generation calls can run in parallel.

Methodology in Plain English

Each training example provides a problem description, a ground-truth solution, and a test suite. The same model plays two roles: a solver that writes code to pass the tests, and a verifier that writes new tests given a problem and a candidate program. Training mixes prompts from the code dataset with verifier prompts built online, and updates a single shared policy with GRPO.

The verifier curriculum has two stages. Stage 1 rewards a generated test only if executing the ground-truth solution on it yields the expected output, and if the test is not an exact copy of an example test from the prompt; the buffer admits executable solver outputs. Stage 2 clears the buffer and admits only executable but incorrect candidates, defined as programs that fail at least one dataset test case. The Stage 2 reward requires a test to pass the ground-truth solution and fail the incorrect candidate, while still excluding example tests. The transition is a hard switch triggered once test-generation accuracy reaches roughly 0.75 (step 200 for the 1.5B model, step 40 for the 7B model).

At inference time, the system samples $N$ candidate code solutions, generates $M$ tests conditioned on each candidate, pools the $N \times M$ tests, executes every candidate on the pooled tests, and selects the candidate with the highest pass-count. In the main experiments $M = 1$.

Training data comes from TACO: 25,433 problems were filtered to 10,605 (fewer than 50 test cases removed), then to 7,918 (problems without submitted solutions removed), then to 6,318 curated instances (at least one Python solution passing all cases). Evaluation uses the TACO validation set (1,000 problems) and LiveCodeBench problems from August 2024 to February 2025. The RL framework is verl, with batch size 128, PPO mini-batch 64, GRPO group size 16, temperature 0.8, and maximum response length 8192; no KL loss is used, and an entropy loss sustains model entropy. Training steps are 200 (Stage 1) and 250 (Stage 2) for 1.5B with a 450-step code baseline, and 40 (Stage 1) and 160 (Stage 2) for 7B with a 200-step code baseline. Reported GPU usage on NVIDIA H100 80GB is 720 hours training and 10/20 hours evaluation for 1.5B, and 960 hours training and 16/32 hours evaluation for 7B, per single run.

Why This Matters

Impact on research. The paper shifts attention from generating more tests to generating better ones, and gives a concrete theoretical link between training objectives and inference-time reliability: the bound $(N-1)\exp(-K(\delta-\alpha)^2/2)$ identifies the number of tests $K$ and the margin $\delta - \alpha$ as the two control knobs, with Stage 1 targeting $\alpha$ and Stage 2 targeting $\delta$. It also provides matched-backbone, matched-budget comparisons (joint vs. code-only vs. test-only RL; joint SFT vs. online RL) that isolate the value of verifier training and of policy-aligned online updates.

Real-world applications:

• Automated grading and judging of code submissions where a hidden reference solution exists but the visible tests are too weak to catch subtle bugs.
• Best-of-$N$ selection in coding assistants, where a locally fine-tuned verifier ranks candidates instead of relying on a separate reward model.
• Unit-test generation for real repositories, where the goal is to construct inputs that break incorrect or regression-prone implementations.
• Data curation and reward modeling for code post-training, by producing execution-verified test cases that can filter training samples.

Industry relevance. The reported token asymmetry (one generated test costs 22–24% of a code candidate for 1.5B, 32–34% for 7B) and the parallelism of independent test-generation calls make test-based selection an economically plausible inference-time scaling strategy, especially where a small fine-tuned model can match or exceed a much larger general baseline.

Future Directions

• Multi-test generation per inference call. The current inference-time pipeline emits only one test per call; generating several at once is more efficient but lacks a reward function that resists reward hacking, since intuitive metrics such as counting correct cases or accuracy rates are gameable.
• Soft stage switching. The present curriculum uses a hard Stage 1 → Stage 2 switch; dynamically adjusting the proportion of the two reward types may improve results, and the authors say they avoided it only because of RL resource demands.
• Relaxing the ground-truth assumption. TCS assumes access to a ground-truth solution $C^{*}$ during post-training (though not at inference); extending the framework to model-based or consistency-based verification would broaden applicability.
• Understanding the self-play synergy. Filtering TCS-generated code with Test-RL tests still underperforms TCS self-verification, which the authors attribute to a synergistic effect where a stronger solver creates harder failure modes for the verifier; characterizing this dynamic more precisely remains open.

Target Audience

Researchers and practitioners in code LLM post-training and reinforcement learning, particularly those working on execution-feedback training, verifier or reward-model learning, unit-test generation, and inference-time scaling for code. It is also relevant to engineers building evaluation or best-of-$N$ selection pipelines for coding assistants, and to readers interested in theory that connects training reward design to inference-time selection reliability. The paper is written at an advanced level and is best read by someone comfortable with RL fine-tuning and execution-based code benchmarks.

Authors’ abstract

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.