AWSがLLM強化学習の「報酬設計」詳細を公開 — スコアが正常に見えたまま学習が止まる「報酬崩壊」をどう防ぐか

8月15日、AWSが「Custom reward functions for multi-turn reinforcement learning with Amazon Nova Forge」と題した記事を公開した。Amazon Nova Forgeを使ったマルチターン強化学習において、カスタム報酬関数をどう設計・実装するかを詳述した内容だ。報酬設計のわずかなズレがモデルに意図しない挙動を静かに教え込むという問題意識を軸に、実装コード・セキュリティ対策・失敗事例まで踏み込んで解説されている。

報酬関数がすべてを決める

強化学習ファインチューニング（RFT）では、報酬関数の設計が学習結果を左右する。わずかにズレた報酬設計は、学習曲線が正常に見えている状態で、モデルに意図しない挙動を静かに教え込む。

Amazon Nova ForgeのマルチターンRFTは、GRPO（Group Relative Policy Optimization）というアルゴリズムを採用している。GRPOは1つの会話プロンプトに対してK個（記事中の実装例ではデフォルト8個）のロールアウト（モデルの応答軌跡）を生成し、それらを報酬でランク付けして、正規化された優位性（アドバンテージ）をもとにモデルの重みを更新する。

ここで重要な原則がある。グループ内でスコアにばらつきが生まれなければ、その報酬成分は勾配に一切貢献しない。全ロールアウトが同じ値を返す報酬成分は、重み付けがどれだけ高くても学習に寄与しないのだ。

マルチターン報酬の3つの構造

単一スカラーの報酬はゲームされやすく、最終ターンだけを評価するスパースな報酬はマルチターンタスクでは学習が収束しにくい。実用的なマルチターン報酬は以下の3種類の信号を組み合わせる。

• エピソードレベルの報酬（結果）: ユニットテスト合格率など、最終的な成果を測る。目標に直結するが、学習初期はほぼゼロになりがち
• ターンレベルの報酬（行動）: 適切なツール呼び出しや確認質問など、望ましい中間行動を評価する
• ペナルティ: 即座のコード推測や応答ループといった失敗モードを明示的に抑制する

実装例：「コードを書く前に質問する」挙動の学習

記事では、Amazon Nova Lite 2.0を対象に、500種類のプログラミングタスクでマルチターンRFTを実施した具体例が示されている。タスクの設計は以下のとおりだ。

• モデルには仕様が不完全なコーディング要求が提示される
• ユーザーシミュレーターが完全仕様を持っており、モデルが質問したときだけ詳細を開示する
• モデルは各ターンで「質問する」か「コードをコミットする」かを選ぶ

「推測してコードを書く→テスト失敗」「質問して仕様を把握する→テスト通過」という構造により、質問してからコミットするという行動をタスク設計で強制している。

4成分の複合報酬

設計上のポイントは2点ある。第一に、**asked_before_codingは正解率に依存しない独立した評価になっている。ただし、最終的なコードコミットを条件とすることで「永遠に質問し続けてコミットしない」という抜け道を塞いでいる。第二に、guessed_immediatelyペナルティによって「推測」戦略を明示的に不利にし、GRPO内のグループで戦略間のばらつきを確保**している。

実際のコードは以下のような形になる。

def asked_before_coding(completion, answer, **kw) -> float:
    msgs = _messages(completion, kw)
    first_q = _first_question_turn(msgs, parser)
    final = _final_code(completion, parser)
    committed = bool(final) and not _is_question(final)
    if first_q == 1 and committed:
        return 1.0          # 最初に質問してコミット（理想）
    if first_q is not None and committed:
        return 0.6          # 後で質問してコミット
    return 0.0              # 質問なし、またはコミットなし

def guessed_immediately(completion, answer, **kw) -> float:
    for m in _assistant_turns(completion, kw):
        code = _code_of(parser.parse(m["content"]))
        return -1.0 if (code and not _is_question(code)) else 0.0
    return 0.0

モデル生成コードの安全な実行

correctness成分はモデルが生成したコードをユニットテストで実行して採点する。RLの探索フェーズで生成されるコードは検証済みでないため、セキュリティ対策が必要だ。記事では以下のアプローチが示されている。

• クレデンシャルやネットワークを生成コードに渡さない（env={"PATH": "/usr/bin"}のみ）
• CPUタイム・メモリ・プロセス数にリソース制限をかける（resource.setrlimit）
• ランダムなnonce（ワンタイムトークン）をマーカーとして使う。これにより、モデルが期待値をstderrに直接書き込んでスコアを偽装することを防ぐ
• 実行されたテスト数を期待値と照合する。モデルが自分で自明なテストを追加してスコアを希薄化することへの対策だ

nonce = secrets.token_hex(8)  # 実行ごとに異なるマーカー
# リソース制限の適用
resource.setrlimit(resource.RLIMIT_CPU, (timeout_s, timeout_s))
resource.setrlimit(resource.RLIMIT_AS, (2 * 1024**3, 2 * 1024**3))  # 2 GB
resource.setrlimit(resource.RLIMIT_NPROC, (64, 64))

報酬崩壊の罠：最も危険な失敗モード

記事が特に強調しているのが報酬崩壊（reward collapse）だ。報酬ハッキング（proxy指標のゲーム）や学習不安定化（KL発散）は損失曲線やトランスクリプトに現れやすい。しかし報酬崩壊は集計スコアが正常に見えたまま、内部でグループ内のばらつきが消滅して学習が停止するという点で、他の失敗モードより検知が難しい。

実際のトレーニングランで発生した事例として、記事では**correctness成分（重み1.0と最も高い）が全ロールアウトで一様にゼロを返し続け、一切の学習信号を生み出していなかったケースが報告されている。タスク難易度の設定が高すぎたため、K個のロールアウトすべてがユニットテストを通過できず、グループ内のスコアが均一なゼロに張り付いた状態だ。集計スコア上は他の成分が機能しているように見えるため、ログを見ただけでは問題に気づきにくい。対策は、metrics_listで各成分のスコア分布を別途ロギングし、成分ごとのグループ内分散をトレーニング中に継続監視する**ことだ。

インフラ構成

Nova ForgeのBYOO（Bring Your Own Orchestration）パスでは、rollout.delegate: trueを設定し、報酬ロジックをAmazon ECS上のコンテナで動かす。コンテナが会話状態・ユーザーシミュレーター・コード実行・検証を管理し、集計報酬スコア（aggregate_reward_score）と成分別スコア（metrics_list）をNova Forgeに返す構成だ。インフラのAWS CDKデプロイを含む詳細はシリーズのPart 1で解説されており、サンプルコードはaws-samples/sample-nova-multi-turn-rl-infraリポジトリで公開されている。

詳細はCustom reward functions for multi-turn reinforcement learning with Amazon Nova Forgeを参照していただきたい。