Entropy-Regularized Rank-Masked Policy Optimization for Test-Time Reinforcement Learning in Code Generation

AI summary

Overview

Research area: Test-time reinforcement learning (TTRL), LLM code generation, policy optimization for large language models.

Technical level: Advanced. The paper assumes familiarity with policy-gradient methods (GRPO/PPO), advantage normalization, and entropy regularization, though its central ideas can be understood without the derivations.

Scope: The paper introduces probe-driven test-time reinforcement learning and a new optimization objective (ERPO) that lets a code-generating LLM improve itself on unlabeled coding problems at deployment time, without access to hidden tests or oracle outputs.

What This Paper Is About

Test-time reinforcement learning works well for math, where sampled solutions can be reduced to a single canonical answer and the majority answer can serve as a pseudo-label. Code breaks this: two programs can be functionally identical yet look completely different, so majority voting on text is meaningless. This paper makes TTRL work for code by having the model generate its own "probe" inputs, executing candidate programs on them, and scoring programs by whether their behavior agrees with the group — then optimizing that noisy signal in a deliberately conservative way.

Key Contributions

1. Probe-driven TTRL for code generation. A framework where the model generates output-free probe inputs directly from the problem statement, executes candidate programs on those shared inputs, and derives a training signal purely from behavioral agreement — no oracle tests, no reference outputs, no canonical answers.

2. Probe Consensus Reward (PCR) and an analysis of its noise structure. PCR scores each candidate by the fraction of probes on which its output falls in the majority set. The paper empirically characterizes PCR's calibration and shows its error is asymmetric: it is a reliable negative signal but an unreliable positive one.

3. Entropy-Regularized Rank-Masked Policy Optimization (ERPO). A policy optimization objective that discards the top half of each candidate group by PCR rank, uses only the lower-ranked samples' signed advantages (which are usually negative) to suppress low-consensus programs, and adds a fixed entropy-ceiling penalty to prevent unbounded entropy growth.

4. Empirical demonstration across in-domain adaptation and zero-shot transfer. ERPO improves both pass@1 and the pass@k frontier on LiveCodeBench, with gains transferring to CodeContests, CodeForces, and TACO, and it outperforms a baseline that has access to ground-truth public test cases.

Main Findings

• PCR is asymmetric, not accurate. On LiveCodeBench, PCR correlates with hidden-test correctness at AUC 0.790 (Qwen3-4B) and 0.750 (Qwen3-8B). Candidates scoring at or below 0.5 fail hidden tests more than 91% of the time — a strong negative signal. But candidates scoring exactly 1.0 are still wrong 44.2% (4B) and 51.9% (8B) of the time, so high consensus is not a correctness certificate.

• Directly optimizing PCR destroys the pass@k frontier. GRPO with PCR rewards picks up a small pass@1 gain on LiveCodeBench but a large pass@16 drop (27.8/30.9 for Qwen3-4B versus 26.0/34.4 for the base model), and the pass@1 gain vanishes entirely on transfer benchmarks. The authors attribute this to overfitting on spurious consensus patterns.

• Using real public tests is not a substitute for adaptation quality. GRPO_Pub, which sees complete inputs and outputs of public test cases, is second-best on pass@1 but shows only marginal pass@16 gains and barely transfers from the 8B model — evidence of reward hacking on the adaptation source.

• Negative-only optimization without entropy control is also unstable. NSR_PCR, which penalizes only below-0.5 samples, improves pass@16 (39.4 on LCB for Qwen3-4B) but barely moves pass@1 (25.7) and eventually suffers entropy explosion and policy collapse, according to the appendix.

• ERPO improves both axes simultaneously. On LiveCodeBench in-domain, ERPO reaches 36.7 pass@1 and 46.3 pass@16 for Qwen3-4B (base: 26.0 / 34.4) and 35.1 / 50.9 for Qwen3-8B. Transfer means across CodeContests, CodeForces, and TACO are 25.4 / 39.3 (4B) and 22.9 / 38.8 (8B), well above every baseline.

• Adaptation beats training-free search on cost and quality. On LiveCodeBench, ERPO's single-sample pass@1 of 36.2 exceeds the base model's pass@32 of 35.9; on CodeContests, 41.5 exceeds the base model's pass@16 of 41.2. Two representative training-free selectors (CodeT-BoN and RM-BoN) stagnate below 30% even at 32 samples, whereas ERPO hits 36.2% with one sample and no search overhead.

• The entropy ceiling is a stability knob, not a performance lever. Ablating H_ceil over {0.2, 0.5, 1.0} shows all three values improve over GRPO_Pub, meaning the ERPO objective itself drives the gains. A ceiling that is too low slows training; one that is too high accelerates early pass@k gains but degrades quickly after peaking. The authors fix H_ceil = 0.5 uniformly across models and datasets.

Methodology in Plain English

Step one — build probes. For each unlabeled coding problem, the model is prompted iteratively to write K test inputs that are format-correct, satisfy the stated constraints, and are structurally different from everything generated so far and from the examples in the problem statement. These probes are generated once per problem and reused for all training steps. Crucially, the model never needs to know the correct outputs — the probes exist only to expose behavior on shared execution points.

Step two — score by consensus. The current policy samples a group of G candidate programs (G = 16) for a problem. Each candidate is executed on the shared probes. For each probe, the outputs are tallied and the majority output (ties credited to all tied outputs) defines the expected behavior. A candidate's PCR score is the fraction of probes on which its output lands in the majority set. Executions that crash or time out get no credit and are excluded from that probe's tally.

Step three — optimize conservatively. Instead of maximizing PCR scores, ERPO computes standard group-normalized advantages (how far each candidate's PCR is above or below the group mean, divided by the group standard deviation), then deletes the top half of the group by rank. The remaining candidates typically have negative advantages, so the gradient step reduces their likelihood. High-PCR candidates receive zero policy-gradient weight — they are never directly reinforced. A separate penalty term activates only when average token-level policy entropy exceeds a fixed ceiling, adding a quadratic cost proportional to the violation. Below the ceiling it contributes nothing.

The intuition: if the group is entirely wrong but agrees, ERPO does not make the shared mistake more likely. If the group is right, the high-consensus programs are simply ignored rather than reinforced. The only action taken is to make low-consensus outliers less probable.

Experimental protocol. Base models are Qwen3-4B and Qwen3-8B in non-thinking mode, trained with verl, vLLM rollouts, and FSDP on eight A100-80GB GPUs (~12 hours per run). Training uses 16 prompts per batch, 16 rollouts per prompt, 150 global steps, learning rate 5e-7 after a 20-step warmup, and 10 generated probes per problem. Two evaluation settings are used: in-domain adaptation on unlabeled LiveCodeBench (or CodeContests) problem statements only, and zero-shot transfer of the resulting checkpoint to three other benchmarks. Hidden tests are never touched during training and are used only post hoc to compute pass@k.

Why This Matters

Research impact. The paper reframes test-time reinforcement learning as a problem of reward asymmetry rather than reward accuracy. Its central claim — that a noisy behavioral signal is more trustworthy as a negative filter than as a positive label — is a general principle that could transfer to any domain where verification is partial but execution is cheap: theorem proving, SQL generation, agentic tool use, simulated robotics, spreadsheet formulas. The paper also challenges the common assumption that pass@1 gains are always worth having, demonstrating empirically that they can come at the cost of the pass@k frontier, which bounds everything downstream inference-time search can extract.

Real-world applications:

• Code assistants and IDE completion. A model adapted this way produces a better single-pass suggestion, which is what users actually see when an autocomplete fires. The paper's LiveCodeBench result — ERPO's pass@1 beating the base model's pass@32 — maps directly to lower latency at equal quality.
• Agentic coding pipelines. Agents that write and iterate on code benefit from a stronger single-step generator, since every extra retry multiplies cost and latency across a multi-turn trajectory.
• Self-improving models on proprietary codebases. A company can adapt a model to its internal style and APIs using only unlabeled problem descriptions plus a sandbox to execute code in, without curating tests or labels for every function.
• On-device and edge code models. Smaller models with better single-sample accuracy reduce the need for expensive test-time sampling, which matters when compute, memory, or power is constrained.
• CI/CD automated repair. Bug-fixing systems that propose one patch and stop benefit from higher pass@1 without needing to run dozens of candidate patches through a full test pipeline.

Industry relevance. The method requires no human-labeled data and no annotated test suites — only a code execution sandbox, which most engineering organizations already have in their CI infrastructure. That makes it deployable in settings where labeled data is the bottleneck. The compute cost (roughly 12 GPU-hours on eight A100s for a full adaptation run) is modest by modern post-training standards, and the inference-time savings compound in production.

Future Directions

• Project-level and industrial-scale code generation. The authors explicitly flag that their evaluation is limited to competitive-programming-style problems and that multi-file, repository-scale code requires substantially more engineering infrastructure for reliable execution, training, and evaluation. Whether probe-based behavioral consensus still works when "correct behavior" spans a whole codebase is open.
• Better probes. The entire method rests on the quality of the generated probe inputs. The paper uses 10 probes per problem generated by a fixed prompt template, and the authors note that probes missing the decisive corner case are a primary source of spurious consensus. Smarter probe generation — adversarial, coverage-guided, or failure-driven — is an obvious lever left largely unexplored.
• Reducing the rank mask's bluntness. The current mask is a hard cutoff at the median of each group. The ablation on mask percentile is relegated to an appendix, and adaptive or confidence-weighted masking is not attempted. Whether a soft mask keyed to PCR calibration could extract more signal from the high-PCr end is unresolved.
• Extending the asymmetry principle beyond code. The finding that low-consensus is a sharply better signal than high-consensus could be tested in other open-vocabulary generation settings where execution or simulation provides behavioral feedback — SQL, shell commands, planning traces, or tool-calling trajectories.

Target Audience

Reinforcement learning researchers working on post-training and test-time adaptation for large language models; ML engineers building code-generation systems who need quality improvements without labeled data or curated test suites; and practitioners of inference-time scaling who care about the pass@k frontier and the cost of sampling many candidates. Readers without a policy-gradient background will follow the conceptual argument — that a noisy reward is safer used as a negative filter — but will need to consult the preliminaries section for the advantage normalization and masking details.

Authors’ abstract

Existing methods for test-time reinforcement learning (TTRL) derive rewards from answer-level self-voting on unlabeled test-time tasks with canonical answers, but this breaks down for code generation because programs cannot be compared by surface form and therefore do not directly provide a usable training signal. To make TTRL applicable to code generation, we propose probe-driven TTRL, which constructs output-free probe inputs from the problem statement, executes candidate programs on these probes, and defines a Probe Consensus Reward (PCR) from the resulting behavioral agreement. PCR provides a behavioral training signal for open-vocabulary programs, but it is not a fully reliable verifier and remains susceptible to reward hacking through spurious consensus. We therefore introduce Entropy-Regularized Rank-Masked Policy Optimization (ERPO), which converts low PCR into conservative negative updates through rank masking and controls policy drift with an entropy ceiling. On coding benchmarks, ERPO substantially improves pass@1 and pass@k in both in-domain adaptation and zero-shot transfer.