Proximal Policy Optimization

Created on August 03, 2024

2024   ·   Reinforcement Learning   PPO   ·   ML/DL Theory

Concept

Proximal Policy Optimization (PPO)는 policy gradient 계열의 Reinforcement Learning algorithm이다.

PPO의 핵심 목적은 policy update 과정에서 policy가 너무 크게 변하는 것을 제한하여 performance collapse를 완화하는 것이다.

Policy gradient 기반 algorithm에서는 policy \(\pi_\theta\)를 parameter \(\theta\)로 표현하고, objective function \(J(\pi_\theta)\)를 증가시키는 방향으로 parameter를 업데이트한다.

\[\theta \leftarrow \theta + \alpha \nabla_\theta J(\pi_\theta)\]

위 수식에서 \(\theta\)는 policy network의 parameter를 의미하고, \(\alpha\)는 learning rate를 의미한다. \(\nabla_\theta J(\pi_\theta)\)는 현재 policy의 objective function을 증가시키는 방향을 나타낸다.

문제는 parameter space에서 작은 변화가 policy space에서도 항상 작은 변화로 대응되지는 않는다는 점이다. 즉, parameter \(\theta\)의 변화량이 작더라도 action probability distribution \(\pi_\theta(a \mid s)\)는 크게 바뀔 수 있다.

이러한 큰 policy 변화는 좋지 않은 trajectory를 생성할 수 있고, 이 trajectory를 기반으로 다시 policy gradient를 계산하면 policy가 더 불안정해질 수 있다. PPO는 이 문제를 해결하기 위해 policy update가 일정한 safe region 안에서 이루어지도록 surrogate objective를 설계한다.

Performance Collapse

Performance collapse는 policy update 이후 agent의 성능이 갑자기 저하되는 현상을 의미한다.

Policy gradient에서는 새로운 policy가 생성한 trajectory를 사용하여 다음 update를 수행한다. 만약 한 번의 update로 policy가 크게 변하면 agent는 이전보다 훨씬 나쁜 trajectory를 생성할 수. 이후 이 나쁜 trajectory를 기반으로 gradient가 계산되면 policy가 더 나빠지는 악순환이 발생할 수 있다.

이 문제는 parameter space와 policy space 사이의 mapping이 동일한 scale로 이루어지지 않기 때문에 발생한다.

Parameter space에서 두 parameter pair의 거리가 같다고 하더라도, 이에 대응되는 policy distribution 사이의 거리는 서로 다를 수 있다.

\[d_\theta(\theta_1, \theta_2) = d_\theta(\theta_2, \theta_3)\] \[d_\theta(\theta_1, \theta_2) = d_\theta(\theta_2, \theta_3) \not\Rightarrow d_\pi(\pi_{\theta_1}, \pi_{\theta_2}) = d_\pi(\pi_{\theta_2}, \pi_{\theta_3})\]

위 수식에서 \(d_\theta\)는 parameter space에서의 distance를 의미하고, \(d_\pi\)는 policy space에서의 distance를 의미한다. 즉, parameter update의 크기를 고정하더라도 실제 policy 변화량은 일정하지 않을 수.

따라서 PPO는 단순히 learning rate \(\alpha\)를 작게 설정하는 방식이 아니라, objective function 자체를 수정하여 policy change를 제한하는 방식을 사용한다.

Relative Policy Performance Identity

PPO의 출발점은 새로운 policy \(\pi'\)가 기존 policy \(\pi\)보다 얼마나 좋아졌는지를 측정하는 것이다.

기존 policy \(\pi\)와 update 이후 policy \(\pi'\)가 있을 때, 두 policy의 objective function 차이는 다음과 같이 표현된다.

\[J(\pi') - J(\pi) = \mathbb{E}_{\tau \sim \pi'} \left[ \sum_{t=0}^{T} \gamma^t A^\pi(s_t, a_t) \right]\]

위 수식에서 \(J(\pi')\)는 새로운 policy의 objective function을 의미하고, \(J(\pi)\)는 기존 policy의 objective function을 의미한다. \(\tau\)는 trajectory를 의미하며, \(\gamma\)는 discount factor를 의미한다. \(A^\pi(s_t, a_t)\)는 기존 policy \(\pi\)를 기준으로 계산한 advantage function이다.

Advantage function은 특정 state-action pair가 현재 policy의 평균적인 행동보다 얼마나 좋은지를 나타낸다.

\[A^\pi(s_t, a_t) = Q^\pi(s_t, a_t) - V^\pi(s_t)\]

위 수식에서 \(Q^\pi(s_t, a_t)\)는 state \(s_t\)에서 action \(a_t\)를 선택했을 때의 action-value function을 의미하고, \(V^\pi(s_t)\)는 state-value function을 의미한다.

만약 \(J(\pi') - J(\pi)\)가 양수라면 새로운 policy \(\pi'\)는 기존 policy \(\pi\)보다 더 좋은 policy라고 해석할 수 있다.

Surrogate Objective

Relative policy performance identity는 새로운 policy \(\pi'\)가 생성한 trajectory에 대한 expectation을 포함한다.

\[J(\pi') - J(\pi) = \mathbb{E}_{\tau \sim \pi'} \left[ \sum_{t=0}^{T} \gamma^t A^\pi(s_t, a_t) \right]\]

하지만 실제 학습에서는 \(\pi'\)를 얻기 전에 \(\pi'\)로부터 trajectory를 sampling할 수 없다. 따라서 기존 policy \(\pi\)로부터 얻은 trajectory를 활용하여 새로운 policy의 성능을 추정할 필요가 있다.

이를 위해 importance sampling을 사용한다.

\[J(\pi') - J(\pi) = \mathbb{E}_{\tau \sim \pi} \left[ \frac{p_{\pi'}(\tau)}{p_{\pi}(\tau)} \sum_{t=0}^{T} \gamma^t A^\pi(s_t, a_t) \right]\]

위 수식에서 \(p_{\pi'}(\tau)\)는 새로운 policy \(\pi'\)가 trajectory \(\tau\)를 생성할 probability를 의미하고, \(p_\pi(\tau)\)는 기존 policy \(\pi\)가 trajectory \(\tau\)를 생성할 probability를 의미한다.

실제 PPO에서는 trajectory 전체의 probability ratio를 직접 사용하기보다, 각 time step에서 action probability ratio를 사용한다.

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

위 수식에서 \(r_t(\pi_\theta)\)는 현재 policy \(\pi_\theta\)와 old policy \(\pi_{\theta_{\text{old}}}\) 사이의 action probability ratio를 의미한다. \(\pi_\theta(a_t \mid s_t)\)는 현재 policy가 state \(s_t\)에서 action \(a_t\)를 선택할 probability이고, \(\pi_{\theta_{\text{old}}}(a_t \mid s_t)\)는 old policy가 동일한 action을 선택할 probability이다.

이를 이용한 conservative policy iteration (CPI) 형태의 surrogate objective는 다음과 같이 표현할 수 있다.

\[J_{\pi_{\text{old}}}^{CPI}(\pi_\theta) = \mathbb{E}_t \left[ r_t(\pi_\theta) A^{\pi_{\text{old}}}(s_t, a_t) \right]\]

이 objective는 old policy로부터 수집한 sample을 사용하면서도, 현재 policy가 해당 action을 얼마나 더 많이 또는 적게 선택하려 하는지를 ratio로 반영한다.

PPO Objective

PPO에는 크게 두 가지 형태가 있다.

첫 번째는 adaptive KL penalty 기반 PPO이고, 두 번째는 clipped surrogate objective 기반 PPO이다.

Adaptive KL penalty 기반 PPO는 objective function에 두 policy 사이의 KL divergence penalty를 추가한다.

\[J^{KLPEN}(\pi_\theta) = \mathbb{E}_t \left[ r_t(\pi_\theta) A^{\pi_{\text{old}}}(s_t, a_t) - \beta KL \left( \pi_\theta(\cdot \mid s_t) \Vert \pi_{\text{old}}(\cdot \mid s_t) \right) \right]\]

위 수식에서 \(KL(\pi_\theta \Vert \pi_{\text{old}})\)은 현재 policy와 old policy 사이의 distribution 차이를 측정한다. \(\beta\)는 KL penalty의 강도를 조절하는 coefficient이다.

Clipped surrogate objective 기반 PPO는 KL penalty를 직접 사용하지 않고, probability ratio \(r_t(\pi_\theta)\)를 일정 범위 안으로 clipping한다.

\[J^{CLIP}(\pi_\theta) = \mathbb{E}_t \left[ \min \left( r_t(\pi_\theta) A^{\pi_{\text{old}}}(s_t, a_t), \text{clip}(r_t(\pi_\theta), 1-\epsilon, 1+\epsilon) A^{\pi_{\text{old}}}(s_t, a_t) \right) \right]\]

위 수식에서 \(\epsilon\)은 clipping range를 정의하는 hyperparameter이다. 일반적으로 \(r_t(\pi_\theta)\)가 \(1-\epsilon\)보다 작아지거나 \(1+\epsilon\)보다 커지면, objective에서 사용되는 ratio가 clipped value로 제한된다.

즉, PPO는 현재 policy가 old policy와 너무 멀어지는 방향으로 update되는 것을 방지한다.

Interpretation of Clipped Objective

Clipped objective의 핵심은 advantage의 부호에 따라 policy update를 다르게 제한하는 것이다.

먼저 \(r_t(\pi_\theta) = 1\)이면 현재 policy와 old policy가 해당 state-action pair에 대해 동일한 action probability를 가진다는 의미이다.

\[r_t(\pi_\theta) = 1 \Rightarrow \pi_\theta(a_t \mid s_t) = \pi_{\text{old}}(a_t \mid s_t)\]

이 경우에는 policy가 아직 변하지 않았다고 볼 수 있다.

만약 \(1-\epsilon \le r_t(\pi_\theta) \le 1+\epsilon\)이면 ratio가 clipping range 안에 있으므로 CPI objective와 동일하게 작동한다.

\[J^{CLIP}(\pi_\theta) = \mathbb{E}_t \left[ r_t(\pi_\theta) A^{\pi_{\text{old}}}(s_t, a_t) \right]\]

이 경우에는 policy update가 safe region 안에 있다고 해석할 수 있다.

Advantage가 양수인 경우, 해당 action은 old policy 기준에서 좋은 action이다. 따라서 policy는 해당 action의 probability를 증가시키는 방향으로 업데이트된다.

하지만 \(r_t(\pi_\theta) > 1+\epsilon\)이면 이미 해당 action probability가 충분히 증가한 상태이다. PPO는 이때 objective 증가를 clipping하여 더 이상의 과도한 증가를 막는다.

\[r_t(\pi_\theta) > 1+\epsilon, \quad A^{\pi_{\text{old}}}(s_t, a_t) > 0\] \[J^{CLIP}(\pi_\theta) = \mathbb{E}_t \left[ (1+\epsilon) A^{\pi_{\text{old}}}(s_t, a_t) \right]\]

반대로 advantage가 음수인 경우, 해당 action은 old policy 기준에서 좋지 않은 action이다. 따라서 policy는 해당 action probability를 감소시키는 방향으로 업데이트된다.

하지만 \(0 < r_t(\pi_\theta) < 1-\epsilon\)이면 이미 해당 action probability가 충분히 감소한 상태이다. PPO는 이때 objective를 clipping하여 더 이상의 과도한 감소를 막는다.

\[0 < r_t(\pi_\theta) < 1-\epsilon, \quad A^{\pi_{\text{old}}}(s_t, a_t) < 0\] \[J^{CLIP}(\pi_\theta) = \mathbb{E}_t \left[ (1-\epsilon) A^{\pi_{\text{old}}}(s_t, a_t) \right]\]

따라서 PPO의 clipping은 좋은 action의 probability를 너무 많이 증가시키는 것과 나쁜 action의 probability를 너무 많이 감소시키는 것을 모두 제한한다.

PPO Algorithm Flow

PPO는 일반적으로 Actor-Critic 구조와 함께 사용된다.

Actor는 policy \(\pi_\theta(a \mid s)\)를 학습하고, Critic은 value function \(V_\phi(s)\)를 학습한다. 여기서 \(\theta\)는 Actor parameter이고, \(\phi\)는 Critic parameter이다.

PPO의 학습 과정은 다음과 같이 정리할 수 있다.

  • Old policy \(\pi_{\theta_{\text{old}}}\)를 사용하여 trajectory를 수집한다.

  • 수집된 trajectory로부터 return 또는 advantage estimate를 계산한다.

  • 현재 policy \(\pi_\theta\)와 old policy \(\pi_{\theta_{\text{old}}}\) 사이의 probability ratio \(r_t(\pi_\theta)\)를 계산한다.

  • Clipped surrogate objective \(J^{CLIP}(\pi_\theta)\)를 최대화하도록 Actor를 업데이트한다.

  • Value loss를 최소화하도록 Critic을 업데이트한다.

  • 일정 횟수의 optimization 이후 현재 policy를 old policy로 갱신한다.

PPO는 old policy로 수집한 trajectory를 여러 epoch 동안 재사용할 수 있다. 이 점에서 PPO는 on-policy 성격을 가지지만, update 과정에서는 old policy와 current policy의 ratio를 사용하는 importance sampling 형태를 포함한다.

PPO and Actor-Critic

PPO는 REINFORCE 또는 Actor-Critic의 objective function을 clipped surrogate objective로 대체하여 적용할 수.

REINFORCE에서는 Monte-Carlo return을 기반으로 policy gradient를 추정한다.

\[\nabla_\theta J(\pi_\theta) = \mathbb{E}_t \left[ \nabla_\theta \log \pi_\theta(a_t \mid s_t) G_t \right]\]

Actor-Critic에서는 return 대신 advantage estimate를 사용한다.

\[\nabla_\theta J(\pi_\theta) = \mathbb{E}_t \left[ \nabla_\theta \log \pi_\theta(a_t \mid s_t) A_t \right]\]

PPO는 이 policy gradient update를 직접 사용하는 대신, 현재 policy와 old policy 사이의 ratio를 포함한 clipped surrogate objective를 사용한다.

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

여기서 \(A_t\)는 advantage estimate를 의미한다. 실제 구현에서는 Generalized Advantage Estimation (GAE)을 사용하여 \(A_t\)를 계산하는 경우가 많다.

PPO는 Actor-Critic 구조의 안정성을 높이는 policy update mechanism으로 이해할 수 있다.

Interpretation

PPO는 policy optimization을 직접적으로 제한하는 algorithm이다.

기존 policy gradient에서는 parameter update 크기를 learning rate로 조절한다. 그러나 parameter space에서의 작은 변화가 policy space에서의 작은 변화로 이어진다는 보장은 없다.

PPO는 이 문제를 action probability ratio \(r_t(\pi_\theta)\)로 다룬다. Ratio가 \(1\)에 가까우면 current policy와 old policy가 비슷하다는 의미이고, ratio가 \(1\)에서 멀어질수록 policy가 크게 바뀌었다는 의미이다.

Clipping range \([1-\epsilon, 1+\epsilon]\)은 policy update의 허용 범위로 해석할 수 있다.

  • \(r_t(\pi_\theta)\)가 clipping range 안에 있으면 policy update를 허용한다.

  • \(r_t(\pi_\theta)\)가 clipping range를 벗어나면 objective의 추가적인 증가를 제한한다.

  • Advantage가 positive이면 좋은 action의 probability 증가를 제한한다.

  • Advantage가 negative이면 나쁜 action의 probability 감소를 제한한다.

이 구조는 policy가 한 번의 update에서 지나치게 멀리 이동하는 것을 막는다. 따라서 PPO는 performance collapse를 줄이고, 더 안정적인 policy optimization을 가능하게 한다.

Summary

PPO는 policy gradient 기반 Reinforcement Learning에서 performance collapse를 완화하기 위해 제안된 algorithm이다.

핵심 문제는 parameter space의 update 크기와 policy space의 변화량이 항상 일치하지 않는다는 점이다. 이로 인해 learning rate만으로는 안정적인 policy update를 보장하기 어렵다.

PPO는 old policy와 current policy 사이의 action probability ratio를 계산하고, 이 ratio를 clipping하여 policy update가 safe region 안에서 이루어지도록 한다.

Clipped surrogate objective는 다음과 같이 정의된다.

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

이 수식에서 \(r_t(\pi_\theta)\)는 current policy와 old policy의 action probability ratio이고, \(A_t\)는 advantage estimate이며, \(\epsilon\)은 clipping range를 의미한다.

결과적으로 PPO는 Actor-Critic 구조에서 policy update를 안정화하는 대표적인 Reinforcement Learning algorithm으로 이해할 수 있다.

References

  • Murphy, Kevin. “Reinforcement learning: an overview.” arXiv preprint arXiv:2412.05265 (2024).
  • Schulman, John, et al. “Proximal policy optimization algorithms.” arXiv preprint arXiv:1707.06347 (2017).