[RL] (Spinning Up) Proof for Don't Let the Past Distract You

 (이 글은 OpenAI Spinning Up의 글을 개인적으로 정리한 내용입니다. 원본)

Extra Material — Spinning Up documentation

이번 글에서는 action이 이전에 얻은 reward에 reinforce되서는 안된다는 것을 증명하고 한다.

먼저 simplest policy gradient에서의 식 중 \(R(\tau)\)를 전개해보면 다음과 같다.

$$ \begin{align} \nabla_{\theta} J(\pi_{\theta}) &= E_{\tau \sim \pi_{\theta}} \bigg[ \sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta} (a_{t} | s_{t}) R(\tau) \bigg] \\ &= E_{\tau \sim \pi_{\theta}} \left[ \sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta} (a_{t} | s_{t}) \sum_{t'=0}^{T} R(s_{t'}, a_{t'}, s_{t'+1}) \right] \\ &= \sum_{t=0}^{T} \sum_{t'=0}^{T} E_{\tau \sim \pi_{\theta}} [ \nabla_{\theta} \log \pi_{\theta} (a_{t} | s_{t}) R(s_{t'}, a_{t'}, s_{t'+1})] \end{align} $$ 

 여기서 다음과 같이 항을 축약해본다. (우항을 \(t, t'\)에 관한 함수로 만드는 것이다.)

$$ E_{\tau \sim \pi_{\theta}} [ f(t, t') ] = E_{\tau \sim \pi_{\theta}} [\nabla_{\theta} \log \pi_{\theta} (a_{t} | s_{t}) R(s_{t'}, a_{t'}, s_{t'+1})] $$

 이제 여기서 \( t' < t \)인 경우, 즉 reinforce된 action 이전에 reward를 받은 케이스에서 위의 식이 0이 되는 것을 보이고자 한다. 이렇게 하면 우리가 원했던 목적을 완전히 증명하게 되는데, 그 이유는 위의 식에서 \(t' < t \) 부분을 제거하게 되면, 아래와 같이

$$ \nabla_{\theta} J(\pi_{\theta}) = E_{\tau \sim \pi_{\theta}} \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] $$

 우리가 원했던 이전에 다뤘던 policy gradient의 reward-to-go 형태만 남기 때문이다.

1. Using the Marginal Distribution

 진행에 앞서, 먼저 expectation \( E_{\tau \sim \pi_{\theta}} [ f(t, t') ] \)를 살펴봐야 한다. 이건 trajectory에 대한 expectation인데, expectation 내에 포함되어 있는 식은 \( s_{t}, a_{t}, s_{t'}, a_{t'}, s_{t'+1} \) 이라는 몇몇 state와 action들로만 구성되어 있다.  그래서 expectation을 계산할 때, 위의 random variable에 대한 marginal distribution만 신경쓰면 된다. 그럼 이렇게 도출할 수 있다.

$$ \begin{align} E_{\tau \sim \pi_{\theta}}[f(t, t')] &= \int_{\tau} P(\tau|\pi_{\theta}) f(t, t') \\ &= \int_{s_{t}, a_{t}, s_{t'}, a_{t'}, s_{t'+1}} P(s_{t}, a_{t}, s_{t'}, a_{t'}, s_{t'+1} | \pi_{\theta})f(t, t') \\ &= E_{s_{t}, a_{t}, s_{t'}, a_{t'}, s_{t'+1}} [f(t, t')] \end{align} $$

 2. Probability Chain Rule

 Joint distribution은 \( P(A, B) = P(B|A)P(A) \) 라는 chain rule of probability을 통해서 conditional probability와 marginal probability에 관해서 계산할 수 있다. 그러면 앞에 잠깐 나왔던 항을 다음과 같이 바꿔볼 수 있다.

$$ P(s_{t}, a_{t}, s_{t'}, a_{t'}, s_{t'+1} | \pi_{\theta}) = P(s_{t}, a_{t} | \pi_{\theta}, s_{t'}, a_{t'}, s_{t'+1}) P(s_{t'}, a_{t'}, s_{t'+1} | \pi_{\theta}) $$

 3. Separating Expectations Over Multiple Random Variables

 만약 두 개의 random variable \(A, B \)에 관해서 expectation을 알고있다면, 이걸 inner expectation과 outer expectation 두 개로 쪼갤 수 있다. 여기서 inner expectation은 outer expectation에서 나온 변수를 상수로 취급하게 된다. 여기서도 probability chain rule을 사용하게 된다. 수학적으로는 다음과 같다.

$$ \begin{align} E_{A, B}[f(A, B)] &= \int_{A, B} P(A, B) f(A, B) \\ &= \int_{A} \int_{B} P(B|A) P(A) f(A, B) \\ &=  \int_{A} P(A) \int_{B} P(B|A) f(A, B) \\ &= \int_{A} P(A) E_{B} \bigg[ f(A, B)|A \bigg] \\ &= E_{A} \bigg[ E_{B} \left[ f(A, B)|A \right] \bigg] \end{align}$$

이제 여기서 \(s_{t}, a_{t}, s_{t'}, a_{t'}, s_{t'+1}\)에 대한 expectation을 구하면 다음과 같이 된다.

$$ \begin{align} E_{\tau \sim \pi_{\theta}}[f(t, t')] &= E_{s_{t}, a_{t}, s_{t'}, a_{t'}, s_{t'+1}}[f(t, t')] \\ &= E_{s_{t'}, a_{t'}, s_{t'+1} \sim \pi_{\theta}} \bigg[ E_{s_{t}, a_{t} \sim \pi_{\theta}} \big[f(t, t') | s_{t'}, a_{t'}, s_{t'+1}\big] \bigg] \end{align} $$

 4. Constants Can Be Pulled Outside of Expectations

만약 expectation 내부의 어떤 항이 expect를 구하는 변수에 대해 상수 취급을 받는다고 하면, 해당 항은 expectation의 외부로 뺄 수 있다. 예를 들어 두 개의 random variable \(A, B\) 에 대한 expectation을 생각해보는데 이번에는 \(f(A, B) = h(A)g(B)\)라고 해보자. 그러면 이전에 구했던 것을 활용해보면 다음과 같다.

$$ \begin{align} E_{A, B}[f(A, B)] &= E_{A}\bigg[ E_{B} \big[f(A, B) | A \big]\bigg] \\ &= E_{A} \bigg[ E_{B} \big[ h(A)g(B) | A \big]\bigg] \\ &= E_{A} \bigg[ h(A) E_{B} \big[ g(B) | A\big] \bigg] \end{align} $$

( \( h(A) \)는 \(B\)에 관한 식이 아니기 때문에 상수로 인정되어 expectation을 구할때 밖으로 뺄 수 있다.)

그러면 앞에서 구했던 expectation 식은 다음과 같이 다시 쓸 수 있게 된다.

$$ \begin{align} E_{\tau \sim \pi_{\theta}} [f(t, t')] &= E_{s_{t'}, a_{t'}, s_{t'+1} \sim \pi_{\theta}} \bigg[ E_{s_{t}, a_{t} \sim \pi_{\theta}} \bigg[ f(t, t') | s_{t'}, a_{t'}, s_{t'+1}\bigg]\bigg] \\ &= E_{s_{t'}, a_{t'}, s_{t'+1} \sim \pi_{\theta}} \bigg[ E_{s_{t}, a_{t} \sim \pi_{\theta}} \bigg[ \nabla_{\theta} \log \pi_{\theta} (a_{t} | s_{t}) R(s_{t'}, a_{t'}, s_{t'+1}) | s_{t'}, a_{t'}, s_{t'+1}\bigg]\bigg] \\ &= E_{s_{t'}, a_{t'}, s_{t'+1} \sim \pi_{\theta}} \bigg[ R( s_{t'}, a_{t'}, s_{t'+1}) E_{s_{t}, a_{t} \sim \pi_{\theta}} \bigg[ \nabla_{\theta} \log \pi_{\theta} (a_{t}|s_{t}) | s_{t'}, a_{t'}, s_{t'+1}\bigg]\bigg] \end{align} $$

 5. Applying the EGLP Lemma.

 마지막 증명은 EGLP lemma에 기반한 것이다. 여기서 innermost expectation에 대해서만 신경쓰면 된다.

$$ E_{s_{t}, a_{t} \sim \pi_{\theta}} \Big[ \nabla_{\theta} \log \pi_{\theta} (a_{t} | s_{t}) | s_{t'}, a_{t'}, s_{t'+1} \Big] = \int_{s_{t}, a_{t}} P(s_{t}, a_{t} | \pi_{\theta}, s_{t'}, a_{t'}, s_{t'+1}) \nabla_{\theta} \log \pi_{\theta} (a_{t} | s_{t}) $$

 이제 여기서 두가지 케이스에 대해서 구분해야 한다. 첫번째 케이스는 \(t' < t\)인데, 이 경우는 action을 취하기전에 reward가 발생한 것이고, 두번째는 \(t' \geq t \) 로써 action을 취한 후에 reward를 고려한 케이스이다.

 Case 1 : Reward Before Action

 만약 \( t' < t \)라면, 현재 policy \(\pi_{\theta}\)에서 나온 action \(a_{t}\)에 대한 conditional probability는 다음과 같다.

$$ P(s_{t}, a_{t} | \pi_{\theta}, s_{t'}, a_{t'}, s_{t'+1}) = \pi_{\theta} (a_{t} | s_{t}) P(s_{t} | \pi_{\theta}, s_{t'}, a_{t'}, s_{t'+1}) $$

 그럼 위에서 구한 innermost expectation은 다음과 같이 쪼갤 수 있다.

$$ \begin{align} E_{s_{t}, a_{t} \sim \pi_{\theta}} \Big[ \nabla_{\theta} \log \pi_{\theta} (a_{t} | s_{t}) | s_{t'}, a_{t'}, s_{t'+1} \Big] &= \int_{s_{t}, a_{t}} P(s_{t}, a_{t} | \pi_{\theta}, s_{t'}. a_{t'}, s_{t'+1}) \nabla_{\theta} \log \pi_{\theta} (a_{t} | s_{t}) \\  &= \int_{s_{t}} P(s_{t} | \pi_{\theta}, s_{t'}, a_{t'}, s_{t'+1}) \int_{a_{t}} \pi_{\theta} (a_{t} | s_{t}) \nabla_{\theta} \log \pi_{\theta} (a_{t} | s_{t}) \\  &= E_{s_{t} \sim \pi_{\theta}} \Big[ E_{a_{t} \sim \pi_{\theta}} \Big[ \nabla_{\theta} \log \pi_{\theta} (a_{t} | s_{t}) | s_{t} \Big] | s_{t'}, a_{t'}, s_{t'+1} \Big] \end{align}$$

이전에 다뤘던 EGLP lemma에서는 다음과 같이

$$ E_{a_{t} \sim \pi_{\theta}} \Big[ \nabla_{\theta} \log \pi_{\theta} (a_{t} | s_{t}) | s_{t} \Big] = 0 $$

이라고 정의했으므로, 결국 \(t' < t\)인 상태에서는 \(E_{\tau \sim \pi_{\theta}} [f(t, t')] = 0 \) 라는 결론을 내릴 수 있게 된다.

 Case 2 : Reward After Action

 그럼 \(t' \geq t\)인 경우는 어떨까? 왜 같은 논리가 성립하지 않을까? 이 경우 action \(a_{t}\)에 대한 conditional probability가 같은 방법으로 나눌 수가 없는데, 그 이유는 지금 이렇게 구하는 과정 자체가 미래에 대해서 conditioning을 하고 있기 때문이다. 이렇게 한번 생각해보자: 만약 매일 아침에 조깅을 할건지, 일찍 출근을 할건지 50 대 50의 확률로 선택한다고 가정해보자. 만약 미래에 일찍 출근하는 것에 대한 조건을 세운다면, 반대로 과거에 조깅을 했을 것에 대한 조건은 어떻게 구할 것인가? 당연히 구하지 못할 것이다. 반대로 과거에 뭔가 결정을 내리기 전에 조건을 걸었을 때, 나중에 조깅을 할 것에 대한 조건은 어덯게 구할 것인가? 다시 50대 50의 확률이 될것이다.

 결국 \(t' \geq t \)인 경우에는 action \(a_{t}\)에 대한 conditional distribution이 \(\pi(a_{t} | s_{t})\)가 아니며, EGLP lemma는 성립하지 않게 된다.