PPO — 近接方策最適化

動機と歴史

問題 ：標準的な方策勾配の更新には、ステップ幅の制約がありません。たった1つの不運なバッチによって、方策がゴミを生成する領域へ押し出される可能性があります \(\rightarrow\) ゴミは低い報酬しか得られない \(\rightarrow\) 次の勾配でさらに悪化する \(\rightarrow\) 回復不能な崩壊に至ります。

解決策の歴史 ：

1. TRPO (Schulman et al. 2015)（2015年）：古い方策と新しい方策の間のKLダイバージェンスを制約します。完全に機能しますが、二次最適化（フィッシャー情報行列、共役勾配法）が必要で、計算コストが高くなります。

2. PPO （2017年）(Schulman et al. 2017)：単純な一次のクリップ付き目的関数によって、同程度の安定性を実現します。実装は10\(\times\)単純で、ほぼ同等に機能し、分散学習にも容易にスケールします。

クリップ付き目的関数

PPOの中核的な革新は、破壊的に大きな方策更新を防ぎながら、実装も簡単に保てるクリップ付き代理目的関数です。

\[ \boxed{L^{\text{CLIP}}(\theta) = \mathbb{E}_t\left[\min\left(r_t(\theta)\hat{A}_t,; \text{clip}(r_t(\theta), 1{-}\epsilon, 1{+}\epsilon)\hat{A}_t\right)\right]} \]

ここで、\(r_t(\theta) = \frac{\pi_\theta(a_t\vert s_t)}{\pi_{\theta_\text{old}}(a_t\vert s_t)}\) は確率比です。

Tip

クリッピングの直感 — 核となる洞察

\(\min\) 演算子は 悲観的な上限 を作ります。

• 良い行動（\(\hat{A} > 0\)） ：その確率を高めたいとします。代理目的 \(r\hat{A}\) は \(r\) が増えるにつれて大きくなります。しかし、clipは \(r = 1 + \epsilon\) で利益に上限を設けます。*「1つの良い例で欲張りにならない」*ためです。

• 悪い行動（\(\hat{A} < 0\)） ：その確率を下げたいとします。\(r\hat{A}\) は \(r\) が減るにつれて改善します。しかし、clipは \(r = 1 - \epsilon\) で利益に上限を設けます。*「1つの悪い例を根拠に、過度に忘却しない」*ためです。

結果として、1回の更新ステップでの方策の変化は最大でも\(\pm\)20%になります。これにより、壊滅的な崩壊と、過信した特化の両方を防ぎます。

PPOの完全な損失

\[ L = L^{\text{CLIP}} - c_1 \underbrace{(V_\theta(s_t) - V^{\text{target}}_t)^2}_{\text{value loss}} + c_2 \underbrace{H[\pi_\theta(\cdot|s_t)]}_{\text{entropy bonus}} \]

• 価値損失 （\(c_1 = 0.1\)）：リターンを予測するようにクリティックを学習します。安定性のため、これもクリップされます。

• エントロピー・ボーナス （\(c_2 = 0.01\)）：決定論的な方策への早すぎる収束を防ぎます。探索に不可欠です。

PPOの勾配と更新則の導出

この節では、RLの目的関数からPPOの更新則に至る数学的な道筋をたどり、クリップ付き代理目的関数がなぜ機能するのかを示します。

ステップ1：RLの目的関数

目標は、方策のもとで期待累積報酬を最大化することです。

\[ J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\left[\sum_{t=0}^T r_t\right] \]

ステップ2：方策勾配定理

方策パラメータに関する \(J(\theta)\) の勾配は次のとおりです。

\[ \boxed{\nabla_\theta J(\theta) = \mathbb{E}_{\pi_\theta}\left[\sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \cdot \hat{A}_t\right]} \]

ここで \(\hat{A}_t\) はアドバンテージ関数です（行動 \(a_t\) が状態 \(s_t\) における平均的な行動と比べてどれだけ優れていたか）。分散を減らすため、完全なリターンをアドバンテージに置き換えています。

ステップ3：オフポリシーデータに対する重点サンプリング

PPOは \(\pi_{\theta_{\text{old}}}\) を使ってデータを収集しますが、更新するのは \(\pi_\theta\) です。この分布の不一致を補正するため、重点サンプリングを適用します。

\[ \nabla_\theta J(\theta) = \mathbb{E}_{\pi_{\theta_{\text{old}}}}\left[\frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{\text{old}}}(a_t|s_t)} \nabla_\theta \log \pi_\theta(a_t|s_t) \cdot \hat{A}_t\right] \]

確率比 \(r_t(\theta) = \frac{\pi_\theta(a_t\vert s_t)}{\pi_{\theta_{\text{old}}}(a_t\vert s_t)}\) を定義します。恒等式 \(\nabla_\theta \log f = \frac{\nabla_\theta f}{f}\) を使うと、次が得られます。

\[ \nabla_\theta J(\theta) = \mathbb{E}_{\pi_{\theta_{\text{old}}}}\left[\nabla_\theta, r_t(\theta) \cdot \hat{A}_t\right] \]

これは、次の 代理目的関数 を最大化することを意味します。

\[ L^{\text{CPI}}(\theta) = \mathbb{E}_t\left[r_t(\theta) \cdot \hat{A}_t\right] \]

ステップ4：制約のない代理目的関数の問題

\(L^{\text{CPI}}\) は有効な目的関数ですが、制約がなければ、1回の勾配ステップで \(r_t(\theta)\) が1.0から大きく離れてしまい、次の問題が起こります。

• 重要度の重みが極端になる \(\rightarrow\) 分散が大きくなる

• 方策が未検証の領域に入る \(\rightarrow\) 報酬モデルが信頼できないスコアを返す

• 壊滅的な崩壊：方策がゴミを生成し、回復できなくなる

TRPOの解決策 ：\(D_{\text{KL}}(\pi_{\theta_{\text{old}}} \\vert \pi_\theta) \leq \delta\) を制約します。二次法が必要で、計算コストが高くなります。

ステップ5：PPOのクリップ付き代理目的関数（一次近似）

PPOは厳密なKL制約を、一次勾配だけを使って同様の挙動を実現する クリップ付き目的関数 に置き換えます。

\[ \boxed{L^{\text{CLIP}}(\theta) = \mathbb{E}_t\left[\min!\left(r_t(\theta)\hat{A}_t,;\text{clip}(r_t(\theta), 1{-}\epsilon, 1{+}\epsilon)\hat{A}_t\right)\right]} \]

勾配の導出 ：

\(L_t = \min(r_t \hat{A}_t,; \bar{r}_t \hat{A}_t)\) とし、\(\bar{r}_t = \text{clip}(r_t, 1{-}\epsilon, 1{+}\epsilon)\) とします。

\[ \nabla_\theta L_t = \begin{cases} \nabla_\theta r_t(\theta) \cdot \hat{A}_t & \text{if } r_t \hat{A}_t < \bar{r}_t \hat{A}_t \text{ (unclipped term is smaller)} \\ 0 & \text{if } r_t \hat{A}_t \geq \bar{r}_t \hat{A}_t \text{ (clipped term is smaller, gradient = 0)} \end{cases} \]

条件を展開すると、次のようになります。

• \(\hat{A}_t > 0\) かつ \(r_t < 1+\epsilon\) の場合 ：勾配は通常どおり流れます。方策は \(\pi_\theta(a_t\vert s_t)\) を増やすよう促されます。

• \(\hat{A}_t > 0\) かつ \(r_t \geq 1+\epsilon\) の場合 ：勾配は ゼロ です。方策はすでに十分に増加しているため、それ以上押し上げません。

• \(\hat{A}_t < 0\) かつ \(r_t > 1-\epsilon\) の場合 ：勾配は通常どおり流れます。方策は \(\pi_\theta(a_t\vert s_t)\) を減らすよう促されます。

• \(\hat{A}_t < 0\) かつ \(r_t \leq 1-\epsilon\) の場合 ：勾配は ゼロ です。方策はすでに十分に減少しているため、それ以上押し下げません。

ステップ6：PPOの完全な更新則

クリップ付き方策損失、価値損失、エントロピー・ボーナスを組み合わせます。

\[ \boxed{\theta_{k+1} = \theta_k + \alpha \cdot \nabla_\theta \left[L^{\text{CLIP}}(\theta) - c_1 L^{\text{VF}}(\theta) + c_2 H[\pi_\theta]\right]} \]

ここで、

\[ \begin{aligned} L^{\text{VF}}(\theta) &= \left(V_\theta(s_t) - V_t^{\text{target}}\right)^2 & &\text{(value function regression loss)} \\ H[\pi_\theta] &= -\sum_a \pi_\theta(a|s_t)\log\pi_\theta(a|s_t) & &\text{(entropy of the policy)} \end{aligned} \]

Tip

まとめ：これが機能する理由

1. 方策勾配定理 は、方策を改善する方向を与えます。

2. 重点サンプリング によって、\(\pi_{\theta_{\text{old}}}\) のデータを複数エポックにわたって再利用できます。

3. クリッピング は重要度の重みが極端になるのを防ぎ、更新を安全に保ちます。

4. min演算子 は、常に（クリップされた項とクリップされていない項の）より保守的な方を選びます。改善に対する悲観的な下限です。

5. 結果 ：一次勾配だけを使い、確率1で単調に改善します。ヘッセ行列も、共役勾配も、ラインサーチも必要ありません。

ロールアウトバッファとロールアウト

PPOでは、データ管理に ロールアウトバッファ と呼ばれる専用の短期ストレージシステムを使います。経験をリプレイバッファに無期限に保存するオフポリシーアルゴリズム（DQNなど）とは異なり、PPOではオンポリシーという数学的制約を満たすため、一時的な構造が必要です。

ロールアウトとは何か？

ロールアウト （軌跡）とは、エージェントが現在の方策を環境内で実行することによって生成される相互作用の系列です。

• プロセス ：エージェントは状態を観測し、行動を選択し、報酬を受け取り、次の状態へ移ります。固定されたステップ数、またはエピソードが終了するまで繰り返します。

• LLM/RLHFの場合 ：ロールアウトは、データセットからプロンプトを取り出し、言語モデルにテキスト終了マーカーに到達するまでトークンを1つずつ生成させ、完全なトークン列を作ることです。各トークンが1つの「ステップ」です。

ロールアウトバッファ

ロールアウトバッファは、ロールアウト段階で収集したすべてのデータを一時的に保存します。生成されたトークン／ステップごとに、次を記録します。

\[ \boxed{\mathcal{B} = \left\{ \left(s_t,; a_t,; \log\pi_{\theta_{\text{old}}}(a_t|s_t),; r_t,; V(s_t)\right) \right\}_{t=1}^{T}} \]

• \(s_t, a_t, r_t\)：状態、取った行動、時刻 \(t\) における報酬。

• \(\log\pi_{\theta_{\text{old}}}(a_t\vert s_t)\)：その行動を、実際にデータを生成した方策のもとで取る対数確率（比の計算に必要）。

• \(V(s_t)\)：価値関数によるベースラインの予測（GAEのアドバンテージ計算に必要）。

ロールアウトバッファのライフサイクル

バッファは、厳密な3段階の時計仕掛けのサイクルで動作します。

1. 収集 ：アクティブな方策が環境と相互作用し、新しい軌跡でバッファを満たします（70Bモデル、batch=128、max_tokens=512の場合、1回のロールアウトあたり最大65Kのトークンレベル遷移）。

2. 学習 ：軌跡全体でGAEアドバンテージを計算します。\(K\) エポック（通常3〜10）にわたってミニバッチの勾配降下を実行し、クリップ付き目的関数で方策の重みを更新します。

3. 消去 ：バッファ全体を 完全に消去 します。PPOはオンポリシーなので、古い方策が生成したデータを次の更新サイクルで安全に再利用することはできません。比 \(r_t(\theta)\) が古くなり、クリッピングの保証が壊れるためです。

Warning

ロールアウトバッファとリプレイバッファの違い

リプレイバッファ （DQN、SAC）：オフポリシー。数百万の遷移を無期限に保存します。ランダムサンプリングを行い、データを多くの更新で再利用します。

ロールアウトバッファ （PPO、GRPO）：オンポリシー。1バッチ分の軌跡を保存します。数エポック使った後、完全に破棄します。サイクルごとに新しいデータが必要です。

これが、PPOに継続的な生成が必要な理由です。更新のたびにバッファが空になり、新しいロールアウトが要求されます。そのため、生成ボトルネック（実時間の60〜70%）が特に深刻になります。

Tip

RLHFにおけるvLLM

RLHF学習では、vLLMは 生成段階 （実時間の60〜70%）に使われます。方策モデルがロールアウトを生成し、その後、報酬モデルがスコアを付けます。主な利点は次のとおりです。

• バッチ生成 ：プロンプトをまたいで256個以上の応答を並列生成します。

• メモリ効率 ：より多くの同時生成を収容できるため \(\rightarrow\)、生成ボトルネック中のGPU利用率が高くなります。

• プレフィックス共有 ：プロンプトあたり \(N=8\) 個の応答を生成する場合（GRPO）、プロンプトのKVは一度だけ計算して8個すべてで共有します。冗長なプリフィルはありません。

• 統合 ：OpenRLHFやTRLなどのフレームワークは、vLLMを生成バックエンドとして使用し、生成ワーカー（vLLM）と学習ワーカー（DeepSpeed/FSDP）を分離します。

RLHFのためのPPO：完全なループ

Note

70BチャットモデルにおけるPPOの具体的なステップ

設定 ：128個のプロンプトのバッチ、Llama-3-70B方策、最大512トークン。

ステップ1 — 生成 ：128個の応答をサンプリングします（temperature=0.7、top-p=0.9）。時間の60%を要します。

ステップ2 — スコアリング ：報酬モデルが各（プロンプト、応答）ペアを採点します。範囲は0.2〜0.95です。

ステップ3 — KL ：トークンごとのKLを計算します：\(\text{KL}_t = \log\pi_\theta(y_t\vert y_{<t}) - \log\pi_\text{ref}(y_t\vert y_{<t})\)。トークン全体での平均KLは通常3〜8です。

ステップ4 — 最終報酬 ：\(R = r_\text{RM} - 0.05 \times \text{mean_KL}\)（最後のトークンにだけ与えます）。

ステップ5 — GAE ：価値ヘッドの予測を使って、各トークン位置の \(\hat{A}_t\) を計算します。アドバンテージをホワイトニングします（平均0、分散1）。

ステップ6 — 更新 ：16個のミニバッチでSGDを4エポック実行します。比を \(\epsilon = 0.2\) でクリップします。勾配ノルムは1.0でクリップします。

結果 ：方策は1ステップあたり\(\sim\)0.005の勝率向上を示します。1万ステップ後には、SFTに対して絶対値で5〜10%改善します。

Warning

LLMのRLにおけるトークン化の落とし穴

トークンごとのKLペナルティとアドバンテージを計算するときは、トークン化によって「ステップ」が決まることを忘れないでください。1つの概念的な行動（たとえば「2024」の出力）が、トークナイザーによって1〜4トークンに分かれることがあります。これにより、次のような微妙な問題が生じます。

• KLの集計 ：同じ意味内容でもトークン化が異なると、トークンごとのKLの合計は異なります（たとえば、まれな単語がより多くのサブワードに分割されると、合計KLペナルティが大きくなります）。

• 信用割当 ：GAEはトークン位置ごとにアドバンテージを割り当てますが、意味的な「意思決定」は複数トークンにまたがることがよくあります。モデルが本当に「決定」するのは単語の最初のトークンだけで、その後のサブワードトークンはほぼ決定論的です。

• 報酬の配置 ：最終トークンにだけ報酬を置くと、それより前のすべてのトークンがGAEを通してクレジットを後方へ伝播させなければなりません。長い応答ほど、シグナルが希薄になります。

緩和策 ：システムによっては、KLを系列長で正規化したり、単語レベルの報酬シェーピングを使ったり、最終トークンではなく意味的な境界で報酬を適用したりします。

詳細な仕組み：ロジットと方策の更新

PPOは、メモリ上で2つの異なるパラメータ状態を管理します。両者は同じニューラルネットワーク構造を共有しますが、最適化中は異なる重みの値を保持します。

Important

中核アーキテクチャ：2つのネットワーク

1. 方策ネットワーク（\(\pi_\theta\)） ：重み \(\theta\) でパラメータ化された、アクティブで稼働中のネットワークです。最適化中のバックプロパゲーションによって継続的に更新されます。

2. 古い方策ネットワーク（\(\pi_{\theta_{\text{old}}}\)） ：重み \(\theta_{\text{old}}\) でパラメータ化された、凍結したスナップショットです。1回の最適化サイクルの間、方策が大きく変化しすぎないよう静的なアンカーとして機能します。

フェーズ1：ロールアウト（データ収集）

データ収集中、エージェントは \(T\) ステップにわたって環境と相互作用します。各時刻 \(t\) では次を行います。

1. 環境が現在の状態／観測 \(s_t\) を返します（LLMの場合：プロンプト＋これまでに生成されたトークン）。

2. 状態 \(s_t\) を現在のネットワークスナップショット（\(\theta_{\text{old}}\)）に通します。

3. ネットワークが正規化前の生の値、すなわち ロジット \(z_{\text{old}}\) を出力します。これは \(\vert V\vert\)（語彙サイズ32K〜128K）の大きさのベクトルです。

4. Softmaxで確率を計算します。

\[ \boxed{P(a \mid s_t) = \text{Softmax}(z_{\text{old}}) = \frac{\exp(z_{\text{old}, a})}{\sum_{j=1}^{|V|} \exp(z_{\text{old}, j})}} \]

5. 行動 \(a_t\)（次のトークン）を \(P(a \mid s_t)\) からサンプリングし、遷移タプル \(\langle s_t, a_t, r_t, s_{t+1} \rangle\) と \(\log \pi_{\theta_{\text{old}}}(a_t \mid s_t)\) をロールアウトバッファに保存します。

Tip

なぜ対数確率を保存するのか？

ロールアウト中に \(\log \pi_{\theta_{\text{old}}}(a_t \mid s_t)\) をスカラーとして保存すると、最適化中に凍結したネットワークを再実行せずに済みます。ミニバッチごとにフルフォワードパスを1回省略できるため、70Bモデルでは大きな効果があります。

フェーズ2：最適化ループ（ミニバッチ更新）

ロールアウトバッファが満杯になると、PPOはミニバッチに対して \(K\) エポック（通常3〜10）を実行します。各勾配ステップでは、保存された状態 \(s_t\) を使って両方の方策のロジットを生成します。

古い方策の評価 （凍結）：

\[ z_{\text{old}} = f(s_t; \theta_{\text{old}}) \quad \longrightarrow \quad \log \pi_{\theta_{\text{old}}}(a_t \mid s_t) = \text{LogSoftmax}(z_{\text{old}})[a_t] \]

実装上のショートカット：再計算する代わりに、ロールアウト時に保存したスカラーを再利用します。

現在の方策の評価 （更新対象）：

\[ z_{\text{new}} = f(s_t; \theta) \quad \longrightarrow \quad \log \pi_\theta(a_t \mid s_t) = \text{LogSoftmax}(z_{\text{new}})[a_t] \]

\(\theta\) は各ミニバッチの勾配ステップ後に更新されるため、\(z_{\text{new}}\) は最適化ループ全体で連続的に変化します。一方、\(z_{\text{old}}\) は完全に静的なままです。

ロジットから確率比へ

PPOの中核となる比は、新しい方策のもとで行動が古い方策よりどれだけ起こりやすいか、または起こりにくいかを測ります。

\[ \boxed{r_t(\theta) = \frac{\pi_\theta(a_t \mid s_t)}{\pi_{\theta_{\text{old}}}(a_t \mid s_t)}} \]

生の確率を割ることによる壊滅的な数値アンダーフロー／オーバーフローを避けるため、計算は 対数空間 で行います。

\[ \begin{aligned} \log \pi_\theta(a_t \mid s_t) &= \text{LogSoftmax}(z_{\text{new}})[a_t] \\ \log \pi_{\theta_{\text{old}}}(a_t \mid s_t) &= \text{LogSoftmax}(z_{\text{old}})[a_t] \end{aligned} \]

差を指数関数に入れることで、比を復元します。

\[ \boxed{r_t(\theta) = \exp!\left(\log \pi_\theta(a_t \mid s_t) - \log \pi_{\theta_{\text{old}}}(a_t \mid s_t)\right)} \]

この比をPPOのクリッピング目的関数に組み込みます。

\[ \boxed{\mathcal{L}^{\text{CLIP}}(\theta) = \hat{\mathbb{E}}_t \left[ \min!\left(r_t(\theta)\hat{A}_t, ;\text{clip}(r_t(\theta),, 1{-}\epsilon,, 1{+}\epsilon),\hat{A}_t\right) \right]} \]

Tip

クリッピングの仕組み

• \(\hat{A}_t > 0\)（良い行動）の場合：比は \(1+\epsilon\) でクリップされ、良い行動を過度に利用できません。

• \(\hat{A}_t < 0\)（悪い行動）の場合：比は \(1-\epsilon\) でクリップされ、悪い行動を過度に罰することができません。

• \(\min(\cdot)\) によって、常により保守的な推定値を選びます。

結果として、信頼領域内で単調に改善し、壊滅的な崩壊は起こりません。

PPOの重みのライフサイクル

PPOの学習フェーズにおける \(\theta\) と \(\theta_{\text{old}}\) の推移。

連続行動空間への拡張

連続行動空間（LLMでは一般的ではありませんが、ロボティクスRLでは重要）では、ネットワークは離散的なロジットではなく分布パラメータを出力します。

• 予測平均ベクトル \(\mu\)

• 予測標準偏差ベクトル \(\sigma\)

対数確率はガウス分布の対数PDFで計算します。

\[ \boxed{\log \pi(a_t \mid s_t) = -\frac{1}{2}\left(\frac{a_t - \mu}{\sigma}\right)^{!2} - \log(\sigma) - \frac{1}{2}\log(2\pi)} \]

続いて比 \(r_t(\theta) = \exp(\log \pi_\theta - \log \pi_{\theta_{\text{old}}})\) を同じように計算し、同じクリッピング目的関数に入力します。

TRLによる実装

HuggingFace TRLライブラリ (Werra et al. 2022) は、LLM向けの主要なRL手法をすべて、プロダクションで利用できる形で実装しています。

from trl import PPOConfig, PPOTrainer, AutoModelForCausalLMWithValueHead
from transformers import AutoTokenizer
from peft import LoraConfig

# Model setup
model = AutoModelForCausalLMWithValueHead.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    torch_dtype=torch.bfloat16, device_map="auto",
    peft_config=LoraConfig(r=64, lora_alpha=16, target_modules=["q_proj","v_proj","k_proj","o_proj"])
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")

# PPO config with all critical hyperparameters
ppo_config = PPOConfig(
    learning_rate=1.5e-6,        # Low LR for stability
    batch_size=128,              # Prompts per step
    mini_batch_size=16,          # Gradient accumulation unit
    ppo_epochs=4,                # Epochs per batch (reuse data)
    gamma=1.0,                   # No discounting (single turn)
    lam=0.95,                    # GAE lambda
    cliprange=0.2,               # PPO epsilon
    cliprange_value=0.2,         # Value function clipping
    vf_coef=0.1,                 # Value loss coefficient
    init_kl_coef=0.05,           # Initial KL penalty
    target_kl=6.0,               # Adaptive KL target
    whiten_rewards=True,         # Normalize advantages
    gradient_accumulation_steps=4,
    max_grad_norm=1.0,
)

ppo_trainer = PPOTrainer(config=ppo_config, model=model, tokenizer=tokenizer,
    dataset=prompt_dataset, data_collator=collator)

# Training loop
for batch in ppo_trainer.dataloader:
    # 1. Generate responses
    query_tensors = batch["input_ids"]
    response_tensors = ppo_trainer.generate(
        query_tensors, max_new_tokens=512, temperature=0.7, top_p=0.9, do_sample=True
    )
    # 2. Score with reward model
    texts = [tokenizer.decode(r, skip_special_tokens=True) for r in response_tensors]
    rewards = [torch.tensor(reward_model.score(q, r)) for q, r in zip(batch["query"], texts)]
    # 3. PPO update (handles KL, GAE, clipping internally)
    stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
    # Monitor: stats["ppo/mean_scores"], stats["ppo/policy/approx_kl"]

重要なハイパーパラメータ