Mastering diverse control tasks through world models(NeurIPS 2025)

느리·2025년 5월 30일

paper-review

목록 보기

14/15

🍦왜 이걸 택했는가~
일단 전통적인 RL에서는 구조가 단순했다. 그런데 Dreamer는 단순히 정책 네트워크만 사용하는 것이 아니라 "세상을 상상하는 뇌"를 만들고 이걸 기반으로 행동을 결정했다.... 는 점이 신기해서!ㅎㅎ

강화학습도 RNN이나 여타 아키텍처를 이용해서 좀 더 실용적이고 범용적으로 만들 수 있겠다... 그런 생각도 함.

넓은 범위에 적용할 수 있는 태스크를 해결하는 범용 알고리즘을 개발하는 것은 AI에서 매우 중요하면서도 도전적인 과제다.
현재 RL 알고리즘은 만들어진 목적에 부합하지만, 새로운 도메인에 적용하기 위해서는 부가적인 인간 전문가와 경험이 필요하다.
논문에서 제시한 Dreamer3은 150가지의 다양한 태스크에서 특화된 방법보다 뛰어난 성능을 보이인다.
= 강화학습을 새로운 문제에 쉽게 적용할 수 있도록 함
정규화, 균형 조정 및 변환에 기반한 다양한 강건성 기술을 통해 이러한 어려움을 극복함
Dreamer는 환경 모델을 배우고 미래 시나리오를 상상함으로서 행동을 개선한다.
별도의 설정 없이 바로 적용했을 때, Dreamer는 현재까지 알려진 바로는 인간 데이터나 커리큘럼 없이 MInecraft에서 처음부터 다이아몬드를 수집하는 최초의 알고리즘이다.
이 성취는 인공지능 분야에서 상당한 난제로 이뤄지며 open world에서 픽셀과 희소한 보상을 통해 원대한 전략을 모색해야 한다.

Learning Algorithm

알고리즘은 World model, Critic, Actor라는 세 가지 신경망(neural network)로 구성된다.
에이전트가 환경과 상호작용하고 있는 동안 구성 요소들은 재생된 경험으로부터 동시에 학습.

DreamerV3는 환경과의 상호작용에서 직접 학습하기 보다는 "World model'을 학습해 미래 시나리오를 상상하면서 정책을 개선하는 모델 기반 강화학습 방식임!

World model Learning

본 논문에서는 세계 모델을 위 그림에서 보는 것처럼 순환 상태 공간 모델(recurrent state-space model)로 구현했다. 우선 encoder는 감각 입력 $x_i$ 를 훈련 시퀀스에서 각 time step $t$ 에 대한 확률적 표현 $z_t$ 로 mapping한다. 그 다음 순환 상태 $h_t$ 를 갖는 시퀀스 모델은 과거 action $a_{t-1}$ 이 주어졌을 때 이러한 표현들의 시퀀스를 예측한다. $h_t$ 와 $z_t$ 의 concatenation은 모델 상태를 형성하고, 이를 통해 보상 $r_t$ 와 episode 지속 flag $c_t \in \{0,1\}$ 를 예측하고 입력값을 재구성해서 정보력이 높은 표현을 보장한다.

일단 world model이 뭔지부터 알아보자. 여기서 말하는 world model이랑 강화 학습 에이전트가 환경 자체를 학습하여 이해하는 모델을 의미한다. 그러니까 사람이 주변 세상을 관찰하고 경험해서 세상의 규칙이나 역학을 배우는 것과 비슷하다. 이 World model은 감각(sensory) 입력을 압축된 형태로 학습하는데, 이떄 RSSM(Recurrent State-Space Model)을 사용한다.

여기서 $x_1, x_2, x_3$ 는 에이전트가 실제 환경에서 관측한 이미지로 게임 화면이나 시각적 입력 등이 있다. 이 관측된 이미지는 잠재 공간(latent space)의 벡터 $z_1, z_2, z_3$ 로 인코딩되는데 이 잠재 벡터는 이미지의 중요한 특징을 압축한 표현이다.

먼저 인코더에서 고차원 센서 입력을 압축된 잠재 표현 (Latent Representation)으로 변환한다.
그럼 이전 행동과 잠재 상태를 바탕으로 다음 잠재 상태의 확률 분포를 예측하며 순환 상태(Recurrent State)를 유지한다.
Dynamic Predictor : 이런 순환 상태만을 바탕으로 다음 잠재 상태의 확률 분포를 예측한다.
보상 예측기(Reward Predictor) : 잠재 상태에서 예상되는 보상을 예측한다
계속 예측기(Continue predictor) : 현재 상태에서 에피소드가 종료될지 여부를 예측한다.
디코더(Decoder) : 잠재 상태에서 원래 센서 입력을 재구성한다. 이는 잠재 표현이 입력의 핵심 정보를 담도록 강제하는 역할을 한다.

[!note] 잠재 상태에서 재구성하는게 왜 핵심 정보를 담도록 하는거지?
데이터 압축 과정에서 정보가 손실될 수 있다. 여기서 재구성 손실이 등장한다. 디코더는 압축된 잠재 상태만을 입력으로 받아 원래의 입력( $x_t$ )과 최대한 유사한 $x_t$ 를 만들려고 시도한다. 재구성 손실은 디코더가 만들어낸 출력과 원래의 입력 사이를 측정하고, 이 차이를 최소화할 수 있도록 모델을 학습시킨다.

==학습 목표==

결국 world model의 학습 목표는 다음과 같은 세 가지 손실을 최소화하는 것이다.

\mathcal{L}(\phi) \doteq E_{q_\phi} \left[ \sum_{t=1}^T (\beta_{\text{pred}} \mathcal{L}_{\text{pred}}(\phi) + \beta_{\text{dyn}} \mathcal{L}_{\text{dyn}}(\phi) + \beta_{\text{rep}} \mathcal{L}_{\text{rep}}(\phi)) \right].

예측 손실, 동역학 손실, 표현 손실. -> 이 부분은 깊게 들어가면 강화학습을 열심히 공부해야 한다...ㅎㅎ

Actor Critic Learning

Actor-Critic은 강화학습 방법론 중 하나로, 미래의 가치를 평가하는 평가자(Critic)과 Critic의 피드백을 바탕으로 행동을 선택하는 행위자(Actor)로 이뤄져 있다.

먼저 왼쪽 아래에 있는 게임 이미지는 실제 환경에서 받아온 관측값 $x_1$ 이다. 이는 encoder를 통해 잠새 표현으로 변환된다. 그렇게 만들어진 잠재 벡터 $z_1$ 는 확률적(latent stochastic) 상태로 "관측으로부터 얻은, 뭐가 있었어??"에 대한 요약 정보다.

그리고 위에 있는 $h_t$ 는 결정적(deterministic) 히든 상태로, 이전 시간의 히든 상태와, 현재 샘플된 $z_t$ 를 RNN에 넣어서 만들어진다.

h_{t+1}=f(h_t, z_t, a_t)

이건 시간에 따라 축적된 정보를 반영하는 메모리 같은 것이라 시간 흐름을 반영한다.
그래서 이 두 개를 합친 것이 바로 $s_t=\{h_t, z_t\}$ 으로, Markov 상태 표현이며 이걸 기반으로 Actor와 Critic이 작동한다.

Critic 학습

먼저 Critic의 경우 World Model이 상상한 상태로부터 예측된 return 분포를 학습한다.

목표! $v_\psi(R_t | s_t)$ 를 학습하라!

= 쉽게 말하면 현재 상태의 "가치(미래 보상의 총합)"를 예측하여 Actor의 학습을 도와주는 것이다.
여기서 사용하는 것이 상상 기반 학습, Trajectory Imagination이다. 실제 환경을 매번 실행하는 것이 아니라(돈이 많이 들고 시간도 들어요!) World model이 상상한 가상 시나리오로 학습을 하는 것이다.

\text{Actor:} \quad a_t \sim \pi_\theta(a_t | s_t) \qquad \text{Critic:} \quad v_\psi(R_t | s_t)

Trajectory 생성
먼저 world model과 actor을 이용해서 상상 속에서 trajectory를 만든다. 여기서 이 trajectory라는 것은 쉽게 말해 상태, 행동, 보상의 총 집합을 말한다. s1 -> a1 -> r1 -> s2 -> a2 -> .... 이런 것!
$\lambda$ -return 계산
애초에 강화학습에서 "지금 당장의 보상"만 보는 것은 말이 안된다. 우리는 미래의 보상까지 일정 비율을 섞어서 계산을 해야 한다. 이걸 $\lambda$ -return 라고 부른다.
Critic은 이렇게 bootstrapped $\lambda$ -return을 이용해서 학습한다.
(솔직히 여기까지 읽었을 때 이거 아키텍처보다 강화학습 비중이 높나 싶었지만 이미 늦었다.)

R^\lambda_t = r_t + \gamma c_t \left[ (1 - \lambda) v_t + \lambda R^\lambda_{t+1} \right]

$v_t$ : critic이 예측한 현재 상태의 기대 보상
$\gamma$ =0.997 : 할인율! -> 이건 미래 보상 감쇠라는 것!

확률 분포로 예측
보상의 형태가 다양할 수 있기 때문에 critic은 보상을 하나의 숫자가 아니라 확률 분포로 에측한다.
그 다음 reaply buffer을 이용해서 실제 경험을 혼합해서 보강하고

[!note] Replay buffer
과거의 경험들(trajectory)을 저장두는 메모리 공간으로 안에는 이렇게 저장이 된다.
$(x_t, a_t, r_t, x_{t+1}, c_t)$
강화학습은 환경을 실행해야 하기 때문에 데이터를 얻는 것이 비싸다. Replay buffer을 이용하면 저장된 과거 경험을 여러 번 반복해서 학습에 이용할 수 있다!

🧩 Dreamer에서는 어떻게 사용?
Dreamer는 world model을 학습할 때 두 가지 종류의 데이터를 사용한다.
1. 실제 환경에서 수집된 경험 - repaly buffer 저장
2. world model이 상상한 trajectory - imagination trajectory

이전 파라미터의 지수 이동 평균 버전을 사용해서 Critic이 자기 자신을 참조하게 하여 예측 안정화를 한다.

Actor 학습

Actor의 경우는 각 시점에서 $s_t$ 를 받아서 행동 $a_t$ 를 샘플링한다.

a_t \sim \pi_\theta (a_t |s_t)

그림에서는 조이스틱이 움직이는 것을 a라고 한다.
그런데 여기서 actor가 주의해야 하는 것은, 너무 똑같은 행동만 반복하지 않는 것이다. 그렇기에 우리는 탐색(exploration)을 한다. 그런데 이 탐험의 적절한 강도를 환경 내에서 보상의 크기(scale)와 빈도(frequency) 모두에 따라 달라진다. 이상적으로는 보상이 드문 환경에서는 더 많이 탐색하고, 보상이 자주 나오거나 가까이에 있을 때는 더 많이 이용하는 것이 좋으며 보상의 절대적인 크기 변화에 탐색 강도가 영향을 받아서는 안된다. 때문에 논문에서는 보상의 크기를 정규화했다.

\dot{L} = - \sum_{t=1}^{T} \frac{\text{sg}(R_t^\lambda - v_\psi(s_t))}{\max(1, S)} \log \pi_\theta(a_t | s_t) + \eta H[\pi_\theta(a_t | s_t)]

솔직히 매우 어렵다...
그냥 보상의 스케일이 환경마다 달라지니까 return 자체를 정규화해서 균형 있게 학습하도록 도와주는 Actor의 손실함수라고 생각하면 됨!

Results

Benchmarks

연속적이고 불연속적인 행동, 시각적 및 저차원 입력, 조밀하고 희소한 보상, 다양한 보상 척도, 2D 및 3D 세계, 절차적 생성을 포함하는 8개의 도메인에서 광범위한 연구를 수행함!

Atari, ProcGen, DMLab, Minecraft, Atari100k, Proprio Control, Visual Control, BSuite

Dreamer는 다양한 도메인에서 PPO보다 훨씬 뛰어난 성능을 보여줬다.

Minecraft

마인크래프트에서 다이아몬드를 수집하는 것은 인공지능 분야에서 오랜 기간 동안 난제였다. (최근에는 MAS를 위한 마인크래프트 벤치마크도 새로 생겼던데, 그런 거 보면 마크는 역시 계속 사용될 것 같음,,,)
Dreamer는 인간 데이터를 사용하지 않고도 마인크래프트에서 처음부터 다이아몬드를 수집하는 최초의 알고리즘이다.

몰랐는데 MineRL이 인간 전문가 궤적(trajectory) 데이터셋을 제공한다고 함.

Ablation

a. Robustness techniques

Dreamer에 도입된 여러 안정화/정규화 기법들을 하나씩 제거한 것! 모두 의미가 있지만 특히 KL objective, return nomalization, symexp twohot loss가 특히 중요함

b. Learning signals (학습 신호의 영향)

world model을 학습할 때 사용하는 두 가지 gradient 신호를 각각 제거했다.
1) 보상과 가치 예측에 대한 gradient 제거 -> 성능 약간 하락
2) 관측값 복원(reconstruction)에 대한 gradient 제거 -> 급락

[!question] 😵‍💫 왜일까?
일단 관측 복원이란 모델이 잠재 상태 표현 ( $h_t, z_t$ )를 기반으로 입력으로 주어졌던 관측 $x_t$ 을 다시 만들어내는 것이었다. 즉, 모델이 풍부하고 일반적인 표현을 학습하도록 강제하고 모델이 환경을 상상하고 예측하는 능력을 키우는데 매우 중요했다.
그럼 여기서 보상/가치 예측 제거는 왜 덜 중요하냐... 할수도 있다.
일단 이건 애초에 Dreamer라는 모델이 무얼 목적으로 하냐가 중요하다. Dreamer는 "보상을 잘 예측하는 모델"이 아니라 "세상을 잘 상상하고 이해하는 모델"이다. world model이 학습한 표현을 기반으로 actor와 critic은 상상 속에서 따로 학습하기 때문에 world model이 보상을 직접 잘 예측하지 않아도 상관없다!

c. Model size scaling (모델 크기 확장)

모델 크기가 커질 수록 성능이 향상됨

d. Replay buffer 크기 영향

충분히 큰 Replay buffer는 더 좋은 성능을 유도함

Conclusion

고정된 하이퍼파라미터로도 다양한 도메인에서 뛰어난 성능을 보이는 범용 강화학습 알고리즘
Dreamer는 150개 이상의 과제에서 탁월한 성능을 보였고, 다양한 데이터 및 연산 자원 조건에서도 안정적으로 학습함으로써, 실제 활용 가능한 강화학습으로의 도약을 보여줌
Dreamer는 추가 조정 없이(out-of-the-box) 사용되어, 마인크래프트에서 다이아몬드를 수집한 최초의 알고리즘
Dreamer는 학습된 월드 모델 기반의 고성능 알고리즘