PaperlayerView exact revision on arXiv

Source-linked AI summary

A^2Agent: Action-Aware Reinforcement Learning for Repository-Level Code Localization Agents

Doyeon Kim, Suyoung Bae, Yumin Lee, Jee-Hyong Lee

arXiv:2608.29831v1cs.CLcs.SE

TL;DR

Repository-level code localization requires agents to identify issue-relevant regions, but sparse trajectory-level signals do not identify effective per-turn actions or ensure discovered regions are committed. A2AGENT combines discovery-and-commit rewards with context-matched action-level advantage estimation. It improves average F1 over SOTA on SWE-Bench Verified and SWE-Bench Pro, while its 4B model surpasses baselines up to 8× larger.

  • Problem

    Existing localization agents rely on sparse trajectory-level signals that cannot distinguish effective per-turn actions or resolve discovery–commit failures.

  • Method

    A2AGENT combines per-turn rewards for discovering and committing gold code regions with action-level advantages estimated by grouping turns sharing exploration context.

  • Results

    A2AGENT improves average F1 over SOTA by 1.58% on SWE-Bench Verified and 8.55% on SWE-Bench Pro, with its 4B model surpassing baselines up to 8× larger.

  • Takeaways & Limitations

    A2AGENT provides action-level supervision for effective exploration and reliable commitment in repository-level code localization.

  • Takeaways & Limitations

    Experiments used small open-source models because of computational resource constraints, leaving optimization for downstream issue resolution as future work.

Abstract

from arXiv · show

Localizing issue-relevant code regions is a critical step in automated software engineering. However, due to their reliance on sparse trajectory-level signals, existing methods cannot identify which per-turn actions are effective and often discover correct code regions during exploration but fail to commit them. To address these limitations, we propose an action-aware reinforcement learning method that combines a per-turn reward sequence rewarding both the discovery and commitment of gold code regions with an action-level advantage estimation scheme that isolates each action's credit by grouping turns sharing the same exploration context. Extensive evaluations show that our method improves the average F1 over the state-of-the-art (SOTA) by 1.58% on SWE-Bench Verified and 8.55% on SWE-Bench Pro, with our 4B model outperforming baselines up to 8x larger. Our code is available at https://github.com/donian00/A2Agent.

1 Introduction

Repository-level code localization is a prerequisite for issue resolution, but existing agents struggle to choose useful actions and to commit code regions they have discovered. A2AGENT addresses both problems with action-aware reinforcement learning and improves average F1 on two benchmarks.

  • Accurate localization identifies the code regions requiring modification and is necessary for effective downstream issue resolution.
  • Growing repository scale and dependency complexity has motivated multi-turn agents that navigate codebases through iterative tool calls.
  • Existing methods cannot distinguish effective from erroneous per-turn actions, leading to redundant tool calls and weak credit assignment.
  • Agents may observe gold code regions during exploration but fail to commit them, while trajectory-level sparse rewards treat these cases like complete failures.
  • A2AGENT combines per-turn rewards for discovering and committing gold regions with action-level advantage estimation matched by exploration context.
  • 1.58% and 8.55% average-F1 gains over SOTA are reported on SWE-Bench Verified and SWE-Bench Pro, respectively.

2 Related Work

Repository-level localization has progressed from retrieval and agent navigation toward post-training, but existing approaches leave per-turn action contribution and final commitment unresolved. A2AGENT uses context-matched action-level estimation alongside reward design and tool-use training.

  • Repository-Level Code Localization: Similarity-based file retrieval is limited by the semantic gap between issue descriptions and code and by difficulty modeling inter-file dependencies.
  • Repository-Level Code Localization: Agent-based methods iteratively invoke tools to navigate repositories, yet identifying useful per-turn actions and committing observed code elements remain unsolved.
  • Agentic Training: SFT-based agentic training struggles to generalize to unfamiliar codebases, while GRPO and GSPO estimate a single trajectory-level advantage from rollout groups.
  • Agentic Training: A2AGENT groups turns sharing the same exploration context to estimate action-level advantages rather than assigning one advantage to every turn.
  • Agentic Training: The training pipeline first fine-tunes tool-use ability and then applies Action-Aware RL to improve localization performance.

3 Preliminary

Repository-level code localization is formulated as a multi-turn decision process in which an agent acts from issue, repository, and interaction-history state. Conventional training observes reward only at the final localization decision, while group-based advantages reflect whole-trajectory performance.

  • Problem Setup: At each turn, the agent observes state s_t and samples action a_t from policy π_θ while interacting with the codebase.
  • Problem Setup: The agent produces a localization decision at final turn T, yielding trajectory τ, and conventional reward evaluates only that final prediction.
  • Group-Relative Advantage Estimation: Group-based methods estimate advantage by subtracting the mean return among trajectories sampled for the same issue, without a separate value function.
  • Group-Relative Advantage Estimation: Because the trajectory-level value is shared across turns, it cannot represent the contribution of individual actions or distinguish different actions under the same state.

4 Proposed Method

A2AGENT formulates repository-level code localization as multi-turn tool-driven decision making and addresses sparse credit assignment with turn-level rewards and action-level advantages. Its design rewards partial discovery and commitment while comparing actions within matched exploration contexts.

  • Task Formulation: A2AGENT treats code localization as a multi-turn process in which tool calls transform interaction history into a final prediction at file, module, and function levels.The action space contains six repository-navigation and retrieval tools, and each action includes a tool name and arguments.
  • Motivation: Trajectory-level signals miss partial discoveries and cannot reliably distinguish effective actions from misleading tool calls.They also fail to address cases where the agent discovers correct regions but does not commit them to its final prediction.
  • Reward Design: The reward design combines a terminal reward for final prediction quality and discovery commitment with turn-level feedback for newly discovered gold locations.The terminal reward includes file-, module-, and function-level F1 plus a Discovery-Commit reward; turn-level rewards preserve signals for partially correct trajectories.
  • Action-Level Advantage Estimation: Action-level advantages compare discounted returns among turns sharing the same preceding exploration history, isolating an action’s contribution relative to alternatives.The state proxy uses preceding K-action sequences, excludes the current action from the grouping key, and relies on deterministic read-only tools for comparable observations.
  • Training Procedure: The training procedure collects multiple trajectories per issue and weights tokens according to clipped turn-level advantages during maximum-likelihood optimization.The objective assigns each token the weight of its corresponding turn, while clipping limits outlier advantages and stabilizes the positive weights.

5 Experiments

A2AGENT is evaluated on held-out repository-level code-localization benchmarks using variable-length predictions and average F1 across file, module, and function granularities. It consistently improves localization, reduces uncommitted discoveries, and lowers inference cost through more targeted search.

  • Experimental Setup: A2AGENT is trained on 2,438 SWE-Gym instances and evaluated on repository-disjoint SWE-Bench Verified and SWE-Bench Pro benchmarks.SWE-Bench Verified contains 500 human-verified instances, while SWE-Bench Pro contains 266 instances.
  • Experimental Setup: The evaluation predicts a variable number of locations and uses instance-level average F1 across file, module, and function granularities.Precision and recall are also reported at each granularity for comparisons across methods with different prediction lengths.
  • Main Results: 1.58% higher average F1 than CodeScout-14B is achieved on SWE-Bench Verified, while the 4B model surpasses baselines up to 8× larger.A2AGENT also consistently outperforms open-source baselines across all three granularities on SWE-Bench Pro.
  • Ablation Study: History Grouping and discovery–commit reward contribute independently and combine for the best result across file, module, and function granularities.History Grouping compares actions under matched exploration contexts, while the discovery–commit reward encourages committing observed gold elements.
  • Further Analysis: Over 40% of baseline failures include a discovered but uncommitted gold location, whereas the discovery–commit reward reduces this gap while improving file-level F1.The Discovery-Commit gap is defined as failures where a gold location was explored but omitted from the final prediction.
  • Further Analysis: History depth H=2 best balances grouping precision and coverage; larger depths sharply reduce coverage and make credit estimation noisier, causing F1 to decline.A2AGENT also reduces average observation length by approximately 40% compared with baselines, yielding the lowest inference time among agent-based methods.

6 Conclusion

A2AGENT combines action-level advantage estimation with a Discovery-Commit reward to improve repository-level code localization and commitment of discovered gold locations. It achieves strong performance while remaining deployable on-premises without external API access or Docker.

  • 6 Conclusion: A2AGENT uses action-level advantages and a Discovery-Commit reward so agents learn which tool calls advance localization and commit discovered gold locations.The method targets both exploration effectiveness and reliable final commitment.
  • 6 Conclusion: A2AGENT outperforms baselines up to 8× larger models in repository-level code localization.The conclusion reports this result without restricting it to a single benchmark.
  • 6 Conclusion: A2AGENT requires no external API access or Docker setup and can run entirely on-premises for source-code privacy requirements.This deployment property is explicitly identified as practical where source-code privacy is required.

Limitations

The experiments were conducted on small open-source models because of computational resource constraints. The method was effective at this scale, while further downstream issue-resolution optimization remains future work.

  • Limitations: Experiments used small open-source models because computational resource constraints prevented evaluation on larger-scale models.This limits the tested model scale rather than reporting a failure of the method.
  • Limitations: Further optimization for downstream issue resolution remains a promising direction despite the method’s effectiveness on that task.The paper reports downstream effectiveness in Table 5 but identifies additional optimization as future work.

A Tool Specifications

A2AGENT uses an existing repository-tool infrastructure spanning graph traversal, lexical and semantic search, entity inspection, summaries, and commit history. Its action-level advantages are reported as stable and approximately zero-centered.

  • Tool Infrastructure: A2AGENT directly adopts prior tool infrastructure so performance gains are attributable to learning-signal design rather than novel tool design.The infrastructure is built on repository graph, BM25, semantic summary, and commit-history indexes.
  • TraverseGraph: TraverseGraph explores repository dependencies from an entity for a bounded number of hops across containment, import, invocation, and inheritance edges.It is the primary entry point for understanding unfamiliar repository structure.
  • SearchCodeSnippet: SearchCodeSnippet uses BM25 keyword or identifier queries to return matching file paths, line ranges, and code snippets.It is most effective when issue identifiers narrow candidates for later entity searches.
  • SearchEntity: SearchEntity retrieves a file, class, or method path together with its raw source code for inspecting narrowed candidates.The tool supports confirmation at function-level localization.
  • Summary Tools: SearchSummary finds semantically similar file summaries, while ViewSummary presents a file’s responsibilities, classes, functions, and dependencies.Together they support candidate narrowing and rapid file-level understanding.
  • SearchCommit: SearchCommit returns matching historical commits with SHAs, messages, dates, and changed files to supplement static analysis with repository-evolution context.Historical co-modification information provides an additional signal for localization.
  • Training Diagnostics: The action-level advantage distribution with H=2 is approximately zero-centered, with µ = −0.006 and σ = 0.404.The paper interprets this distribution as indicating a stable and unbiased baseline for action-level credit assignment.

C Benchmark Details

The benchmark section defines SWE-Bench-based localization evaluation, describes baseline and tool configurations, and presents the A2AGENT training procedure. Evaluation targets are extracted at file, module, and function granularities, while A2AGENT retains the shared tool infrastructure used by existing methods.

  • Benchmark Definition: SWE-Bench instances pair real-world GitHub issues with corresponding pull requests from popular open-source Python repositories.For localization, gold targets are extracted from the gold patch at file, module, and function granularities.
  • Experimental Controls: A2AGENT uses the same tool infrastructure as existing methods, focusing the benchmark comparison on training methodology rather than tool-set novelty.The tool configurations are summarized in Table 7.
  • Evaluation Benchmarks: SWE-Bench Verified is a human-validated 500-instance subset, whereas SWE-Bench Pro is a harder 266-instance benchmark spanning GPL-licensed and commercial repositories.SWE-Bench Pro is described as targeting contamination mitigation and long-horizon tasks.
  • Baseline Methods: Agentless uses a fixed three-stage prompting pipeline, while CoSIL, OrcaLoca, and LocAgent employ progressively different graph, scheduling, pruning, or indexing strategies.These baselines differ in inference organization and repository-navigation mechanisms.
  • Training Procedure: Algorithm 1 initializes from an SFT policy, samples N rollouts per query, computes discounted returns and action-level advantages, and optimizes the policy.The procedure uses history depth H, discount factor γ, and advantage weights ω0:H.
  • Agentic Training Baselines: RepoSearcher trains localization agents through two-stage narrowing and teacher-based ToolTrain distillation, while RepoNavigator and CodeScout use GRPO and GSPO respectively.RepoNavigator uses a Jump tool; CodeScout uses a Unix terminal and LocalizationFinish.

E Evaluation Protocol

The evaluation predicts variable numbers of locations and measures instance-level average F1 across file, module, and function granularities. Precision and recall are also reported to compare systems producing different numbers of candidate locations, but repository entities outside submitted candidates are excluded from precision.

  • Instance-level average F1 is computed between predicted and ground-truth locations across file, module, and function granularities.
  • Precision uses the submitted candidate set as its denominator, excluding all other repository entities and leaving no negative class.

F Implementation Details

Implementation uses repository-wide static indexes, teacher SFT warm-up, and iterative advantage-weighted training, with controlled rollout and optimization settings. Case studies contrast A2AGENT’s successful adaptive localization with SFT-only failure on an unseen codebase.

  • Training infrastructure: Training uses four NVIDIA RTX Pro 6000 GPUs with PyTorch, Python 3.12, HuggingFace Transformers, DeepSpeed ZeRO-3, and vLLM.
  • Training procedure: Teacher SFT warm-up filters 1,410 trajectories matching at least one gold location at each granularity.SFT uses learning rate 2×10^-4 and sequence length 36K, computing loss only on assistant turns.
  • Rollout settings: Rollouts use temperature 1.0, top-p 0.8, repetition penalty 1.05, at most 20 tool calls, and a 40K context limit.Evaluation results use a single greedy-decoding run.
  • Training procedure: Advantage-weighted training uses environment-derived token weights rather than a clipped PPO or GRPO policy-gradient objective.The procedure repeatedly collects rollouts from the current policy and trains with advantage-weighted maximum likelihood.
  • Training procedure: The method performs up to three rollout-training iterations, with β=1.5 and κ=5.0 selected by validation F1.
  • Repository indexing: Static dependency-graph and BM25 indexes are built once per repository snapshot and reused across inference rollouts.The read-only index is shared across rollouts without execution-environment or state-restoration requirements.
  • Case study: A2AGENT successfully localizes the gold location on sphinx-doc/sphinx#10449, whereas SFT-only training fails to reach it.The case study attributes the contrast to A2AGENT adapting its search strategy while SFT imitates teacher trajectories.
Loading 2608.29831v1…