Reinforcement learning method improves code generation tasks at test time
Entropy-Regularized Rank-Masked Policy Optimization for Test-Time Reinforcement Learning in Code Generation
Machine LearningComputation and Language
Summary
Generating computer programs is hard because small changes can make big differences, so usual techniques to check answers don't work well. The authors created a way to test programs by running them on special test inputs made from the problem description and seeing if they behave similarly. They use this behavior as a kind of reward to teach the system better. To avoid tricks and incorrect learning, they use a method that carefully adjusts training updates and keeps the program's guessing patterns diverse. Their approach improves the quality of generated code in several tests.
What this means in practice
- •For software developers: Improve automated code synthesis tools by adapting generation policies using behavioral signals from running code on custom test inputs.
- •For devops teams: Enhance CI/CD pipelines with adaptive code generation that refines outputs during deployment using reinforcement learning from test behaviors.
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.