Reinforcement Learning for Code Optimization

The Core Problem: Why Timing Breaks Code RL

Reinforcement learning has cracked code correctness—models generate programs, test them, and get rewarded for passing test cases. Extending this to code optimization (making programs faster) sounds trivial: just add execution time to the reward. It isn't. The authors show that raw timing signals are too noisy, too sparse, and too unstable for standard RL algorithms to learn from. Models trained this way generate solutions barely faster than baselines, and more often fail correctness entirely.

The breakthrough insight is that execution time needs to be treated as a learnable signal, not raw data fed directly into the reward. The authors restructure the entire pipeline across three dimensions.

How the Three-Stage Fix Works

Stage 1: Better Testing.
The authors build DMC-Optim, a benchmark with large-scale optimization tests and a carefully calibrated sandbox. This isn't just about having more test cases—it's about creating a measurement environment where timing is reproducible enough to ground learning. Standard benchmarks lack the scale and rigor needed for RL to extract a consistent optimization signal.

Stage 2: Learnable Rewards.
Rather than feeding wall-clock time directly into the loss function, the authors compose correctness and speed through an offline simulator. This simulator predicts which code configurations are most promising before the model commits to them. Think of it as a learned heuristic that filters the reward landscape, reducing noise and sparsity by focusing on decisions that matter.

Stage 3: Algorithm Adaptation.
They modify GRPO (a policy gradient method) and its evaluation to handle sparse, noisy timing data. Standard RL algorithms assume dense, clean rewards; timing has neither. The adaptation keeps the model from thrashing on measurement noise.

Results That Matter

On DMC-Optim with Qwen 2.5 7B, the approach lifts strict top-50% pass@1 from 18.0% to 31.3%. For the larger CWM 32B model, it jumps from 30.7% to 50.4%—a near-doubling. At stricter evaluation thresholds (top-30%), CWM improves by 125% relative to baseline. Crucially, pure correctness scores are preserved, so the model doesn't sacrifice passing cases to chase speed.

When the sandbox is degraded (simulating noisy real-world measurement), the optimization-aware RL reaches 100–200% improvement over standard baselines. On the LCB benchmark, CWM wins 83% of median-sample speed comparisons against the non-optimized baseline.

The gap to human performance remains real: the optimized model reaches about 14% of the human rate for complexity-class improvements (humans achieve 28%). But the gains are substantial and systematic.

Why This Matters

Code optimization is harder than code correctness because the signal is weak—most test cases either pass or fail, but speed improvements are continuous and noisy. This work shows that RL can still work at optimization when you treat execution time not as raw ground truth but as a learned, structured signal. The three-stage pipeline is a template: better data environments, reward shaping, and algorithm hardening unlock RL in domains where simple reward specifications fail.