FlashREINFORCEとは？
FlashREINFORCEは、エージェント型言語モデルの非同期学習向けに設計された、クリティック不要の強化学習フレームワークです。可変長の推論、ツール呼び出し、環境とのインタラクションによって軌跡が不規則なタイミングで完了する、長期的な処理を想定しています。
各プロンプトに対して複数の兄弟ロールアウトを生成する代わりに、FlashREINFORCEはプロンプトごとに1つのロールアウトを使用します。そのため、B本の軌跡を含むバッチではB個の異なるプロンプトを扱うことができ、同じプロンプトへの複数の応答を待たずに、完了した軌跡を学習器へ送ることができます。
このリポジトリには、スタンドアロンのPyTorch参照損失、CPUでのオプティマイザー更新例、テスト、R1およびQwen2.5-Math向けのMoltランチャー、推論、Pythonツール、Mixture-of-Expertsモデル、ALFWorld向けの実験設定が含まれています。完全な非同期学習ランチャーではMoltを使用します。
強化学習はREINFORCEであるべきです。
なぜ単一ロールアウトを使うのか？
グループ相対的手法は、同じプロンプトから複数のロールアウトを必要とします。非同期エージェント処理では、この要件によって同期の障壁が生じ、固定されたロールアウト予算で扱える異なるプロンプトの数が減る可能性があります。
単一ロールアウト学習では兄弟ロールアウトが不要になりますが、次の3つの安定性に関する問題が生じます。
- プロンプトレベルのベースラインを構築するための兄弟グループが存在しません。
- 非同期に収集された軌跡は、古い方策で生成されている可能性があります。
- すべてのトークンにバッチ内で同じ重みを与えると、長い失敗軌跡が更新を支配する可能性があります。
FlashREINFORCEは、One-Batch REINFORCE、Sequence Trust Region、Sample-Mean Optimizationによってこれらの問題に対処します。
主要コンポーネント
1. One-Batch REINFORCE
新しい各バッチについて、学習器は独立したプロンプト間でスカラー報酬を中心化します。軌跡iの報酬をR_iとすると、そのアドバンテージは軌跡報酬からバッチ平均を引いた値です。
mean_reward = sum(R_j for j in batch) / B
A_i = R_i - mean_reward平均を上回る軌跡には正のアドバンテージが、下回る軌跡には負のアドバンテージが与えられます。これにより、学習されたクリティックなしで符号付きのフィードバックが得られます。学習器はバッチに対してオプティマイザー更新を1回実行し、その後バッチを破棄します。
2. Sequence Trust Region
すべての軌跡には、サンプリングされた各トークンを実際に生成した行動方策の確率を保持させる必要があります。学習器は現在の方策の確率を再計算し、トークンの重要度比を求めます。
rho_i_t = exp(current_log_probability_i_t
- behavior_log_probability_i_t)この比は、保存された履歴におけるサンプリング済み行動の分布を補正します。ただし、履歴自体が古い行動方策から生成されたという事実は補正しないため、FlashREINFORCEは軌跡全体にわたる累積ドリフトも評価します。
サンプリングされた各行動について、行動方策の確率pと現在の確率qからベルヌーイKLの近似値を計算します。
d_i_t = p * log(p / q)
+ (1 - p) * log((1 - p) / (1 - q))トークン値は軌跡全体で平均されます。軌跡の平均近似値が設定された信頼しきい値以下の場合にのみ、その軌跡が受け入れられます。
mean_drift_i = sum(d_i_t for t in trajectory_i) / T_i
mask_i = 1 if mean_drift_i <= delta else 0拒否された軌跡はトークン損失に寄与しません。したがって、信頼性の判定は個々のトークンではなく、保存された履歴全体に適用されます。
3. Sample-Mean Optimization
トークンレベルのバッチ平均では、長い軌跡ほど自動的に大きな影響力を持ちます。長い失敗ロールアウトに有用な中間ステップが含まれていても、最終的に負の報酬で終わる場合には、これは有害になる可能性があります。
FlashREINFORCEはまず各軌跡内でトークンの寄与を平均し、次に軌跡間で平均します。
trajectory_term_i = (
mask_i * advantage_i / trajectory_length_i
* sum(importance_ratio_i_t for each token t)
)
objective = sum(trajectory_term_i for i in batch) / Bその結果、応答の長さに関係なく、各軌跡の重みは1/Bになります。実際の損失では、重要度比をdetachして使用します。基本手法では、学習されたクリティック、比のクリッピング、参照モデルのフォワードパスは使用しません。
インストールとローカルセットアップ
FlashREINFORCEリポジトリをクローンまたはダウンロードし、ルートディレクトリに移動して、テスト依存関係とともに編集可能モードでインストールします。
pip install -e '.[test]'続いて、付属のCPU例を実行します。
python examples/toy_update.pyこの例では、参照用の目的関数がCPU上でオプティマイザー更新を生成できることを確認します。これは機能検証用の例であり、論文のベンチマーク結果を再現するものではありません。
最後に、テストスイートを実行します。
python -m pytest -q例とテストが正常に完了すれば、ローカルの参照実装とその依存関係が利用可能であることを確認できます。
参照実装を確認する
中心となる参照目的関数はflashreinforce/loss.pyにあります。これは論文とREADMEで説明されている次の処理を網羅しています。
- バッチ全体での報酬の中心化。
- トークンレベルの重要度サンプリング。
- シーケンスレベルの信頼性に基づく受け入れ。
- サンプル平均最適化。
- エントロピーに基づく失敗トークンのフィルタリング（オプション）。
まずexamples/toy_update.pyで、参照損失がオプティマイザー更新にどのように組み込まれるかを確認してください。次にflashreinforce/loss.pyを調べ、例のテンソルとマスクが目的関数の各処理段階にどのように対応しているかを確認します。
CPUの例は、完全な分散ロールアウト基盤を必要としないため、最も安全な基本利用パスです。
1回限りの非同期更新を理解する
本番環境でのFlashREINFORCE更新は、次の6段階で行われます。
- 収集: ロールアウトワーカーが、プロンプトごとに完了した軌跡を1つずつ独立して送信します。各軌跡には報酬と、トークンのサンプリングに実際に使用された行動方策の確率が含まれます。
- バッチ化: 学習器は、次のB個の完了済み軌跡を取得し、それらの報酬を中心化して符号付きアドバンテージを得ます。
- 補正: 学習器は現在の方策の対数確率を再計算し、トークンの重要度比を構成します。
- 信頼性判定: 各軌跡についてサンプリングされた行動の平均KL近似値を計算し、その値が設定されたしきい値内にある場合に限り、軌跡全体を受け入れます。
- 最適化: 受け入れた各軌跡内で平均を取り、バッチ全体でさらに平均を取って、オプティマイザーを1ステップ実行します。
- 破棄: 収集したバッチを追加の学習器更新のために再利用しません。
最後の破棄ステップは、この手法の重要な一部です。FlashREINFORCEでは、同じ行動スナップショットから複数の逐次ミニバッチ更新を行うのではなく、新たに収集したバッチごとに、完全なバッチ更新を1回だけ実行します。
Moltの学習設定をプレビューする
リポジトリには、Moltの非同期実装を通じて学習を開始するためのscripts/train_molt.pyが含まれています。完全な学習を割り当てる前に、ドライランモードを使用して学習フラグを確認してください。
python scripts/train_molt.py --dry-runGPUのセットアップ、データ準備、学習コマンド、論文とコードの対応関係、既知の再現上の制限については、docs/training.mdを参照してください。examples/README.mdには、タスクプリセット、アブレーション、機能チェック付きのネイティブトレーナー起動方法が記載されています。
R1およびQwen2.5-Mathの起動スクリプトはリポジトリで固定されています。ツール利用とALFWorldの設定には、互換性のあるエージェントと環境も必要です。そのため、これらの実験を開始する前に、実行時要件を確認してください。
各軌跡で保持すべきデータ
この更新は、生成中に取得した情報に依存します。概念的な学習フローでは少なくとも、完了したトークン列、スカラー報酬、サンプリングされた各行動に対応する行動確率が必要です。その後、学習器は保存された行動と履歴を評価して、対応する現在の方策の確率を取得します。
記録されたサンプリング確率を、後から再構成した確率に置き換えないでください。後からモデルを評価しても、推論エンジンがサンプリング時に実際に使用した確率と一致する保証はありません。
高度な安定性向上のヒント
サンプリングのサポートを維持する
厳密な重要度補正では、対象方策が行動方策のサポート外に確率質量を割り当てないことを前提とします。過度に積極的なtop-kやtop-pの切り捨ては、この前提に違反する可能性があります。そのため、サンプリング設定はサポート要件を考慮して選択してください。
トークン重要度サンプリングを軌跡全体の補正とみなさない
トークンの重要度比は、保存された履歴を条件とするサンプリング済み行動を補正しますが、履歴自体の分布は補正しません。行動方策と学習器の累積乖離が大きすぎる軌跡を除外するには、Sequence Trust Regionを使用してください。
更新を新鮮に保つ
収集したバッチごとにオプティマイザー更新を1回実行し、その後でバッチを破棄してください。古いバッチを複数回の学習器更新に再利用すると、文書化された1回限りの手法から外れ、方策の不一致が増大する可能性があります。
軌跡レベルの重み付けを維持する
バッチ内の軌跡を平均する前に、各軌跡内のトークン寄与を平均してください。有効なトークンをすべて平坦化してバッチ全体で平均すると、応答長による乗数が再び導入され、長い軌跡に不均衡な重みが与えられます。
信頼マスクを正しく解釈する
Sequence Trust Regionは、軌跡全体に対して1つの受け入れ判定を行います。軌跡がしきい値を超えた場合、その軌跡のすべてのトークン損失がその更新から除外されます。
検証と再現を分ける
CPUの例で検証できるのは、オプティマイザー更新だけです。研究結果を再現するには、学習ガイドとexamplesガイドに記載されたGPU、データ、モデル、エージェント、環境、非同期ランタイムの設定が必要です。
報告されている適用分野
報告されている実験は、数学的推論、マルチターンのPythonツール利用、混合専門家モデルの学習、インタラクティブなALFWorldでの意思決定を対象としています。READMEには、DeepSeek-R1-Distill-Qwen-1.5Bで6,000回の更新まで安定した長い思考連鎖の学習、Qwen2.5-Math-1.5Bでの数学的推論の結果、Qwen2.5-7B-InstructおよびQwen3-30B-A3BでのPythonツール実験、Qwen2.5-7B-InstructでのALFWorld実験が報告されています。
これらの結果はフレームワークの想定される適用範囲を示すものですが、小規模なCPU更新例で同様の結果が得られるとは限りません。
役立つリソース
まとめ
FlashREINFORCEは、単一ロールアウトによるプロンプトカバレッジと、コンパクトなクリティック不要の更新を組み合わせます。バッチを中心とした報酬が符号付きアドバンテージを提供し、トークン重要度サンプリングが古いアクションに対処し、Sequence Trust Regionが蓄積する軌跡のドリフトを制御します。また、サンプル平均最適化により、応答の長さが軌跡の重みを決めることを防ぎます。
まずは、編集可能なインストール、CPU更新の例、テストから始めてください。これらが動作したら、--dry-runを使ってMoltランチャーをプレビューし、トレーニングガイドに従ってモデル固有のGPUおよびデータ準備を行ってください。このプロジェクトはApache License 2.0の下で公開されています。