MLエンジニアへの道 #44 - ポリシー勾配法 続

Last Edited: 3/1/2025

この記事では、強化学習における高度なポリシー勾配法を紹介します。

ML

前回の記事では、ポリシー勾配定理と2つのオンポリシーポリシー勾配法(REINFORCEとActor-Critic)を紹介しました。 これらは理論上、状態-行動価値法よりも優れた収束保証を持っていますが、非i.i.d.データや探索の困難さにより、収束に苦労することがあります。 この記事では、より高度なポリシー勾配法でこれらの問題に対処していきます。

A3C

Actor-Critic法の収束に時間がかかる場合、並行処理を使って収束を速める方法があります。具体的には、ターゲットポリシーを持つグローバルActor-Criticを用意し、 CPUコアごとに、行動ポリシーに基づいて行動し勾配を計算するActor-Criticをワーカーとして実行します。ワーカーが終了状態に達するか最大タイムステップに達したら、 収集した経験を使って勾配を計算し、グローバルターゲットポリシーと批評家のパラメータを更新し、その更新されたパラメータをワーカーの行動ポリシーと批評家にコピーして探索を継続させます。

このプロセスは特定の総タイムステップ数に達するまで続けられます。複数のワーカーを実行している場合、あるワーカーが他のワーカーより早く終了することがあります。 そのような場合、他のワーカーが終了するのを待たずに、他のワーカーがまだ探索している間にグローバル探索家と批評家のパラメータを非同期に更新することで、CPUをフル活用できます。 また、ここでの批評家は損失に対してQ(s,a,w)Q(s, a, w)ではなくV(s,w)V(s, w)を出力するアドバンテージを使用できます。 複数のワーカーが実行するローカル探索家と批評家からの勾配を使って非同期にパラメータを更新することで、 グローバルアドバンテージ探索家と批評家を訓練するこの方法は、Asynchronous Advantage Actor-Critic(A3C)と呼ばれています。

A3Cは、複数のワーカーが並列にバッチをサンプリングするミニバッチ勾配更新を実行するか、小規模のリプレイバッファを使用すると解釈できます。 Machine Learning with Phil(2022)によるこちらのリンクされた動画で、 アルゴリズムの詳細な説明とA3CのPyTorch実装を確認できます。(前回の記事のActor-Critic法との主な違いは、torch.multiprocessingの使用と勾配の蓄積です。) 動画で示されているように、A3Cはターゲットポリシーと行動ポリシーが同じであると仮定し、オンポリシーのポリシー勾配を計算しますが、これは良い仮定ではないかもしれません。 また、ハイパーパラメータに敏感で、学習結果に実行ごとのばらつきが観察されます。

A2C

OpenAIの研究者たちはA3Cを分析し、CPU利用のために設定された非同期更新がモデルのパフォーマンスを向上させるという証拠を見つけられませんでした。 そこで彼らはA3Cから非同期性を取り除き、グローバルターゲットActor-Criticとワーカーの間に重みの更新を同期できるコーディネーターを導入したA2Cを考案しました。 パラメータを同時に更新することで、ワーカー内のすべてのActor-CriticモデルがグローバルターゲットActor-Criticと同じ最新のパラメータセットを使用することを保証し、 A2Cを真にオンポリシーにして、潜在的により速い収束を達成できます。

A2Cは、ミニバッチをバッチに結合してパラメータを更新すると解釈でき、GPUをより効率的に利用できます。実際、OpenAIの研究者たちは、より大きなポリシーを使用する場合、 より良いパフォーマンスとより速い収束を観察すると主張しています。A2Cを実装する際には、セマフォや、コーディネーターがすべてのワーカーの探索終了を待つことを可能にする代替アプローチの使用を検討するとよいでしょう。

オフポリシーポリシー勾配

A3Cは行動ポリシーの導入を反映していませんが、行動ポリシーβ(as)\beta(a|s)を導入することで、報酬関数J(θ)J(\theta)とポリシー勾配θJ(θ)\nabla_{\theta} J(\theta)を以下のように書き直すことができます。

J(θ)=sSdβ(s)Vπθ(s)=sSdβ(s)aAπθ(as)Qπθ(s,a)ddθJ(θ)=xSdβ(s)aAddθπθ(as)Qπθ(s,a)+πθ(as)ddθQπθ(s,a) J(\theta) = \sum_{s \in S} d_{\beta}(s) V_{\pi_{\theta}}(s) = \sum_{s \in S} d_{\beta}(s) \sum_{a \in A} \pi_{\theta}(a | s) Q_{\pi_{\theta}}(s, a) \\ \frac{d}{d \theta} J(\theta) = \sum_{x \in S} d_{\beta}(s) \sum_{a \in A} \frac{d}{d \theta} \pi_{\theta}(a | s) Q_{\pi_{\theta}}(s, a) + \pi_{\theta}(a | s) \frac{d}{d \theta}Q_{\pi_{\theta}}(s, a)

上記の勾配は積の法則を用いて展開されています。Degrisら(2013年)による証明を使用すると、θQπθ(s,a)\nabla_\theta Q_{\pi_\theta}(s, a)を無視した近似勾配を使用しても、 真の局所最小値への収束が保証されることがわかります。これにより、勾配計算を以下のように簡略化できます。

ddθJ(θ)=xSdβ(s)aAddθπθ(as)Qπθ(s,a)+πθ(as)ddθQπθ(s,a)=xSdβ(s)aAddθπθ(as)Qπθ(s,a)=xSdβ(s)aAβ(as)πθ(as)β(as)Qπθ(s,a)ddθπθ(as)πθ(as)=Esdβ,aβ[πθ(as)β(as)Qπθ(s,a)ddθln(πθ(as))] \frac{d}{d \theta} J(\theta) = \sum_{x \in S} d_{\beta}(s) \sum_{a \in A} \frac{d}{d \theta} \pi_{\theta}(a | s) Q_{\pi_{\theta}}(s, a) + \pi_{\theta}(a | s) \frac{d}{d \theta}Q_{\pi_{\theta}}(s, a)\\ = \sum_{x \in S} d_{\beta}(s) \sum_{a \in A} \frac{d}{d \theta} \pi_{\theta}(a | s) Q_{\pi_{\theta}}(s, a) \\ = \sum_{x \in S} d_{\beta}(s) \sum_{a \in A} \beta(a | s) \frac{\pi_{\theta}(a | s)}{\beta(a | s)} Q_{\pi_{\theta}}(s, a) \frac{\frac{d}{d \theta} \pi_{\theta}(a | s)}{\pi_{\theta}(a | s)} \\ = \text{E}_{s \sim d_{\beta}, a \sim \beta} [\frac{\pi_{\theta}(a | s)}{\beta(a | s)} Q_{\pi_{\theta}}(s, a) \frac{d}{d \theta} \ln(\pi_{\theta}(a | s))]

ここで、オフポリシーポリシー勾配に必要な唯一の調整は、重要度サンプリング比πθ(as)β(as)\frac{\pi_\theta(a | s)}{\beta(a | s)}を重みとした加重平均を使用するだけで良いことがわかります。

TRPO

A3Cでは、ワーカーの行動ポリシーはターゲットポリシーと同じであると仮定されていますが、本当は異なるため、より良いパラメータ更新のために損失とその勾配を計算する際に考慮する必要があります。 ワーカーの行動ポリシーπθold\pi_{\theta_{\text{old}}}を導入すると、上記を使用して目的関数の勾配θJ(θ)\nabla_\theta J(\theta)を以下のように書くことができます。

ddθJ(θ)=Eθold[πθ(as)πθold(as)Qπθ(s,a)ddθln(πθ(as))] \frac{d}{d \theta} J(\theta) = \text{E}_{\theta_{\text{old}}} [\frac{\pi_{\theta}(a | s)}{\pi_{\theta_{\text{old}}}(a | s)} Q_{\pi_{\theta}}(s, a) \frac{d}{d \theta} \ln(\pi_{\theta}(a | s))]

ワーカーは行動ポリシーを用いて獲得したQ値を使用して目的関数を計算しており、行動ポリシーがターゲットポリシーと合理的に類似していると仮定できるため、Qπθold(s,a)Q_{\pi_{\theta_{\text{old}}}}(s, a)を使用して目的関数を近似できます。

ddθJ(θ)Eθold[πθ(as)πθold(as)Qπθold(s,a)ddθln(πθ(as))] \frac{d}{d \theta} J(\theta) \approx \text{E}_{\theta_{\text{old}}} [\frac{\pi_{\theta}(a | s)}{\pi_{\theta_{\text{old}}}(a | s)} Q_{\pi_{\theta_{\text{old}}}}(s, a) \frac{d}{d \theta} \ln(\pi_{\theta}(a | s))]

上記の表現は、自然対数の導関数を展開することでさらに導出できます。

ddθJ(θ)Eθold[πθ(as)πθold(as)Qπθold(s,a)ddθ(πθ(as))πθ(as)]=Eθold[ddθ(πθ(as))πθold(as)Qπθold(s,a)]J(θ)Eθold[πθ(as)πθold(as)Qπθold(s,a)] \frac{d}{d \theta} J(\theta) \approx \text{E}_{\theta_{\text{old}}} [\frac{\pi_{\theta}(a | s)}{\pi_{\theta_{\text{old}}}(a | s)} Q_{\pi_{\theta_{\text{old}}}}(s, a) \frac{\frac{d}{d \theta}(\pi_{\theta}(a | s))}{\pi_{\theta}(a | s)}] \\ = \text{E}_{\theta_{\text{old}}} [\frac{\frac{d}{d \theta}(\pi_{\theta}(a | s))}{\pi_{\theta_{\text{old}}}(a | s)} Q_{\pi_{\theta_{\text{old}}}}(s, a)] \\ J(\theta) \approx \text{E}_{\theta_{\text{old}}} [\frac{\pi_{\theta}(a | s)}{\pi_{\theta_{\text{old}}}(a | s)} Q_{\pi_{\theta_{\text{old}}}}(s, a)]

学習を安定させるためにアドバンテージを使用し、r(θ)=πθ(as)πθold(as)r(\theta) = \frac{\pi_\theta(a | s)}{\pi_{\theta_{\text{old}}}(a | s)}を用いて表現を簡略化すると、最終的な代理目的関数に到達できます。

J(θ)Eθold[r(θ)Aπθold(s,a)] J(\theta) \approx \text{E}_{\theta_{\text{old}}} [r(\theta) A_{\pi_{\theta_{\text{old}}}}(s, a)]

代理目的関数はθθold\theta \approx \theta_{\text{old}}を前提としており、これは同期アプローチを使用するか、 各パラメータ更新のステップサイズ(学習率)が小さいことを確保することで維持できます。 Trust Region Policy Optimization(TRPO) は、 KLダイバージェンスに制約を課すことでステップサイズを小さく保ちます。

Esdπθold[DKL(πθold(.s)πθ(.s))]δ \text{E}_{s \sim d_{\pi_{\theta_{\text{old}}}}} [D_{\text{KL}}(\pi_{\theta_{\text{old}}}(.|s)||\pi_{\theta}(.|s))] \leq \delta

ここで、π(.s)\pi(.|s)はアクション分布を表し、δ\deltaはハイパーパラメータです。近似勾配は共役勾配法を使用して効率的に解くことができ、 制約を満たす十分なステップサイズはバックトラッキングライン検索によって決定できます(残念ながらこれらの手法はこの記事の範囲外です)。 適切かつ慎重なパラメータ更新により、計算コストは素朴なA3Cよりも高くなりますが、論文で証明されているように、 単調なポリシー改善保証を伴うより安定した学習が可能になります。

PPO

制約を適用する代わりに、ハイパーパラメータβ\betaでスケーリングされたKLダイバージェンスペナルティを加える(勾配上昇法を実行しているため減算する)ことができます。 これにより、従来の最適化アルゴリズムを使用し、簡単に実装することができます。

J(θ)Eθold[r(θ)Aπθold(s,a)]βEsdπθold[DKL(πθold(.s)πθ(.s))] J(\theta) \approx \text{E}_{\theta_{\text{old}}} [r(\theta) A_{\pi_{\theta_{\text{old}}}}(s, a)] - \beta \text{E}_{s \sim d_{\pi_{\theta_{\text{old}}}}} [D_{\text{KL}}(\pi_{\theta_{\text{old}}}(.|s)||\pi_{\theta}(.|s))]

しかし、固定されたβ\betaを選択しても、適切な信頼領域を促進することはできません。そのため、KLダイバージェンスが目標のKLダイバージェンスよりも小さいか大きいかによってβ\betaを調整することができます。 具体的には、KLダイバージェンスが目標の3分の2未満の場合はβ\betaを半分にして大きな更新を許可し、KLダイバージェンスが目標の1.5倍以上の場合はβ\betaを2倍にして更新にペナルティを与えることができます。 これまで、πθ\pi_{\theta}πθold\pi_{\theta_{\text{old}}}がどれだけ異なるかの尺度としてKLダイバージェンスを使用してきましたが、r(θ)r(\theta)の値が1に近いほど、2つのポリシーはより類似しているため、 r(θ)r(\theta)も使用することができます。したがって、r(θ)r(\theta)を特定の範囲内にクリッピングして大きなステップサイズを避け、信頼領域と同様の効果を達成できると考えることができます。

J(θ)Eθold[clip(r(θ),1ϵ,1+ϵ)Aπθold(s,a)] J(\theta) \approx \text{E}_{\theta_{\text{old}}} [\text{clip}(r(\theta), 1 - \epsilon, 1 + \epsilon) A_{\pi_{\theta_{\text{old}}}}(s, a)]

ここで、ϵ\epsilonは更新の閾値を設定するためのハイパーパラメータです。r(θ)r(\theta)がこれらの境界の外にある場合、ポリシーの更新を許可しません。 しかし、上記のアプローチは理想的ではありません。なぜなら、負のアドバンテージ(r(θ)A<(1+ϵ)Ar(\theta) A < (1 + \epsilon)Aの場合)では比率を1+ϵ1 + \epsilonでクリッピングし、 正のアドバンテージ(r(θ)A<(1ϵ)Ar(\theta)A < (1 - \epsilon)Aの場合)では比率を1ϵ1 - \epsilonでクリッピングすることになるからです。 より保守的な代理目的関数をもたらす場合にのみ比率をクリッピングするために、以下のように上記を修正することができます。

J(θ)Eθold[min(r(θ)Aπθold(s,a),clip(r(θ),1ϵ,1+ϵ)Aπθold(s,a))] J(\theta) \approx \text{E}_{\theta_{\text{old}}} [\min(r(\theta)A_{\pi_{\theta_{\text{old}}}}(s, a), \text{clip}(r(\theta), 1 - \epsilon, 1 + \epsilon) A_{\pi_{\theta_{\text{old}}}}(s, a))]

上記を代理目的関数として適用することで、より単純な計算でKLダイバージェンス制約やペナルティを適用するのと同じ効果を得ることができます。 アドバンテージ批評家のための二乗誤差と探索を促進するポリシーのエントロピーボーナスHπθ(s)H_{\pi_{\theta}}(s)を組み合わせると、 近位ポリシー最適化(PPO) の目的関数JPPO(θ)J^{\text{PPO}}(\theta)に到達することができます。これはより単純な計算でTRPOと同等の結果を達成します。

Hπθ(s)=aπθ(as)log(πθ(as))JCLIP(θ)=Eθold[min(r(θ)Aπθold(s,a),clip(r(θ),1ϵ,1+ϵ)Aπθold(s,a))]JPPO(θ)=E[JCLIP(θ)c1(V(s)Vθ(s))2+c2Hπθ(s)] H_{\pi_{\theta}}(s) = - \sum_{a} \pi_{\theta}(a | s) \log(\pi_{\theta}(a | s)) \\ J^{\text{CLIP}}(\theta) = \text{E}_{\theta_{\text{old}}} [\min(r(\theta)A_{\pi_{\theta_{\text{old}}}}(s, a), \text{clip}(r(\theta), 1 - \epsilon, 1 + \epsilon) A_{\pi_{\theta_{\text{old}}}}(s, a))] \\ J^{\text{PPO}}(\theta) = \text{E}[J^{\text{CLIP}}(\theta) - c_1 (V_{*}(s) - V_{\theta}(s))^2 + c_2 H_{\pi_{\theta}}(s)]

ここで、c1c_1c2c_2は批評家の損失とエントロピーボーナスが全体の目的関数に与える影響を調整するハイパーパラメータです。 エントロピーボーナスHπθ(s)H_{\pi_{\theta}}(s)はポリシーが一様分布に近いときに大きくなり、より多くの探索を促します。 この目的関数に対して勾配上昇法を実行するため、二乗誤差を減算することが重要です。 (PyTorchの実装例はMachine Learning with Phil(2021)によるこちらの動画で確認できます。)

PPOのシンプルさと優れたパフォーマンスにより、様々な複雑な強化学習問題(Dota 2のプレイ、ロボットの制御など)で使用され、 顕著な成功を示しています。また、教師あり学習による次のトークン予測の後、人間のフィードバックを報酬として使用し(reinforcement learning form human feedbackまたはRLHFと呼ばれる)、 人間が好む回答を生成するようにChatGPTなどのLLMを微調整するためにも使用されています。最近世界を席巻したDeepSeek-R1でさえ、その訓練全体にGRPOと呼ばれるPPOの変種を使用しており、 PPOの比類ない能力と多大な可能性を示しています。

結論

この記事では、いくつかの高度なActor-Critic手法、A3C、A2C、TRPO、PPOを紹介しました。A2Cはこの中で真にオンポリシーである唯一のアルゴリズムですが、 紹介されたすべての手法は一般的にオンポリシーのポリシー勾配法と考えられています。しかし、決定論的ポリシーを学習するDPGのようなすべてのオフポリシー手法では、 重要度サンプリング比を使用して目的関数を修正することができます。DPG、SAC、TD3など、他にも多くのポリシー勾配法がありますので、 興味がある方はぜひチェックしてみてください。

リソース