MuZero: 범용 인공지능에 한발 다가서다

김성준
CURG
Published in
15 min readJan 30, 2021

김성준, Virtual Machine & Optimization Lab in Seoul Nat’l University

본 포스팅에서는 Schrittwieser, Julian, et al. “Mastering atari, go, chess and shogi by planning with a learned model.” Nature 588.7839 (2020): 604–609.[1]을 다룬다.

뮤제로는 게임의 규칙을 스스로 학습하며, 다양한 도메인에 적용 가능하다

들어가며
몇 년 전 겨울, 대학 동아리 지인들과 3박4일 여행을 갔을 때의 이야기다. 다같이 숙소에 짐을 풀고 TV를 보며 잡담을 나누는데, 채널을 돌리던 와중 우연히 바둑TV 채널을 시청하게 되었다. 다같이 대국을 바라보며 즐거운 시간을 보냈는데, 화면을 자세히 보니 우측 상단에 생소한 막대 그래프가 표시되어 있었다. 알고 보니 여러 종류의 바둑 인공지능 프로그램이 두 기사의 승률을 실시간으로 예측하여 그래프로 표현한 것이었다. 구글 딥마인드의 알파고(AlphaGo)가 바둑을 평정하며 세상을 놀라게 하기 전까지는 경험해보지 못했던, 매우 낯선 모습이었다.

이처럼 알파고의 등장은 인공지능 역사에 한 획을 그은 사건으로 남아 있다. 그 파급력 때문인지 이세돌 九단과의 대국 이후 어언 5년이 지난 지금도 많은 사람들은 인공지능 하면 알파고를 먼저 떠올린다.

구글 딥마인드의 알파고가 바둑을 정복한 뒤로도, 딥마인드는 기존의 알파고 알고리즘을 계승하는 연구를 꾸준히 진행해왔다. 가장 먼저 등장한 알파고 제로(AlphaGo Zero)는 사전에 인간의 바둑 기보를 바탕으로 알파고 신경망을 학습시키지 않고도 셀프 플레이만으로 점차 강해지는 바둑 인공지능을 가능케 했다. ‘제로’라는 네이밍에서 알 수 있듯이, 게임의 규칙만 주어진다면 무에서부터 시작하여 어떠한 외부 지식 없이 스스로 전략을 창조해나가는 인공지능이 탄생한 것이다. 이를 약간 변형하여, 바둑뿐만 아니라 체스 및 쇼기로도 도메인을 확장한 알파제로(AlphaZero)는 체스의 Stockfish, 쇼기의 Elmo 등 기존 인공지능의 성능을 압도하는 등 모델이 명확한 게임의 전 영역을 제패하였다. 여기에 스타크래프트2 그랜드 마스터에 오른 알파스타(AlphaStar)는 덤이다.

하지만 딥마인드는 여기서 만족하지 않았다. 기존에 초점을 맞추던 보드게임에서 시야를 넓혀, 그간 모델 기반 접근법으로는 저조한 성능을 보여왔던 아타리 게임에 도전장을 내민 것이다. 그들의 새로운 시도는 성공적이었으며, OpenAI Gym에서 제공하는 아타리 게임 도메인에서 현존하는 인공지능의 성능을 뛰어넘었다. 더욱 놀라운 점은, 아타리에서 SOTA(state-of-the-art)급의 성능을 보여주면서, 동일한 인공지능을 바둑, 체스 및 쇼기에 그대로 적용하더라도 알파제로와 비견되거나 그 이상의 성능을 나타낸다는 점이다. 이처럼 지난해 12월 네이처에 발표된 인공지능 알고리즘이 서로 별개로 다루어져왔던 도메인을 하나로 통합하였으니, 바로 이번 글에서 소개할 뮤제로 알고리즘이다.

뮤제로는 각 도메인에서 기존 인공지능 대비 뛰어난 성능을 보여준다

그러면 본격적으로 뮤제로가 어떻게 동작하는지, 선행 연구였던 알파제로와는 어떤 차별점을 가지는가를 살펴보자.

뮤제로는 어떻게 동작하는가?
뮤제로 알고리즘은 크게 세 단계를 반복적으로 수행함으로써 내부 신경망을 학습시킨다.

  • 환경으로부터 관측이 주어졌을 때, 플래닝을 통해 정책(policy)을 얻는다.
  • 정책으로부터 다음 행동을 샘플링 및 수행한다.
  • 에피소드에서 얻어진 경로(trajectory)로부터 내부 신경망을 학습시킨다.

이로부터 학습된 뮤제로 신경망은 플래닝에 필요한 환경의 세 요소를 모델링하는데 활용된다.

  • 가치(value): 현재 상태가 얼마나 좋은가?
  • 정책(policy): 무슨 행동이 가장 이로울까?
  • 보상(reward): 최근 한 행동이 얼마나 좋았을까?

그러면 뮤제로 알고리즘의 각 단계를 하나씩 찬찬히 살펴보도록 하자.

  1. 플래닝(Planning)

플래닝이란, 현재 상태가 주어졌을 때 어떤 행동을 취할 것인가에 대한 정책을 찾기 위한 과정이라고 이해할 수 있다. 알파고 계보에서와 마찬가지로, 뮤제로 또한 몬테카를로 트리 탐색(MCTS, Monte-Carlo Tree Search)을 반복함으로써 플래닝을 수행한다. 몬테카를로 트리의 각 노드는 특정 상태(state)를 의미하며, 엣지는 특정 행동(action)에 대응된다. 임의의 부모 노드에서 엣지를 따라 자식 노드로 향하는 과정은, 이전 상태에서 어떤 행동에 의해 다음 상태로 상태 전이(state transition)가 일어나는 과정으로 볼 수 있다.

뮤제로 알고리즘에서는 몬테카를로 트리 탐색을 통해 플래닝을 수행한다

뮤제로에서의 플래닝에는 세 가지 함수가 활용된다. 이들 중 예측 함수 및 동역학 함수는 신경망으로 이루어져 있으며, 학습이 진행되면서 내부 파라미터가 점차 튜닝된다.

1–1. 표현 함수 (representation function)
먼저, 환경으로부터 관측되는 정보를 신경망의 입력으로 사용할 수 없으므로, 이를 적절히 인코딩하여 현재 상태로 변환해주는 표현 함수 h를 거친다. 이렇게 얻어진 상태는 MCTS의 루트 노드로 사용된다.

representation function

1–2. 예측 함수 (prediction function)
트리의 각 노드는 특정 상태를 표현하는데, 이를 예측 함수 f의 입력으로 넣으면 1) 현재 이 상태가 얼마만큼의 가치를 지니는가를 추정하는 가치 v와, 2) 어떤 행동을 선택하는 것이 가장 바람직할 것인가를 추정하는 정책 p를 출력으로 내보낸다.

prediction function

여기까지는 선행 연구인 알파제로에서의 플래닝과 별다른 차이가 없다. 하지만 뮤제로에서는 상태 전이에 있어 동역학 함수를 도입함으로써 알파제로와 차별점을 두었으며, 이로부터 보다 다양한 도메인에 적용 가능한 일반화된 모델을 생성해낼 수 있었다.

1–3. 동역학 함수 (dynamics function)
현재 상태와 수행할 행동이 주어졌을 때, 기존 알파제로는 모델에 기반하여 다음 상태를 계산해내는 일종의 시뮬레이터로서 동작한다. 바둑을 예로 들면, 알파제로는 특정 공간에 착수하면 내 색상의 돌이 해당 위치에 생성된다는 규칙을 알고 있으며, 일련의 규칙에 따라 다음 상태가 어떻게 정의되는가를 내부적으로 정확히 계산할 수 있다는 의미다.

뮤제로에서는 이러한 결정론적 시뮬레이터 방식을 사용하지 않는다. 대신 이전 상태 s_(t-1)와 수행할 행동 a_t이 입력으로 주어졌을 때, 이로부터 1) 다음 상태 s_t 및 2) 행동에 대한 보상 r_t을 추론하는 동역학 함수 g를 사용한다.

dynamics function

다음 상태가 신경망의 출력으로부터 얻어지기 때문에, 시뮬레이터와는 달리 얻어진 상태 정보로부터 실제 환경에서의 관측을 복원해낼 수 없다. 언뜻 보면 정확한 상태 정보를 뭉개버리는 등 학습에 악영향을 미치는 것처럼 보인다. 하지만 반대로 생각하면, 특정 상태에서 주어지는 다양한 정보 중 중요한 정보만을 추출해내 모델링할 수 있다는 의미이기도 하며, 바로 그렇기에 시각적으로 복잡한 도메인인 아타리에서 성공적으로 동작할 수 있었다. (벽돌깨기 게임에서, 과연 벽돌의 색상이나 주변 이펙트가 학습에 중요할까?) 더 중요한 특징으로, 규칙이 명확한 보드게임과는 달리, 현실 세계에서의 많은 문제는 정확한 규칙 및 모델을 찾기가 불가능에 가깝기 때문에, 동역학 함수의 도입은 보다 일반화된 알고리즘으로의 도약을 위해 필수적이라고 볼 수 있다.

위의 세 가지 함수가 정의되었다면 MCTS에 기반한 시뮬레이션을 반복적으로 수행할 수 있다. 각 시뮬레이션마다 선택, 확장 및 백업이 순차적으로 이루어진다.

  • 선택 (Selection)

자식 노드가 여러 개 있을 때, 어떤 노드를 방문할 것인가는 PUCT(probabilistic upper confidence tree)에 따라 결정된다. PUCT는 착취 항(exploitation term)과 탐사 항(exploration term)으로 구분할 수 있다. 착취 항은 가치가 높다고 판단되는 노드를 방문하게끔 유도하며, 탐사 항은 방문 횟수가 상대적으로 적은 노드를 방문하게끔 유도한다. 학습 초기에는 루트 노드의 정책이 한쪽으로 치우쳐질 우려가 있으므로, 루트 노드의 정책값에 디리클레 노이즈(Dirichlet noise)를 추가함으로써 초기 탐사를 장려한다. 이러한 자식 노드의 선택은 리프 노드에 도달할 때까지 반복된다.

PUCT는 착취 항과 탐사 항의 조합으로 구성된다
  • 확장 (Expansion)

리프 노드, 리프 노드의 부모 노드 및 이 둘을 잇는 엣지를 입력으로 하여, 예측 함수 및 동역학 함수가 MCTS 확장을 수행한다. 확장 결과 리프 노드 아래 자식 노드들이 추가되며, 이들이 새로운 리프 노드가 된다.

  • 백업 (Backup)

리프 노드부터 루트 노드까지 거슬러 올라가며 백업을 수행한다. 백업 과정에는 1) 노드 방문 횟수 1만큼 증가, 2) 노드의 평균 가치 조정, 3) 엣지에서의 보상 조정 등이 포함된다.

시뮬레이션을 여러번 수행하여 루트 노드에서의 자식 노드 방문 횟수가 어느 정도 누적되면, 이는 다음 단계인 행동하기 단계에서 행동을 샘플링하는데 활용된다.

2. 행동하기 (acting)

플래닝 단계에서는 MCTS를 수행했을 때의 루트 노드의 자식 노드 방문 횟수에 대한 정보를 얻었다. 행동하기 단계에서는 이를 기반으로 특정 관측에 대응되는 정책, 즉 각 행동을 선택할 확률에 대한 분포를 얻어내며, 얻어진 확률 분포로부터 실제 환경에서 수행할 행동을 샘플링 및 이를 수행한다.

플래닝으로부터 얻어진 정보를 활용하여 다음 행동을 샘플링한다

다음 행동을 샘플링하는 과정에서, 플래닝 단계에서의 자식 노드 방문 횟수뿐만 아니라 지식 증류에서도 소개되었던 온도(temperature)라는 파라미터가 추가적으로 사용된다. 온도가 무한대로 설정되면 확률 분포는 균등 분포를 따르며, 온도가 0으로 수렴하면 방문 횟수가 가장 많은 노드에 대응되는 행동이 선택되게끔 설계되어 있다.

각 자식 노드의 방문 횟수와 온도로부터 정책값이 계산된다.

논문에 기재된 메소드를 살펴보면, 학습 초기에는 온도를 1.0으로 설정했다가, 중간에 0.5를 거쳐 나중에는 0.25까지 온도를 낮춘다. 이는 학습이 진행될수록 다음 행동을 더욱 그리디하게 선택하도록 유도함을 의미한다.

이러한 샘플링은 각 에피소드가 끝날 때까지 반복적으로 수행되며, 한 에피소드가 끝나면 해당 에피소드에서의 경로 데이터를 리플레이 버퍼에 저장하여 추후 학습 단계에 사용한다.

3. 신경망 학습 (training)

이전 경로를 unroll하면서 각 신경망에 대한 학습을 수행한다.

각 에피소드로부터 얻어진 경로에는 매 스텝에서의 환경으로부터의 관측값, 선택한 행동 및 행동에 따른 보상 등 유의미한 데이터가 담겨있다. 앞서 플래닝에 필요한 환경의 세 요소를 모델링하기 위하여 뮤제로 신경망을 학습시킨다는 말을 기억하는가? 경로에 저장된 값들은 이들 요소에 대한 타겟으로 사용된다.

  • 예측 함수로부터 추론된 정책 ↔ MCTS로부터 얻어진 정책
  • 예측 함수로부터 추론된 가치 ↔ 감가율 γ를 적용한 n-step 보상의 총합 (=리턴)
  • 동역학 함수로부터 추론된 보상 ↔ 환경으로부터 관측된 보상

이들 세 요소에 대한 손실 함수와 L2 정규화를 하나로 묶으면 전체 손실 함수에 대한 수식을 유도해낼 수 있다.

세 종류의 손실 함수 및 L2 정규화로부터 전체 손실 함수가 계산된다

주목할 만한 사항으로, 가치 및 보상에 대한 손실 함수로 보드게임 도메인에서는 평균 제곱 오차를, 아타리 도메인에서는 크로스 엔트로피 오차를 사용했을 때 학습이 안정적이었다고 한다. 정책에 대한 손실 함수는 두 도메인 모두 크로스 엔트로피 오차를 사용했다고 한다. 그 밖에도, 보드게임 도메인에서는 감가율을 1로 설정함으로써 보드게임 승패에 따른 리턴이 온전히 반영될 수 있게끔 설계하였다.

실험 결과
뮤제로 알고리즘이 보드게임과 아타리 도메인을 아우르는 만큼, 이들 각각에 대한 성능 평가가 이루어졌다. 먼저 바둑 도메인에 대한 결과를 선행 연구였던 알파제로와 비교했을 때, 20개의 ResNet 블록을 사용하는 알파제로 대비 16개의 ResNet 블록만을 사용한 뮤제로의 성능이 오히려 더 뛰어났다고 한다. 저자들은 뮤제로가 MCTS에서의 연산을 캐싱할 수 있었으며(caching its computation), 동역학 모델의 추가적인 응용이 포지션에 대한 더 깊은 이해를 동반했다고 해석한다.

(a) 바둑 도메인에서의 뮤제로와 알파제로 ELO, (b) 아타리 도메인에서의 뮤제로와 R2D2 평균 보상

이전까지 아타리 도메인에서 SOTA였던 모델 프리 방식의 R2D2 알고리즘과 비교했을 때, 57개의 아타리 중 42개에 대하여 모델 기반 방식으로 동작하는 뮤제로가 더 높은 성능을 보였다.

아타리 도메인에서 뮤제로가 타 알고리즘의 성능을 압도했다

기존 실험에 추가로, 딥마인드에서는 적은 양의 데이터만으로 모델 학습이 가능케 하는 뮤제로 재분석(MuZero Reanalyze) 방법론을 도입하여 그 성능을 평가하였다. 기존 실험에서는 모델 학습에 200억 프레임을 사용하였는데, 전체 데이터의 100분의 1 수준인 2억 프레임만을 전체 학습 과정의 10%로 할당하고, 나머지 90%는 과거에 이미 사용되었던 경로 데이터를 재사용하는 방식이다. 놀랍게도 하이퍼 파라미터만 약간 수정해주면 적은 학습 데이터만으로도 오버피팅 없이 원활한 모델 학습이 가능했으며 타 알고리즘들의 성능을 압도했다고 한다.

마지막으로, Ms. Pac-Man이라는 도메인에서 뮤제로의 목적 함수를 모델 프리 Q-learning 기반의 목적 함수로 대체하여 분석한 결과 뮤제로의 평균 보상이 월등히 높음을 확인할 수 있었다고 한다. R2D2는 Q-learning에 기반을 두기 때문에 베이스라인으로 사용되었다. 그 밖에도, 학습 단계에서 매 행동마다 가능한 행동 공간보다 적은 횟수의 시뮬레이션(6 또는 7)만을 수행하는 상황에서도 효과적으로 정책을 수립하는 등 충분히 괜찮은 성능을 보인다는 점이 주목할 만하다.

(c) Ms. Pac-Man 도메인에서의 뮤제로와 Q-learning, R2D2 평균 보상, (d) 학습 단계에서 매 행동마다 수행된 시뮬레이션 횟수에 따른 평균 보상

결론
본 포스팅에서 다룬 뮤제로는 상당히 폭넓은 도메인에 적용 가능한 범용 알고리즘이며, 그 성능 또한 기존 SOTA를 뛰어넘었다. 본문에서 다루었던 도메인 이외에도, 뮤제로는 새로운 동영상 압축 기술의 개발과 YouTube의 비용을 절감할 수 있는 새로운 동영상 인코딩 방법을 찾는 데 이미 실용화되고 있다고 하며[2], 실제로 약 5% 정도의 비디오 압축률 향상이 있었다고 한다. 이전 포스팅에서 다루었던 알파폴드(AlphaFold)가 단백질 접힘 문제에만 초점을 맞추었다면, 뮤제로는 보다 범용적으로 활용될 수 있는 인공지능이라는 점에서 차별점이 있다.

알파고의 역사를 살펴보면, 시간이 지날수록 다룰 수 있는 도메인은 확장되면서 내부적으로 사용되는 알고리즘은 점차 간결해진다는 느낌을 받는다. 그러면서도 더 적은 학습 데이터만으로 더 빠른 시간 내에 더 나은 성능을 지닌 에이전트를 생성해낸다. 또다른 범용 인공지능으로의 발전 가능성으로 주목받는 OpenAI의 GPT-3에서도 풀고자 하는 문제에 대해 몇 개의 예시만 보고 태스크에 적응하여 문제를 푸는 few shot learning만으로 높은 성능을 얻어냈다는 점을 생각해보면, 범용 인공지능 개발에 있어 양질의 데이터 및 알고리즘의 정교함보다는 인간의 직관을 설계할 수 있는 직관(?)이 더 중요할지도 모른다는 것이 필자의 개인적인 생각이다.

--

--