[CS285] 9. Advanced Policy Gradients

💡 이 글은 『2024 PseudoLab 전반기 강화학습팀』으로 진행되었으며, CS285 Fall 2023를 따라 정리했습니다.

1. Introduction

지금까지 공부한 policy gradient와 policy iteration은 유사한 점이 있다. 두 알고리즘 모두 advantage estimator $A^ \pi (s, a)$를 계산하고, 이를 이용하여 policy update를 하는 두 단계를 반복한다. 두 알고리즘의 차이점이라면 policy gradient는 policy를 gradient를 사용해 soft update를 하고, policy iteration은 argmax를 사용해 hard update를 한다는 점이다. 이러한 관점에서 policy gradient는 policy iteration의 soften version이라고 볼 수 있다. 따라서 만약 advantage estimator가 완벽하지 않다면, soft policy change가 더 reasonable하고, 따라서 policy gradient가 잘 작동하는 것이다. 이를 수학적으로 나타내보자.

수학적으로 policy gradient를 policy iteration과 다음과 같이 연결지을 수 있다.

\[J(\theta ^ \prime) - J(\theta) = \mathbb{E}_ {\tau \sim p_ {\theta ^ \prime} (\tau)} \left[ \sum _ t \gamma ^ t A ^ {\pi _ \theta} (s_t, a_t) \right]\]

해석하면 objective의 improvement $J(\theta ^ \prime) - J(\theta)$는 new policy에서의 trajectory distribution 하에서 previous policy의 advantage를 계산한 값과 같다. 즉, advantage를 old policy에서 계산하고, 이를 바탕으로 new policy에 대하여 update하는 것인데, 이는 policy iteration의 방식과 동일하다. 하지만 policy iteration과는 달리 policy distribution이 soft하고, 따라서 위 식은 policy gradient를 policy iteration처럼 표현한 식이라고 볼 수 있다. 위 식의 증명은 아래를 참고하면 간단하다. 아래에서 $\mathbb{E}_ {s_0 \sim p(s_0)} \left[V^ {\pi _ \theta} (s_0) \right] = \mathbb{E}_ {\tau \sim p_ {\theta^ \prime} (\tau)} \left[ V^ {\pi _ \theta} (s_0) \right]$인 이유는 initial state marginal이 항상 똑같기 때문이다.

우변의 식을 $\theta^ \prime$에 대하여 최대화하는 것은 결국 $J(\theta ^ \prime)$를 최대화하는 것과 같다. 우변의 값을 계산하기 위해 알 수 없는 $\theta ^ \prime$을 제거하자. 이를 위해 importance sampling을 사용하면 아래와 같다.

그래도 $s_t \sim p_ {\theta ^ \prime} (s_t)$를 얻을 수 없다는 문제가 있다. 따라서 어쩔 수 없이 아래 근사를 사용한다.

이 근사가 충분히 참이 되기 위해서는 $p_ {\theta ^ \prime} (s_t)$와 $p_ {\theta} (s_t)$가 비슷해야 한다. 이는 policy $\pi _ {\theta ^ \prime}$가 $\pi _ \theta$와 비슷하다는 것을 의미한다. 따라서 그러한 제한 조건 하에서는 위 식 $\bar{A} (\theta^ \prime)$를 최대화하는 것이 $J(\theta ^ \prime)$를 최대화하는 것과 같다. 그렇다면 이러한 제한 조건을 수학적으로 나타내고, 그 조건 하에서 위 식을 최대화하는 방법을 알아보자.

2. Bounding the Distribution Change

임의의 분포 $\pi_ \theta$에 대하여 두 분포 $\pi_ \theta$와 $\pi_ {\theta ^ \prime}$가 비슷하다는 것은 $\vert \pi_ \theta (a_t \vert s_t) - \pi_ {\theta ^ \prime} (a_t \vert s_t) \vert \leq \epsilon$이라는 식으로 나타낼 수 있다. 이때 $p_ {\theta ^ \prime} (s_t)$를 다음과 같이 나타낼 수 있다.

\[p_ {\theta ^ \prime} (s_t) = (1 - \epsilon) ^ t p_ {\theta} (s_t) + (1 - (1 - \epsilon) ^ t) p_ {\text{mistake}} (s_t)\]

따라서 두 분포 $p_ {\theta ^ \prime} (s_t)$와 $p_ {\theta} (s_t)$의 total variation divergence(TV divergence)는 다음과 같이 나타낼 수 있다.

\[\begin{aligned} \vert p_ {\theta ^ \prime} (s_t) - p_ {\theta} (s_t) \vert &= (1 - (1-\epsilon) ^ t) \vert p_ {\text{mistake}} (s_t) - p_ {\theta} (s_t) \vert \\ &\leq 2 (1 - (1 - \epsilon) ^ t) \leq 2 \epsilon t \\ \end{aligned}\]

이제 objective value의 bound를 계산해보자. 먼저 임의의 함수 $f(s_t)$에 대하여 다음 부등식이 성립하는 것을 쉽게 알 수 있다.

이를 이용하여 objective value의 bound를 계산하면 다음과 같다.

마치 ELBO를 계산하는 것과 같이, 좌변의 true objective value를 최대화하는 것은 우변의 bound를 최대화하는 것과 같다. 즉, $\vert \pi_ \theta (a_t \vert s_t) - \pi_ {\theta ^ \prime} (a_t \vert s_t) \vert \leq \epsilon$이라는 constraint 하에서 bound를 maximize하면 true objective를 maximize하는 것과 같은 효과를 낼 수 있다. 이를 아래와 같이 나타낼 수 있다.

3. Policy Gradients with Constraints

실제로는 TV divergence 대신 KL divergence를 사용하여 $D_ {KL} (\pi_ {\theta ^ \prime} \Vert \pi_ \theta) \leq \epsilon$이라는 constraint를 사용한다. 이러한 constraint를 강제하기 위하여 일종의 regularization loss를 사용한다. 최적화 측면에서는 Lagrangian을 사용하여 constraint를 부여하는 것으로 볼 수 있고, KL divergence의 값에 따라 Lagrangian multiplier를 조절한다. 이를 dual gradient descent라고 한다. 이 방법은 guided policy search나 PPO에서 볼 수 있다. 물론 $\lambda$를 휴리스틱하게 설정하는 것도 가능하다.

그러나 이 방법은 계산량이 많아 불완전하게 사용되는 경우가 많고, 대신 여기에서는 natural gradient에 대해 소개한다.

4. Natural Gradient

Natural gradient는 또 다른 방법으로 constraint를 부여하는 방법이다. 지금까지의 최적화 목표를 다시 보면 다음과 같다.

여기서 objective $\bar{A} (\theta ^ \prime)$를 최적화하기 위하여 해당 함수의 Taylor 1차 근사에 대한 최적화 값을 찾는 것으로 문제를 근사한다. 즉, 특정 constraint 하에서 linear approximation을 최적화하는 것은 closed-form solution이 존재하는 간단한 문제이고, 이러한 근사는 상당히 정확한 편이다.

이 문제의 closed-form solution은 다음과 같이 natural gradient로 나타낼 수 있다. 이때 $\mathbf{F}$는 Fisher-information matrix로 sample을 사용하여 $\mathbb{E}_ {\pi_ \theta} \left[\nabla_ \theta \log \pi_ \theta (a \vert s) \nabla_ \theta \log \pi_ \theta (a \vert s) ^ T \right]$로 추정할 수 있다. Natural gradient에 대한 설명은 CS285 Lecture 5 강의를 참고하자.

\[\theta^ \prime = \theta + \alpha \mathbf{F} ^ {-1} \nabla_ \theta J (\theta) \quad \alpha = \sqrt{\frac{2 \epsilon}{\nabla_ \theta J(\theta) ^ T \mathbf{F} \nabla_ \theta J( \theta)}}\]

추가로 TRPO(Trust Region Policy Optimization)를 통해 Fisher-information matrix를 더 효율적으로 계산하는 방법도 있다.

CS285 카테고리 내 다른 글 보러가기