Domain-Adaptable Reinforcement Learning for Code Generation with Dense Rewards

2026-05-20

AI Summary

This work addresses the challenges large language models face in generating code that is correct, safe, and adheres to domain-specific constraints—such as physical limitations in robotic systems. The authors propose a reinforcement learning fine-tuning framework based on Proximal Policy Optimization (PPO), incorporating a dense reward mechanism that jointly evaluates syntax, functional correctness, coding style, safety, and sim-to-real executability. To enable effective credit assignment from program execution outcomes back to individual tokens, they introduce a token-level reward mapping strategy. The approach facilitates cross-domain adaptive code generation, achieving a 19% improvement in pass@1 on the MBPP benchmark and reducing execution failure rates by 51% on RoboEval, thereby significantly enhancing both functional correctness and domain-specific adaptability of generated code.

Abstract

Large language models show strong potential for automated code generation, but lack guarantees for correctness, quality, safety, and domain-specific constraints. For instance in robotics, where code generation is increasingly being used for planning and executing actions, awareness of the environment and physical constraints is critical. To facilitate the adaption of code-generating LLMs to diverse requirements, including domain-specific ones, we present a reinforcement learning framework that fine-tunes pre-trained LLMs using proximal policy optimization. Our customizable execution-aware reward formula captures and optimizes syntax, functional correctness, code style, security, and simulator executability. A token-level reward mapping mechanism enables effective credit assignment from execution outcomes to generated tokens. The framework is evaluated on general-purpose code generation (MBPP/MBPP+) and robotic program synthesis (RoboEval). The results show substantial improvements in functional correctness and simulator executability, including an absolute pass@1 increase of 19% on MBPP and a reduction in execution failures by 51% on RoboEval. These findings demonstrate that structured reinforcement learning can effectively align language models to correct program generation and domain-specific requirements.

Problem

• code generation
• domain-specific constraints
• functional correctness
• large language models
• reinforcement learning

Innovation

• reinforcement learning
• code generation
• dense rewards
• domain adaptation
• token-level reward