Mutation Guided Reinforcement Learning for Improving Test Case Generation

Project PUMA is a framework that combines Reinforcement Learning (RL) with Mutation Testing to enhance test case generation for large language models (LLMs).

Current Approach

To evaluate effectiveness, we generate multiple buggy variants (mutants) of a correct program using mutation testing tools such as mutmut (only for python). Each generated test case is executed against both correct and mutated code using a Python execution framework using subprocess. A test is considered effective if it passes on the correct code and fails on a sufficient number of mutated versions.

We are currently implementing a PPO (Proximal Policy Optimization) loop using the TRL library, where:

• The policy model generates test cases
• The value model estimates expected reward
• Rewards are assigned based on mutation detection performance
• Fixing the main script to generate dataset and prompt the LLM all in one go.

Key Features

• Syntax-Aware Mutation Engine: Utilizes libcst and mutmut operators to inject faults at the Concrete Syntax Tree (CST) level, ensuring all mutants are syntactically valid Python code.
• RLHF with PPO: Implements a full Reinforcement Learning from Human Feedback (RLHF) loop to optimize model completions based on execution feedback.
• Comprehensive Dashboards: Generates visual reports (results_dashboard.png) to track model performance, pass rates, and training stability.
• Benchmark Support: Native integration with the HumanEval dataset for standardized performance tracking.

Tech Stack

• Core Logic: Python 3.10+
• Datasets: HumanEval+, leetcode-in-c github repository (planning to use for assembly)
• Deep Learning: PyTorch, HuggingFace Transformers, Accelerate
• Reinforcement Learning: TorchRL, TensorDict, TRL(planning to use)
• Static Analysis: LibCST, Mutmut
• Visualization: Matplotlib

Installation

1. Clone the repository:

   git clone https://github.com/your-username/ProjectPUMA.git
   cd ProjectPUMA

2. Set up a virtual environment:

   python -m venv .venv
   source .venv/bin/activate  # On Windows: .venv\Scripts\activate

3. Install dependencies:

   pip install -r requirements.txt

Usage

Prepare the data

python3 src/main copy.py --output_dir ./test_run

This is an incomplate main script which can be used for generating and mutating HumanEval dataset.

Running the mutation analysis pipeline

python3 src/main.py

After loading data to ../data/source and ../data/mutations directories, the script willgenerate unit test cases for each code prompt in batches of 32 using the CodeLlama model, run the tests using python's unittest module by appending the generated tests to the original content in the temporary files, categorize the results of each test run as passed, failed, or timed out, save the results in a json file in ../logs/results.json, and generate a comprehensive dashboard in ../logs/results_dashboard.png.

Results

The project generates a comprehensive dashboard in logs/results_dashboard.png. This dashboard provides:

• Overall Test Status Distribution: Percentage of mutants passed, failed, or timed out.
• Pass Rate by File Type: Granular breakdown of model performance across different code categories.