tool-integrated-rl-repo-search

name	tool-integrated-rl-repo-search
title	Tool-Integrated RL for Repository Bug Localization
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2508.03012
keywords	["reinforcement-learning","code-search","bug-localization","tool-use"]
description	Two-stage post-training framework combining rejection-sampled SFT and RL for LLM-guided repository code search and issue localization.

Tool-Integrated RL for Repo Deep Search

This framework teaches LLMs to locate buggy code in repositories through structured tool use and reinforcement learning. The two-stage approach combines supervised fine-tuning on high-quality trajectories with RL-based exploration, enabling models to learn sophisticated repository navigation strategies.

Core Concept

Developers spend significant time locating bugs in code repositories—a task requiring systematic exploration of file structure, function definitions, and import relationships. Rather than hoping LLMs can solve this through in-context learning alone, this framework explicitly trains them to use navigation tools effectively. The key insight: start with supervised learning of basic tool use, then deepen understanding through RL rewards that measure accuracy and ranking quality.

Architecture Overview

• Lightweight RepoSearcher Tools: Six minimalist tools (GetRepoStructure, GetImportOfFile, SearchClass, SearchFunction, SearchClassMethod, Exit)
• Stage 1 (SFT): Rejection sampling: keep only trajectories that correctly identify buggy code for supervised fine-tuning
• Stage 2 (RL): Reward-based policy optimization using nDCG@k metric to balance correctness and ranking quality
• Ranking-Aware Rewards: Unlike binary correctness, rewards account for ranking of candidate functions

Implementation Steps

Step 1: Define RepoSearcher Tool Interface

Create a minimal but complete tool set for repository navigation:

from abc import ABC, abstractmethod
from typing import List, Dict, Tuple
import os

class RepoSearcherTool(ABC):
    """Base class for repository search tools."""
    @abstractmethod
    def execute(self, repo_path: str, query: str) -> str:
        pass

class GetRepoStructure(RepoSearcherTool):
    """Retrieve directory structure and top-level files."""
    def execute(self, repo_path: str, query: str = None) -> str:
        """
        Returns: Formatted tree view of repo structure.
        """
        tree = []
        for root, dirs, files in os.walk(repo_path):
            # Limit depth to avoid overwhelming output
            depth = root.replace(repo_path, '').count(os.sep)
            if depth > 2:
                continue

            indent = '  ' * depth
            tree.append(f"{indent}{os.path.basename(root)}/")
            subindent = '  ' * (depth + 1)
            for file in files[:]:  
                tree.append()

         .join(tree)

 ():
    
     () -> :
        
        full_path = os.path.join(repo_path, file_path)
        imports = []

        :
             (full_path, )  f:
                 line  f:
                     line.strip().startswith()  line.strip().startswith():
                        imports.append(line.strip())
         FileNotFoundError:
             

         .join(imports[:])  

 ():
    
     () -> :
        
        matches = []
         root, dirs, files  os.walk(repo_path):
             file  files:
                  file.endswith():
                    
                filepath = os.path.join(root, file)

                :
                     (filepath, )  f:
                         line_num, line  (f, ):
                               line:
                                rel_path = os.path.relpath(filepath, repo_path)
                                matches.append()
                :
                    

         .join(matches[:])  matches  

 ():
    
     () -> :
        
        matches = []
         root, dirs, files  os.walk(repo_path):
             file  files:
                  file.endswith():
                    
                filepath = os.path.join(root, file)

                :
                     (filepath, )  f:
                         line_num, line  (f, ):
                               line:
                                rel_path = os.path.relpath(filepath, repo_path)
                                matches.append()
                :
                    

         .join(matches[:])  matches  

 ():
    
     () -> :
        
        class_name, method_name = query.split()
        
        
         

 ():
    
     () -> :
         

 :
    
     ():
        .repo_path = repo_path
        .tools = {
            : GetRepoStructure(),
            : GetImportOfFile(),
            : SearchClass(),
            : SearchFunction(),
            : SearchClassMethod(),
            : Exit(),
        }

     () -> :
        
         tool_name   .tools:
             
         .tools[tool_name].execute(.repo_path, query)

Step 2: Implement Stage 1 - Rejection-Sampled SFT

Generate and filter high-quality training trajectories:

from dataclasses import dataclass
from typing import List, Tuple

@dataclass
class Trajectory:
    """Represents a sequence of tool calls and responses."""
    steps: List[Tuple[str, str, str]]  # (tool_name, input, response)
    final_answer: str
    is_correct: bool
    ground_truth: str

class SFTDataCollector:
    """
    Generate trajectories using base LLM, keep only correct ones.
    Rejection sampling: sample trajectories and filter by correctness.
    """
    def __init__(self, model, repo_searcher, num_samples=10):
        self.model = model
        self.repo_searcher = repo_searcher
        self.num_samples = num_samples

    def generate_trajectories(self, issue_description: str, ground_truth_functions: List[str]) -> List[Trajectory]:
        """
        Sample multiple trajectories and filter for correctness.
        Only keep trajectories that identify ground truth as top prediction.
        """
        trajectories = []

        for _ in range(self.num_samples):
            # Initialize trajectory
            steps = []
            current_state = f"Issue: "

            
             step  ():  
                
                prompt = 

                response = .model.generate(prompt)
                tool_name, tool_input = ._parse_response(response)

                 tool_name == :
                    

                
                tool_result = .repo_searcher.execute_tool(tool_name, tool_input)
                steps.append((tool_name, tool_input, tool_result))
                current_state += 

            
            final_prompt = 
            final_answer = .model.generate(final_prompt)
            predicted_functions = ._extract_functions(final_answer)

            
            is_correct = (gt  predicted_functions  gt  ground_truth_functions)

            trajectory = Trajectory(
                steps=steps,
                final_answer=final_answer,
                is_correct=is_correct,
                ground_truth=.join(ground_truth_functions)
            )
            trajectories.append(trajectory)

        
        correct_trajectories = [t  t  trajectories  t.is_correct]
         correct_trajectories

     () -> [, ]:
        
        lines = response.split()
        tool_name = 
        tool_input = 

         line  lines:
             line.startswith():
                tool_name = line.replace(, ).strip()
             line.startswith():
                tool_input = line.replace(, ).strip()

         tool_name, tool_input

     () -> []:
        
        
         re
        functions = re.findall(, text)
         functions  [text.split()[].strip()]

 ():
    
    collector = SFTDataCollector(model, repo_searcher, num_samples=)

    sft_data = []
     issue, ground_truth_funcs  issues_and_functions:
        trajectories = collector.generate_trajectories(issue, ground_truth_funcs)
         trajectory  trajectories:
            
            sft_data.append({
                : ,
                : trajectory.steps,
                : trajectory.final_answer
            })

     sft_data

Step 3: Implement Stage 2 - RL with nDCG Rewards

Train using ranking-aware rewards:

import numpy as np
from typing import List

def compute_ndcg_at_k(predictions: List[str], ground_truth: List[str], k: int = 5) -> float:
    """
    Compute nDCG@k metric: Normalized Discounted Cumulative Gain.
    Rewards both correctness and ranking quality.

    A correct function at rank 1 gets higher reward than at rank 5.
    """
    # Binary relevance: 1 if function is ground truth, 0 otherwise
    relevance = [1.0 if pred in ground_truth else 0.0 for pred in predictions[:k]]

    # Discounted cumulative gain
    dcg = sum(rel / np.log2(i + 2) for i, rel in enumerate(relevance))

    # Ideal DCG: perfect ranking would have all ground truth first
    ideal_relevance = [1.0] * min(k, len(ground_truth))
    idcg = sum(rel / np.log2(i + 2) for i, rel in enumerate(ideal_relevance))

    # Normalized
    ndcg = dcg / idcg if idcg > 0 else 0.0
    return ndcg

class RLTrainer:
    """
    Train model using RL with nDCG rewards.
    Reinforcement learning explores beyond supervised trajectories.
    """
     ():
        .model = model
        .repo_searcher = repo_searcher
        .learning_rate = learning_rate

     () -> :
        
        ndcg = compute_ndcg_at_k(predictions, ground_truth, k=)
        length_penalty =  * trajectory_length  

        reward = ndcg - length_penalty
         reward

     () -> [, [Trajectory]]:
        
        trajectories = []

         _  ():  
            steps = []
            current_state = 

            
             step  ():
                prompt = 
                tool_name, tool_input = ._sample_action(prompt)

                 tool_name == :
                    

                tool_result = .repo_searcher.execute_tool(tool_name, tool_input)
                steps.append((tool_name, tool_input, tool_result))
                current_state += 

            
            final_answer = .model.generate()
            predictions = ._extract_functions(final_answer)

            
            reward = .compute_reward(predictions, ground_truth, (steps))

            trajectory = Trajectory(
                steps=steps,
                final_answer=final_answer,
                is_correct=(p  ground_truth  p  predictions),
                ground_truth=.join(ground_truth)
            )
            trajectory.reward = reward
            trajectories.append(trajectory)

         (t.reward  t  trajectories) / (trajectories), trajectories

     () -> [, ]:
        
        
        
         , 

     () -> []:
        
        

     ():
        
         epoch  (num_epochs):
            total_reward = 

             issue, ground_truth  issues_and_ground_truth:
                episode_reward, trajectories = .run_episode(issue, ground_truth)
                total_reward += episode_reward

                
                 trajectory  trajectories:
                    advantage = trajectory.reward - total_reward  
                    
                    ._update_policy(trajectory, advantage)

            ()

     ():
        
        
        

Step 4: Evaluate on Benchmarks

Measure performance improvement:

def evaluate_model(model, repo_searcher, test_issues, test_ground_truth) -> Dict:
    """Evaluate model on test set, measuring localization accuracy."""
    predictions = []

    for issue in test_issues:
        # Run model once (no sampling)
        trajectory = []
        current_state = f"Issue: {issue}"

        for _ in range(10):
            tool_name, tool_input = model.act(current_state)
            if tool_name == 'Exit':
                break
            result = repo_searcher.execute_tool(tool_name, tool_input)
            trajectory.append((tool_name, tool_input, result))
            current_state += f"\n{tool_name}({tool_input})"

        final_answer = model.generate(f"{current_state}\nFunctions?")
        predictions.append(model.extract_functions(final_answer))

    # Compute metrics
    ndcg_scores = [
        compute_ndcg_at_k(pred, truth, k=5)
        for pred, truth in zip(predictions, test_ground_truth)
    ]

    return {
        'mean_ndcg@5': np.mean(ndcg_scores),
        'success_rate': sum(1 for pred, truth in zip(predictions, test_ground_truth)
                           if any(p in truth for p in pred)) / len(predictions),
    }

Practical Guidance

When to Use:

• Bug localization in large codebases (>100K lines)
• Scenarios where ground truth function labels are available for training
• Applications emphasizing systematic code exploration over pattern matching
• Cases where ranking quality matters (top-k function suggestions)

When NOT to Use:

• Small codebases where simple string search suffices
• Scenarios with limited labeled training data (<50 examples)
• Real-time systems requiring <100ms latency (requires full trajectory generation)
• Domains without clear hierarchical structure (unstructured data)

Hyperparameters:

Parameter	Default	Impact
num_sft_samples	10	Trajectories sampled per issue in Stage 1; higher = better SFT data quality
sft_reject_threshold	1.0	Keep only perfectly correct trajectories; lower = more training data
rl_episodes_per_issue	3	Trajectories sampled per issue in RL; balance exploration vs. computation
ndcg_k	5	Top-k functions evaluated in reward; match use case requirements
length_penalty	0.01	Penalize long trajectories; higher = encourage efficiency

Reference

Paper: Tool-integrated RL for Repo Deep Search (2508.03012)

• Two-stage training: SFT + RL with ranking-aware rewards
• nDCG@k metric balances correctness and ranking quality
• 32B ToolTrain model matches Claude-3.7-Sonnet on function-level localization