“速通” PPO - Nagi-ovo

プロキシマルポリシー最適化

ついにここ数年の NLP 分野で比較的注目を集めている RL アルゴリズムの一つに到達しました。

On-Policy アルゴリズムでは、データ収集に使用するポリシーとトレーニングに使用するポリシーが同じであるため、データを一度使用すると捨てなければならず、再度データを収集する必要があり、トレーニング速度が非常に遅くなります。

PPO の背後にある直感#

PPO の理念は、各トレーニングサイクルでのポリシーの変更を制限することで、ポリシーのトレーニングの安定性を向上させることです：急激なポリシーの更新を避けることです。

Screenshot 2024-10-11 at 13.53.20

これは二つの理由からです：

この分野の経験に基づくと、トレーニング中の小さなポリシー更新は最適解に収束する可能性が高いです。
ポリシー更新において、過度に大きなステップは「崖から落ちる」（悪いポリシーを得る）可能性があり、回復には長い時間がかかり、元のレベルに戻れないこともあります。

クリップされた代理目的関数#

振り返り：ポリシー目的関数#

私たちの目標は、勾配上昇（または勾配下降の負関数）を取ることで、エージェントがより高い報酬をもたらす行動を選択し、負の効果をもたらす可能性のある動作を避けることです。

L^{PG}(\theta) = \mathbb{E}_t \left[ \log \pi_\theta(a_t | s_t) * A_t \right]

$\log \pi_\theta(a_t | s_t)$ ：状態 $s_t$ において行動 $a_t$ を選択する対数確率で、現在のポリシーでこの行動を取る確率がどれほど大きいかを示します。
$A_t$ ：アドバンテージ関数（Advantage）で、 $A > 0$ の場合、この行動は現在の状態で他の可能な行動よりも良いことを示します；逆に、 $A < 0$ の場合は劣ります。
しかし、古典的な PG 手法には一つの問題があります：ポリシー更新のステップサイズの選択が非常に重要です。

ステップサイズが小さすぎると、トレーニングプロセスが非常に遅くなります；
ステップサイズが大きすぎると、トレーニング中の変動が大きくなり、トレーニングが不安定になる可能性があります。

そこで、PPOは新しいアプローチ、クリップされた代理目的関数を提案しました。これはポリシーの変化の範囲を制限することで、ポリシー更新が過度に急激にならないようにし、トレーニングプロセスの安定性を保ちます。

この新しい目的関数は次のようになります：

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

比率関数#

ここでの重要な部分は比率関数 $r_t(\theta)$ で、これは現在のポリシーと以前のポリシーの間の行動確率の比率を示します：

r_t(\theta) = \frac{\pi_\theta(a_t | s_t)}{\pi_{\theta_{\text{old}}}(a_t | s_t)}

比率は現在のポリシーと古いポリシーの偏差の程度を反映します：

$r_{t(\theta)}\in (0, 1)$ の場合、現在のポリシーではこの行動を選択する確率が減少したことを示します。
$r_t(\theta) > 1$ の場合、現在のポリシーでは行動 $a_t$ が以前よりも選ばれる可能性が高いことを示します。

未クリップ部分#

式の未クリップ部分は次の通りです：

L^{CPI}(\theta) = \hat{\mathbb{E}}_t \left[ \frac{\pi_\theta(a_t | s_t)}{\pi_{\theta_{\text{old}}}(a_t | s_t)} \hat{A}_t \right] = \hat{\mathbb{E}}_t \left[ r_t(\theta) \hat{A}_t \right]

未クリップの目的関数では、 $\ r_t(\theta)$ がアドバンテージ値 $\hat{A}_t$ に直接掛けられます。もし行動 $a_t$ が現在のポリシーの下で古いポリシーよりも優れている（つまりアドバンテージ値 $\hat{A}_t > 0$ ）場合、私たちはその行動を推奨し、逆にその影響を弱めます。これは標準的なポリシー勾配最適化の方向です。

しかし、前述のように、制約のないポリシー更新はトレーニングの不安定性を引き起こす可能性があります。比率 $r_t(\theta)$ が 1 を大きく超えると、ポリシー更新が過度になり、トレーニングプロセスで収束が難しくなります。

この時、PPO はクリッピングポリシーを導入し、比率の範囲をクリップします。

クリップ部分#

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

ここで、 $\min$ 操作の導入が見られます。比率 $r_t(\theta)$ が設定された閾値 $[1 - \epsilon, 1 + \epsilon]$ を超えた場合、クリッピング操作は比率をこの範囲内に制限します。これにより、ポリシー更新が過度になるのを防ぎます。

クリップ比率関数は次の通りです：

\text{clip}\left( r_t(\theta), 1 - \epsilon, 1 + \epsilon \right)

これは、比率 $r_t(\theta)$ が設定された範囲（元の論文では $\epsilon = 0.2$ ）を超えると、 $[0.8, 1.2]$ の間にクリップされ、ポリシー更新の安定性が保証されることを意味します。私たちはクリップされた値と未クリップの値の間の最小値を取ります。これにより、最終的な目的関数が過度に楽観的にならず、より保守的な推定に向かうことが保証されます。

可視化#

Pasted image 20241013022437

まず、私たちはクリップされた目的と未クリップの目的の間の最小値を取ることを思い出してください。

状況 1 と 2：比率が範囲内#

この二つの状況では、クリッピングは行われません。ポリシーは $A_t$ の正負に応じて適切に更新されます。これは PPO の理想的な状態であり、すべてが期待通りに進行します。

状況 1： $A_t > 0$ かつ $p_t(\theta) \in [1 - \epsilon, 1 + \epsilon]$
- アドバンテージ関数 $A_t$ が正であることは、この行動が期待よりも良いことを意味します。
- $p_t(\theta)$ がこの範囲内にあることは、ポリシーの変化が小さいことを示し、私たちはこの行動を奨励したいため、クリッピングは行いません。
- 結果：目的関数は正であり、勾配更新はポリシーをさらにこの行動に偏らせます。
状況 2： $A_t < 0$ かつ $p_t(\theta) \in [1 - \epsilon, 1 + \epsilon]$
- アドバンテージ関数が負であることは、この行動が期待よりも悪いことを意味します。
- 同様に、比率が範囲内にあるため、クリッピングは行いません。私たちはこの行動の実行を減らしたいと考えています。
- 結果：目的関数は負であり、勾配更新はポリシーをこの行動から遠ざけます。

状況 3 と 4：比率が範囲を下回る#

ここでの比率は、現在のポリシーが古いポリシーよりもこの行動の確率を過小評価していることを示しています。何が起こるでしょうか？

状況 3： $A_t > 0$ かつ $p_t(\theta) < 1 - \epsilon$
- 行動は良好です（アドバンテージ関数が正）ですが、新しいポリシーはこの行動の確率を低く見積もっています。
- クリッピングは行わず、この優れた行動の確率を 増加させたい ため、勾配が強く更新を促進します。
- 結果：目的関数は正であり、勾配はこの行動を奨励します。
状況 4： $A_t < 0$ かつ $p_t(\theta) < 1 - \epsilon$
- 行動は非常に悪い（アドバンテージ関数が負）ですが、ポリシーはこの行動の確率を減少させています。
- しかし、クリッピングを行います。なぜなら、確率がすでに $1 - \epsilon$ を下回っており、さらに減少させることは過度な罰を与え、トレーニングを不安定にする可能性があるからです。
- 結果：目的関数はクリップされ、勾配は更新されず、この行動の確率は下限に保たれます。

状況 5 と 6：比率が範囲を超える#

ここでは、ポリシーが行動に対して過度に自信を持っていることを示しています。これは、新しいポリシーがこの行動の実行確率を過大評価していることを意味します。

状況 5： $A_t > 0$ かつ $p_t(\theta) > 1 + \epsilon$
- 行動は良好ですが（アドバンテージ関数が正）、新しいポリシーはその実行確率を過大評価しています。
- クリッピングを行います。なぜなら、ポリシーがこの行動に過度に偏ることを望まないからです。 $A_t$ が正であっても、ポリシーの更新幅を制限する必要があります。
- 結果：目的関数はクリップされ、勾配は更新されず、ポリシーの変化幅が制限されます。
状況 6： $A_t < 0$ かつ $p_t(\theta) > 1 + \epsilon$
- 行動は悪いですが、ポリシーはその実行確率を高めています。これは明らかに望ましくありません。
- この時、比率はすでに範囲を超えているため、クリッピングは行いません。目的関数は負であり、勾配はポリシーをこの悪い行動から遠ざけます。
- 結果：目的関数は負であり、勾配はポリシーをこの行動から遠ざけます。

なぜクリッピングされた場合、勾配は 0 になるのか？#

理由は、比率 $r_t(\theta)$ が $1 - ϵ$ または $1 + ϵ$ にクリップされると、導関数は比率 $r_t(\theta)$ にアドバンテージ $A_t$ の導関数を掛けたものではなく、 $(1 - ϵ)A_t$ または $(1 + ϵ)A_t$ の導関数になりますが、これらの二つの表現の導関数は 0 です。

まとめ#

まとめると、PPO の目標は、クリップされた代理目的を通じて、現在のポリシーと古いポリシーの間の変化の範囲を制限することです。私たちは、確率比が $[1 - ϵ, 1 + ϵ]$ の範囲を超えるインセンティブを取り除きました。なぜなら、一度比率がその範囲を超えると、勾配は 0 になり、ポリシーの更新が停止するからです。

PPO の更新プロセスでは、ポリシーを更新するのは次の二つの状況のみです：

比率 $r_t(\theta)$ が $[1 - ϵ, 1 + ϵ]$ の範囲内にあるとき。
比率が範囲外であるが、アドバンテージ関数が比率をその範囲に近づけるとき。

最後に復習すると、PPO のクリップされた代理目的損失は三つの部分から構成されています：

クリップされた代理目的関数：ポリシー更新の変化の範囲を制限します。
価値損失関数：価値関数の平均二乗誤差を最小化するために使用されます。
エントロピーボーナス：十分な探索を維持するために使用され、ポリシーが早期に局所最適に陥るのを防ぎます。

これら三つの部分が組み合わさって、PPO が安定してポリシーを更新し、十分な探索性を保持できるようにします。

コード実装#

さて、コードの観点から PPO の実装を深く理解してみましょう。cleanrl のppo.pyの最も重要な部分に焦点を当て、その動作原理を簡潔に説明します。

1. ポリシーネットワークと価値ネットワークの構造#

class Agent(nn.Module):
    def __init__(self, envs):
        super().__init__()
        # クリティックネットワーク：状態を価値にマッピング（神経ネットワークを使用して状態の良し悪しを推定）
        self.critic = nn.Sequential(
            layer_init(nn.Linear(np.array(envs.single_observation_space.shape).prod(), 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, 1), std=1.0), # 最後の層の初期化は学習の安定性に非常に重要です
        )
        # アクターネットワーク：状態を行動確率にマッピング（ポリシーの神経ネットワーク）
        self.actor = nn.Sequential(
            layer_init(nn.Linear(np.array(envs.single_observation_space.shape).prod(), 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, envs.single_action_space.n), std=0.01), # 小さな標準偏差は初期ポリシーがほぼ均一分布に近いことを保証します
        )

これは典型的な二重ネットワークアーキテクチャです：

アクター（ポリシーネットワーク）は行動の確率分布を出力します
クリティック（価値ネットワーク）は状態の価値を予測します
二つのネットワークはシンプルな二層 MLP 構造（64-64）を採用しています
正交初期化（orthogonal initialization）を使用してトレーニングの安定性を助けます

2. GAE（一般化アドバンテージ推定）の実装#

# GAE計算：逆方向にアドバンテージ関数とリターン値を計算
with torch.no_grad():
    next_value = agent.get_value(next_obs).reshape(1, -1)
    advantages = torch.zeros_like(rewards).to(device)
    lastgaelam = 0
    for t in reversed(range(args.num_steps)):
        # GAE（一般化アドバンテージ推定）の優雅な実装
        # 時系列差分誤差の指数加重和として理解できます
        delta = rewards[t] + args.gamma * nextvalues * nextnonterminal - values[t]
        advantages[t] = lastgaelam = delta + args.gamma * args.gae_lambda * nextnonterminal * lastgaelam
    returns = advantages + values  # リターン値 = アドバンテージ関数 + 価値推定

このコードは GAE の再帰的計算プロセスを示しています：

後ろから前に TD 誤差（デルタ）を計算します
これらの TD 誤差を指数加重の方法で累積します
gamma と lambda のハイパーパラメータは価値推定のバイアス - バリアンスのトレードオフを制御します

3. PPO のコア損失関数の計算#

# PPOのコア：ポリシーの変化を防ぎつつポリシーを改善します
_, newlogprob, entropy, newvalue = agent.get_action_and_value(b_obs[mb_inds], b_actions.long()[mb_inds])
ratio = (newlogprob - b_logprobs[mb_inds]).exp()  # 重要性サンプリング比率

# 有名なPPO-Clip目的関数
pg_loss1 = -mb_advantages * ratio
pg_loss2 = -mb_advantages * torch.clamp(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
pg_loss = torch.max(pg_loss1, pg_loss2).mean()  # 悲観的（最悪のケース）ポリシー損失

# 価値関数損失もクリッピングを使用して旧予測に近づけます
if args.clip_vloss:
    v_loss_unclipped = (newvalue - b_returns[mb_inds]) ** 2
    v_clipped = b_values[mb_inds] + torch.clamp(
        newvalue - b_values[mb_inds],
        -args.clip_coef,
        args.clip_coef,
    )
    v_loss_clipped = (v_clipped - b_returns[mb_inds]) ** 2
    v_loss = 0.5 * torch.max(v_loss_unclipped, v_loss_clipped).mean()

# 総合損失関数：ポリシー損失、価値損失、探索のためのエントロピー報酬項を組み合わせます
loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef

ここでは前述の三つの重要なコンポーネントが実装されています：

クリップされた代理目的は max 操作を使用してクリッピングを実現します
価値関数損失もクリッピングメカニズムを使用しています（これは OpenAI の実装の特徴です）
エントロピーボーナスは探索を奨励するために使用されます

最近、プラットフォームの ipfs がクラッシュしたため、長い間保留されていましたが、この投稿の後、Huggingface DeepRL シリーズが正式に完結します～