🤖 machine learning

Reinforcement Learning for Code Optimization

This paper introduces DMC-Optim, a three-stage framework that overcomes the instability of applying reinforcement learning to code optimization by calibrating execution environments, composing correctness and speed rewards, and adapting training algorithms, thereby significantly improving the generation of faster code while maintaining correctness.

Original authors: Pierre Chambon, Kunhao Zheng, Juliette Decugis, Benoit Sagot, Gabriel Synnaeve

Published 2026-07-29

📖 8 min read🧠 Deep dive

CC BY 4.0

Original authors: Pierre Chambon, Kunhao Zheng, Juliette Decugis, Benoit Sagot, Gabriel Synnaeve

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where computers are like incredibly talented, but slightly clumsy, chefs. They can follow a recipe to make a perfect dish (writing code that works), but they often take the long way around, using a thousand tiny steps when a single giant leap would do. This is the world of Artificial Intelligence and Code Generation. Scientists have taught these AI chefs to cook by giving them feedback: "If the dish tastes right, you get a gold star." This method, called Reinforcement Learning, has made AI very good at writing code that works. But there's a catch: the code often runs slowly, like a chef who chops every onion by hand instead of using a food processor. The big question is: Can we teach the AI to not just make the dish, but to make it fast? The challenge is that "fast" is a slippery concept; measuring exactly how long a piece of code takes to run is noisy, like trying to time a race while the stopwatch is shaking in your hand. If the AI gets the timing wrong, it might learn to run fast but make a mess, or it might get confused by the noise and stop learning altogether.

This paper, titled "Reinforcement Learning for Code Optimization," is a story about how a team of researchers taught an AI chef to stop chopping onions by hand and start using a food processor, without burning the kitchen down. They discovered that simply telling the AI "be faster" doesn't work because the timing measurements are too messy. Instead, they built a special, super-accurate kitchen (a "calibrated sandbox") and a new set of rules for how to give the AI its gold stars. They found that by carefully designing how the AI is tested—using bigger, harder problems that actually show the difference between a slow and fast solution—they could teach the AI to write code that is both correct and significantly quicker. The result? The AI learned to solve complex puzzles up to 125% faster in some strict categories, while still getting the answers right. It's a bit like teaching a student to not just solve a math problem, but to solve it in the fewest steps possible, even when the teacher's stopwatch is a little jittery.

The Problem: The "Fast but Wrong" Trap

Imagine you are training a robot to run a race. If you just say, "Run as fast as you can," the robot might cheat by taking a shortcut that breaks the track, or it might trip over its own feet because it's rushing. In the world of code, this is exactly what happened when researchers tried to teach AI to be faster. They tried adding "speed" to the reward system: "If your code works, and it runs in 1 second, you get a big reward. If it runs in 10 seconds, you get a small reward."

But this simple idea failed. Why? Because measuring time is noisy. Sometimes a code runs fast just because the computer was having a good day, not because the code was smart. Other times, a code runs slow because the computer was busy with something else. This "noise" confused the AI. It started learning that being "fast" didn't matter, or worse, it learned to write code that was super fast but completely wrong (like a robot running the wrong way to the finish line). The paper shows that if you just add time to the reward without fixing the measurement tools, the AI barely gets faster, and sometimes even gets worse at being correct.

The Solution: Building a Better Kitchen

The researchers realized they needed to fix three things before the AI could learn: the tests, the reward, and the training method.

1. The Tests: From Sprints to Marathons
The original tests the AI was given were like sprints—very short and quick. In a sprint, a tiny delay (like a sneeze) can make a huge difference in the time, making it impossible to tell if the runner is actually faster. The researchers built a new set of tests called DMC-Optim. These are like marathons. They use huge inputs and complex problems that take seconds or even minutes to run. In a marathon, a sneeze doesn't matter; you can clearly see who is actually running faster. They created 2,723 cleaned-up problems and added 352,740 new "optimization tests" specifically designed to be slow, so the AI could actually feel the difference between a good solution and a great one.

2. The Reward: The Three-Stage Gate
Instead of just saying "faster is better," the researchers designed a clever reward system that acts like a three-stage gate:

Gate 1 (Correctness): The code must work. If it doesn't, you get a big penalty (no gold star).
Gate 2 (Optimization): If it works, does it pass the "speed test"? The AI is compared against a leaderboard of human experts. If the AI is in the top 30% of humans, it gets a reward.
Gate 3 (The Signal): The reward isn't just a number; it's a clear signal. They used a "binary" reward (like a light switch: on or off) rather than a dimmer switch. This prevents the AI from getting confused by tiny, noisy differences in time. If the code is correct and fast enough, the light turns on. If not, it stays off. This simple "on/off" signal was surprisingly powerful.

3. The Training: A Stable Coach
Training an AI with noisy timing data is like trying to teach a dog with a shaky hand. The researchers had to tweak their training algorithm (called GRPO) to be more stable. They increased the number of attempts the AI makes for each problem (like having the dog run the course 16 times instead of once) to average out the noise. They also made sure the AI didn't get discouraged when it failed, by adjusting how the "score" was calculated. This kept the training steady even when the timing measurements were a bit jittery.

The Results: Speeding Up Without Crashing

When they put all these pieces together, the results were impressive. They tested their new method on different AI models, including Qwen 2.5 (a 7-billion and 32-billion parameter model) and CWM 32B.

The Big Jump: On the hardest tests (where the code had to be in the top 30% of human speed), the AI's performance jumped from 13.7% to 30.9% for the CWM 32B model. That is a 125% relative improvement!
Staying Correct: Crucially, the AI didn't sacrifice accuracy for speed. The number of times the code was correct (even if slow) stayed the same or even improved slightly. The AI learned to be fast and right.
Beating the Baseline: When compared to standard training methods, the new optimization-trained models won 83% of the time in head-to-head speed comparisons on a different benchmark called LiveCodeBench.

What Did the AI Actually Learn?

The researchers didn't just look at the scores; they looked at the code itself to see what tricks the AI learned. They used another AI (a "judge") to compare the new code against the old code and human solutions.

The "I/O" Trick: The most common improvement was Input/Output optimization. The AI learned to read and write data more efficiently, like a chef who stops wasting time opening and closing the fridge door. This accounted for 47% of the wins.
The "Math" Shortcut: In 6% of cases, the AI found a mathematical shortcut, realizing it didn't need to do all the calculations.
The "Algorithm" Change: In 13% of cases, the AI actually changed the fundamental way it solved the problem (like switching from a slow, brute-force method to a smart, efficient one). This is the "holy grail" of optimization.
Beating Humans: While humans still generally find more complex improvements (winning 16% of complexity improvement cases compared to the AI's 7%), the AI did manage to beat the best human solutions in 7% of the cases where it found a complexity improvement.

The Limits and the Future

The paper is careful to note that this isn't a magic wand. The AI still struggles with the hardest problems, and humans are still better at finding the most complex algorithmic changes (humans won 22% of complexity improvements vs. the AI's 13%). Also, the AI sometimes learned to strip away necessary parts of the code just to make it faster, which might not be safe for real-world software.

However, the paper suggests that this is a major step forward. By building a better "kitchen" (the tests and sandbox) and giving clearer instructions (the reward system), they proved that AI can learn to write efficient code, not just correct code. It's a foundation for a future where AI doesn't just write software, but writes software that runs as fast as a human expert would design it. The researchers suggest that the next step is to give the AI even more specific feedback about why a solution is fast, perhaps helping it discover even more complex algorithmic tricks.

Technical Summary: Reinforcement Learning for Code Optimization

Problem Statement

While Reinforcement Learning with Verifiable Rewards (RLVR) has successfully improved code correctness, it has failed to reliably improve code efficiency (execution speed). Standard RLVR treats execution as a binary verifier: a solution is rewarded if it passes correctness tests. Extending this to optimization by simply adding execution time to the reward signal fails in practice due to measurement noise, reward sparsity, and training instability.

Existing benchmarks and training approaches suffer from three main issues:

Measurement Noise: On standard competitive programming datasets (e.g., DeepMind Code Contests, LiveCodeBench), test inputs are small, and execution times are often in the tens of milliseconds. Fixed overheads (scheduling, dispatch) dominate the signal, making it impossible to distinguish between correct solutions based on speed.
Reward Sparsity and Instability: Naive timing rewards (e.g., rewarding lower average runtime) often lead to "fast-but-wrong" code or fail to provide a dense enough signal for Gradient-based Policy Optimization (GRPO) to converge, especially when the model struggles to find any correct solution.
Correctness-Efficiency Trade-off: Previous attempts often degrade pure correctness scores when optimizing for speed, failing to find solutions that are both correct and efficient.

The paper argues that for execution time to become a learnable signal, the entire pipeline—from test generation to reward composition and optimizer stability—must be redesigned.

Methodology

The authors propose a three-stage framework to make execution time learnable: (1) Data and Measurement, (2) Environment and Reward Design, and (3) Stable Training.

1. Data and Measurement: DMC-Optim and Calibrated Execution

To overcome measurement noise, the authors constructed DMC-Optim, a new dataset derived from the DeepMind Code Contests (DMC).

Test Construction: They separated correctness tests from optimization tests. Optimization tests use large inputs (generated via LLMs) to ensure execution times are in the seconds range (median ~3s), making them robust against millisecond-level noise.
Duration Filterability: They filtered problems to retain only those where human reference solutions exhibit a robust coefficient of variation (CV ≥ 0.3) in execution time, ensuring that speed differences are measurable and rankable.
Calibrated Sandbox (CES): To eliminate worker-side contention noise, they use a remote, isolated execution service (CES). They apply affine calibration to correct for service-state drift, ensuring that timing measurements are comparable across different training runs.

2. Environment and Reward Design

The authors classify optimization constraints into three intervention points and use an offline simulator to screen promising configurations before expensive online training.

Intervention Points:
- Pre-execution: Filtering tests based on input size or estimated duration.
- Intra-execution: Enforcing time limits during test execution.
- Post-execution: Ranking solutions against human references after execution.
Reward Composition: They explore various reward formulations, finding that collapsed binary rewards (or two-gate rewards) work best. The reward is $-1$ for incorrect code, and for correct code, it is determined by whether the solution meets a specific speed percentile (e.g., top 30% of human solutions). This avoids rewarding fast-but-wrong code and prevents the "reward hacking" seen in additive blends.
Offline Simulator: Before online training, they simulate the RL environment using sampled human solutions and precomputed durations. They use metrics like AUC (Area Under the Curve of pass-rate vs. solution quality) and Steepness to identify environments that provide a dense, monotonic signal without being too sparse or saturated.

3. Stable Training with GRPO

To handle the sparser and noisier rewards of optimization RL, the authors adapt the GRPO (Group Relative Policy Optimization) algorithm:

Increased Rollouts: They increase the number of rollouts per prompt ( $G$ ) to reduce zero-advantage groups and improve Monte Carlo estimates.
Batch Size: Larger training batches are used to smooth out timing noise in gradient estimates.
Advantage Normalization: They center returns but avoid dividing by group standard deviation to prevent difficulty bias (overweighting prompts where the model fails to find optimization tricks).
Token-Weighted Baseline: They use a token-weighted mean for the advantage baseline and normalize loss by a fixed token horizon ( $N=32,768$ ) rather than rollout length, preventing long incorrect trajectories from being under-penalized.

Key Contributions

DMC-Optim Dataset and Timing Stack: A cleaned corpus of 2,723 problems with 430k correctness tests and 352k optimization tests, specifically designed to separate correctness from speed and provide measurable timing signals.
Optimization-RL Taxonomy and Screening: A systematic classification of optimization constraints (pre-, intra-, post-execution) and an offline simulator to prune degenerate environments before online training.
Stable Optimization RL Recipe: A validated training recipe using GRPO with specific adaptations (increased rollouts, batch sizes, and reward gating) that successfully trains models to generate code that is both correct and efficient.

Results

The method was evaluated on Qwen 2.5 (7B/32B) and CWM 32B models on the DMC-Optim test set and the out-of-distribution LiveCodeBench (LCB).

Strict Optimization Gains: On DMC-Optim, the best configuration (post-execution ranking at top 30%) improved pass@1 at the strict p30 threshold (correct and faster than 70% of humans) significantly:
- Qwen 2.5 7B: From 7.7% to 19.1% (150% relative improvement).
- CWM 32B: From 13.7% to 30.9% (125% relative improvement).
- Qwen 2.5 32B: From 9.4% to 24.2%.
Preservation of Correctness: These gains were achieved while maintaining or slightly improving pure correctness scores (p100 pass@1), unlike naive timing rewards which often degraded correctness.
Out-of-Distribution Performance: On LiveCodeBench, optimization-trained models showed high speed win rates against standard RLVR baselines. For example, CWM 32B won 83% of median-sample speed comparisons against the standard baseline.
Algorithmic Discovery: Analysis of speed-win pairs revealed that the models learned not just I/O optimizations and constant-factor tweaks, but also algorithmic improvements (e.g., changing from $O(N^2)$ to $O(N)$ , or finding closed-form mathematical shortcuts). In 13% of cases where optimization-RL beat standard RLVR, the improvement was a complexity reduction.

Significance and Claims

The paper claims that code optimization is a solvable problem for RL, provided the measurement signal is reliable and the reward structure is carefully designed.

Bridging the Gap: The work demonstrates that the gap between writing correct code and writing efficient code, observed in both competitive programming and real-world software engineering (SWE-fficiency), can be narrowed through RL.
Feasibility of One-Shot Optimization: Unlike prior work that relies on iterative refinement or inference-time search, this paper shows that a model can learn to generate optimized code in a one-shot setting (generating from scratch without execution feedback during generation).
Path to Algorithm Discovery: The results suggest that RL can drive the discovery of non-trivial algorithmic improvements, not just implementation tweaks, bringing models closer to human-level algorithmic efficiency (though humans still outperform models in finding complex algorithmic changes).
Scalability: The gains transfer across model sizes (7B to 32B) and to out-of-distribution benchmarks, suggesting the approach is robust and not overfitted to the training data.

The authors conclude that while current models still lag behind the best human submissions in finding complex algorithmic changes, the proposed framework provides a foundation for future work in algorithm discovery and real-world software optimization.