Mastering Atari with Discrete World Models
이산 월드 모델로 Atari 마스터하기
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, Jimmy Ba (2021)
이산 잠재 표현(categorical latent)을 도입한 세계 모델과 상상 속 액터-크리틱 학습을 결합하여, 모델 기반 강화학습으로서 최초로 Atari 벤치마크에서 모델 프리 알고리즘(Rainbow)의 성능에 필적하는 성과를 달성했다.
배경
DreamerV1은 연속 제어(DMC)에서 뛰어난 성과를 보였지만, Atari와 같은 이산 행동 공간과 시각적으로 복잡한 환경에서는 모델 프리 알고리즘에 미치지 못했다. 특히 Atari 55개 게임 벤치마크에서 모델 기반 RL이 모델 프리의 최고 수준(Rainbow DQN)에 도달하지 못하는 것은 해당 분야의 장기적 과제였다. 가우시안 잠재 공간의 표현력과 이산 행동 공간에서의 정책 학습의 한계가 주요 병목이었다.
핵심 아이디어
DreamerV2의 핵심 혁신은 세계 모델의 잠재 상태를 가우시안에서 범주형 분포(categorical distribution)로 전환한 것이다. 32개의 범주형 변수, 각각 32개 클래스를 가지는 잠재 상태는 총 32^32의 조합 가능성을 제공하여 복잡한 환경을 효과적으로 표현한다. 이산 잠재 변수의 그래디언트는 straight-through 추정기로 전파한다. 추가적으로 KL 밸런싱이라는 기법을 도입하여 사전(prior)과 사후(posterior) 분포 간의 KL 발산 학습 속도를 비대칭적으로 조절함으로써, 세계 모델이 더 정보가 풍부한 잠재 표현을 학습하도록 유도한다.
방법론
RSSM의 확률론적 부분을 32x32 범주형 분포로 대체하고, straight-through 그래디언트로 역전파한다. KL 손실을 α:β = 0.8:0.2로 비대칭 분배하여 사후분포가 사전분포를 따라가도록 유도한다. 상상 속에서 Reinforce + straight-through 혼합 그래디언트로 액터를 학습하고, 이산 행동 공간에서는 straight-through을 통한 원-핫 행동 샘플링을 사용한다.
주요 결과
Atari 55개 게임에서 인간 정규화 중앙값 점수 209.2%를 달성하여, 모델 프리 알고리즘 Rainbow(223%)에 근접하며 IQN(218%)과 동등한 수준에 도달했다. 이는 모델 기반 RL이 Atari에서 모델 프리와 경쟁적 성능을 달성한 최초의 사례이다. 동시에 DMC 연속 제어 벤치마크에서도 DreamerV1을 상회했다.
임팩트
DreamerV2는 모델 기반 RL의 가능성에 대한 인식을 근본적으로 전환하여, '모델 기반은 모델 프리에 미치지 못한다'는 통념을 깨뜨렸다. 이산 잠재 표현은 이후 DreamerV3, IRIS 등에서 채택되었으며, 범주형 세계 모델은 비디오 예측과 세계 시뮬레이터 연구에도 영향을 미쳤다.