[강화학습 기초지식] Part 3: Intro to Policy Optimization (2)

본 포스트는 OpenAI에서 공개한 강화학습 교육자료인 Spinning Up의 도움을 받아, 강화학습(Reinforcement Learning, RL)에 대한 기초개념을 정리해보고자 제작한 시리즈의 일환입니다.

1)기초 지식 2)주요 논문, 3)최신 논문 의 순서로 시리즈를 정리하고자 합니다.

아래 링크에서 더 상세한 내용을 찾아볼 수 있습니다.

Part 3: Intro to Policy Optimization - Spinning Up documentation

In this section, we'll discuss the mathematical foundations of policy optimization algorithms, and connect the material to sample code. We will cover three key results in the theory of policy gradients: In the end, we'll tie those results together and describe the advantage-based expression for the policy gradient-the version we use in our Vanilla Policy Gradient implementation.

https://spinningup.openai.com/en/latest/spinningup/rl_intro3.html

Expected Grad-Log-Prob Lemma

Policy optimization의 근간이 될 policy gradient

이번 포스트에선, 지난번에 이어 policy gradient에 있어 굉장히 폭넓게 쓰이는 Expected Grad-Log-Prob(EGLP) lemma부터 시작하여, policy gradient에 대해 좀 더 깊이 살펴보겠습니다. 지난 포스트에서 가볍게 다뤄본 적 있으니 어렵게 느껴지진 않을 것입니다.

EGLP Lemma: Suppose that $P_\theta$ is a parametrized probability distribution over a random variable, $x$. Then:

$\underset{x\sim P_\theta}{\mathbb{E}}[\nabla_\theta\log P_\theta(x)]=0.$

Grad-log-probability를 random variable $x$에 대해 평균을 취했을 때 결국 0으로 수렴한다는 내용입니다.

Proof: Recall that all probability distributions are normalized:

$\int_x P_\theta(x)=1.$

Take the gradient of both sides of the normalization condition:

$\nabla_\theta\int_x P_\theta(x)=\nabla_\theta1=0.$

Use the log derivative trick to get:

$\begin{aligned} 0 &= \nabla_\theta\int_xP_\theta(x) \\ &=\int_x \nabla_\theta P_\theta(x) \\ &=\int_x P_\theta(x) \nabla_\theta \log P_\theta(x) \\ \therefore 0&=\underset{x \sim P_\theta}{\mathbb{E}}[\nabla_\theta\log P_\theta(x)]. \end{aligned}$

Don't Let the Past Distract You

지난 포스트에서 살펴본 policy gradient의 식을 다시 되짚어봅시다.

$\nabla_\theta J(\pi_\theta)=\underset{\tau\sim\pi_\theta}{\mathbb{E}} \left[ \sum_{t=0}^T \nabla_\theta \log\pi_\theta(a_t|s_t) \cdot R(\tau) \right]$

매번 gradient들을 계산할 때마다, log-prob값들에게 $R(\tau)$라는 weight가 곱해집니다. 가만히 곱씹어보면, $\tau$는 trajectory 전체 구간을 의미하는데, 이게 의미가 있을까요? 우리는 agent가 저지른 action을 해당 action로 말미암은 결과에 의거하여 강화합니다. 즉, 해당 action을 저지른 이후에 획득한 reward만이 의미가 있습니다. 때문에 이번 소단원의 제목이 Don't Let the Past Distract You 인 것입니다. 이러한 직관적인 해석은 수학적으로도 증명이 완료되었으며, 위 식을 아래와 같이 변형할 수 있습니다.

$\nabla_\theta J(\pi_\theta)=\underset{\tau\sim\pi_\theta}{\mathbb{E}} \left[ \sum_{t=0}^T \nabla_\theta \log\pi_\theta(a_t|s_t) \sum_{t'=t}^T R(s_{t'},a_{t'},s_{t'+1}) \right]$

이처럼 action이 취해진 이후의 reward로만 action을 강화시키는 방법을 reward-to-go policy gradient라 합니다. 위 식에서 $R(\tau)$ 대신 대입된 $\sum_{t'=t}^T R(s_{t'},a_{t'},s_{t'+1})$를 reward-to-go라 지칭하기 때문입니다.

Implementing Reward-to-Go Policy Gradient

Reward-to-go policy gradient를 PyTorch를 이용해 구현해볼까요? 이번에도 Spinning Up Github에서 코드를 차근차근 살펴보겠습니다. 지난 포스트와 거의 바뀐 것이 없는 것을 알 수 있습니다. 달라진 점은 아래와 같습니다.

# 2_rtg_pg.py, line 17:

def reward_to_go(rews):
    n = len(rews) # 해당 trajectory의 전체 길이를 의미함.
    rtgs = np.zeros_like(rews)
    for i in reversed(range(n)):
        rtgs[i] = rews[i] + (rtgs[i+1] if i+1 < n else 0)
		'''
		rtgs[n-1] = rews[n-1] + 0
		rtgs[n-2] = rews[n-2] + rews[n-1]
		rtgs[n-3] = rews[n-3] + rews[n-2] + rews[n-1]
		...
		rtgs[0] = rews[0] + rews[1] + ... + rews[n-2] + rews[n-1]
		'''
    return rtgs

reward_to_go라는 함수를 새로 추가하고, 아래와 같이 수정하면 끝입니다.

# 1_simple_pg.py, line 91:

		            # the weight for each logprob(a|s) is R(tau)
		            batch_weights += [ep_ret] * ep_len


# 2_rtg_pg.py, line 98:

								# the weight for each logprob(a_t|s_t) is reward-to-go from t
                batch_weights += list(reward_to_go(ep_rews))

Baselines in Policy Gradients

위에서 배운 EGLP lemma의 식을 강화학습의 관점에서 변형해봅시다.

$\underset{x\sim P_\theta}{\mathbb{E}}[\nabla_\theta\log P_\theta(x)]=0.$

였던 식에 $x$대신 $a_t$를, $P_\theta$대신 $\pi_\theta$를 대입하면 아래와 같습니다.

$\underset{a_t\sim \pi_\theta}{\mathbb{E}}[\nabla_\theta\log \pi_\theta(a_t|s_t)]=0.$

여기에 $s_t$에 의존하는 임의의 함수 $b(s_t)$를 곱해줘도 여전히 방정식은 성립하겠죠.

$\underset{a_t\sim \pi_\theta}{\mathbb{E}}[\nabla_\theta\log \pi_\theta(a_t|s_t)\cdot b(s_t)]=0.$

이 값은 0이기에, 위에서 구한 policy gradient식의 양변에 더하거나 빼도 방정식을 유지할 수 있을 것입니다.

$\nabla_\theta J(\pi_\theta)=\underset{\tau\sim\pi_\theta}{\mathbb{E}} \left[ \sum_{t=0}^T \nabla_\theta \log\pi_\theta(a_t|s_t) \left( \sum_{t'=t}^T R(s_{t'},a_{t'},s_{t'+1}) -b(s_t) \right) \right]$

이 함수 $b$를 우리는 baseline이라 일컫습니다. Baseline에 다양한 종류의 함수를 대입할 수 있지만, on-policy value function $V^\pi(s_t)$를 대입하는 것이 가장 보편적입니다. 그렇다면 우리는 왜 굳이 더하나마나한 값을 $\nabla_\theta J(\pi_\theta)$에 추가하는걸까요? Baseline은 $\nabla_\theta J(\pi_\theta)$의 sample estimate의 분산을 줄이는데 탁월한 효과가 있기 때문입니다. 분산을 줄임으로서 policy gradient의 update과정이 빨라지고 안정적이게 됩니다.

실제 baseline을 구현할 때, $V^\pi(s_t)$의 정확한 값을 계산하는 것은 불가능하기에, 우리는 neural network로 근사된 근사값 $V_\phi(s_t)$를 구합니다. 이는 policy $\pi$가 update될 때마다 동시에 update되는 값으로, 가장 최신 policy에 의거하여 update됩니다. $V_\phi$의 학습을 위해 VPG, TRPO, PPO, A2C 등 다양한 알고리즘들이 이용되며, 아래의 mean-squared-error를 최소화시키는 방향으로 학습이 이루어집니다. 여기서 $k$는 epoch를 의미합니다.

$\phi_k=\argmin_\phi \underset{s_t,\hat{R}_t \sim \pi_k}{\mathbb{E}} \left[ \left( V_\phi(s_t)-\hat{R}_t \right) ^2\right]$

Other Forms of the Policy Gradient

지금까지 배운 내용을 요약, 정리해볼까요? 아래의 기본 form에서 weight $\Phi_t$를 다르게하여 reward-to-go policy gradient와 baseline을 구현할 수 있습니다. 세 방법 모두 다른 분산을 가짐에도 불구하고 동일한 기대값을 가집니다(EGLP lemma).

General form of the policy gradient: $\nabla_\theta J(\pi_\theta)=\underset{\tau\sim\pi_\theta}{\mathbb{E}} \left[ \sum_{t=0}^T \nabla_\theta \log\pi_\theta(a_t|s_t) \cdot \Phi_t \right]$

where weight $\Phi_t$ could be any of $\Phi_t = R(\tau)$, or $\Phi_t = \sum_{t'=t}^T R(s_{t'},a_{t'},s_{t'+1})$, or $\Phi_t = \sum_{t'=t}^T R(s_{t'},a_{t'},s_{t'+1}) -b(s_t)$.

그 외에도 알아놓으면 좋을 두 가지의 format을 소개하겠습니다.

On-policy action-value function: $\Phi_t = Q^{\pi_\theta}(s_t,a_t)$

Advantage function: $\Phi_t = A^{\pi_\theta}(s_t,a_t) = Q^{\pi_\theta}(s_t,a_t) - V^{\pi_\theta}(s_t)$

On-policy action-value function를 사용할 경우의 타당성 증명은 link에서 볼 수 있으며, advantage function을 쓰는 경우엔 value function $V^{\pi_\theta}(s_t)$를 baseline으로 쓰는 것이기에 증명은 동일합니다.

Spinning Up에선 더욱 다양한 weight의 적용과 응용방안을 공부하기 위해 Generalized Advantage Estimation (GAE)를 살펴볼 것을 강력하게 추천합니다.

Reference

[1] R. S. Sutton, A. G. Barto, et al., Introduction to reinforcement learning, vol. 135. MIT press Cambridge, 1998.

[2] OpenAI. Spinning Up, [Online]. Available: https://spinningup.openai.com/