マルチエージェント REINFORCE

REINFORCE は、サンプリングされた報酬を使用してポリシーを直接最適化するポリシー勾配手法のクラスです。これらの手法は、シンプルさと効率性から、LLM の微調整によく使用されます。例としては、GRPO、Dr. GRPO、RLOO、ReMax、TreeRPO、および REINFORCE++ などがあります。 REINFORCE は、複数の LLM エージェントが各ターンで応答し、それらの共同応答に共有された報酬を与えるマルチエージェント設定に拡張できます。

すべて MAREINFORCE trainers 1回の学習ループを共有します。各 rollout ノードにおいて、すべてのエージェントは $$ K $$ 応答をサンプリングし、 trainer インデックスに一致する共同応答を1つの共有された報酬で評価し、各エージェントは結果の共有された利点に基づいて自身の対数確率を重み付けしてからバッファされたオプティマイザステップを実行.

MAREINFORCE の学習フロー各共同サンプルの獲得は、共有された報酬を与え、それぞれのエージェントのポリシー勾配を重み付けします。

MAREINFORCE#

マルチエージェント REINFORCE (MAREINFORCE) のポリシー勾配のサロゲート、利点正規化または参照正規化の前に、次のように表すことができます:

J(\theta_i) = \mathbb{E}_{\mathbf{o}_0 \sim \mathcal{D}, \mathbf{h}_t \sim \boldsymbol{\pi}_{\boldsymbol{\theta}}} \Bigg[\sum_{t=0}^{H-1} R_t \cdot \log \pi_{\theta_i}(a_{i,t}\mid h_{i,t})\Bigg],

ここで、

\( R_t \)

はターン

\( t \)

の報酬であり、

\( H \)

は時間 (つまり、対話ターンの数) です。初期のデータセット

\mathcal{D}

と、ポリシー

\boldsymbol{\pi}_{\boldsymbol{\theta}}

に従うすべての episodes の共同行動履歴からの期待値が計算されます。

REINFORCE の手法は、価値推定のために critic モデルを使用しません。そのポリシー勾配の推定は、確率的なサンプリングと長期的な信用割り当てにより、高い分散を持つ可能性があります。この分散を軽減するための一般的なアプローチには、アクションに依存しないベースラインを使用するか、より多くのサンプルを平均することを含みます。たとえば、各共同履歴 $\mathbf{h}_t$ における $$ K $$ のサンプルを平均することが挙げられます。

J(\theta_i) = \mathbb{E}_{\mathbf{o}_0 \sim \mathcal{D}, \mathbf{h}_t \sim \boldsymbol{\pi}_{\boldsymbol{\theta}}} \Bigg[\frac{1}{K}\sum_{k=1}^{K} \sum_{t=0}^{H-1} \left(R^{k}_t - b(\mathbf{h}_t)\right) \cdot \log \pi_{\theta_i}(a^{k}_{i,t}\mid h_{i,t})\Bigg],

ここで、ベースラインの

b(\mathbf{h}_t)

は行動に依存しません。

MAGRPO#

マルチエージェントグループ相対的ポリシー最適化 (MAGRPO) は、MAREINFORCE のインスタンスであり、GRPO にインスパイアされており、グループ平均ベースラインは $$ K $$ のサンプルの平均リターンです。デフォルトの aligned の共同モードでは、アドバンテージ正規化前のサロゲートは次のとおりです:

J(\theta_i) = \mathbb{E}_{\mathbf{o}_0 \sim \mathcal{D}, \mathbf{h}_t \sim \boldsymbol{\pi}_{\boldsymbol{\theta}}} \Bigg[\frac{1}{K}\sum_{k=1}^{K}\sum_{t=0}^{H-1} \left(R^{k}_t - \frac{1}{K}\sum_{l=1}^{K}R^{l}_t\right) \cdot \log \pi_{\theta_i}(a^{k}_{i,t}\mid h_{i,t})\Bigg].

含まれているサンプルのグループ平均は、上記の一般的な式で想定されているアクションに依存しないベースラインではありません。マルチターンでのトレーニングでは、各ブランチのリターンは、その子ノードの即時報酬と平均リターンであり、割引は適用されません。 cross モードでは、各エージェントの完了は、その完了を含むすべての共同組み合わせのリターンの合計を受け取ります。

デフォルトでは、advantage_normalization=True は各ノードの生成グループ内で、集団の標準偏差を使用して利点を中心化およびスケーリングし、最小の分母として 1e-6 を使用します。参照 KL の正規化が有効になっている場合、サンプリングされた参照ペナルティは、ベースラインと正規化の適用前に返却値から差し引かれます。

現在の MAGRPOTrainer 損失スライスには、最終的な完了 token が含まれていないため、その実装されたシーケンススコアは、上記のサロゲートにおける完全なシーケンスの対数確率とは異なります。標準的な設定は joint_mode="cross" です。承認された "crossed" のエイリアスは、カルテシアンな組み合わせを生成しますが、同じ完了ごとの集約を適用しません。

MAGRPOConfig パラメータ

num_agents: 設定されたエージェントの数、これは trainer の設定と一致する必要があります
num_turns: episode ごとのターン数
num_train_epochs: epochs の学習回数
agent_learning_rate: 学習率
logging_steps: ステップごとのログ記録間隔
num_generations: 各エージェントに対してプロンプトごとにサンプリングする世代数
max_new_tokens: 生成する最大 tokens の数
temperature: サンプリング用の温度
top_p: Top-p をサンプリングに使用
top_k: Top-k をサンプリングに使用
discount: 保持された設定フィールド; 現在の MAGRPO 実装は、この値を使いません
joint_mode: 共同行動の構成 (インデックス対応の aligned、カルマン積の cross)
early_termination_threshold: 学習ブランチは、ノードの平均即時報酬がこの閾値を超えると拡張を停止します (デフォルト -0.2; None は早期終了を無効化)
rollout_buffer_size: 各エージェントごとに更新する前にバッファリングするノードサンプルの数
train_batch_size: 各最適化ステップあたりのノードサンプルの数; デフォルトは rollout_buffer_size
advantage_normalization: 各ノードの生成グループ内で利潤を中央揃えおよびスケーリングするかどうか (デフォルト True)
eval_interval: 評価は、epoch の最初のバッチを含む、N で割り切れるバッチインデックスで実行されます; 0 は定期評価を無効化
eval_num_samples: 評価実行あたりのサンプルの数
eval_batch_size: DataLoader バッチサイズ評価
external_prompt_passthrough: マルチターン学習で外部プロンプトを直接使用するかどうか
advantage_mode: ベースラインモード (mean, max, rloo, raw)
parallel_training: エージェント実行モード (none または mp)
agent_devices: エージェントのデバイス配置; parallel_training="mp" の使用には明示的に必要
reference_kl_enabled: froze-reference KL ペナルティを有効にするかどうか (デフォルト False)
reference_kl_coef: 参照ペナルティの重み (デフォルト 0.1)
reference_devices: 参照モデルのオプションなデバイス配置

MAGRPOTrainer パラメータ

agent_model または agents: 同種エージェント用のモデル識別子文字列、またはモデル識別子のリストまたはプリロードされたモデル; プリロードされたモデルは agents を通じて提供
num_agents: agent_model から作成されるエージェントの数 (デフォルト 2), args と独立して構成。agents のみ提供された場合、その長さがエージェント数を決定
tokenizer: 共有の tokenizer または 1 つ/エージェント; モデル識別子についてはオプション、プリロードされたモデルについては必須
reward_func: 浮動小数点数のリストを返す関数 (必須)
reward_processor: 報酬に適用するオプションなプロセッサ (例: スケーリング)
formatters: プロンプトフォーマット用の単一関数またはエージェントごとに関数リスト
args: MAGRPOConfig のインスタンス (オプション)
train_dataset: train() を呼び出す際に必要なデータセット
eval_dataset: 評価用データセット (任意)
model_config: モデル構成辞書 (任意)
wandb_config: Weights & Biases のログ構成 (任意)
external_transition: ターン間の遷移を提供する関数; num_turns > 1 (デフォルトは 2) を使用する場合に必要
eval_logger: 任意タスクメトリクスコールバック、eval_logger(agent_completions_turns=..., batch_items=..., prompts=...) で呼び出されます; 評価された各アイテムごとに 1 つのメトリクス辞書を返し、eval_aggregator と共に使用
eval_aggregator: 任意コールバック、eval_aggregator(metrics, num_turns=...) で呼び出されます; 戻される turn_k/... のメトリクスは eval/ でログされます。評価用コールバックは両方とも提供する必要があります

MAGRPO はサンプルを rollout バッファに収集し、更新時にログ確率を再計算しますが、重要サンプリングや比率クリッピングは行いません。したがって、サンプルは、同じバッファまたは rollout ツリーからの以前の更新よりも古い可能性があります。trainer に critic モデルがないため、value_clip_range は適用されません。

trainer は、固定された学習バッチサイズ 1 を使用し、num_generations=2 および train_batch_size=rollout_buffer_size を含む、少なくとも DataLoader のサイズが必要です。デフォルトでは、各エージェントは、1 つの最適化ステップでバッファされたノードサンプルの損失を平均します。より小さなミニバッチは、複数の最適化ステップを生成します。バッファの一部は、各 epoch の末尾でフラッシュされます。

その他のバリエーション#

CoMLRL は異なるベースラインを持つ他の MAREINFORCE バリエーションも提供します:

MARLOO: マルチエージェント REINFORCE 1 つのノードのサンプルグループ内の他の世代の平均リターン、つまり他のエージェントの平均リターンをベースラインとして使用します。

J(\theta_i) = \mathbb{E}_{\mathbf{o}_0 \sim \mathcal{D}, \mathbf{h}_t \sim \boldsymbol{\pi}_{\boldsymbol{\theta}}} \Bigg[\frac{1}{K}\sum_{k=1}^{K}\sum_{t=0}^{H-1} \left(R^{k}_t - \frac{1}{K-1}\sum_{l=1, l\neq k}^{K}R^{l}_t\right) \cdot \log \pi_{\theta_i}(a^{k}_{i,t}\mid h_{i,t})\Bigg].

MAReMax: 複数エージェントの REINFORCE、グループ最大値 (MAReMaxConfig および MAReMaxTrainer) を使用。ベースラインは、ノードでの最大サンプリングされたグループの戻り値です。この実装では、別個の貪欲なベースライン応答を生成しません。

J(\theta_i) = \mathbb{E}_{\mathbf{o}_0 \sim \mathcal{D}, \mathbf{h}_t \sim \boldsymbol{\pi}_{\boldsymbol{\theta}}} \Bigg[\frac{1}{K}\sum_{k=1}^{K}\sum_{t=0}^{H-1} \left(R^{k}_t - \mathrm{max}_l\, R^l_t \right) \cdot \log \pi_{\theta_i}(a^{k}_{i,t}\mid h_{i,t})\Bigg].

MAREINFORCETrainer, MARLOOTrainer, および MAReMaxTrainer は、以下から派生 comlrl.trainers.reinforce.MAGRPOTrainer. それらの対応する構成クラスは継承 MAGRPOConfig, デフォルト値で advantage_mode 値は raw, rloo, および max, それぞれ。それらは同じ trainer インターフェースと構成フィールド.

4 つの trainers は、共有された戻り値から引かれるベースラインのみが異なります。1 つのノード ( $$ K=4 $$ ) での共同応答の場合、アニメーションは正規化前の利点の比較を示しています: MAREINFORCE は元の戻り値を保持し、REINFORCE のようにします; MAGRPO はグループの平均を差し引きます (GRPO のように); MARLOO は他の $$ K-1 $$ の戻り値の平均を差し引きます (RLOO のように)、これは MAGRPO の利点を $$ K/(K-1) $$ でスケーリングします; および MAReMax はグループの最大値を差し引くため、どの応答にも正の利点はありません。

MAREINFORCE 系のベースライン各行は同じ共有された報酬から異なるベースラインを差し引きます; 影付きのバーは正規化前の利点を表します。

これらの変数の式は、正規化前の利点を記述しています。すべての変数は advantage_normalization=True を継承し、元の戻り値と最大ベースラインの利点を再中心化します。標準偏差のフロアと数値精度を除き、4 つのモードは同じ正規化された利点を提供します。ユーザーは、上記のベースラインの違い (元の戻り値の推定値を含む MAREINFORCE) を維持するために、advantage_normalization=False を設定できます。