뮤제로(MuZero) 강화학습 알고리즘 - 개요 및 알고리즘 구조 (1)

개요

알파고로 유명한 딥마인드사에서 2020년 즈음 발표한 강화학습 알고리즘인 뮤제로(MuZero)를 리뷰하려 한다. 
논문 제목 : 'Mastering Atari, Go, chess and shogi by planning with a learned model'
supplementary information : https://www.nature.com/articles/s41586-020-03051-4 페이지 아래쪽

뮤제로는 알파제로에서 한단계 더 발전된 강화학습 알고리즘으로, 알파제로는 보드게임(바둑, 체스류)만 학습이 가능했다면 뮤제로는 아타리(atari)게임 및 보드게임 둘다 학습이 가능한 구조이다. 알고리즘 설명에 앞서, 뮤제로 알고리즘의 중요한 키워드인 dynamics에 대한 이해를 돕기 위해 아타리 게임 하나를 소개한다. 아래 링크는 lunar lander라는 게임이다. 링크에서 직접 해볼 수 있다.


해보면 상당히 어렵다. 그리고 어떻게 하는건지 파악하기까지 시간이 꽤나 걸린다. 알 수 있는건 방향키 4개를 쓸 수 있다는 것 뿐이다. 직접 해보면서 무슨 키를 누르면 무슨 일이 일어나는지, 우주선이 어느 방향으로 가는지 감각을 익혀야 한다. 이러한 감각을 dynamics라고 설명할 수 있다.

알고리즘이 게임을 학습하려면 이러한 dynamics를 알고 있어야 한다. 바둑의 경우 바둑판에 돌을 두면 그 다음은 돌이 추가된 모양의 바둑판이 나타날 것이다. 그런데 이걸 알고 있는 것은 사람이 바둑이 뭔지를 대략 이해하고 있기 때문인 것이지 당연한 것은 아니다. 위의 lunar lander 게임의 경우는 어떤 상황에서 어떤 키를 누르면 그 다음이 어떻게 될지 정확히 알기가 힘들다. 사람은 감각적으로 학습하겠지만 그걸 정확히 파악해서 알고리즘에게 알려주는 것은 어려운 일이다. 

뮤제로는 이러한 dynamics를 스스로 학습할 수 있는 model-based RL 알고리즘이다. 때문에 적용 가능 범위가 훨씬 넓어졌다고 할 수 있다. 영상 rate control 관련에 뮤제로를 사용한 연구도 존재한다. 
(알파고~알파제로는 이러한 dynamics 학습 능력이 없는 model-free 알고리즘이다. 그럼에도 학습이 가능했던 이유는 dynamics 정보를 바둑, 체스 simulator가 완벽하게 제공을 했기 때문이다.)

뮤제로 알고리즘 구조

뮤제로 알고리즘의 바탕이 되는 알고리즘은 알파제로 알고리즘이다. 알파제로 알고리즘에 dynamics를 학습할 수 있는 구조 및 이를 뒷받침하는 여러 기법들을 추가한 형태이다. 뮤제로 알고리즘은 크게 3가지의 function(뉴럴 네트워크)에 기반해 동작한다. Representation function, dynamics function, prediction function, 3개의 네트워크가 각자 역할을 하는 MuZero model이 구성된다. 
MuZero model

오른쪽 μ(뮤)가 뮤제로 모델이고, 왼쪽이 그를 구성하는 3개의 네트워크들이다. 위 식은 뮤제로 모델이 o(observation, 실제 화면)들로부터 p(policy), v(value), r(reward)를 예측해낼 수 있고, 이에 따라 미래에 가능한 a(action)들을 계획할 수 있다는 의미이다. Hypothetical future trajectories를 lookahead search(MCTS)로 그려낼 수 있다.

Representation function
가장 먼저, 뮤제로 알고리즘에서 state는 있는 그 자체로 표현되지 않고 hidden state로 변환되어서 계산이 된다. 
(Deepmind official은 아니고 공개된 pseudo code에 기반해서 오픈소스 구현한 werner-duvaud의 github에서 제공한 그림.)

알고리즘의 시작 단계에 이전 32프레임에 해당하는 observation(+action)들을 hidden state로 바꾸는 작업이 필요하고 이를 representation function이 담당한다. Resnet 형태의 평범한 backbone으로 구현이 된다. (논문에서는 convolutional and residual architecture as AlphaZero, 16 residual block, 256 hidden planes). Stacked observation들을 6x6에 맞게 downsample하고 backbone에 inference한다. 함수 표현으로는 아래와 같다.
Representation function of muzero algorithm

Prediction function
Prediction function은 hidden state로부터 p,v (policy, value) prediction을 하는 역할이다. 이 부분은 알파제로와 거의 동일하다. 

Hidden state를 backbone에 통과시키고 policy head, value head를 통해 p,v를 predict한다. 함수 표현으로는 아래와 같다.
Prediction function of muzero algorithm

Dynamics function
Dynamics function은 Hidden state의 state transition과 그에 따른 reward prediction을 담당한다. 
선택된 action이 image(one-hot plane)의 형태로 hidden state의 채널에 concat된다. (Atari의 경우 18개의 action이 가능하고, 이게 6x6의 plane에 적절하게 배치된다.) Concat된 채널들의 채널 수를 하나 줄이고, backbone을 통과시킨다. 이 결과는 다음 hidden state로 쓰이고, 이 값으로부터 reward prediction도 실행된다. 함수 표현으로는 아래와 같다.
Dynamics function of muzero algorithm

이렇게 3개의 뉴럴 네트워크로 구현되는 function을 가지고 기존 observation들로부터 p,v,r을 predict할 수 있게 되었고 이에 기반해 lookahead search(MCTS) 및 self-play를 할 수 있게 된다.
  

Self-play and training cycle

알고리즘을 학습시키기 위해서는 학습시킬 데이터를 모으기 위한 self-play 과정이 필요하다. 아직 완전히 학습되지 않은 알고리즘 스스로 게임을 진행하고, 어떨 때 어떤 행동을 선택했고 어떤 결과를 얻었는지 등의 정보를 기록한다. 게임이 끝나면 저장된 정보를 꺼내와 training에 사용해서 성능을 개선한다. 성능이 개선된 알고리즘을 가지고 다시 self-play를 하고 저장하고 training하는 것을 계속 반복한다. 이러한 사이클을 계속 반복하면 알고리즘은 높은 성능을 가지게 된다.

다음 글에서 self-play, mcts, training등의 과정에 대해 추가로 설명하려 한다.


댓글

가장 많이 본 글

구글 람다(LaMDA)란? - 구글의 언어 모델

알파고 강화학습 원리

버텍스 AI란? - 구글 인공지능 플랫폼

카타고와 바둑 두어보기

뉴럴 네트워크란?

블로그 글 목록

뉴럴 네트워크를 학습시키는 방법