ACORN: Cost-Aware RL Routing for LLM-Driven Software Development

ACORN is a cost-aware model-routing framework for software engineering tasks. Given a query and a pool of candidate LLMs, it learns to select the smallest model that can solve the query while balancing correctness and inference cost.

The repository contains the code and processed data required to reproduce the paper experiments on code translation, code repair, and code generation. It intentionally excludes model weights, intermediate checkpoints, experiment logs, discarded runs, and third-party IRT-Router code.

ACORN has three components:

• Processed routing dataset. Each input is paired with the smallest candidate model that succeeds and a cost-sensitive reward vector.
• Cost-aware router training. Standard SFT and Fixed-Cost baselines train the router with alternative objectives.
• Adaptive RL-Agent training. ACORN first initializes an RL-Agent with supervision and then alternates router updates with adaptive penalty-score updates.

Source Code

Environment

Create an isolated Python environment and install the dependencies:

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Alternatively, create the Conda environment supplied in environment.yml:

conda env create -f environment.yml
conda activate acorn-routing

The experiments require a CUDA-capable PyTorch installation and local Hugging Face-compatible checkpoints for the router backbones. Model weights are not redistributed in this repository.

Dataset

The processed paper splits are included under data/processed_dataset/:

data/processed_dataset/
├── train.jsonl  # 9,862 samples
├── valid.jsonl  # 1,187 samples
└── test.jsonl   # 1,197 samples

See data/README.md for the record schema. The split checksums are recorded in data/SHA256SUMS.

Training

Set MODEL_PATH to a local router backbone before running a training script. The scripts expose their hyperparameters as environment variables, with the Qwen-0.5B paper settings as defaults.

# Standard Classification baseline
MODEL_PATH=/path/to/router-backbone bash scripts/train_sft.sh

# Fixed-Cost baseline
MODEL_PATH=/path/to/router-backbone bash scripts/train_fixed_cost.sh

# ACORN: RL-Agent initialization + alternating training
MODEL_PATH=/path/to/router-backbone bash scripts/train_acorn.sh

For the random-initialization ablation, run src/training/train_acorn.py without --rlagent_path, using the same data and optimization settings as the pre-trained configuration.

The paper's controlled initialization comparison is available as a single entry point:

MODEL_PATH=/path/to/router-backbone bash scripts/run_initialization_ablation.sh

Evaluation

Evaluate a trained router on the paper test split:

MODEL_PATH=/path/to/trained/router bash scripts/evaluate.sh

Evaluation logs can be aggregated into routing and task-level summaries:

RESULTS_ROOT=/path/to/evaluation-results bash scripts/summarize_results.sh

RQ3 and RQ4 Analyses

The paper-specific analysis scripts are organized by research question:

# Generate deterministic RQ3 training subsets.
bash scripts/create_rq3_subsets.sh

# Run the RQ3 subset and equal-budget experiments.
MODEL_PATH=/path/to/router-backbone bash scripts/run_rq3_sample_efficiency.sh
MODEL_PATH=/path/to/router-backbone bash scripts/run_rq3_training_dynamics.sh

# Inspect the available analysis entry points.
python experiments/rq3/plot_sample_efficiency.py --help
python experiments/rq3/plot_training_dynamics.py --help
python experiments/rq4/export_rlagent_scores.py --help
python experiments/rq4/select_case_studies.py --help
python experiments/rq4/plot_case_studies.py --help

# Export and analyze RQ4 RL-Agent scores.
MODEL_PATH=/path/to/trained/router RLAGENT_PATH=/path/to/rl_agent_final.pt \
  bash scripts/run_rq4_interpretability.sh

Repository Layout

.
├── data/                    # Processed paper dataset
├── experiments/             # RQ3, RQ4, and empirical-study analyses
├── scripts/                 # Reproducible local training and evaluation entry points
├── src/
│   ├── training/            # SFT, Fixed-Cost, RL-Agent, and ACORN trainers
│   ├── evaluation/          # Inference and correctness/cost evaluation
│   ├── results/             # Paper-result aggregation utilities
│   └── utils/               # Shared data, prompt, and RL-Agent utilities
├── requirements.txt
└── README.md

Reproducibility Notes

• Use the same candidate model pool and decoding configuration as the paper when reproducing reported FLOPs.
• Training, tokenization caches, checkpoints, and evaluation outputs are generated under outputs/ and ignored by Git.
• The RQ3 scripts use seed 42 and nested 10%, 25%, 50%, 75%, and 100% training subsets.

Citation

Citation details will be added with the final paper release.