Improving code generation with smarter test-time learning rewards
Entropy-Regularized Rank-Masked Policy Optimization for Test-Time Reinforcement Learning in Code Generation
Machine LearningComputation and Language
Summary
Training AI to write code is tricky because you can't easily compare correct answers by just looking at the programs themselves. To solve this, the authors created a method that tests code on specially designed inputs and checks if different versions behave the same way. This agreement gives a reward signal to help the AI learn better. Since this signal can sometimes be misleading, the authors include a way to be cautious about updates and keep the AI's learning stable. Their approach shows better results when the AI adapts to new coding tasks or tackles tasks it has never seen before.
reinforcement learningcode generationtest-time learningreward hackingpolicy optimizationprobe inputsbehavioral agreemententropy regularizationrank maskingpass@k
Authors
Jiacheng Xu, Feng Chen, Xiuneng Xu, Bo An
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.