[RL] CS285 - Reducing Variance

(해당 포스트는 UC Berkeley 에서 진행된 CS285: Deep Reinforcement Learning, Decision Making and Control를 요약한 내용이며, 그림들은 강의 장표에서 발췌한 내용입니다.)

CS 285

Reducing variance

이전 포스트에서 소개했다시피, policy gradient 알고리즘은 high variance로 인해서 practical problem에 적용하기 어려웠다. 그러면 이 variance를 줄일 수 있는 방법이 있을까? 다시 policy gradient의 핵심 부분을 살펴본다.

$$ \nabla_{\theta}J(\theta) \approx \frac{1}{N} \sum_{i=1}^N \big( \sum_{t=1}^T \nabla_{\theta} \log \pi_{\theta} (a_{i, t} \vert s_{i, t}) \big) \big(\sum_{t=1}^T r(s_{i, t}, a_{i, t})\big) $$

여기에서 causality라는 개념을 적용했다. (나는 사실 causality라는 개념을 signal processing 공부할때 처음 배웠던 내용인데, 찾아보면 굳이 주제에 한정짓지 않은 보편적인 개념인것 같다.) 이 causality는 시간적으로 현재에 발생한 일이 과거에 영향을 끼치지 않는다는 개념이고, RL에 한정지어서 설명하자면 time \(t'\)에서 수행되었던 policy \(\pi\)는 \(t'\)시점 이전의 time \(t\), 즉 \(t < t'\)때의 reward에 영향을 미치지 않는다는 것을 나타낸다. 어떻게 보면 causality 특성이 markov property와 비슷하다고 느낄 수 있지만, markov property는 state 관점에서 미래에 발생한 state가 과거의 state와 독립적이라는 것을 설명하는 요소이기 때문에 완전히 다르다. 이전의 partial observability에서 나왔던 것처럼 어떤 process를 따르냐에 따라서 markov property가 성립할 수도, 혹은 성립하지 않을 수도 있지만, causality는 항상 참인 명제이다. (강의에 나온것처럼 시간을 여행할 수 있다던가 과거의 action에 영향을 줄 수 있는 case라면 설명이 다르겠지만...)

그래서 위의 수식을 다시 괄호를 정리해서 보면, 아래와 같이 되는데

$$ \nabla_{\theta}J(\theta) \approx \frac{1}{N} \sum_{i=1}^N  \sum_{t=1}^T \nabla_{\theta} \log \pi_{\theta} (a_{i, t} \vert s_{i, t}) \big(\sum_{t'=1}^T r(s_{i, t'}, a_{i, t'})\big) $$

이를 causality 측면에서 보면 timestep \(t\)가 policy와 reward에 모두 영향을 주기 때문에 현재의 policy가 과거의 reward에 영향을 주는 case가 발생하게 된다. 그래서 이 causality를 적용해서 과거의 reward가 영향을 받지 않게끔 reward term을 수정하면 아래와 같이 된다.

$$ \nabla_{\theta}J(\theta) \approx \frac{1}{N} \sum_{i=1}^N  \sum_{t=1}^T \nabla_{\theta} \log \pi_{\theta} (a_{i, t} \vert s_{i, t}) \big(\sum_{t'=t}^T r(s_{i, t'}, a_{i, t'})\big) $$

물론 이렇게 하면 sample수가 달라졌기 때문에 기존의 estimator에 영향이 생기지만, 그래도 전체적으로 unbiased한 성향은 유지할 수 있다. 그리고 전체적으로 sample 수가 줄어들었기 때문에 variance도 같이 줄어든 효과를 가져온다. 참고로 이렇게 causality로 인해서 변화된 부분인

$$ \sum_{t'=t}^T r(s_{i, t'}, a_{i, t'}) $$

을 reward-to-go라고 하고(아마 내가 번역했던 책에서는 이동 보상이라고 표현했던 것 같다..), 수집한 reward 중 policy가 과거에 얻은 reward를 제외한 나머지 reward의 총합을 나타내고, 수식상에서는 보통 \(\hat{Q}_{i, t}\)로 표현된다. 그래서 이를 policy gradient에 반영한 식은 아래와 같이 된다.

$$ \nabla_{\theta}J(\theta) \approx \frac{1}{N} \sum_{i=1}^N  \sum_{t=1}^T \nabla_{\theta} \log \pi_{\theta} (a_{i, t} \vert s_{i, t}) \hat{Q}_{i, t} $$

Baselines

이밖에도 variance를 낮출 수 있는 방법이 있는데, 그 중 하나가 이번에 소개할 baseline 기법이다. policy gradient 수식을 가져오면,

$$ \nabla_{\theta} J(\theta) \approx \frac{1}{N} \sum_{i=1}^N \log p_{\theta}(\tau) r(\tau) $$

와 같은데, 이 식은 기본적으로 이전 포스트에서 설명했던 것처럼 good trajectory의 log probability를 높이고, bad trajectory의 log probability를 낮추는 것으로 동작한다. 그리고 일반적으로 이 경우는 good trajectory가 positive reward를 가지고, bad trajectory는 negative reward를 가지는 것을 말한다. 그런데 만약 전체 reward 분포가 이렇게 positive / negative가 아닌, 그냥 positive로만 되어 있는 경우는 어떨까? 예를 들어서 good trajectory는 매우 큰 reward를 가지고, bad trajectory는 매우 작은 reward를 가지게끔 되면, 결과적으로 모든 trajectory에 대한 log-probability가 증가하는 방향으로 바뀐다. 이를 해결할 수 있게 어떤 기준점, 즉 baseline을 정하고, reward term에 이 baseline을 반영해서 0을 기준으로 positive/negative로 나눠질 수 있게 바꾸는 것이다.

$$ \nabla_{\theta} J(\theta) \approx \frac{1}{N} \sum_{i=1}^N \log p_{\theta}(\tau) [r(\tau) - b] $$

여기서 \(b = \frac{1}{N}\sum_{i=1}^N r(\tau) \)를 말한다. 이렇게 하면, 평균 보상값보다 좋은 trajectory는 log probability가 커지고 평균 보상값보다 나쁜 trajectory는 낮아지는 효과가 생기면서 결과적으로 variance를 낮추는 효과를 보여주게 된다. 

그림 1. Baseline을 적용하면서 log probability의 변화

그러면 결과적으로 baseline을 빼주면 기존의 policy gradient에서 새로운 term인 \(\mathbb{E}[\nabla_{\theta} \log p_{\theta}(\tau)b]\)를 빼준 셈인데, 이렇게 해도 학습에 영향이 없을까? 이때 이전에 사용했던 log-derivative trick을 다시 활용할 수 있다.

$$ p_{\theta} (\tau) \nabla_{\theta} \log p_{\theta}(\tau) = \nabla_{\theta} p_{\theta}(\tau) $$

결국 동일한 trick을 써서 정리를 해보면 아래와 같이 된다.

$$ \mathbb{E}[\nabla_{\theta} \log p_{\theta}(\tau) b ] = \int p_{\theta} (\tau) \nabla_{\theta} \log p_{\theta} (\tau) b d\tau = \int \nabla_{\theta} p_{\theta} (\tau) b d\tau = b \nabla_{\theta} \int p_{\theta}(\tau) d\tau $$

그런데 마지막 적분은 아마 아는 사람은 알겠지만, 확률에 대한 적분이므로 총합이 1이다. 이 상수에 대한 \(\nabla_{\theta}\)를 구한 것이므로 최종적으로는 0이라는 결과가 나온다. 즉, 이렇게 추가된 expectation term은 전체 policy gradient를 계산하는데 전혀 영향을 끼치지 않는다. 이 결과는 trajectory를 infinite하다고 생각했을 때 0인 것이지, finite sample이라고 가정하면 0이 아닌데, 이런 특성을 통해서 policy gradient가 unbiased하면서 variance를 낮추는 효과를 계속 누릴 수 있다. 물론 지금 이 부분에 다루는 것처럼 trajectory의 평균을 baseline으로 삼는 것이 항상 좋은 것은 아니지만, 그래도 적어도 variance를 낮추는 기대 효과를 가져올 수 있다.

Analyzing variance

그러면 지금까지 계속 high variance를 낮출 수 있는 방법에 대해서 소개했는데, 그러면 이 variance라는 것을 표현할 수 있을까? 사실 이게 가능하다면 이 variance를 낮출 수 잇도록 optimal baseline을 설계할 수 있는 근거가 생기게 된다. 가장 기본적인 내용이지만 variance를 구하는 식은 아래와 같다.

$$ \text{Var}[x] = \mathbb{E}[x^2] - \mathbb{E}[x]^2 $$

위의 식에 맞춰서 policy gradient의 variance는 다음과 같이 표현할 수 있다.

$$ \text{Var} = \mathbb{E}_{\tau \sim p_{\theta}(\tau)}[(\nabla_{\theta} \log p_{\theta}(\tau)(r(\tau) - b))^2] - \mathbb{E}_{\tau \sim p_{\theta}(\tau)}[\nabla_{\theta} \log p_{\theta}(\tau)(r(\tau) - b)]]^2 $$

사실 두번째 term은 바로 앞에서 소개했던 것처럼 \(b\)를 추가한 영향이 없으므로 단순히 \(\mathbb{E}_{\tau \sim p_{\theta}(\tau)}[\nabla_{\theta} \log p_{\theta}(\tau) r(\tau)] \)로 축약할 수 있다. 다시 말해 이 값은 \(b\)와 상관없는 수식이다. 그러면 이 전체 variance를 \(b\)에 관해서 미분을 취하면 아래와 같이 정리된다.(참고로 \(g(\tau)\)는 \(\nabla_{\theta} \log p_{\theta}(\tau)\)를 줄인 것이다.)

$$ \begin{aligned} \frac{d\text{Var}}{db} = \frac{d}{db} \mathbb{E}[g(\tau)^2(r(\tau) - b)^2]  &= \frac{d}{db}\big(\mathbb{E}[g(\tau)^2 r(\tau)^2]  - 2 \mathbb{E}[g(\tau)^2r(\tau)b] + b^2 \mathbb{E}[g(\tau)^2]\big) \\ &= -2 \mathbb{E}[g(\tau)^2 r(\tau)] + 2b\mathbb{E}[g(\tau)^2] = 0 \end{aligned} $$

그러면 첫번째 term에 대한 식을 쭉 풀어서 해를 찾으면 되는데, 중간에 풀어보면 \(b\)와 관련된 수식이 미분을 통해서 없어지고, 최종적으로 \(b\)에 대한 수식으로 정리할 수 있다.

$$ b = \frac{\mathbb{E}[g(\tau)^2 r(\tau)]}{\mathbb{E}[g(\tau)^2]} $$

이렇게 optimal baseline을 구할 수 있는데, 이 baseline을 구성하는 요소가 앞에서 언급한 grad-log-pi와 reward term이다. 그래서 policy를 구성하는 parameter가 100개이면, gradient하면서 각 parameter에 대한 baseline이 다 생성이 되고, 이 optimal baseline은 결국 기대 reward인데, gradient에 의해서 weight가 가해진 형태임을 알 수 있다.