안녕하세요. 오랜만에 World Model로 다시 돌아왔습니다. 이전에 Dreamder V1을 리뷰해봤는데요. Dreamder V3의 경우, 그 구조를 유지하면서도 굉장히 다양한 도메인에서 높은 성능을 유지하였습니다. 그럼 시작해보겠습니다.
<Introduction>
강화학습은 Go, Dota, LLM 개선 등에는 쓰였지만, 실제로는 영역이 바뀌면 알고리즘 설정을 다시 맞추어야 한다는 문제가 크다는 것을 논문은 시작하는데요. PPO는 널리 쓰이는 표준 알고리즘이지만, continuous control, discrete action, sparse reward, image input, spatial environment, board game 등에서는 task별 전문화된 알고리즘이 자주 사용됩니다.
즉, “새로운 과제마다 전문가가 알고리즘을 다시 튜닝해야 한다면 강화학습의 실제 적용이 어렵다” 는 점이 문제였는데요. 이 문제를 해결하기 위해 본 연구에서는 고정 하이퍼파라미터로 여러 영역을 학습하는 DreamerV3를 제안하였습니다.
실제로 Figure 1에서 Dreamer가 특정 영역 전용 알고리즘이 아니라 여러 영역에서 작동한다는 것을 보여줍니다. Atari, ProcGen, DMLab, Minecraft 등에서 Dreamer와 비교 알고리즘의 성능을 요약한 것을 볼 수 있는데, Dreamer가 고정된 하이퍼파라미터로로 PPO 및 다수 전문화된 알고리즘보다 높은 성능을 기록한 것을 볼 수 있습니다.
Figure 2에서는 본 연구에서 얼마나 다양한 도메인을 가져갔는지 예시 화면을 볼 수 있습니다. Figure를 보시게되면 실험 영역이 굉장히 다양하다는 것 뿐만 아니라, Dreamer가 로봇 제어, 비디오게임, 3D 미로, Minecraft 같은 서로 다른 시각 구조를 가진 환경에서 평가되는 것 또한 확인할 수 있습니다.
<Learning algorithm>
그럼, 본 논문의 방법론에 해당하는 Learning algorithm에 대해서 설명해드리고자 합니다. DreamerV3는 3개의 신경망으로 구성되는데요.
- 세계 모델(world model)
- 비평가(critic)
- 행위자(actor)
이렇게 3개의 신경망으로 구성됩니다.
세계 모델은 가능한 행동의 결과를 예측하고, 비평가는 그 결과의 가치를 판단하며, 행위자는 가장 가치 있는 결과로 가는 행동을 선택하게 됩니다. 세 구성 요소는 환경 상호작용 중 replated experience에서 동시에 학습하게 됩니다.
<전체 학습 파이프라인>
각 신경망에 대한 설명에 들어가기 전에 전체 학습 파이프라인에 대해서 간단하게 설명하고자 합니다. Dreamer는 실제 환경에서 얻은 겅험을 replayu buffer에 저장하고, 그 경험을 다시 꺼내 세계 모델을 학습합니다. 세계 모델은 실제 관측을 latent stste로 압축하고, 그 state 안에서 미래를 예측합니다. 그 다음 행위자와 비평가는 실제 환경이 아니라 세계 모델이 예측한 상상 속의 궤적 위에서 학습을 진행하는데요.
이를 쉽게 말하면, 에이전트는 실제 환경에서 경험을 조금씩 모으고, 그 경험으로 머리속 연습장을 만들고, 그 이후에는 그 연습장에서 여러 미래를 상상해보고, 어떤 행동이 더 좋은지를 배운다고 보시면 되겠습니다.
학습은 크게 네 단계로 볼 수 있다.
- 환경에서 관측, 행동, 보상, 계속 여부를 수집한다.
- replay buffer에서 sequence batch를 꺼내 RSSM(이는 뒤에서 설명하도록 하겠습니다)을 학습한다.
- 재생된 입력의 표현에서 시작해 세계 모델과 행위자가 상상 궤적을 만든다.
- 비평가는 lambda-return을 맞히고, 행위자는 정규화된 반환과 엔트로피 항을 이용해 행동 분포를 개선한다.
이렇게 보니 얼추 학습이 어떻게 진행되는지 이해가 되실까요? 그럼 세부적으로 어떻게 각 신경망이 구성되는지, 학습하는지 보도록 하겠습니다.
<World model learning>
세계 모델은 센서 입력을 압축된 표현으로 바꾸고, 가능한 행동에 따른 미래 표현과 보상을 예측합니다. 이를 Recurrent State-Space Model(RSSM)으로 구현하였는데요. 인코더는 입력 $x_t$를 확률적 표현 $z_t$로 바꾸고, sequence model은 과거 상태와 행동 $a_{t-1}$로 순환 상태 $h_t$를 만듭니다. 그리고 $h_t$와 $z_t$를 합친 모델 상태에서 입력, 보상, 에피소드 계속 여부를 예측합니다.
여기서 중요한 점은 세계 모델이 단순히 보상만 맞히는 것이 아니라 input reconstruction도 함께 수행한다는 점입니다. 표현 $z_t$가 과제 보상에만 맞춘 정보가 아니라, 환경을 이해하는 데 필요한 풍부한 정보를 담도록 만들게 합니다.
Figure 4에서는 DMLab 미로와 사족 보행 로봇에서 5개 context image를 보고 45 frame 뒤까지 예측한 결과를 보여줍니다. 모델은 중간 실제 이미지를 보지 않고 주어진 행동 시퀀스만으로 미래 화면을 예측하는 것을 볼 수 있는데, 세계 모델이 단기 픽셀만 복원 하는 것이 아니라 긴 시간의 환경 구조를 학습한다는 점을 확인할 수 있습니다.
수식 1의 경우, RSSM의 순차 모델(sequence model), 인코더(encoder), 동역학 예측기(dynamics predictor), 보상 예측기(reward predictor), 계속 예측기(continue predictor), 디코더(decoder)를 정의합니다.
$h_t$는 순환 상태, $z_t$는 확률적 잠재 표현, $a_t$는 행동, $r_t$는 보상, $c_t$는 에피소드 계속 여부, $x_t$는 관측 입력을 의미합니다. 모델 상태는 $h_t$와 $z_t$의 결합으로 구성되고, 이 상태가 행위자와 비평가의 입력이 됩니다. 세계 모델이 어떤 정보를 기억하고 무엇을 예측하는지 명확히 위해서 해당 수식이 사용되게 됩니다.
해당 수식은 세계 모델 전체 loss를 정의하는데요. $L_{\mathrm{pred}}$는 예측 loss, $L_{\mathrm{dyn}}$은 동역학 loss, $L_{\mathrm{rep}}$은 표현 loss를 의미합니다. 본 연구에서는 $\beta_{\mathrm{pred}}=1, \beta_{\mathrm{dyn}}=1, \beta_{\mathrm{rep}}=0.1$을 사용했다고 합니다. 세계 모델 파라미터 $\phi$는 sqeuence batch 전체에서 end-to-end로 최적화되며, 입력 재구성, 미래 표현 예측, 예측 가능한 표현 학습을 동시에 하기 위해 해당 수식을 사용합니다.
해당 수식은 예측 loss, 동역학 loss, 표현 loss의 구체적 형태를 정의한다. 예측 loss는 입력,보상, 계속 여부를 맞히고, 동역학 loss는 예측 분포가 인코더 표현을 맞히게 하며, 표현 loss는 표현이 예측 가능해지게 합니다. 표현이 입력 정보를 담으면서도 미래 예측에 적합하도록 균형을 맞추기 위해 해당 수식을 사용하는데요. stop-gradient와 free bits를 사용하며, KL 손실은 1 nat 아래에서는 잘라 학습 부담을 줄였다고 합니다.
또한, 본 연구에서 설명하길, 이전 세계 모델이 환경의 시각 복잡도에 따라 표현 손실 스케일을 다르게 조정해야 했다고 설명하는데요. DreamerV3는 free bits와 작은 표현 loss를 함께 사용해 이 문제를 줄입니다. 또한 벡터 관측에는 symlog 변환을 적용해 큰 입력값과 큰 reconstruction gradients를 막았습니다.
추가로, 논문에서는 인코더와 동역학 예측기의 categorical distribution을 99%의 신경망 출력과 1%의 uniform distribution의 혼합으로 만든다고 설명한다. 이 기법은 1% unimix라고 볼 수 있는데요. 분포가 완전히 deterministic이 되어 KL loss가 튀는 문제를 막기 위한 장치입니다. 이러한 장치를 넣는 이유는 해당 알고리즘이 여러 도메인에서 안정적으로 사용되기 위해서임을 명심하시면 좋을 것 같습니다.
<Critic learning>
비평가는 모델 상태 $s_t=\{h_t,z_t\}$에서 앞으로 얻을 반환된 것의 분포를 예측하는데요. 행위자와 비평가는 실제 환경에서 직접 긴 roll-out을 하지 않고, 세계 모델이 만든 추상 궤적에서 학습합니다. 비평가의 역할을 쉽게 말하면, “지금 이 상태가 앞으로 얼마나 좋은 결과로 이어질까”를 평가하는 채점자라고 생각하시면 이해가 좀더 빠르실 것 같습니다.
해당 수식은 행위자인 Actor와 비평가인 Critic의 기본 형태를 정의합니다. $\pi_\theta$는 행동 분포, $v_\psi$는 반환 분포, $s_t$는 모델 상태를 의미합니다. 정책 학습과 value 학습이 같은 latent state 위에서 이루어진다는 점을 보이기 위해 사용하며, 환경 상호작용 때는 lookahead planning을 따로 하지 않고 행위자 네트워크에서 행동을 샘플링합니다.
해당 수식은 비평가 loss와 bootstrap $\lambda$ return(부트스트랩 람다 반환)을 정의하고, $\gamma=0.997$은 할인율, $c_t$는 계속 여부, $\lambda$는 직접 보상과 미래 가치 예측의 혼합 정도를 의미합니다. 상상 지평보다 먼 미래 보상을 고려하기 위해 해당 수식을 사용하는데요. 비평가는 반환을 정규분포가 아니라 지수 간격의 bin 위의 categorical distribution로 예측합니다.
본 연구에서는 비평가가 자신의 예측에 의존하는 목표를 회귀하므로 학습이 불안정할 수 있다고 하는데요. 이를 줄이기 위해 비평가를 자기 파라미터의 지수이동평균 출력쪽으로 정규화합니다. 여기서 추가로, 학습 초기의 보상 예측기와 비평가 출력의 가중치를 0으로 초기화하여 초기부터 큰 보상이나 큰 가치를 할루시네이션하지 않게 합니다. 정말 디테일한 장치들이 많이 들어가있죠…?
여기서 끝이 아닙니다… 비평가는 상상 궤적뿐 아니라 replay buffer에서 샘플링한 실제 궤적에도 loss를 적용합니다. 본 연구에서는 상상 궤적의 loss 스케일을 $\beta_{\mathrm{val}}=1$, 재생 궤적 비평가 loss 스케일을 $\beta_{\mathrm{repval}}=0.3$으로 두었는데요. 이렇게 설정한 이유는 보상 예측이 어려운 환경에서 실제 보상 시퀀스가 가치 예측을 보조하도록 하기 위해서 설정하였다고 합니다.
<Actor learning>
행위자는 반환을 최대화하면서 entropy regularizer를 통해 탐색합니다. 본 연구에서는 보상 크기와 보상 빈도가 환경마다 다르기 때문에 고정 엔트로피 계수 $\eta=3\times10^{-4}$를 쓰려면 반환 스케일을 안정적으로 맞춰야한다고 하였는데요. 그래서 Dreamer는 반환을 5번째와 95번째 백분위수 차이로 정규화하였습니다. 단, 희소 보상에서 작은 반환 차이를 과하게 키우지 않도록 분모에 $\max(1,S)$를 사용하였습니다.
해당 수식은 반환 정규화에 쓰는 스케일 S를 정의합니다. 95번째 백분위수와 5번째 백분위수의 차이를 지수이동평균으로 부드럽게 만들며, 이상치에도 흔들리지 않는 반환 범위를 얻기 위해 사용됩니다. 해당 수식을 사용하여 희소 보상에서 표준편차 정규화가 보상을 과도하게 키우는 문제를 피할 수 있습니다.
<Robust predictions>
본 연구에서 목표하는 바는 introduction에서 말씀드렸지만, 다양한 도메인에서 동작하는 일반적인 알고리즘을 만드는 것인데요. 그런데 입력 재구성, 보상 예측, 반환 예측의 목표값 스케일이 환경마다 크게 다를 수 있습니다. 큰 값을 제곱 loss로 직접 예측하면 발산할 수 있고, 절댓값 loss나 Huber loss는 학습을 정체시킬 수 있다고 본 논문의 연구자들은 판단하였는데요. 그래서 DreamerV3는 symlog, symexp, symexp twohot loss를 사용하였습니다. symlog와 symexp twohot은 목표값 크기가 달라도 gradient 크기가 과도하게 달라지지 않도록 만드는 장치라고 보시면 되겠습니다.
해당 수식으로 목표값을 symlog로 압축한 뒤 예측하고, 읽을 때 symexp로 복원합니다. f(x,\theta)는 신경망 출력, y는 목표값, \hat y는 복원된 예측값을 의미합니다. 큰 목표값이 gradient를 폭발시키는 문제를 줄이기 위해 사용됩니다.
해당 수식은 symlog와 symexp을 정의한 것인데요. 부호는 유지하고 크기만 로그 또는 지수 방식으로 변환합니다. 보상과 반환은 음수와 양수를 모두 가질 수 있기 때문입니다. 원점 근처에서는 거의 항등 함수처럼 작동해 작은 값 학습을 해치지 않습니다.
해당 수식은 신경망이 지수 간격 bin 위의 확률분포를 출력하고, 그 평균을 예측값으로 읽는 것을 의미하는데요. B는 bin 위치, softmax는 각 bin의 확률을 의미합니다. 반환이나 보상처럼 확률적이고 스케일이 큰 목표를 안정적으로 예측하기 위해 사용하며, 확률 평균이 bin 사이에 있을 수 있으므로 연속값 예측이 가능합니다.
twohot target에 대한 onehot encoding entropy를 나타냅니다. twohot은 목표값에 가장 가까운 두 bin에 가중치를 나누어 주는 인코딩입니다. 목표값 크기와 gradient 크기를 분리하기 위해 사용됩니다
<Results>
드디어, 실험파트입니다. Dreamer는 일반성을 평가하기 위해서, 8개 영역, 150개 이상의 task에서 평가됩니다. 이 실험들은 continuous action, discrete action, visual input, low-dimensional vector inputd, ense reward, sparse reward, reward scale, procedural generation을 포함합니다.
<Benchmarks>
비교 대상은 두 종류인데요. 첫째, 모든 영역에서 고정 하이퍼파라미터를 쓰는 강한 PPO이고, 둘째, 각 벤치마크에 맞춰 설계되거나 튜닝된 전문화된 알고리즘입니다. Figure 1을 통해서 각 벤치마크에서의 성능을 볼 수 있는데, Dreamer가 다양한 전문화된 알고리즘과 비교되며, 특히 PPO보다 모든 영역에서 크게 높은 성능을 보이는 것을 확인할 수 있습니다. 각 벤치마크는 서로 다른 입력,행동,보상 구조를 갖는데요. 다만 Dreamer의 경우 세계 모델과 여러 기법 덕분에 환경마다 하이퍼파라미터를 새로 맞추지 않아도 됩니다.
Figure 5는 Iron Ingot, Iron Pickaxe, Diamond를 발견한 학습 에이전트 비율을 보여주는데요. 평균 반환만으로는 실제로 다이아몬드까지 도달했는지 알기 어렵기 때문에, 후반 아이템 발견 여부를 따로 보여주는데, 비교된 알고리즘들은 Iron Pickaxe까지는 어느 정도 진행하지만, Dreamer만 다이아몬드를 발견했음을 볼 수 있습니다. 사람의 개입 없이 다이아몬드를 발견한 알고리즘은 DreamerV3가 최초라고 합니다.
Figure 6에서는 Robutness 기법 제거, Learning signals 제거, model size 확장, replay ratio 확장을 비교하였습니다. KL objective, return normalization, symexp twohot regression, symlog가 평균 성능에 기여하는 것을 볼 수 있습니다. 또한 reconstruction loss가 표현 학습에서 특히 중요한 것으로 분석할 수 있습니다.
<Atari>
Atari 벤치마크는 57개 Atari 2600 게임으로 구성되고, bugets은 200M 프레임인데요. 논문에서는 sticky action 시뮬레이터 설정을 사용하였다고 합니다. Dreamer가 MuZero보다 적은 계산 자원으로 더 높은 성능을 보이는 것을 확인할 수 있는데요. Table 6에서 Dreamer의 gamer mean이 3381% 비해서 MuZero는 각각 3054%, PPO는 892% 것을 볼 수 있습니다. 이 결과는 Dreamer가 이미지 입력과 이산 행동을 가진 오래된 게임 벤치마크에서도 강하다는 점을 보여주는데요. 다만 Table 6을 task 별로 보면 Dreamer가 모든 Atari 게임에서 항상 최고는 아님을 알 수 있습니다.
<ProcGen>
ProcGen은 16개 게임으로 구성되며, randomized level에서 visual distraction을 통해 에이전트의 강건성과 일반화를 평가합니다. 논문에서는 hard difficulty와 unlimited level 설정을 사용하고, 50M 프레임에서 비교하였다고 합니다.
Table 7에서 normalized mean은 Dreamer 66.01, PPG 64.89를 기록한 것을 볼 수 있는데요. Dreamer가 튜닝된 전문 알고리즘 PPG와 비슷하거나 약간 높은 평균 성능을 보이고, Rainbow와 PPO를 능가하는 것을 볼 수 있습니다. 마찬가지로, task별로 보면 Dreamer가 일부 task에서는 잘하고, 일부 task에서는 성능이 낮은 것을 볼 수 있습니다. 즉, “Dreamer가 모든 게임을 압도한다”라기보다, 절차적 일반화 환경에서도 고정 설정으로 전문화된 알고리즘 수준의 평균 성능을 낸다는 의미를 가진다고 이해하시면 좋을 것 같습니다.
<DMLab>
DMLab은 30개 3D 과제로 구성되며, spatial reasoning과 temporal reasoning을 요구하는 task입니다.
Table 8에서 human mean capped는 Dreamer 71.4%, PPO 35.9%, 100M IMPALA 31.0%, 1B IMPALA 66.3%, 10B IMPALA 85.1%, 10B R2D2+ 85.4%를 보여주는데요. 논문에서는 Dreamer가 100M step 만으로 1B step IMPALA보다 높은 성능을 보이는데요. 해당 결과는 Dreamer의 데이터 효율성이 있음을 보여주는 결과이기는 하나 10B step까지 학습한 IMPALA나 R2D2+와 비교하면 Dreamer보다 높은 성능도 존재합니다. 따라서 해당 결과들을 통해서 항상 데이터 효율성이 좋다가 아니라, 적은 환경 상호작용으로도 높은 성능을 얻을 수도 있음을 보였다고 받아들여주시면 좋겠습니다.
이렇게 DreamerV3를 다뤄봤습니다. DreamerV1에 비해서 그렇게까지 구조가 많이 변하지 않은 것 같은데 강건화 기법을 정말 잘 적용하여 이렇게까지 성능을 끌어올린 것이 참 대단한 논문인 것 같습니다. 지금까지 읽어주셔서 감사합니다.
