Reinforcement Learning for Code Optimization
Abstract
RL for code correctness is now established: have the model generate a program, run it against hidden test cases, and reward solutions that pass. Extending this to code optimization seems straightforward: just add execution time to the reward. But in practice, once timing drives the reward, small problems in measurement noise, reward sparsity, or GRPO instability overwhelm the signal and make RL fail: generated solutions are barely faster, and more of them can fail. We make execution time learnable through three stages: (1) how code is tested, by building DMC-Optim with large optimization tests and a calibrated sandbox; (2) how speed is turned into reward, by composing correctness and speed in the RL environment and using an offline simulator to predict the most promising configurations; and (3) how the model learns from that reward, by adapting GRPO and evaluation to the sparser, noisier timed-execution setting. On DMC-Optim, the strongest optimization-aware configurations improve strict top-decile pass@1 from 18.1% to 27.7% on Qwen 2.5 7B and from 28.7% to 42.4% on CWM 32B, while preserving most of the pure-correctness score. When the timing sandbox is degraded, robust optimization RL reaches up to roughly 130% improvement over standard RLVR. On LCB, CWM 32B wins up to 70.3% of best-sample speed comparisons against standard RLVR; relative to the fastest correct human submissions per problem, it reaches about half the human rate of complexity-class improvements (14% vs. 28%).