Teaching LLMs to Write Clean Code: Quality-Aware Feedback via Reinforcement Learning Training

This repository contains the code to replicate experiments of our study "Teaching LLMs to Write Clean Code: Quality-Aware Feedback via Reinforcement Learning Training", submitted to [TBD].

It contains the datasets, training scripts, inference scripts, analysis pipelines, and results of experiments of code generated by large language models aligned with Supervised Fine-Tuning (SFT), Direct Preference Optimization (DPO), and Proximal Policy Optimization (PPO).

---

Environment Setup

Two conda environments are required:

• myEnv (Python 3.9.21) — used for all training, inference, and quality analysis steps.
• myPyTEnv (Python 3.7.16) — used exclusively for python-taint (PyT) static analysis.

Create both environments by running from the repository root:

bash setup_envs.sh

After the environments are created, replace the following files in the installed trl package with the modified versions provided in the assets/ folder:

• trl/trainer/ppo_trainer.py
• trl/trainer/utils.py

---

Datasets

The repository includes three datasets covering different training stages:

• SFT dataset — located in datasets/sft/. Cleaned version of the Stack dataset used to supervised fine-tune the base models on code generation tasks, split into training, validation, and test sets. Each instance has the following structure:

  {
    "instruction": "<function signature with a docstring describing its functionality>",
    "output": "<reference code>"
  }

• DPO dataset — located in datasets/dpo/. Extended PoisonPy dataset used to train models with Direct Preference Optimization, split into training and test sets. Each instance has the following structure:

  {
    "prompt": "<natural language prompt>",
    "chosen": "<higher quality code>",
    "rejected": "<lower quality code>"
  }

• PPO dataset — located in datasets/ppo/. Derived from the Emergent Misalignment dataset and processed through a multi-tool static analysis pipeline (Semgrep, PyT, DeVAIC) to filter and label secure and insecure code samples. Split into training, validation, and test sets. Each instance has the following structure:

  {
    "messages": [
      { "role": "user", "content": "<natural language prompt>" },
      { "role": "assistant", "content": "<reference code>" }
    ]
  }

---

Training

Training scripts are located in training/. For each of the three training techniques and each of the four models (CodeGPT 125M, CodeGen 350M, QwenCoder 0.5B, DeepSeek 1.3B), a dedicated runnable script is provided:

• Supervised Fine-Tuning (SFT) — scripts in training/sft/, one per model (run_sft_CodeGPT.py, run_sft_CodeGen.py, run_sft_QwenCoder.py, run_sft_DeepSeek.py). Checkpoints are saved to models/sft/<Model>-finetuned/.

  torchrun --nproc_per_node=<N> training/sft/run_sft_<Model>.py

• Direct Preference Optimization (DPO) — scripts in training/dpo/, one per model (run_dpo_CodeGPT.py, run_dpo_CodeGen.py, run_dpo_QwenCoder.py, run_dpo_DeepSeek.py). Use --finetuned yes (default) to start from the SFT checkpoint, or --finetuned no to train directly from the pretrained base model. Models are saved to models/sft_dpo/ or models/dpo/ respectively.

  torchrun --nproc_per_node=<N> training/dpo/run_dpo_<Model>.py [--finetuned yes|no] [--model_path <path>]

• Proximal Policy Optimization (PPO) — scripts in training/ppo/, one per model (run_ppo_CodeGPT.py, run_ppo_CodeGen.py, run_ppo_QwenCoder.py, run_ppo_DeepSeek.py). The --metric argument controls the reward function and must be one of: bertscore, codebleu, edit, crystalbleu, pylint, semgrep, custom_pylint, custom_semgrep. Use --finetuned yes (default) for PPO on top of the SFT model, or --finetuned no for PPO directly from the pretrained model. Models are saved to models/sft_ppo/ or models/ppo/ respectively. The reward functions shared across all PPO scripts are defined in training/ppo/reward_functions.py.

  accelerate launch --num_processes <N> training/ppo/run_ppo_<Model>.py --metric <metric> [--finetuned yes|no]

---

Inference

A single script handles inference for all models and training types:

bash inference/run_inference.sh <model_path> <mode>

• <model_path>: path to the model folder (training type and metric are extracted automatically from the folder name).
• <mode>: dpo runs a single inference pass on the DPO test set; ppo runs two passes, once on the secure test set and once on the insecure test set.

Examples:

# PPO model
bash inference/run_inference.sh models/ppo/CodeGen-ppo-custom_semgrep ppo

# DPO model
bash inference/run_inference.sh models/dpo/DeepSeek-dpo dpo

Inference results are saved to results/inference/<train_type>/.

---

Analysis

Quality Analysis

Static analysis of the generated code is performed using Semgrep across all configured rule sets. The pipeline is controlled by two scripts located in results/scripts/:

• quality_analyze_code.py — takes a .jsonl inference file, extracts the generated code, writes it to temporary .py files, and runs Semgrep, saving the results in batched .json files under results/scripts/quality_outputs/.
• quality_process_results.py — processes the Semgrep output, computes per-function defect rates, error rates, and issue breakdowns by category and severity. A convenience shell script runs the full quality pipeline in a single command:

bash results/scripts/run_quality_analysis.sh results/inference/<train_type>/<file>.jsonl

This automatically cleans the inference file, runs the Semgrep analysis, and processes the results.

Correctness Analysis

Functional correctness of the generated code is measured using Edit Distance (ED). Run the correctness analysis script as follows:

python results/scripts/correctness_similarity_metric.py <path_to_file>.jsonl

Statistical Analysis and Output Characterization

Three additional scripts, located in results/scripts/, support the analyses reported in the paper. All of them read the post-processed generations from results/inference_cleaned/ (searched recursively, following the naming convention {Model}_{config}_inference[_test_{secure|insecure}][_{reward}]_cleaned.jsonl) and write their outputs to results/:

• characterize_outputs.py — characterizes the model outputs: decomposes the error category into empty outputs, prompt-template echo, truncation-suspect generations, and other syntax errors, and compares output lengths and empty/trivial-output rates between quality-only and composite PPO rewards (Mann–Whitney U test). Note that parseability is assessed with Python's ast module, a stricter criterion than the tolerant parsing of the evaluation pipeline, so error totals are upper bounds of the error rates reported in the paper.
• stats_analysis.py — runs the paired statistical analysis of the paper: McNemar's test with Wilson 95% confidence intervals on the per-prompt clean/non-clean labels, and the Wilcoxon signed-rank test with Cliff's delta on per-sample Edit Distance, over the pre-specified family of 16 comparisons, with Holm correction applied within each metric family. It reads the per-sample quality labels from results/scripts/quality_labels.csv, reclassifies empty outputs as errors prior to the analysis (same procedure as the paper), and validates labels and Edit Distance values against the paper's figures and Table 3 before testing.
• run_bandit.py — re-analyzes the outputs of the Semgrep-rewarded configurations (R_s, R_cs) and their baselines with Bandit, and reports the clean-rate deltas against each baseline to verify that the observed trends hold under an independent tool.

python results/scripts/characterize_outputs.py

python results/scripts/stats_analysis.py    # requires results/scripts/quality_labels.csv

python results/scripts/run_bandit.py

Results

The results/ folder contains the collected results across all models and training configurations:

• quality_results.xlsx: aggregated quality metrics (clean rate, defect rate, error rate, issue counts by category and severity) for all evaluated models.
• correctness_results.xlsx: aggregated reference-similarity metric (Edit Distance) for all evaluated models.
• inventory.csv: list of the loaded generation files with their (model, configuration, test set) mapping and row counts.
• error_breakdown.csv: per-configuration decomposition of the error category (empty, prompt-template echo, truncation-suspect, syntax error), as percentages of the total.
• length_by_reward.csv: output length statistics and empty/trivial-output rates of the PPO configurations, grouped by reward type (quality-only vs. composite).
• summary.txt: human-readable summary of the error breakdown and length analyses, including the Mann–Whitney length comparisons.
• stats_quality.csv: paired McNemar tests on clean-function rates (Wilson 95% confidence intervals, discordant pairs, raw and Holm-adjusted p-values) for the 16 pre-specified comparisons.
• stats_ed.csv: paired Wilcoxon signed-rank tests on per-sample Edit Distance (Cliff's delta effect sizes, raw and Holm-adjusted p-values) for the same comparisons.
• bandit_reanalysis.csv: clean/defective/error rates of the Semgrep-rewarded configurations and their baselines, re-assessed with Bandit.
• bandit_trend_check.csv: clean-rate deltas (configuration minus baseline) under Bandit, supporting the cross-tool validity check discussed in the paper.

---