ACER: Sample Efficient Actor-Critic With Experience Replay
ACER: Sample Efficient Actor-Critic With Experience Replay
제목에서도 볼 수 있듯이, 딥마인드에서 나온 Sample Efficient Actor-Critic With Experience Replay 는 Actor-Critic method에 Experience Replay를 접목시켜 데이터 효율성을 높인 새로운 강화학습 알고리즘을 제안하는 논문입니다. A3C의 off-policy 버전이라고 생각하셔도 됩니다.
논문 내용을 요약하면 다음과 같습니다.
-
Experience Replay를 도입해서 Sample efficiency를 향상시켰다.
-
Experience Replay를 사용하기 위해 그래디언트 계산에 off-policy correction을 추가했다. Importance sampling을 사용할 것인데 그냥 사용하면 bounded 되지 않은 importance weight 값이 여러번 곱해져 variance가 너무 커질 수 있으니, Off-policy Actor-Critic 에서 소개된 근사식을 사용해 그래디언트를 계산한다.
[Eq. 1]$ g^{marg} = \mathbb{E}{x_t \sim \beta, a_t \sim \mu} [\rho_t \nabla{\theta}\log \pi_\theta (a_t|x_t)Q^\pi(x_t, a_t)] $ -
식 (1)을 계산하기 위해선 $ Q^\pi $ 값을 알아야 하므로 Retrace($\lambda$ ) 라는 알고리즘을 사용해 $ Q^\pi$ 를 추정한다.
-
식 (1)의 importance weight $ \rho_t$ 도 너무 커질 수 있으니 clipping 해준다. 그런데 bias term을 적절히 추가해 식(1)에 대해 unbiased estimator임을 유지시킨다.
-
TRPO를 적용시키는데, 업데이트 후 policy가 현재 policy가 아닌 이때까지 policy의 moving average와 멀어지지 않게 적용시킨다. 단 policy parameter space가 아니라 최종 distribution의 statistics의 space에 대해 적용한다.
※ Action space가 continuous 한 경우 몇가지 다른 점이 생기는데, 이는 나중에 다루도록 하겠습니다.
그럼 이제 각 부분을 차례대로 자세히 살펴보겠습니다. 지금은 요약문이 잘 이해되지 않더라도 설명을 보고 나면 이해하기 훨씬 쉬워질 것입니다!
Gradient Estimation
Policy Gradient 기반의 강화학습 알고리즘의 목표는 policy의 성능을 증가시키는 방향의 그래디언트를 구해 policy를 업데이트하는 것입니다. 보통 성능은 (discounted) 보상의 합으로 정의됩니다.
$ \pi_\theta(a|x) $를 “$\theta $ 를 매개변수로 사용하는 policy $ \pi $ 를 따를 경우 현재 상태가 $ x $일 때 행동 a를 취할 확률”로 정의했을 때, Policy Gradient Theorem에 따르면 성능을 올리기 위해선
[Eq. 2]
와 같이 $ g $ 를 계산한 후 $\theta $를 $g$ 에 따라 갱신하면 됩니다. $ A^\pi $ 는 일반적으로 쓰이는 action-value 함수, 또는 discounted return(discounted 보상의 합), temporal difference ($r_t + \gamma V^\pi (x_{t+1})-V^\pi(x_t))$) 중 아무 것이나 넣어도 성립합니다. 다만 어떤 걸 선택하느냐에 따라 function approximation을 사용했을 때 bias와 variance에서 정도의 차이가 생기게 됩니다.
이제 그래디언트 계산식을 off-policy에 맞게 수정해봅시다. 위의 그래디언트 계산식은 데이터를 생성하는 behavior policy와 학습이 진행되는 대상인 target policy가 일치할 때만 target policy의 성능을 증가시키는 방향의 그래디언트임이 보장되는데, behavior policy가 target policy와 다를 때도 계산된 그래디언트가 target policy의 성능을 증가시키는 방향의 그래디언트가 되게 계산식을 수정한다는 뜻입니다.
On-policy 알고리즘을 Off-policy로 만들기 가장 쉬운 방법은 importance sampling을 사용하는 것입니다.
$ \pi $ 를 target policy, $ \mu $를 behavior policy, $ \rho_t = \frac{\pi(a_t|x_t)}{\mu (a_t|x_t)} $ 로 정의한 후
[Eq. 3]
로 계산하면 $ \hat{g}^{imp}$ 는 unbiased estimator가 됩니다. 즉 $ \hat{g}^{imp}$의 기댓값은 원래 구하려고 했던 그래디언트의 기댓값과 동일합니다.이는 각 trajectory에 대한 그래디언트를 계산할 때, behavior policy에서 뽑힌 trajectory가 target policy에서 발생했을 확률을 곱해 보정하는 것으로 이해할 수 있습니다.
그런데 $ \pi $ 든 $ \mu $ 든 확률값이기 때문에 $ \rho_t $ 의 값에는 제한이 없고, 심지어 이런 $ \rho_t $가 계속해서 곱해진 값을 사용하기 때문에 식 (3)을 그대로 사용하면 variance가 너무 커서 학습이 불안정합니다. 그렇다고 $ \rho_t$ 의 값을 임의로 clipping 등을 적용해 범위를 제한시키면 bias가 너무 커져 학습이 제대로 진행되지 않습니다.
대신 요약문에서도 언급했듯이, Off-policy Actor-Critic 에서 소개된 그래디언트 근사식을 사용하면 이런 문제를 우회할 수 있습니다.
[Eq. 4]
당연하지만 $ Q_\pi(x_t, a_t)$ 부분은 $ \theta $ 에 대한 그래디언트 계산 대상이 아닙니다.
식 (4)를 살펴보면 $ \rho_t $ 에 대한 곱셈항이 사라지고 그냥 값을 추정하기만 하면 되는 것을 알 수 있습니다. 또한 그래디언트 계산에 실제로 관측한 보상값인 $r_t$를 사용한 식 (3)과 달리, 그래디언트를 계산하기 위해 $ Q^\pi $ 를 필요로 하는 것을 알 수 있습니다.
Off-policy Actor-Critic 논문에서는 $R_t^\lambda = r_t + (1-\lambda)\gamma V(x_{t+1})+\lambda \gamma \rho_{t+1} R_{t+1}^\lambda$ 라는 재귀식을 통해 $ Q^\pi $ 를 계산하는데, 보시다시피 재귀적으로 계산하며 $\rho_t $ 를 계속 곱해주기 때문에 학습이 불안정해질 수 있습니다. 그래서 ACER에서는 Retrace라는 알고리즘을 통해 $ Q^\pi $ 를 계산합니다.
Q function Estimation
Safe and Efficient Off-Policy Reinforcement Learning 에서 소개된 Retrace($ \lambda$)는 off-policy 세팅에서 $ Q $ 함수를 계산하게 해주는 알고리즘입니다. 단순히 $ \rho = \frac{\pi(a_t|x_t)}{\mu(a_t|x_t)} $ 대신 $\rho = \lambda \min(1, \frac{\pi(a_t|x_t)}{\mu(a_t|x_t)}) $ 를 쓰는 것만으로 임의의 $\pi, \mu $에 대해 수렴함도 보장하면서 variance도 더 낮게 됩니다. Off-policy correction에는 tree backup이나 $ Q(\lambda$) 등의 다른 방법도 있지만 discrete 케이스에서는 Retrace의 성능이 가장 좋았다고 합니다. ACER에서는 $ \lambda = 1 $ 로 사용합니다.
Retrace에 따라 $ Q $ 함수를 업데이트하는 경우 target value는 아래와 같이 계산됩니다.
[Eq. 5]
식을 살펴보면 $ Q^{ret} $ 는 재귀적으로 계산됨을 알 수 있습니다. 꼭 한 에피소드 전체 trajectory를 가져와서 계산할 필요는 없고, $ x_0 $ 부터 $x_k$ 까지의 부분 trajectory를 대상으로 계산할 때 $ x_k $ 가 terminal이면 $ Q^{ret}(x_k, a_k) = 0$, 그렇지 않을 경우 $Q^{ret}(x_k, a_k) = V(x_k) $1 로 초기값을 주고 재귀적으로 나머지 상태에 대해 계산하면 됩니다.
즉 $ Q $ 함수를 파라미터 $ \theta_v$ 를 사용해 근사하는 경우, $ Q_{\theta_v}$ 를 업데이트하기 위해선
- 식 (5)에서 $ Q $ 대신 $Q_{\theta_v}$ 를 사용해 $ Q^{ret}$를 계산한 다음
- $Q_{\theta_v}가 $ $Q^{ret}$ 에 가까워지게 MSE를 사용해 업데이트하면 됩니다.
결국 그래디언트 값 $ (Q^{ret}(x_t, a_t) - Q_{\theta_v}(x_t, a_t) \nabla_{\theta_v} Q_{\theta_v} (x_t, a_t)) $ 을 따라 업데이트를 진행하는 것입니다.
Return-based기 때문에 critic의 학습이 빠를 뿐만 아니라, multi-step return을 사용하기 때문에 식 (4)에서 그래디언트를 구할 때 $ Q^\pi $ 값으로 사용되어 bias도 상당히 줄여줍니다.
Importance Weight Truncation with Bias Correction
식 (3)에서 식 (4)로 넘어오며 importance weight에 대한 곱셉항을 제거했지만, $ \rho_t $ 가 unbounded라는 사실은 변함이 없기 때문에 여전히 학습이 불안정해질 수 있는 요소가 남아 있습니다. 혹시 헷갈릴까 첨언하자면 윗 문단에서 진행한 importance weight의 clipping은 $ Q $ 함수 학습에 사용하는 weight를 대상으로 진행한 것입니다. 이 문단에서 다루는 importance weight는 gradient estimation에 사용되는 weight를 말합니다.
하지만 식 (4)에 간단한 변형을 가해 importance weight가 bounded 된 unbiased gradient estimator2를 만들 수 있습니다! $ \bar{\rho_t} = \min (c, \rho_t) $ 라고 정의한 뒤, 그래디언트 계산식을 아래와 같이 두 항으로 나눠주면 됩니다.
[Eq. 6]
$ c \le \rho_t $ 인 경우 위의 식에서 오른쪽 항은 0이 되므로 식 (4)와 식 (6)은 동일합니다.
$ c > \rho_t $ 인 경우도 살펴보겠습니다.
[Eq. 7]
가 성립합니다. 여기서 $ c > \rho_t $ 이므로 $ \bar{\rho_t} = c $ 이기 때문에 식 (7)의 왼쪽 항은 식 (6)의 왼쪽 항과 동일합니다.
이제 식 (7)의 오른쪽 항과 식 (6)의 오른쪽 항이 같음을 보이겠습니다.
[Eq. 8]
식 (6)의 오른쪽 항은 $ c > \rho_t $ 일 때만 활성화되므로, 위아래에 $ \pi(a|x_t)$ 를 곱해 $ \mu $ 에 대한 기댓값을 $ \pi $ 에 대한 기댓값으로 바꾸니 식 (6)과 식 (7)이 같음을 볼 수 있었습니다.
이제 $ \bar{\rho_t} $ 의 값은 최대 $ c $가 되고, $ \frac{\rho_t(a)-c}{\rho_t(a)}$ 또한 최대 1이기 때문에 importance weight에 관련된 값이 모두 bounded 되어 학습의 불안정성 요소가 제거된 것입니다. 대신 그 대가로 behavior policy인 $ \mu $가 아니라 $ \pi $를 따르는 $a $ 에 대해 $ Q, \pi $에 관한 기댓값을 계산하게 되었는데, 다행히 $ a $ 의 확률 분포는 parametrized policy $ \pi $를 계산해 알 수 있고 $ Q $ 또한 뉴럴넷으로 근사해두었기 때문에 별도의 샘플링 없이 식 (6)을 사용할 수 있습니다.
다만 식 (6)에서 왼쪽 항은 실제로 경험한 trajectory에 대해 계산하기 때문에 $ Q $ 함수를 $ Q^{ret} $으로 추정해 사용할 수 있는 반면, 오른쪽 항은 실제 trajectory에서 겪어보지 못한 행동을 한다고 가정하기 때문에 $ Q^{ret} $를 계산할 수 없어 $ Q_{\theta_v} $로 값을 사용하게 됩니다.
[Eq. 9]
여기에 베이스라인으로 value 함수 $ V $ 를 사용해 variance를 더 줄일 수 있습니다. $ Q_{\theta_v}$ 와 $ \pi $ 를 알기 때문에 기댓값으로 $ V_{\theta_v} $ 를 계산해 사용할 수 있습니다.
최종 식은 다음과 같습니다.
[Eq. 10]
Efficient TRPO
TRPO 를 간단하게 설명하면, 한 번 업데이트 했을 때 policy가 지나치게 많이 바뀌지 않게 제약하는 알고리즘입니다. Learning Rate를 낮출 경우 policy parameter $\theta$ 의 변화량은 작게 유지되겠지만, $ \theta $ 의 변화가 작다고 해서 최종적으로 계산된 행동 별 확률값의 차이도 작다는 보장이 없기 때문에 parameter space에서의 변화량이 아니라 policy space의 변화량을 직접적으로 제약하는 방식입니다.
보통 TRPO를 적용하면 parameter $ \theta $ 에 대해 업데이트 직전과 직후 policy 값의 차이를 제약하는 방식으로 적용합니다. 그런데 ACER에서는 특이하게도 업데이트 직후의 policy가 직전 policy가 아닌 이때까지의 moving average parameter의 policy와 차이나지 않게 제약합니다.
또 다른 점 하나는 파라미터 $ \theta $에 대해 직접 제약을 걸지 않고 distribution statistics $ \phi_\theta $ 에 대해 건다는 것입니다. 뉴럴넷의 파라미터를 $ \theta $라고 했을 때, 주어진 상태 $ x $에 대해 각 행동에 대한 확률 값을 계산하기 위해선 $ \theta $에 관한 함수인 $ \phi_\theta$ 로 distributuion statistics $ \phi_\theta(x)$ 를 계산한 후, 이를 사용해 distributuion function $ f $ 를 만들게 됩니다.
행동이 왼쪽/오른쪽만 있는 discrete action space를 예로 들면, 뉴럴넷의 모든 파라미터를 통칭하는 것이 $ \theta $ 이고, 뉴럴넷 자체는 categorical distributuion의 statistics인 logits를 계산하는 함수 $ \phi_\theta$가 되고, 주어진 상태 $ x $ 에 대해 계산한 logits 값은 $ \phi_\theta(x) $ 가 됩니다. 그리고 logits에 softmax를 씌운 categorical distribution이 policy $ \pi $가 될 것이고, statistics을 사용해 실제 확률을 계산하는 함수인 softmax가 $ f $ 인 것입니다. 따라서 $ x $는 $\phi_\theta $와 $ f $를 거쳐 $ \pi $ 로 계산되게 됩니다.
그런데 그래디언트 계산은 $ \theta $ 에 대해 하는데 제약은 그 중간과정인 $\phi_\theta $ 에 대해 걸까요? 그 이유는 TRPO를 적용하는 방식이 parameter space의 그래디언트 값을 수정하는 방식으로 이루어지기 때문입니다. 원래 목표를 달성하기 위한 그래디언트를 구해둔 뒤 그래디언트 성분 중 policy간 차이를 키우는 방향을 제거하는 방법으로 동작하는데, 만약 $ \theta $ 에 대해 그래디언트를 바로 구한다면 한 데이터 샘플 당 (가능한 행동의 수 + 1)번 그래디언트를 계산해 후처리해야 하기 때문에 계산량이 너무 늘어나게 됩니다. 따라서 $ \theta $ 보다 훨씬 차원이 적고 computational graph 상에서 $ \pi $와 가까운 $ \phi_\theta $ 에 대해 제약을 걸어 계산량에서 이득을 보는 것입니다. 실제로 $ \theta $는 수백만에서 수천만까지의 차원을 가질 수 있지만 $ \phi_\theta $는 많아봤자 가능한 행동 개수의 몇 배이기 때문에 큰 이득을 볼 수 있습니다. Section의 제목에 efficient가 들어가는 이유도 이 때문입니다.
수식은 다음과 같습니다. 여기서 사용한 $ \hat{g}^{acer} $ 는 $ \theta $에 대해 그래디언트를 계산한 식 (9)와 달리 $ \phi_\theta $ 에 대해 그래디언트를 계산한 결과입니다.
[Eq. 11]
직관적으로 해석하면 $ k = \nabla_{\phi_\theta(x_t) }D_{KL}[f(\cdot|\phi_{\theta_a}(x_t)) ||f(\cdot|\phi_{\theta}(x_t))]$ 라고 했을 때 $ k $ 방향 성분이 일정 이하인 벡터 중 $ \hat{g}^{acer} $ 에 가장 가까운 벡터 $ z $ 를 찾아 $ \hat{g}^{acer}$ 를 대체할 새로운 그래디언트로 사용하는 것입니다. $ k $ 가 의미하는 것은 moving average policy와 현재 policy 간의 KL Divergence를 늘리는 방향의 그래디언트이기 때문에, 결국 $ z $ 는 원래 사용하려고 했던 그래디언트 $ \hat{g}^{acer}$ 에서 KL Divergence를 늘리는 성분을 제거한 그래디언트가 됩니다.
또 최적화 목표가 $ z $ 에 대한 단순한 이차항이고, constraint도 $ z $ 에 대한 일차부등식이기 때문에 해당 최적화 문제는 Quadratic Progamming이 됩니다. 즉 KKT 조건을 만족하기 때문에 다음과 같이 쉽게 해를 찾을 수 있습니다.
[Eq. 12]
이제 $ \hat{g}^{acer} $ 대신 $ z^* $ 를 사용해 파라미터 $ \theta $ 를 업데이트하면 됩니다.
그런데 이 방식에 장점만 있는 것은 아닙니다. $\theta$ 대신 $ \phi_\theta$에 대한 그래디언트를 사용해 계산량에서 이득을 얻은 대신, 그만큼 제약 조건의 정밀도가 희생되게 됩니다. 원래 증가시키려고 했던 성능에 대한 함수를 $ J $ 라고 부를 때, 뉴럴넷 파라미터 $ \theta $ 의 업데이트에 쓰이는 그래디언트는 $ \frac{\partial{J}}{\partial \theta}$ 입니다. 그런데 위의 변형된 TRPO는 $ \frac{\partial J}{\partial\theta} = \frac{\partial J}{\partial \phi_\theta} \cdot \frac{\partial \phi_\theta}{\partial\theta}$ 중 앞쪽항, 즉 $ \frac{\partial J}{\partial \phi_\theta} $ 에 대해서만 제약을 걸기 때문에 실제로 $ \theta $ 에 대해 업데이트가 이루어졌을 때 KL Divergence가 $ \theta $의 변화량에 얼마나 민감하고 얼마나 바뀌는지는 완벽히 통제할 수 없게 됩니다. 다만 이를 사용하지 않는 경우 계산량이 너무 크기 때문에 사용하는 편이 훨씬 이득이게 됩니다.
구현
일반적인 딥러닝 알고리즘과 달리 ACER는 그래디언트를 직접 구하고 수정한 뒤 적용하기 때문에 처음 접하는 경우 구현이 까다로울 수 있습니다. 여기서는 헷갈리거나 고민하기 쉬운 몇몇 부분만 짚어 보겠습니다.
$ Q^{ret} $ 를 사용할 때는 $ Q^{ret} $으로 그래디언트가 흐르지 않게 해야 합니다. 이는 DQN 등과 동일한 부분으로, target value이기 때문에 $ Q^{ret} $를 계산하는 경로 자체는 최적화 대상이 아니기 때문에 PyTorch의 detach 함수나 with no_grad() 구문, 또는 TensorFlow의 stop_gradient 함수 등으로 그래디언트 계산을 위한 computational graph에서 제외시켜야 합니다.
또한 계산하는 그래디언트가 목표 함수를 감소시키는 방향이 아니라 증가시키는 방향이기 때문에, 라이브러리에서 기본 제공하는 optimizer를 사용하는 경우 그래디언트에 -1을 곱해서 넘겨줘야 합니다. 보통 loss 함수에 대해 그래디언트를 계산하고 loss 함수를 감소시키는 방향으로 적용시켜야 하기 때문에 주어진 그래디언트의 반대 방향으로 파라미터를 업데이트하기 때문입니다.
주의할 점은 성능 함수를 최대화시키는 방향의 그래디언트를 계산한 후, TRPO를 먼저 적용한 다음 -1을 곱해서 optimizer에 넘겨줘야 한다는 점입니다. 이 순서를 지켜야 적용되는 그래디언트에서 KL Divergence를 증가시키는 성분이 제거되기 때문입니다.
그리고 주어진 output에 대한 특정 input의 그래디언트를 구할 일이 많은데, 이는 PyTorch의 torch.autograd.grad 함수를 사용하거나 TensorFlow의 tf.gradients 함수로 구할 수 있습니다. 이를 잘 활용하면 배치 단위로 그래디언트를 구할 수 있는데, statistics 자체도 각 샘플마다 값이 나오기 때문에 배치 단위의 statistics에 대한 배치 단위의 목표값의 그래디언트를 구하면 GPU를 최대한 활용하며 반복문 없이 그래디언트를 구할 수 있게 됩니다.
트리키한 부분은 Bias Correction term에서 식 (10)에서 $ \mathbb{E}_{a \sim \pi} $ 항 내의 그래디언트를 구하는 것입니다. 한 샘플 당 여러 행동에 대해 그래디언트를 각각 구한 다음, 행동값에 따라 다른 후처리를 해준 뒤 더해야 하기 때문입니다. 한 파라미터를 사용해 계산한 여러 output이 있을 때, 여러 output의 해당 파라미터에 대한 그래디언트를 동시에 계산하면 각 output에 대한 그래디언트가 나오는 것이 아니라 그래디언트의 합이 나오기 때문입니다.
그런데 여기서 우리가 사용하는 후처리는 전부 곱셈 뿐이므로, 행동에 따라 다른 곱해야 하는 factor 들을 미리 구해두고, 해당 factor 들에는 그래디언트가 흐르지 않게 처리한 뒤 미리 곱한 뒤 각 output에 대한 그래디언트를 구하면 각 행동에 대해 그래디언트를 구한 후 후처리를 해 더한 것과 같은 결과를 얻을 수 있습니다.
수식으로 표현하면 다음과 같습니다. $ a $ 를 행동, $ s $ 를 상태, $ x $ 를 파라미터라고 두면
$ \sum_{a \in A} [g(a|s) * \frac{df(a|s)}{dx}]$ 을 구해야 하는데, $ \frac{df(·|s)}{dx} $를 계산하면 $ \frac{df(a|s)}{dx}$ 가 모든 a에 대해 나오는 것이 아니라 $ \sum_{a \in A} [\frac{df(a|s)}{dx}]$ 이 나오기 때문에 $ a $ 에 따라 서로 다른 후처리를 할 수 없는 것입니다.
그래서 원래는 a에 대한 for loop 등 반복문을 사용해 각 행동당 한번씩 torch.autograd.grad를 통해
$ \frac{df(a|s)}{dx}$ 를 각각의 a에 대해 계산하고
거기에 $ g(a|s)$ 을 곱한 것을 다 더하는 식으로 계산해야 합니다.
그런데 후처리가 그냥 곱셈이기 때문에 그래디언트 계산에서 후처리 항을 배제하는 경우, 즉 $\frac{dg(a|s)}{dx} = 0$을 만족시키면 $\frac{d(g(a|s) * f(a|s))}{dx} = g(a|s) * \frac{df(a|s)}{dx}$ 이므로 $ h(a|s) = g(a|s) * f(a|s) $로 정의하면
$$ \frac{dh(·|s)}{dx} = \sum_{a \in A} [\frac{dh(a|s)}{dx}] = \sum_{a \in A} [\frac{d(g(a|s) * f(a|s))}{dx}] = \sum_{a \in A} [g(a|s) * \frac{df(a|s)}{dx}] $ 가 성립해 배치 단위 연산이 가능해지게 됩니다. 실제로 OpenAI의 baselines 구현도 이 방식을 사용하고 있습니다.
========================================