Policy Gradient를 처음부터 이해하기

TL;DR

• Policy gradient의 목표는 기대 return $J(\theta )$를 최대화하는 것이다.
• log-derivative trick을 쓰면 gradient를 trajectory 분포 위의 기댓값 형태로 바꿀 수 있다.
• baseline과 advantage estimator를 쓰면 bias 없이 variance를 줄이며 더 안정적으로 학습할 수 있다.

---

이 글의 목표

강화학습이나 RLHF(Reinforcement Learning from Human Feedback)를 보다 보면 아래 식을 자주 만나게 된다.

$${\nabla }_{\theta }J(\theta )=\mathbb{E}\left[\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)A^{\pi }(s_t,a_t)\right]$$

RLHF, PPO, DPO, GRPO는 이름도 다르고 수식도 달라 보이지만, 결국 공통으로 묻는 질문은 같다.

어떻게 하면 좋은 응답의 확률을 높일 수 있을까?

이 글은 그 질문에 대한 가장 기본적인 답인 policy gradient를 처음부터 따라가며 이해하는 노트다. 목표는 policy gradient 식을 외우는 것이 아니라, 왜 그런 모양이 되는지를 이해하는 것이다. 그 과정속에서 수학적 엄밀성이 떨어질 수도 있겠지만, 이해를 우선으로 두도록 하자.

이 글에서는 어떻게 위 식으로 바뀌는지 하나씩 따라가본다.

---

왜 Policy Gradient부터 봐야 하나

RLHF, PPO, DPO, GRPO는 서로 다른 방법처럼 보이지만 큰 흐름에서는 모두 “좋은 응답을 더 자주 생성하게 만드는 방법”이라는 공통점을 가진다.

• RLHF는 reward model이 준 점수를 바탕으로 policy를 업데이트한다.
• PPO는 policy gradient를 더 안정적으로 업데이트하는 방법이다.
• DPO는 RLHF의 일부 단계를 직접 최적화 문제로 다시 쓴 것이다.
• GRPO는 다시 policy optimization 쪽으로 돌아오되, 상대 비교를 더 적극적으로 활용한다.

즉, policy gradient를 이해하면 이후의 RLHF 계열 알고리즘이 훨씬 덜 낯설어진다.

---

기호 정의 및 연습

수식 설명에 앞서, 사용할 기호는 다음과 같이 정의한다.

• $\theta$ (세타): 모델 파라미터
• $\pi$ (파이): 정책(policy)
• $\tau$ (타우): trajectory(하나의 전체 경험)
• $\mathbb{E}$ (기댓값): expectation, 확률을 고려한 평균
• $\sum$ (시그마): 합
• $\prod$ (곱): 여러 항의 곱
• ${\nabla }_{\theta }$ (나블라 세타): 세타에 대한 gradient
• $\log$ (로그): 보통 자연로그
• $\mid$ (given, 바): “~가 주어졌을 때”
• $\sim$ (“~를 따른다”): 샘플링 또는 분포를 따른다
• $s$ (state): 상태
• $a$ (action): 행동
• $r$ (reward): 보상

그렇다면 이제 다음 수식은 어떻게 읽히는가?

$${\pi }_{\theta }(a\mid s)$$

상태 $s$가 주어졌을 때, 정책 ${\pi }_{\theta }$가 행동 $a$를 선택할 확률(또는 확률밀도)

---

평균과 기댓값

평균과 기댓값은 너무 익숙한 개념이지만, 한번 짚고 넘어가자. 한 줄로 요약하면 다음과 같다.

평균은 이미 관측한 데이터의 요약이고,
기댓값은 확률분포가 만들어내는 장기적인 평균이다.

평균: 이미 관측한 데이터의 중심

숫자 세 개 $2,4,10$ 이 있다고 하자. 이들의 평균은 $\frac{2+4+10}{3}$ 이다. 이 값은 이미 손에 들어온 데이터를 똑같이 한 번씩 반영해서 계산한 것이다. 즉, 평균은 관측된 표본(sample)의 중심을 요약한다.

이를 일반화하면, $n$개의 관측값 $x_1,\dots ,x_n$의 평균은 다음으로 정의된다.

$$\bar{x}=\frac{1}{n}\sum _{i=1}^n{x}_i$$

기댓값: 확률을 고려한 평균

이번에는 어떤 값이 나올 확률이 서로 다르다고 하자.

• 1이 나올 확률: 0.8
• 5가 나올 확률: 0.2

이때 확률까지 고려한 평균은 $0.8\cdot 1+0.2\cdot 5=1.8$ 이다. 이 값이 바로 기댓값(expected value) 이다. 즉, 기댓값은 단순한 평균이 아니라 확률을 가중치로 둔 평균이다.

이를 일반화하면, 이산형 확률변수 $X$의 기댓값은 다음으로 정의된다.

$$\mathbb{E}[X]=\sum _{x}xP(X=x)$$

연속형 확률변수라면 합 대신 적분을 사용하여

$$\mathbb{E}[X]={\int }_{-\infty }^{\infty }xf(x)dx$$

로 쓴다.

그렇다면 다음 수식은 어떻게 읽히는가?

$$\mathbb{E}[X]=\sum _{x}xP(X=x)$$

가능한 모든 값 $x$에 대해, 그 값이 나올 확률 $P(X=x)$와 값 자체 $x$를 곱해서 모두 더한 것

---

강화학습에서의 목표: 기대 Return 최대화

강화학습에서는 한 번의 실행 결과를 trajectory라고 부르며, 이는 다음과 같이 $s_t$(시점 $t$의 상태), $a_t$(시점 $t$의 행동)로 구성된다.

$$\tau =(s_0,a_0,s_1,a_1,\dots ,s_T,a_T,s_{T+1})$$

그리고 이 trajectory $\tau$ 전체의 점수를 return 이라고 한다.

$$R(\tau )=\sum _{t=0}^T{\gamma }^t{r}_t$$

• $R(\tau )$: trajectory $\tau$의 전체 return
• $r_t$: 시점 $t$의 보상(reward)
• $\gamma$ : discount factor

이때, 감마 $\gamma$는 미래 보상을 얼마나 중요하게 볼지 정한다.

• $\gamma =1$: 미래 보상도 동일하게 중요
• $\gamma <1$: 먼 미래의 보상은 조금 덜 중요

이제 정책 ${\pi }_{\theta }$의 목표는 기대 return이 큰 trajectory를 더 자주 생성하도록 하는 것이다.

$$J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[R(\tau )]$$

여기서 $\tau \sim {\pi }_{\theta }$ 는 현재 정책 아래에서 trajectory가 샘플링된다는 뜻의 약식 표기다.

조금 더 엄밀하게 쓰면, 정책 ${\pi }_{\theta }$ 와 환경의 상태 전이 확률이 함께 결정하는 trajectory 분포를 $p_{\theta }(\tau )$ 라고 둘 수 있고, 이때

$$J(\theta )={\mathbb{E}}_{\tau \sim p_{\theta }}[R(\tau )]$$

와 같이 쓸 수 있다. 즉, 이후에 등장할 $p_{\theta }(\tau )$ 는 정책 ${\pi }_{\theta }$ 와 환경의 상태 전이 확률이 함께 결정하는 trajectory distribution이다.

즉, policy optimization 은 기대 return $J(\theta )$를 최대화하는 문제와 같다. 실제 학습에서는 이 기댓값을 여러 rollout의 표본평균으로 추정한다.

---

그래도 와닿지 않는다면?

기댓값이 아직 추상적으로 느껴진다면, 정의를 그대로 풀어 써보면 된다.

$$J(\theta )={\mathbb{E}}_{\tau \sim p_{\theta }}[R(\tau )]$$

이 식은 이산형(discrete)인 경우 다음과 같이 쓸 수 있다.

$$J(\theta )=\sum _{\tau }{p}_{\theta }(\tau )R(\tau )$$

• $p_{\theta }(\tau )$: 현재 정책과 환경 전이 아래에서 trajectory $\tau$가 생성될 확률
• $R(\tau )$: 그 trajectory의 return

그렇다면 이제 이 식은 어떻게 읽히는가?

가능한 모든 trajectory에 대해, “그 trajectory가 나올 확률 × 그 trajectory의 return”을 전부 더한 것이 정책의 기대 return이다.

---

우리의 목표는 최대화

이제 다시 한번 목표를 상기시켜보자. 우리는 $J(\theta )$를 최대화하고 싶다. 주어진 목표함수는

$$J(\theta )=\sum _{\tau }{p}_{\theta }(\tau )R(\tau )$$

이므로, $\theta$에 대해 미분하면

$$\begin{array}{rl}{\nabla }_{\theta }J(\theta )& ={\nabla }_{\theta }\sum _{\tau }{p}_{\theta }(\tau )R(\tau )\\ & =\sum _{\tau }{\nabla }_{\theta }(p_{\theta }(\tau )R(\tau ))\\ & =\sum _{\tau }R(\tau ){\nabla }_{\theta }{p}_{\theta }(\tau )\end{array}$$

를 얻는다. 여기서 $R(\tau )$는 “주어진 trajectory의 return”으로 보고, $\theta$에 직접 의존하지 않는다고 둔다. 따라서 미분은 $p_{\theta }(\tau )$에만 걸리고, 다음과 같이 정리 가능하다.

$${\nabla }_{\theta }J(\theta )=\sum _{\tau }R(\tau ){\nabla }_{\theta }{p}_{\theta }(\tau )$$

---

Log-Derivative Trick

이제 로그미분 트릭을 사용해 식을 더 다루기 쉬운 형태로 바꿔보자. 먼저 다음 항등식을 떠올리자.

$${\nabla }_{\theta }\log f(\theta )=\frac{{\nabla }_{\theta }f(\theta )}{f(\theta )}$$

이를 $f(\theta )=p_{\theta }(\tau )$에 적용하면

$${\nabla }_{\theta }\log p_{\theta }(\tau )=\frac{{\nabla }_{\theta }{p}_{\theta }(\tau )}{p_{\theta }(\tau )}$$

이고, 양변에 $p_{\theta }(\tau )$를 곱하면

$${\nabla }_{\theta }{p}_{\theta }(\tau )=p_{\theta }(\tau ){\nabla }_{\theta }\log p_{\theta }(\tau )$$

를 얻는다. 이제 이 식을 앞에서 얻은

$${\nabla }_{\theta }J(\theta )=\sum _{\tau }R(\tau ){\nabla }_{\theta }{p}_{\theta }(\tau )$$

에 대입하면,

$$\begin{array}{rl}{\nabla }_{\theta }J(\theta )& =\sum _{\tau }R(\tau ){\nabla }_{\theta }{p}_{\theta }(\tau )\\ & =\sum _{\tau }R(\tau )p_{\theta }(\tau ){\nabla }_{\theta }\log p_{\theta }(\tau )\\ & =\sum _{\tau }{p}_{\theta }(\tau ){\nabla }_{\theta }\log p_{\theta }(\tau )R(\tau )\end{array}$$

를 얻는다.

이제 식이 다시 기댓값의 형태로 보인다. 왜냐하면 일반적으로 어떤 함수 $f(\tau )$에 대해

$${\mathbb{E}}_{\tau \sim p_{\theta }}[f(\tau )]=\sum _{\tau }{p}_{\theta }(\tau )f(\tau )$$

이기 때문이다.

따라서 위 식은 다음처럼 다시 쓸 수 있다.

$${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim p_{\theta }}\left[{\nabla }_{\theta }\log p_{\theta }(\tau )R(\tau )\right]$$

이 식은 어떻게 읽히는가?

policy gradient는 trajectory별 return과 trajectory log-probability gradient의 곱을, 현재 정책이 만드는 trajectory 분포 위에서 평균낸 것이다.

즉, 평균적으로 보았을 때 return이 큰 trajectory일수록 그 trajectory를 만들어낸 선택들의 확률을 높이는 방향으로 업데이트 가 이루어진다.

이 식이 policy gradient 유도의 첫 번째 핵심식이다.

---

Trajectory 확률은 어떻게 생겼을까?

이제 $\log p_{\theta }(\tau )$를 더 구체적으로 보자. trajectory의 확률은 보통 다음과 같이 쓸 수 있다.

$$p_{\theta }(\tau )={\rho }_0(s_0)\prod _{t=0}^T{\pi }_{\theta }(a_t\mid s_t)P(s_{t+1}\mid s_t,a_t)$$

• ${\rho }_0(s_0)$: 초기 상태 $s_0$가 나올 확률
• ${\pi }_{\theta }(a_t\mid s_t)$: 정책이 상태 $s_t$에서 행동 $a_t$를 선택할 확률
• $P(s_{t+1}\mid s_t,a_t)$: 환경이 다음 상태 $s_{t+1}$로 전이될 확률

즉, 하나의 trajectory가 나올 확률은

1. 처음 상태가 어떻게 시작되는지
2. 각 시점에서 정책이 어떤 행동을 선택하는지
3. 그 행동 이후 환경이 어떻게 다음 상태로 전이되는지 를 모두 곱한 것으로 볼 수 있다.

---

로그: 곱을 합으로 바꾸기

이제 이 식을 미분하기 쉬운 형태로 바꾸기 위해 로그를 취해보자. 로그의 기본 성질은 다음과 같다.

$$\log (ab)=\log a+\log b$$

이를 이용하면 trajectory 확률의 로그는 다음과 같이 쓸 수 있다.

$$\log p_{\theta }(\tau )=\log {\rho }_0(s_0)+\sum _{t=0}^T\log {\pi }_{\theta }(a_t\mid s_t)+\sum _{t=0}^T\log P(s_{t+1}\mid s_t,a_t)$$

이제 $\theta$에 대해 미분하면

$${\nabla }_{\theta }\log p_{\theta }(\tau )={\nabla }_{\theta }\log {\rho }_0(s_0)+\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)+\sum _{t=0}^T{\nabla }_{\theta }\log P(s_{t+1}\mid s_t,a_t)$$

를 얻는다. 표준적인 model-free setting에서는 (1) 초기 상태 분포 ${\rho }_0$는 $\theta$와 무관하고, (2) 환경 전이 $P$도 $\theta$와 무관하다고 둔다.

따라서 앞의 두 항은 0이 되고, 결국 남는 것은 다음과 같다.

$${\nabla }_{\theta }\log p_{\theta }(\tau )=\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)$$

이를 앞에서 얻은 식에 다시 대입하면, 다음과 같다.

$${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim p_{\theta }}\left[\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)R(\tau )\right]$$

이것이 REINFORCE의 가장 기본적인 형태다.

이 식은 어떻게 읽히는가?

어떤 trajectory의 결과가 좋았다면, 그 trajectory 안에서 선택된 행동들의 확률을 높이는 방향으로 업데이트한다.

즉,

• 좋은 결과를 낸 행동은 더 자주 하게 만들고
• 나쁜 결과를 낸 행동은 덜 자주 하게 만든다.

---

그런데 왜 전체 Return을 다 곱할까?

위 식은 맞지만, 분산이 큰 형태다. 과거 보상은 현재 행동 $a_t$가 바꿀 수 없는 값 이다. 그런데도 전체 return $R(\tau )$를 그대로 곱하면, 현재 행동과 무관한 변동까지 gradient estimator에 함께 들어간다. 그 결과 기댓값은 같더라도 샘플마다 업데이트가 더 크게 흔들린다. 따라서 시점 $t$의 행동은 그 이후에 얻는 보상으로 평가하는 편이 더 자연스럽다.

$$G_t=\sum _{t^{\prime }=t}^T{\gamma }^{t^{\prime }-t}{r}_{t^{\prime }}$$

이처럼 시점 $t$ 이후의 보상만 합친 값을 reward-to-go라고 부른다. 그러면 policy gradient는 다음과 같이 쓸 수 있다.

$${\nabla }_{\theta }J(\theta )=\mathbb{E}\left[\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)G_t\right]$$

이 식은 어떻게 읽히는가?

시점 $t$의 행동은 그 이후의 결과로 평가하자.

---

Baseline은 왜 빼도 될까?

이제 policy gradient는 흔히 다음처럼 쓰기도 한다.

$${\nabla }_{\theta }J(\theta )=\mathbb{E}\left[\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)(G_t-b(s_t))\right]$$

여기서 $b(s_t)$는 baseline이다. 상태 $s_t$에서의 기준점이라고 보면 된다.

baseline을 빼는 이유는 분산을 줄이기 위해서다. 직관적으로는 “그 행동의 절대 점수”보다 “평균보다 얼마나 더 좋았는가”를 보는 쪽이 업데이트를 더 안정적으로 만든다.

그런데 baseline을 빼면 gradient 자체가 바뀌는 것처럼 보인다. 핵심은, action에 의존하지 않는 baseline 항의 기댓값은 0 이라는 점이다.

먼저 아래 식을 보자.

$$\sum _a{\pi }_{\theta }(a\mid s){\nabla }_{\theta }\log {\pi }_{\theta }(a\mid s)=0$$

왜냐하면

$${\nabla }_{\theta }\log {\pi }_{\theta }(a\mid s)=\frac{{\nabla }_{\theta }{\pi }_{\theta }(a\mid s)}{{\pi }_{\theta }(a\mid s)}$$

이므로,

$${\pi }_{\theta }(a\mid s){\nabla }_{\theta }\log {\pi }_{\theta }(a\mid s)={\nabla }_{\theta }{\pi }_{\theta }(a\mid s)$$

가 되고, 따라서

$$\sum _a{\pi }_{\theta }(a\mid s){\nabla }_{\theta }\log {\pi }_{\theta }(a\mid s)=\sum _a{\nabla }_{\theta }{\pi }_{\theta }(a\mid s)={\nabla }_{\theta }\sum _a{\pi }_{\theta }(a\mid s)={\nabla }_{\theta }1=0$$

가 된다.

그러므로 baseline 항의 조건부 기댓값은

$${\mathbb{E}}_{a_t\sim {\pi }_{\theta }(\cdot \mid s_t)}\left[{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)b(s_t)\right]=b(s_t)\sum _a{\pi }_{\theta }(a\mid s_t){\nabla }_{\theta }\log {\pi }_{\theta }(a\mid s_t)=0$$

이 된다.

즉, action에 의존하지 않는 baseline은 평균적으로 사라진다. 따라서 baseline을 빼도 bias는 생기지 않고, 대신 variance는 줄일 수 있다.

---

Baseline을 Value로 두면 Advantage의 추정치가 된다

baseline의 가장 대표적인 선택은 상태가치 함수다.

$$b(s_t)=V^{\pi }(s_t)$$

여기서 상태가치 함수와 행동가치 함수는 각각 다음처럼 정의된다.

$$V^{\pi }(s_t)={\mathbb{E}}_{\pi }[G_t\mid s_t],Q^{\pi }(s_t,a_t)={\mathbb{E}}_{\pi }[G_t\mid s_t,a_t]$$

즉,

• $V^{\pi }(s_t)$는 상태 $s_t$에 도착했을 때, 이후 정책 $\pi$를 따라가면 평균적으로 얼마나 좋은가
• $Q^{\pi }(s_t,a_t)$는 상태 $s_t$에서 행동 $a_t$를 한 뒤, 이후 정책 $\pi$를 따라가면 평균적으로 얼마나 좋은가

를 나타낸다.

이 둘의 차이가 바로 advantage다.

$$A^{\pi }(s_t,a_t)=Q^{\pi }(s_t,a_t)-V^{\pi }(s_t)$$

이 식은 다음처럼 읽을 수 있다.

상태 $s_t$에서 행동 $a_t$가, 그 상태의 평균적인 선택보다 얼마나 더 좋았는가

그런데 실제 rollout에서 우리가 직접 관측하는 것은 $Q^{\pi }(s_t,a_t)$ 자체가 아니라, 그 시점 이후 실제로 얻은 return인 $G_t$다.

$$G_t=\sum _{t^{\prime }=t}^T{\gamma }^{t^{\prime }-t}{r}_{t^{\prime }}$$

중요한 점은 $G_t$가 $Q^{\pi }(s_t,a_t)$와 항상 같은 값은 아니라는 것이다. $G_t$는 하나의 rollout에서 나온 샘플값이고, $Q^{\pi }(s_t,a_t)$는 그런 샘플들의 조건부 기댓값이다. 이를 수식으로 쓰면

$${\mathbb{E}}_{\pi }[G_t\mid s_t,a_t]=Q^{\pi }(s_t,a_t)$$

이다.

따라서 baseline을 상태가치 함수로 두면

$$G_t-V^{\pi }(s_t)$$

는 advantage 그 자체라기보다는, advantage의 샘플 추정치로 해석하는 것이 정확하다. 실제로 그 조건부 기댓값은

$$\begin{array}{rlrl}\mathbb{E}[G_t-V^{\pi }(s_t)\mid s_t,a_t]& =\mathbb{E}[G_t\mid s_t,a_t]-V^{\pi }(s_t)& & \text{(조건부 기댓값의 선형성)}\\ & =Q^{\pi }(s_t,a_t)-V^{\pi }(s_t)& & \text{(}{Q}^{\pi }\text{의 정의)}\\ & =A^{\pi }(s_t,a_t)& & \text{(}{A}^{\pi }\text{의 정의)}\end{array}$$

가 된다.

즉, $G_t-V^{\pi }(s_t)$는 샘플마다 흔들릴 수는 있지만, 평균적으로는 정확히 advantage를 맞춘다. 그래서 실제 학습에서는 이 값을 사용해 “이 행동이 평균보다 얼마나 더 좋았는가”를 샘플 단위에서 추정한다고 볼 수 있다.

이 관점에서 policy gradient는 다음처럼 쓸 수 있다.

$${\nabla }_{\theta }J(\theta )=\mathbb{E}\left[\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t){\hat{A}}_t\right]$$

여기서 ${\hat{A}}_t$는 advantage의 추정치이며, 가장 단순한 형태로는 다음처럼 둘 수 있다.

$${\hat{A}}_t=G_t-V^{\pi }(s_t)$$

이 식은 직관적으로 어떻게 읽히는가?

평균보다 더 좋았던 행동이면 그 확률을 높이고, 평균보다 덜 좋았던 행동이면 그 확률을 낮춘다.

실제에서는 보통 $V^{\pi }$를 정확히 알 수 없기 때문에, 이를 근사하는 value function $V_{\varphi }(s_t)$를 학습해서 사용한다. 또한 더 안정적인 학습을 위해 TD error나 GAE 같은 더 정교한 advantage estimator를 쓰기도 한다. 정리하면, baseline을 상태가치 함수로 두면 policy gradient의 기댓값은 유지하면서도 분산을 줄일 수 있고, 이때 $G_t-V^{\pi }(s_t)$는 advantage의 가장 기본적인 샘플 추정치로 이해할 수 있다.

---

LLM과 RLHF 문맥으로 다시 읽기

이제 이 식을 언어모델 쪽으로 다시 해석해 보자. 전통적인 강화학습에서는 상태 $s$, 행동 $a$ 를 다루지만, LLM에서는 이를 $x$ :prompt, $y$:completion 로 생각할 수 있다.

즉, ${\pi }_{\theta }(y\mid x)$ 는 “프롬프트 $x$가 주어졌을 때 응답 $y$를 생성할 확률”이다.

엄밀하게 보면, LLM을 강화학습으로 해석할 때는 completion 전체를 한 번에 내는 것이 아니라 토큰 단위의 순차적 의사결정 과정 으로 보는 쪽이 더 표준적이다.

즉, 상태(state)는 “지금까지 생성한 prefix”이고, 행동(action)은 “다음 토큰을 무엇으로 생성할 것인가”에 해당한다.

다만 RLHF를 처음 이해하는 단계에서는 이를 prompt $x$와 completion $y$ 수준으로 추상화해도 큰 흐름을 이해하는 데는 충분하다.

그러면 policy gradient도 같은 구조로 읽을 수 있다.

좋은 응답이면 그 응답의 확률을 높이고,
나쁜 응답이면 그 응답의 확률을 낮춘다.

RLHF는 결국 이 아이디어를 reward model, KL penalty, PPO 같은 장치와 결합한 것이다. 즉, RLHF는 policy gradient를 언어모델 post-training에 맞게 이식한 것 이라고 보는 편이 이해하기 쉽다.

---

다시 정리

policy gradient의 전개는 다음 흐름으로 이해할 수 있다.

1. 정책의 목표는 기대 return을 최대화하는 것이다.

$$J(\theta )={\mathbb{E}}_{\tau \sim p_{\theta }}[R(\tau )]$$

즉, 정책이 만들어내는 trajectory들의 기대 return이 커지도록 $\theta$를 업데이트하고 싶다.

2. 기댓값은 확률가중합으로 풀어쓸 수 있다.

$$J(\theta )=\sum _{\tau }{p}_{\theta }(\tau )R(\tau )$$

즉, 가능한 모든 trajectory에 대해 “그 trajectory가 나올 확률 $\times$ 그 trajectory의 return”을 더한 값이 정책의 기대 return다.

3. 이제 이를 $\theta$에 대해 미분한다.

$${\nabla }_{\theta }J(\theta )=\sum _{\tau }R(\tau ){\nabla }_{\theta }{p}_{\theta }(\tau )$$

여기서는 표준적인 model-free setting에서 $R(\tau )$가 $\theta$에 직접 의존하지 않는다고 둔다.

4. 로그미분 트릭을 적용하면 gradient를 더 다루기 쉬운 형태로 바꿀 수 있다.

$${\nabla }_{\theta }{p}_{\theta }(\tau )=p_{\theta }(\tau ){\nabla }_{\theta }\log p_{\theta }(\tau )$$

따라서 다음식을 얻을 수 있다.

$${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim p_{\theta }}\left[{\nabla }_{\theta }\log p_{\theta }(\tau )R(\tau )\right]$$

5. trajectory 확률의 로그를 전개하면, 결국 정책 항만 남는다.

$${\nabla }_{\theta }\log p_{\theta }(\tau )=\sum _t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)$$

따라서

$${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim p_{\theta }}\left[\sum _t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)R(\tau )\right]$$

가 된다. 이것이 REINFORCE의 가장 기본적인 형태다.

6. 하지만 시점 $t$의 행동을 평가할 때 전체 return을 그대로 쓰면 불필요한 분산이 커질 수 있다.

그래서 시점 $t$ 이후의 return만 사용하는 reward-to-go를 도입한다.

$$G_t=\sum _{t^{\prime }=t}^T{\gamma }^{t^{\prime }-t}{r}_{t^{\prime }}$$

그러면 식은 다음과 같이 쓸 수 있다.

$${\nabla }_{\theta }J(\theta )=\mathbb{E}\left[\sum _t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)G_t\right]$$

7. 여기에 baseline을 빼면 기댓값은 유지하면서 분산을 더 줄일 수 있다.

$${\nabla }_{\theta }J(\theta )=\mathbb{E}\left[\sum _t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)(G_t-b(s_t))\right]$$

특히 baseline으로 상태가치 함수를 사용하면

$$b(s_t)=V^{\pi }(s_t)$$

이고, 이때

$$G_t-V^{\pi }(s_t)$$

는 advantage의 샘플 기반 추정치로 볼 수 있다. 이상적인 advantage는

$$A^{\pi }(s_t,a_t)=Q^{\pi }(s_t,a_t)-V^{\pi }(s_t)$$

이며, 실제 학습에서는 보통

$${\hat{A}}_t\approx G_t-V^{\pi }(s_t)$$

또는 그 변형을 사용한다. 따라서 policy gradient는 최종적으로

$${\nabla }_{\theta }J(\theta )=\mathbb{E}\left[\sum _t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t){\hat{A}}_t\right]$$

처럼 이해할 수 있다.

즉, policy gradient의 핵심은 다음 한 문장으로 요약된다.

평균보다 더 좋았던 행동이면 그 확률을 높이고, 평균보다 덜 좋았던 행동이면 그 확률을 낮춘다.

---

다음 글

다음 글에서는 REINFORCE와 baseline 를 살펴본다.

policy gradient의 기댓값이 실제 rollout 샘플을 이용한 업데이트 규칙으로 어떻게 바뀌는지, 그리고 왜 이 방식이 high variance 문제를 갖는지 정리할 예정이다.

---

시리즈 메모

• 01. Policy Gradient를 처음부터 이해하기
• 02. REINFORCE와 baseline, 그리고 RLOO
• 03. PPO는 왜 필요한가
• 04. SFT 다음에 왜 RLHF가 필요했는가
• 05. RLHF objective: reward model + KL penalty + policy update
• 06. RLHF에서 DPO로 가는 수식 전개
• 07. PPO에서 GRPO로

📎 Reference

References

• RLHF Book — Reinforcement Learning (i.e. Policy Gradient Algorithms)
• Ronald J. Williams (1992) — Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning
• Richard S. Sutton, David A. McAllester, Satinder P. Singh, Yishay Mansour (1999) — Policy Gradient Methods for Reinforcement Learning with Function Approximation
• OpenAI Spinning Up — Part 3: Intro to Policy Optimization
• OpenAI Spinning Up — Vanilla Policy Gradient