技術概要：密な報酬を伴うコード生成のためのドメイン適応型強化学習

問題定義

大規模言語モデル（LLM）は、自動化されたコード生成において大きな可能性を示していますが、正確性、品質、安全性、およびドメイン固有の制約への準拠については保証が欠如していることが頻繁にあります。検証可能な報酬を用いた強化学習（RLVR）は、実行ベースのフィードバックを用いてコード生成を最適化する手法として登場しましたが、現在の実装はスパースなシーケンスレベルの報酬に大きく依存しています。これらの標準的なアプローチでは、最終的なコードシーケンスにエラーが生じると、そのペナルティがシーケンス全体に適用され、失敗に寄与した特定のトークンに関する詳細なフィードバックを提供できません。さらに、既存の手法は、ロボット工学においてコードがコンパイルされるだけでなく、ハードウェア上で安全に機能する必要があるような、物理的な安全性や環境認識といったドメイン固有の制約を組み込むことにしばしば苦労しています。

手法

著者らは、密なトークンレベルの報酬帰属メカニズムを通じて報酬のスパース性に対処する、**近接方策最適化（PPO）**に基づく統合的なファインチューニングフレームワークを提案します。このフレームワークは、3 つのフェーズからなる反復ループとして動作します。

1. ロールアウトフェーズ: 方策モデル（$\pi_\theta$）がコードをトークン単位で生成します。生成中に、logits プロセッサ（Syncode）が各トークンの構文の正しさを評価し、即時フィードバックを提供します。
2. 評価フェーズ: シーケンスが完了すると、複数のコンポーネントを集約することで包括的な報酬が計算されます。
   • 構文の正しさ（$R_{sync}$）: 言語の文法に違反するトークンをフラグ付けします。
   • コードスタイルと脆弱性（$R_{lint}$）: Ruff リンターを使用してエラーやセキュリティ上の問題を検出します。
   • 分布正則化（$R_{KL}$）: 破滅的な方策のドリフトを防ぐために、参照モデルからの乖離に対してペナルティを課します。
   • タスク固有の報酬（$R_{opti}$）: ユニットテストの合格率、データフローグラフ（DFG）の一致、シミュレータのフィードバック（例：ロボット工学向けの RoboSim）などを含むカスタマイズ可能なシグナル。
   • 価値モデル（$V_\phi$）は、期待報酬を推定してアドバンテージを計算します。
3. 最適化フェーズ: 方策モデルと価値モデルは、計算されたアドバンテージと複合報酬を用いて更新されます。

主要な革新：密なトークンレベル帰属
エンド・オブ・シーケンス（EOS）でのみ報酬を割り当てる標準的な PPO 実装（例：trl）とは異なり、このフレームワークはスパースなシーケンスレベルのシグナルを密なトークンレベルのシグナルに変換します。

• 正確な帰属が可能である報酬（例：リンターチェック、構文エラー）については、違反を責任のある特定のトークンにマッピングします。
• 正確な帰属が困難な報酬（例：ユニットテストの結果）については、報酬またはペナルティを生成されたコードトークン全体に均等に分配し、非コードトークンにはシグナルを与えません。
  これにより、各時間ステップで連続的なフィードバックシグナル（$R_t$）が生成され、きめ細かいクレジット割り当てが可能になります。

主要な貢献

1. 統合 PPO フレームワーク: 構文制約、静的解析、実行結果、シミュレータフィードバックを単一の報酬ループに統合するファインチューニングアーキテクチャ。
2. 密なトークンレベル報酬帰属: スパースなシーケンスレベルのフィードバックを局所的なトレーニングシグナルに変換するメカニズムにより、モデルがコードのどの部分が成功または失敗に寄与したかを学習できるようにします。
3. ドメイン適応性: 汎用 Python プログラミングとロボットプログラム合成という 2 つの異なるドメインにおけるフレームワークの有効性の実証。コンパクトなモデルが、正確性と実行性の両面で大幅な改善を達成できることを示しています。

実験結果

このフレームワークは、汎用コード生成用のMBPP/MBPP+ベンチマークと、ロボットタスク用のROBOEVALベンチマークにおいて、Qwen2.5-Coder-1.5Bモデルを使用して評価されました。

• 汎用コード生成:
  • MBPPにおいて、RL によって洗練されたモデルは0.653 の pass@1を達成し、ベースモデル（0.460）に対して**+19.3%**の絶対的な改善を示しました。
  • MBPP+において、モデルは 0.413 から0.556（+14.3%）に改善し、エッジケースに対する堅牢性の向上を示しています。
• ロボットプログラム合成:
  • ベースモデルは 80 タスクのうち 77 タスクで実行可能なコードを生成できず（主に Python レベルのエラー）、シミュレータ内での成功した実行は 0 でした。
  • RL ファインチューニング後、Python レベルのエラーは 11 に減少し、14 タスクが正常に解決されました。
  • これは実行失敗の 51% 削減を意味し、出力を主に実行不可能なコードから、環境を認識し実行可能なプログラムへと転換しました。
  • 特筆すべきは、ファインチューニングされた 1.5B モデルが、より大きな 7B モデルに適用された ROBO-INSTRUCT 手法よりもはるかに大きな相対的な性能向上を示し、コンパクトなモデルと大規模モデルの間のギャップを狭めたことです。

意義と主張

本論文は、構造化された強化学習、特に密な実行認識報酬を備えたものは、LLM をプログラムの正確性とドメイン固有の要件に整合させるための実用的かつ拡張可能なアプローチであると主張しています。

著者らは、このアプローチが表面的なパターンマッチングを超えて、構造的に整合性があり意味的に一貫したプログラム生成を促進すると強調しています。多段階のフィードバックを最適化ループに直接統合することで、このフレームワークは機能的な正確性と実行可能なロボット動作の両方への反復的な整合を可能にします。結果は、蒸留ベースの手法の計算オーバーヘッドやスパースな報酬シグナルの制限なしに、コンパクトな言語モデルでも正確性と実行性の大幅な改善を達成できることを示唆しています。

本研究は、報酬における長さバイアスの対処や、より広範な実行ベースのメカニズムの探求など、今後の研究が必要であるとしつつも、提案された手法が、汎用的なコード生成と、ロボット工学のような特殊なドメインで必要とされる厳格な安全性および物理的制約との間のギャップを効果的に埋めていると結論付けています。