Contrastive Code Graph Embeddings for Reinforcement Learning-Based Automated Code Refactoring

ICLR 2026 Conference Submission25618 Authors

Keywords: Code Refactoring

Abstract: We propose a novel reinforcement learning (RL) framework for automated code refactoring that uses contrastive pre-trained code graph embeddings to overcome the limitations of the traditional heuristic-based reward functions. The key challenge is balancing the implementation of syntactic improvements - while maintaining the semantics of the code being refactored - something that necessarily requires the existing RL approaches to accomplish and that most often do last year because of the handcrafted nature of their metrics. Our approach presents a syntax-guided contrastive encoder that acquires structural invariant representations of code graphs by relating structurally augmented variants under a self-supervised objective. These embeddings are then combined with standard measures of code quality in a composite reward function, allowing the RL agent to reason about both low-level changes to the syntactic structure as well as high-level changes in the semantic structure. The policy network itself, which takes the form of a graph attention network, runs on the joint representation space directly, which models dependency on the context on the code structure.

Primary Area: transfer learning, meta learning, and lifelong learning

Submission Number: 25618