Scalable trust-region method for deep reinforcement learning using Kronecker-factored approximation 논문의 세미나가 있었다.


전체적인 내용은 actor-critic 방식에 피셔 정보 행렬을 근사하기 위한 K-FAC 최적화를 적용한 second-order method를 통해 훈련속도와 sample efficiency를 올리며 computation cost를 줄이는 ACKTR라는 새로운 강화학습 방식의 제시였다.


의사결정 모델이나 네트워크 구조 등의 수정 등이 아닌 단지 최적화 과정에서의 새로운 개념 제시만으로도 얼마나 획기적으로 강화학습의 성능을 향상시킬수 있는지 잘 드러난 논문이었다. 일반적인 형태의 논문 리뷰는 아니지만, 가장 중요하지만 가장 이해하기 힘들었던 최적화 과정이 어떤 식으로 제시됐는지에 대해 중점적으로 되짚고 당시의 세미나 발표를 돌아보며 작성하고자 한다.


Natural gradient using Kronecker-factored approximation의 개념은 다음과 같다. Steepest descent(second-order) 방식에선 non-convex형태의 손실함수 $\mathcal{J}(\theta)$를 업데이트하기 위해 $\mathcal{J}(\theta+\Delta\theta)$를 최소화하는 방향으로 $\Delta\theta$를 업데이트한다. 이때 사용되는 norm은 $\lVert{x}\rVert_B=(x^TBx)^{\frac{1}{2}}$ 과 같은 형태에 $\lVert\Delta\theta\rVert_B<1$ 라는 제약 조건을 가진다.


이 때 파라미터의 업데이트 값은 다음과 같은 비례관계를 따른다.


$\Delta\theta\propto -B^{-1}\nabla_\theta\mathcal{J}$


만약 $B$가 단위행렬 $I$라면, 위에서 설명한 norm은 Euclidean norm방식과 같으며, $\nabla_\theta\mathcal{J}$는 일반적인 gradient descent의 standard gradient와 같다. 그러나 단위행렬일 경우 업데이트는 온전히 파라미터 $\theta$에 영향을 받을 수 밖에 없다. 이는 모델의 파라미터 선택에 임의성을 부여함으로 적절하지 않으며, 최적화 경로에 영향을 주게 된다.


natural gradient(second-order) 방식은 이 $B$ 행렬에 KL divergence의 2차 미분 근사치인 피셔 정보 행렬 $F$를 사용하며, 이는 파라미터 $\theta$와 독립성이 보장된다.


이 부분에 대한 이해가 부족한 상태에서 교수님이 따로 첨삭을 해주셨는데, 피셔 정보 행렬은 (모든 파라미터) x (모든 파라미터) 크기의 행렬로 정보의 분포를 유추하는 개념이며, 이는 파라미터의 독립성이 보장되지 않는 gradient descent와 다르게 파라미터별로 업데이트되는 정도가 달라지게 되는 효과를 볼 수 있다고 한다.


문제는 파라미터 업데이트가 $\Delta\theta\propto -F^{-1}\nabla_\theta\mathcal{J}$와 같이 이루어진다면, 피셔 정보 행렬의 역행렬을 구해야 하고, 피셔 정보 행렬의 거대한 크기와 역행렬을 구하는 과정을 고려했을 때 이는 엄청난 computation cost가 요구된다는 것이다.(단순 미분 횟수로만 생각해 봤을 때도, 1차 미분만 필요한 first-order 방식과 2차 미분까지 필요한 second-order 방식이 속도 차이가 심한 것은 당연하다) 지금까지 기계학습에서 second-order method가 사용되지 않았던 이유이다. 여기서 K-FAC(Kronecker-factored approximate curvature)이 등장한다.


앞서 피셔 정보 행렬이 분포의 차이를 나타내는 KL divergence의 2차 미분과 같다고 했다. 이를 분포와 파라미터의 관계로 나타내면 다음과 같다.


$F_\ell=\mathbb{E}[\mathrm{vec}{\nabla_WL}\mathrm{vec}{\nabla_WL}^\mathsf{T}]$


$L={\log}p(y\lvert x)$은 신경망에서 입력에 대한 출력의 분포에 대한 log-likelihood, $W{\in}\mathbb{R}^{C_{out}\times{C_{in}}}$는 가중치 행렬을 뜻하고 $\ell$은 레이어를 뜻한다.


입력 활성화 함수가 $a \in \mathbb{R}^{C_{in}}$, 활성화 적용 행렬이 $s=Wa$일때, 가중치의 기울기는 $\nabla_WL = (\nabla_sL)a^\mathsf{T}$로 주어진다. 이때 피셔 정보 행렬의 근사를 다음과 같이 유도할 수 있다.


$F_\ell=\mathbb{E}[\mathrm{vec}{\nabla_WL}\mathrm{vec}{\nabla_WL}^\mathsf{T}]=\mathbb{E}[aa^{\mathsf{T}}\bigotimes\nabla_sL(\nabla_sL)^{\mathsf{T}} \ \quad \ \approx\mathbb{E}[aa^\mathsf{T}]\bigotimes\mathbb{E}[\nabla_sL(\nabla_sL)^{\mathsf{T}}]:=A\bigotimes S:=\hat F_\ell$


하나의 피셔 정보 행렬이 $A\bigotimes S$라는 두 행렬의 크로네커 곱으로 나눠졌다. 여기서 크로네커 곱의 기본 성질인 $(P\bigotimes Q)^{-1}=P^{-1}\bigotimes Q^{-1}$, $(P \bigotimes Q)\mathrm{vec}(T)=PTQ^{\mathsf{T}}$를 적용할 수 있다.


$\mathrm{vec}(\Delta W) = \hat F_\ell^{-1}\mathrm{vec}{\nabla_W\mathcal{J}} = \mathrm{vec}(A^{-1}\nabla_W\mathcal{J}S^{-1})$


결과적으로 하나의 거대한 행렬의 역행렬을 구하는 연산에서 두개의 작은 행렬의 역행렬을 구하는 연산으로 식이 간략화됐고, 이는 오직 $W$의 크기와 비슷한 수준의 computation cost만 요구된다.


그렇다면 이 Natural gradient를 어떻게 강화학습에 적용할 수 있을까? 강화학습에서 policy의 log-likelihood인 $\log{\pi}_\theta (a\lvert s_t)$가 $L=\log p(y\lvert x)$과 동일하고, 파라미터 $\theta$가 가중치 행렬 $W$와 동일하다고 보면 다음과 같은 식을 얻을 수 있다.


$F=\mathbb{E}_{p(\tau)}[\nabla_\theta \log\pi(a_t{\lvert}s_t)(\nabla_\theta \log\pi(a_t{\lvert}s_t))^{\mathsf{T}}]$


$p(s_0)=\prod_{t=0}^{T}\pi(a_t\lvert s_t)p(s_{t+1}\lvert s_t,a_t)$로 주어졌을 때, $p(\tau)$는 궤적의 분포를 나타내고, 실제로 훈련중에 수집된 궤적에 대한 기대치를 근사한다.


그러나 위 식은 policy의 분포에 대한 내용이다. 여기서 critic에 natural gradient를 적용하기 위해 이 문제를 최소제곱 근사 문제로 생각해볼 수 있다. 최소제곱 근사 문제에서 second-order알고리즘은 일반적으로 가우스-뉴턴 행렬 $G:=\mathbb{E}[J^TJ]$로 곡률을 근사한다.($J$는 자코비안 행렬이다.) 링크된 논문에 따르면, 가우스-뉴턴 행렬은 가우시안 관찰 모델 하에 피셔 정보 행렬과 동일하다. 이는 K-FAC을 critic에 적용할 수 있는 근거가 된다. critic의 출력 $v$가 $\sigma = 1$인 정규 분포 $p(v\lvert{s_t})\sim\mathcal{N}(v;V(s_t),\sigma^2)$를 따른다고 가정하면, ciritc에 적용되는 피셔 정보 행렬이 정규 출력 분포를 따르는 것으로 정의할 수 있다.


actor-critic의 출력은 각각이 완전히 다른 네트워크가 아닌 하나의 lower-layer가 각각 actor와 critic으로 나누어지는 형식이다. actor와 critic의 독립성을 유지하는 동시에 둘 간의 불일치를 방지한다. 즉, 둘 사이의 분포가 일치하므로, $p(a,v\lvert{s})=\pi(a\lvert{s})p(v\lvert{s})$에 따라 피셔 정보 행렬을 구성할 수 있다. 최종적으로 피셔 정보 행렬에 K-FAC을 근사하게 되면 다음과 같은 식을 얻을 수 있다.


$\mathbb{E}_{p(\tau)}[\nabla\log p(a, v{\lvert}s)\nabla\log p(a,v{\lvert}s)^T]$


추가적으로 이 방식에서 step-size는 신뢰 구간 방식으로 자동으로 결정된다.


$\min(\eta_\mathrm{max}, \sqrt{\frac{2\delta}{\Delta\theta^{\mathsf{T}}\hat F\Delta\theta}})$


$\eta_\mathrm{max}$는 러닝 레이트 최대값, $\delta$는 신뢰구간 범위 하이퍼 파라미터이다. actor와 critic간의 불일치 발생시엔 다양한 세트의 $\eta_\mathrm{max}$와 $\delta$를 튜닝한다. 그러나 앞서 말했듯이 동일한 lower-layer의 공유 덕분에 둘 간의 일치성이 보장되기에 단 하나의 세트만 튜닝하면 충분하다.


해당 최적화 기법을 이용해 이산적인 환경(Atari)과 연속적인 환경(MuJoCo)에서 다양한 성능 향상 비교 실험을 진행했는데, 종합적으로 다음과 같은 성과들을 얻을 수 있었다.


  • sample efficiency 성능 A2C, TRPO방식에 비해 2~3배 향상 (batch-size를 키워도 훈련이 효과적으로 진행됨)
  • computation cost 획기적으로 감소 (second-order 방식임에도 first-order 방식을 적용한 A2C에 비해 고작 25%정도의 비용 증가)
  • 연속적인 실험 환경에서 입력을 raw-pixel형태로 줘도 성공적으로 훈련이 가능


발표 때 기초적인 부분에 대한 설명과 슬라이드가 생략된 것과, 지금 이 글을 쓰고있는 순간보다 훨씬 부족한 상태의 이해도로 발표를 진행한 것에 대한 아쉬움이 많이 남는 세미나였다. 확실히 지금까지 봤던 논문들에 비해 어려워서 발표가 우왕좌왕한 부분이 있었던 것 같다.


현재 Tron게임의 강화학습 에이전트 연구에 이 ACKTR 방식을 적용해보는 중인데, DQN, A3C등으로도 유의미한 성능 향상을 끌어내지 못했으나 이 ACKTR방식은 강력한 Minimax 상대로도 비등비등한 승률을 기록하는 정도로 성능이 잘 나오는 중이다. 확실히 좋은 방법을 제시한 좋은 논문이다.