[논문 리뷰] (GAE) HIGH-DIMENSIONAL CONTINUOUS CONTROL USING GENERALIZED ADVANTAGE ESTIMATION

"High-Dimensional Continuous Control Using Generalized Advantage Estimation"
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, Pieter Abbeel

이 논문에서는 허용 가능한 수준의 편향을 유지하면서 정책 기울기 추정치의 편차를 줄이기 위해 Generalized Advantage Estimation (GAE) 개념과 value functions와 trust region optimization를 사용하여 학습에 필요한 샘플 수를 크게 줄임으로써 신경망으로 모델링된 정책을 보다 안정적으로 개선할 수 있다고 주장합니다.

즉, GAE 개념을 TRPO 알고리즘에 적용하여 일반적으로 많은 수의 샘플이 필요하다는 점과 들어오는 데이터의 비정형성에도 불구하고 안정적이고 지속적인 개선을 얻는 것이 어렵다는 Policy Gradient Method의 문제를 개선합니다.

TRPO 알고리즘은 다음 편에 다루도록 하고, 이번에는 GAE에 대해 정리하려 합니다.

N-steps Return

지난 CS285_HW2 풀이 중 Actor-Critic 알고리즘과 Monte Carlo를 이용한 Advantage Estimation을 비교해보았습니다.

critic의 advantage는 neural network baseline을 이용해서 variance를 낮출 수 있는 반면 bias는 높습니다. 그 이유는 만약 estimated value is wrong인 경우 즉, 학습이 올바르지 못한 경우는 bias가 증가하기 때문입니다.

반면, MC는 sampling data를 이용해 직접 추정하기 때문에 unbiased 이지만, single-sample estimation의 한계로 variance가 높습니다.

이 두 특성을 적절히 이용하기 위해 n-step return 방법을 소개했습니다.

n값을 조정하여 추정의 편향과 분산 사이의 균형을 조절할 수 있습니다. n이 증가할수록 편향은 감소하고 분산은 증가합니다.

여기에서 더 나아가서 모든 n-step을 고려해서 Advantage을 estimation해보는 건 어떨까? 각 n-step에 대한 advantage를 exponential weighted sumation 즉, 가중평균을 이용해 Advantage estimation을 결정하는 것이다.

Generalized Advantage Estimation (GAE)

V를 approximate value function라 한다면,

으로 정의되며, 이는 $V$의 TD 잔차로, 할인 $\gamma$를 사용합니다. 주목할 점은 $\delta^V_t$가 행동 $a_t$의 이점을 추정하는 것으로 간주될 수 있다는 것입니다. 실제로, 올바른 가치 함수 $V = V_{\pi,\gamma}$를 가지고 있다면, 그것은 $\gamma$-정의된 이점 추정기이며, 사실 $A_{\pi,\gamma}$의 편향되지 않은 추정기입니다:

그러나 이 추정기는 $V = V_{\pi,\gamma}$일 때만 $\gamma$-정의됩니다. 그렇지 않으면 편향된 정책 기울기 추정치를 제공합니다.

이제 k개의 이러한 $\delta$ 항의 합을 고려해 보겠습니다. 이는 다음과 같이 표현됩니다

$\hat{A}_{(k),t}$가 근본적으로 k-단계 반환의 추정치와 $-V(s_t)$라는 기준선 항을 포함한다고 볼 수 있습니다. $\delta V_t = \hat{A}_{(1),t}$와 유사하게, $\hat{A}_{(k),t}$는 이점 함수의 추정기로 간주될 수 있으며, $V = V_{\pi,\gamma}$일 때만 $\gamma$-정의됩니다.

그러나 k가 $\infty$로 갈수록 편향은 일반적으로 작아집니다. 왜냐하면 $\gamma^k V(s_{t+k})$항이 점점 더 많이 할인되기 때문이며, ( -V(s_t) ) 항은 편향에 영향을 미치지 않기 때문입니다. 즉, k를$\infty$로 취하면 다음과 같은 식을 얻습니다:

이는 단순히 경험적 반환에서 가치 함수 기준선을 뺀 것입니다.

exponentially weighted sum

GAE는 이전에 설명된 k-단계 추정기의 지수 가중 평균으로 정의된다.

이는 $1 - \lambda$ 계수로 조정된 $\delta V_t$의 각 k-단계 합의 계수로 이루어진 합입니다. 이것은 다음과 같이 간단한 수식으로 나타낼 수 있습니다

이를 요약하면 다음과 같습니다.

위의 수식에서 볼 수 있듯이, GAE는 벨만 잔차 항의 할인된 합을 포함하고 있으며, 이는 return function를 수정한 MDP에서의 return과 유사합니다. 이러한 구성은 TD(λ) (Sutton & Barto, 1998)를 정의하는 데 사용된 구성과 유사하지만, 여기서는 Value 함수가 아닌 Advantage 함수를 추정합니다.

이점 추정식에는 $\lambda = 0$과 $\lambda = 1$인 두 가지 주목할 만한 특별한 경우가 있다.

• $\text{GAE}(\gamma, 0): \hat{A}_t := \delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$
• $\text{GAE}(\gamma, 1): \hat{A}_t := \sum_{l=0}^{\infty} \gamma^l \delta_{t+l} = \sum_{l=0}^{\infty} \gamma^l (r_{t+l} - V(s_t))$

$\text{GAE}(\gamma, 1)$는 $V$의 정확성에 관계없이 $\gamma$에 의해 정의되지만, 높은 분산 때문에 실용적인 문제가 발생할 수 있습니다. 반면에 $\text{GAE}(\gamma, 0)$는 $V = V_{\pi,\gamma}$일 때 $\gamma$에 의해 정의되며, 그렇지 않으면 편향을 유발하지만 일반적으로 분산이 훨씬 낮습니다. $\lambda$ 값이 0과 1 사이일 때, $\text{GAE}(\gamma, \lambda)$는 편향과 분산 사이에서 절충을 이루며, $\lambda$ 매개변수에 의해 제어됩니다.

즉, $\gamma$는 주로 가치 함수 $V^{\pi,\gamma}$의 스케일을 결정하며, 이는 $\lambda$에 의존하지 않습니다. $\gamma$를 1보다 작게 설정하면 정책 기울기 추정에 편향을 도입하게 되며, 이는 가치 함수의 정확도에 상관없이 발생합니다. 반면, $\lambda$를 1보다 작게 설정하면 가치 함수가 부정확할 때만 편향을 도입합니다. 실제로는 $\lambda$가 $\gamma$보다 훨씬 적은 편향을 도입하기 때문에, 정확한 가치 함수에서 $\lambda$의 값이 더 낮은 것이 효과적입니다.

결과적으로, GAE를 사용하여 discounted policy gradient $g^\gamma$의 편향된 추정치를 구성할 수 있습니다. 이는 아래 수식으로 나타낼 수 있습니다.

여기서 $\hat{A}_t^{GAE(\gamma,\lambda)}$는 GAE로 계산된 이점 추정치입니다. 이 식은 λ가 1일 때 등식이 성립합니다. 이를 통해 보다 안정적이고 효율적인 학습이 가능하도록 도와주며, 복잡한 환경에서도 강화 학습 알고리즘이 효과적으로 작동할 수 있도록 합니다.

GAE 구현

구현은 CS285 HW2 과제 중 간소화된 GAE를 이용해 Hopper-v4 환경에 적용해보았다. 단, GAE를 더 간단하게 구현하기 위해 재귀적으로 표현하면,

    def estimate_advantage(self, obs: np.ndarray, rews_list: np.ndarray, q_values: np.ndarray, terminals: np.ndarray):
        if self.nn_baseline:
            values_unnormalized = self.actor.run_baseline_prediction(obs)
            assert values_unnormalized.ndim == q_values.ndim

            values = (values_unnormalized - np.mean(values_unnormalized)) / (np.std(values_unnormalized) + 1e-8)
            values = values * (np.std(q_values) + 1e-8) + np.mean(q_values)

            if self.gae_lambda is not None:
                values = np.append(values, [0])
                rews = np.concatenate(rews_list)

                batch_size = obs.shape[0]
                advantages = np.zeros(batch_size + 1)

                for i in reversed(range(batch_size)):
                    if terminals[i]:
                        delta = rews[i] - values[i]
                        advantages[i] = delta
                    else:
                        delta = rews[i] + self.gamma * values[i+1] - values[i]
                        advantages[i] = delta + self.gamma * self.gae_lambda * advantages[i+1]

                advantages = advantages[:-1]

            else:
                advantages = q_values - values

        else:
            advantages = q_values.copy()

        if self.standardize_advantages:
            advantages = (advantages - np.mean(advantages)) / (np.std(advantages) + 1e-8)

        return advantages

여기서 'values = np.append(values, [0])'를 해주는 이유는 재귀적으로 GAE를 구현했기 때문에 $t$시점에서 Advantage estimation을 구하려면 $t+1$ 시점의 값이 필요하기 때문입니다. 계산을 간단하게 하기 위해 0을 추가하려 인덱싱 오류를 방지하고 계산을 용이하게 하기 위함입니다.

GAE를 구현하는 핵심 파트는 'for i in reversed(range(batch_size)):' 반복문입니다. 여기에서는 시퀀스 끝에서부터 역순으로 계산하게 됩니다. 여기서 각 스텝에 대한 처리 방법을 예시를 들면,

• Batch Size가 3인 경우
  • Step 3 (i=2): 가장 마지막 타임스텝입니다. 여기서는 $\delta = \text{rews}[2] - \text{values}[2]$를 계산하고, $\text{advantages}[2] = \delta$를 설정합니다. 마지막 타임스텝이므로 $A_{t+1}$은 0입니다.
  • Step 2 (i=1): 두 번째 타임스텝입니다. $\delta = \text{rews}[1] + \gamma \cdot \text{values}[2] - \text{values}[1]$을 계산하고, $\text{advantages}[1] = \delta + \gamma \lambda \cdot \text{advantages}[2]$를 설정합니다.
  • Step 1 (i=0): 첫 번째 타임스텝입니다. $\delta = \text{rews}[0] + \gamma \cdot \text{values}[1] - \text{values}[0]$을 계산하고, $\text{advantages}[0] = \delta + \gamma \lambda \cdot \text{advantages}[1]$를 설정합니다.

이러한 방식으로 각 타임스텝에 대한 이점을 계산하며, 마지막으로 더미 값을 제거하기 위해 마지막 이점 값을 제외합니다. 이 과정을 통해 각 상태-행동 쌍의 이점을 추정할 수 있으며, 이를 통해 정책을 개선하는 데 사용될 수 있습니다.