Replication package for the paper "Beyond the Need for Speed: Energy-Aware Code Generation via Simulation-Guided Reinforcement Learning" (Saurabhsingh Rajput and Tushar Sharma).
Code models strictly prioritize functional correctness, leaving software energy efficiency as an unoptimized byproduct. Training models to generate energy-efficient code requires reproducible feedback at scale, which physical hardware measurement cannot reliably provide due to variance.
In this paper, we replace hardware profiling with a deterministic architectural simulation harness to build Green Tea, a corpus of 3.5 million evaluations across 1,474 C++ problems. We train an energy-aware code model via supervised fine-tuning on energy-contrastive pairs, followed by closed-loop reinforcement learning (GRPO) using simulation-in-the-loop feedback. To rigorously evaluate deployment readiness, we introduce the Correctness-Adjusted Reduction in Energy Total (CARET), a metric that explicitly penalizes code that sacrifices functionality for efficiency.
On 143 held-out problems, our simulation-in-the-loop pipeline achieves 12.63% CARET, nearly tripling the gain of fine-tuning alone, and beats the energy efficiency of human-expert references on 58.4% of its valid outputs. Our analysis also exposes the IPC trap: standard throughput proxies like Instructions-Per-Cycle (IPC) actively misrank true energy efficiency on 67.8% of problems, demonstrating the necessity of direct energy simulation. By releasing our dataset and infrastructure, we bypass the 263,000 CPU-hours required for reproduction, empowering the community to train energy-efficient code generation models.
We train code models to optimize a program's energy directly, using deterministic architectural simulation (Sniper/McPAT) as a scalable, reproducible energy signal, in two stages: energy-contrastive supervised fine-tuning (Energy-SFT) followed by a simulation-in-the-loop reinforcement-learning stage (GRPO). This repository contains the training and evaluation code, the analysis scripts that regenerate every table and figure, the pre-computed simulation results needed for that analysis, and the held-out benchmark.
| Artifact | Location |
|---|---|
| Code, analysis scripts, pre-computed sim results, 143-problem benchmark | This repository |
| Green Tea dataset (raw: 3,507,435 simulations, execution master, training pairs) | Zenodo: 10.5281/zenodo.21210099 |
| Trained model checkpoints (Energy-SFT, GRPO, cross-family) | Hugging Face: huggingface.co/saurabh-singh-rajput (see Models) |
The raw dataset and model weights are hosted externally because of their size; everything required to reproduce the paper's numbers and figures is included here.
All trained checkpoints are on Hugging Face under huggingface.co/saurabh-singh-rajput. Merged models load directly with transformers; adapters are LoRA weights (peft) that reference the listed base.
| Model | Kind | CARET |
|---|---|---|
| green-tea-qwen2.5-coder-14b-grpo (headline) | GRPO adapter | 12.63% |
| green-tea-qwen2.5-coder-14b-energy-sft | merged | 4.45% |
| green-tea-qwen2.5-coder-14b-runtime-sft | adapter | 7.64% |
| green-tea-qwen2.5-coder-14b-edp-sft | adapter | 11.02% |
| green-tea-qwen2.5-coder-7b-energy-sft | adapter | 8.71% |
| green-tea-qwen2.5-coder-0.5b-energy-sft | adapter | 4.07% |
| green-tea-llama-3.1-8b-energy-sft | merged | 5.97% |
| green-tea-llama-3.1-8b-grpo | GRPO adapter | cross-family |
| green-tea-deepseek-coder-6.7b-energy-sft | merged | 1.29% |
| green-tea-deepseek-coder-6.7b-grpo | GRPO adapter | 5.99% |
green-tea/
├── config.env # set PROJECT_ROOT, then `source config.env`
├── requirements.txt # Python dependencies
├── energy_data_collection/ # Sniper/McPAT energy-simulation harness
├── finetuning/ # Energy-SFT + GRPO training and evaluation
│ ├── sft_train_trl.py # supervised fine-tuning
│ ├── grpo_trainer_vllm.py # simulation-in-the-loop GRPO
│ ├── evaluation.py # generation + Sniper evaluation
│ └── data/ # pre-computed sim_results (see "Data")
├── analysis/ # scripts that regenerate all tables/figures
│ ├── rq1/ rq2/ rq3/ rq4/ # one directory per research question
│ └── *.py # dataset, RAPL, correlation analyses
├── evaluation/ # benchmark comparison utilities
├── validation/ # McPAT-vs-RAPL calibration and validation
└── sniper/ # Sniper simulator (git submodule)
- Clone with the Sniper submodule:
git clone --recurse-submodules git@github.com:SMART-Dal/green-tea.git cd green-tea - Install Python dependencies (Python 3.10):
pip install -r requirements.txt
- Configure paths:
# edit PROJECT_ROOT in config.env to your clone location, then: source config.env
- The analysis scripts run directly against the pre-computed data in this repository and do not require Sniper. Sniper/McPAT is only needed to regenerate energy simulations from scratch (see
energy_data_collection/README.md).
All analysis scripts resolve paths relative to the repository root and read the pre-computed simulation results included here.
python3 analysis/rq2/rq2.py # main results and cross-family table; energy distribution and decomposition figures
python3 analysis/rq3/rq3.py # initialization-by-reward ablation, pass@k, robustness, optimization-level sweep
python3 analysis/rq4/rq4.py # transformation taxonomy, Stuart-Maxwell shift, per-category compiler analysis
python3 analysis/rq1/rq1.py # within-problem energy landscape and IPC-trap analysis
python3 analysis/training_plots.py # training-dynamics and hero-progression figurespython3 analysis/rapl_validation_analyze.py analysis/rapl_results.jsonl --pairs analysis/rapl_validation_pairs.jsonl \
--rapl-scanf analysis/rapl_scanf_results.jsonl --pairs-scanf analysis/rapl_scanf_pairs.jsonl # Table 1: RAPL cross-check
python3 analysis/robustness_audit.py --dataset <zenodo green-tea-dataset dir> --testcases <PIE merged_test_cases dir>
# Table 7 (all CARET conventions), paired tests, ghost-execution label sensitivity (analysis/matched_labels.jsonl),
# single-input reward fidelity, and the RQ1 variability figures; run on a compute node (re-executes programs natively)analysis/frequency_sensitivity/: re-simulation of test pairs at the EPYC 9554P's 3.1 GHz base clock (Section 4.2).analysis/library_case_study/: functions copied verbatim from widely used open-source projects, their standalone harnesses with test vectors, generations, and complete evaluation results, in three settings (systems_internal_data,systems_stdin,algorithm_libraries; select withCASE_SET).
analysis/rq2/rq2.py reads the cross-family simulation results from data/cross_family_sim_results/ (override with the GREEN_TEA_DATA environment variable). The RQ1 landscape and dataset-statistics scripts additionally require the full execution corpus from Zenodo (completed_samples.jsonl, PIE_Dataset/).
finetuning/data/*_sim_results/andfinetuning/data/new_models/*/sim_results/: per-model Sniper evaluation results used to compute the reported energy-reduction rates.data/cross_family_sim_results/: cross-family evaluation results (seven models).finetuning/data/sft_energy_r32/,sft_energy_r128/: LoRA-rank robustness results.finetuning/data/{sft,grpo}_pass_at_k/: pass@k results.finetuning/data/sft_pairs_test.jsonl: the 143-problem held-out benchmark.analysis/*.jsonl: aggregated training histories, per-solution metrics, and RAPL validation records.
PIE_Dataset/withexecution_master.jsonland per-problem batches: the raw energy-simulation corpus.completed_samples.jsonl: all 3,507,435 per-execution energy measurements.finetuning/data/grpo_train.jsonl,sft_pairs_train.jsonl: the energy-contrastive training pairs.
Scripts under finetuning/slurm/, energy_data_collection/slurm_execution_master.sh, and similar .sh files are SLURM job templates. Adapt the --account, --mail-user, module load lines, and any cluster paths to your environment before use.
@misc{rajput2026greentea,
title = {Beyond the Need for Speed: Energy-Aware Code Generation via Simulation-Guided Reinforcement Learning},
author = {Rajput, Saurabhsingh and Sharma, Tushar},
year = {2026},
note = {Preprint}
}Released under the MIT License (see LICENSE). The Sniper submodule and the PIE/CodeNet-derived data retain their original licenses.