Domain-Adaptable Reinforcement Learning for Code Generation with Dense Rewards

TL;DR

This paper introduces a reinforcement learning framework that fine-tunes pre-trained language models for domain-specific code generation, improving correctness and executability in robotics and general programming tasks.

Contribution

It presents a customizable, execution-aware reinforcement learning approach with token-level credit assignment to adapt language models to diverse code generation domains.

Findings

19% increase in pass@1 on MBPP/MBPP+ benchmarks

51% reduction in execution failures on RoboEval

Substantial improvements in functional correctness and simulator executability

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.