Diffusion for World Modeling: Visual Details Matter in Atari

https://arxiv.org/abs/2405.12399

Diffusion for World Modeling: Visual Details Matter in Atari

 

초록

월드 모델은 강화 학습 에이전트를 안전하고 샘플 효율적인 방식으로 훈련하는 유망한 접근 방식입니다. 최근의 월드 모델들은 주로 이산 잠재 변수 시퀀스를 사용하여 환경 동역학을 모델링하고 있습니다. 하지만 이러한 압축된 이산 표현은 강화 학습에 중요한 시각적 세부 사항들을 무시할 수 있습니다. 동시에, 확산 모델(diffusion models)은 이미지 생성의 주요 접근 방식이 되었으며, 기존의 이산 잠재 변수를 모델링하는 방법들에 도전하고 있습니다. 이러한 패러다임의 변화에 영감을 받아, 우리는 다이아몬드(DIAMOND, DIffusion As a Model Of eNvironment Dreams)라는 확산 월드 모델에서 훈련된 강화 학습 에이전트를 소개합니다. 우리는 확산을 월드 모델링에 적합하게 만들기 위해 필요한 주요 설계 선택들을 분석하고, 개선된 시각적 세부 사항이 에이전트의 성능 향상으로 어떻게 이어질 수 있는지 보여줍니다. 다이아몬드는 경쟁적인 Atari 100k 벤치마크에서 평균 인간 정규화 점수 1.46을 달성하며, 월드 모델에서 전적으로 훈련된 에이전트들 중 새로운 최고 성과를 기록했습니다. 또한, 다이아몬드의 확산 월드 모델이 정적 Counter-Strike: Global Offensive 게임 플레이에 훈련하여 상호작용이 가능한 신경망 게임 엔진으로 단독 실행될 수 있음을 보여줍니다. 월드 모델링에서 확산 연구의 미래를 촉진하기 위해, 우리는 우리의 코드, 에이전트, 비디오 및 플레이 가능한 월드 모델들을 https://diamond-wm.github.io 에서 공개합니다.

 

1. 서론

환경의 생성 모델, 또는 "월드 모델"(Ha와 Schmidhuber, 2018)은 일반 에이전트가 환경에 대해 계획하고 추론할 수 있도록 돕는 중요한 구성 요소로 점점 중요해지고 있습니다(LeCun, 2022). 강화 학습(RL)은 최근 몇 년간 다양한 성공 사례를 보여주었지만(Silver et al., 2016; Degrave et al., 2022; Ouyang et al., 2022), 샘플 비효율성으로 인해 실제 응용이 제한되는 것이 잘 알려져 있습니다. 월드 모델은 다양한 환경에서 강화 학습 에이전트를 훈련하는 데 있어 큰 샘플 효율성 향상을 보여주었으며(Hafner et al., 2023; Schrittwieser et al., 2020), 이를 통해 실제 환경에서의 경험을 통한 학습을 가능하게 합니다(Wu et al., 2023).

 

그림 1: 시간에 따른 다이아몬드의 상상 전개 상단 행은 학습된 확산 월드 모델 𝐃θ에서 에이전트의 정책 πϕ가 일련의 행동을 취하는 과정을 묘사합니다. 환경 시간 t는 수평 축을 따라 흐르고, 수직 축은 denoising 시간 τ가 𝒯에서 0으로 거꾸로 흐르는 것을 나타냅니다. 구체적으로, (깨끗한) 과거 관측값 𝐱<t0와 행동 a<t를 주어 초기 노이즈 샘플 𝐱t𝒯에서 시작하여, 역 노이징 과정을 𝐃θ를 반복 호출하여 시뮬레이션하고, 다음 (깨끗한) 관측값 𝐱t0을 얻습니다. 상상 과정은 자귀회귀적이며, 예측된 관측값 𝐱t0과 정책에 의해 수행된 행동 at는 다음 시간 단계의 조건화에 포함됩니다. 이 과정의 애니메이션 시각화는 https://diamond-wm.github.io에서 확인할 수 있습니다.

 

최근의 월드 모델링 방법들(Hafner et al., 2021; Micheli et al., 2023; Robine et al., 2023; Hafner et al., 2023; Zhang et al., 2023)은 주로 환경 동역학을 이산 잠재 변수들의 시퀀스로 모델링하고 있습니다. 잠재 공간의 이산화는 다중 단계 시간 수평선에서 누적되는 오류를 피하는 데 도움이 됩니다. 그러나 이러한 인코딩은 정보를 잃게 하여, 일반성과 재구성 품질의 손실로 이어질 수 있습니다. 이는 자율주행 차량 훈련과 같이 과제에 필요한 정보가 명확하지 않은 더 현실적인 시나리오에서는 문제가 될 수 있습니다(Hu et al., 2023). 예를 들어, 시각 입력에서 교통 신호등이나 멀리 있는 보행자와 같은 작은 세부 사항들이 에이전트의 정책에 영향을 미칠 수 있습니다. 이산 잠재 변수의 수를 늘리면 이러한 손실 압축 문제를 완화할 수 있지만, 그만큼 계산 비용이 증가합니다(Micheli et al., 2023).

 

-----

먼저 이산 잠재 변수(discrete latent variables)의 개념을 간단한 예시로 설명하겠습니다:

자율주행 차량이 도로를 주행하는 상황을 생각해봅시다:

 

연속적 표현의 경우:

# 실제 환경의 연속적인 상태
position = 23.7 meters
speed = 45.3 km/h
traffic_light_distance = 156.8 meters
pedestrian_position = (187.2, 3.4) meters

이산화된 표현의 경우:

# 이산화된 상태 표현
position_bucket = 2        # 0-30m는 버킷 2
speed_bucket = 4          # 40-50km/h는 버킷 4
traffic_light = "far"     # 거리를 near/medium/far로 구분
pedestrian = "present"    # 보행자 존재 여부만 표현

장점:

1. 오류 누적 방지
   • 연속 값으로 여러 단계를 예측할 때: 23.7m → 24.1m → 24.6m → 25.2m → ...
   • 작은 예측 오차가 시간이 지날수록 커질 수 있음
   • 이산값으로는: 버킷2 → 버킷2 → 버킷3 과 같이 오차가 누적되지 않음
2. 계산 효율성
   • 유한한 수의 상태만 고려하면 되므로 계산이 단순화됨

단점:

1. 정보 손실
   • 실제 도로에서 보행자까지의 정확한 거리(187.2m)가 아닌 단순히 "있다/없다"로만 표현됨
   • 특히 멀리 있는 작은 물체(예: 원거리의 보행자나 교통 표지판)의 세부 정보가 손실될 수 있음

-----

 

동시에, 확산 모델(Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020)은 고해상도 이미지 생성에서 지배적인 패러다임이 되었으며(Rombach et al., 2022; Podell et al., 2023), 이 모델이 노이즈 과정을 역으로 학습하는 방식은 이산 토큰을 모델링하는 기존의 방법들(Esser et al., 2021; Ramesh et al., 2021; Chang et al., 2023)을 도전하게 만들며, 월드 모델링에서 이산화를 줄일 수 있는 유망한 대안을 제공합니다. 또한, 확산 모델은 조건화를 쉽게 할 수 있고, 모드 붕괴 없이 복잡한 다중 모달 분포를 유연하게 모델링할 수 있는 것으로 알려져 있습니다. 이러한 특성들은 월드 모델링에 매우 중요한데, 조건화를 잘 따르는 것은 월드 모델이 에이전트의 행동을 더 잘 반영하게 하여 더 신뢰할 수 있는 신용 할당을 가능하게 하고, 다중 모달 분포를 모델링하는 것은 에이전트에게 더 다양한 학습 시나리오를 제공할 수 있기 때문입니다.

 

이러한 특성에 영감을 받아, 우리는 다이아몬드(DIAMOND, DIffusion As a Model Of eNvironment Dreams)라는 확산 월드 모델에서 훈련된 강화 학습 에이전트를 제안합니다. 우리의 확산 월드 모델이 긴 시간 수평선에서도 효율적이고 안정적이기 위해서는 신중한 설계 선택이 필요하며, 우리는 이러한 선택의 중요성을 정성적으로 분석합니다. 다이아몬드는 잘 알려진 Atari 100k 벤치마크에서 평균 인간 정규화 점수 1.46을 달성하며, 월드 모델에서 전적으로 훈련된 에이전트들 중 새로운 최고 성과를 기록했습니다. 추가적으로, 이미지 공간에서 작동하는 것은 우리의 확산 월드 모델이 환경을 대체할 수 있는 드롭인 방식으로 활용될 수 있게 하며, 이를 통해 월드 모델 및 에이전트의 행동에 대한 더 깊은 통찰을 제공합니다. 특히, 일부 게임에서 향상된 성능은 중요한 시각적 세부 사항을 더 잘 모델링한 결과임을 발견했습니다. 우리의 월드 모델이 단독으로 효과적임을 추가로 증명하기 위해, 우리는 다이아몬드의 확산 월드 모델을 87시간의 정적 Counter-Strike: Global Offensive(CSGO) 게임 플레이(Pearce와 Zhu, 2022)에 훈련시켜, 인기 있는 게임 맵인 Dust II에 대해 상호작용이 가능한 신경망 게임 엔진을 생성했습니다. 우리는 우리의 코드, 에이전트 및 플레이 가능한 월드 모델을 https://diamond-wm.github.io에서 공개합니다.

 

2. 사전 지식

2.1 강화 학습과 월드 모델

우리는 환경을 일반적인 부분적으로 관측 가능한 마코프 결정 과정(POMDP) (Sutton and Barto, 2018)으로 모델링합니다. 이 과정은 (𝒮, 𝒜, 𝒪, T, R, O, γ)로 표현되며, 여기서 𝒮는 상태의 집합, 𝒜는 이산 행동의 집합, 𝒪는 이미지 관측의 집합입니다. 전이 함수 T: 𝒮 × 𝒜 × 𝒮 → [0, 1]는 환경의 동역학 p(𝐬ₜ₊₁ ∣ 𝐬ₜ, 𝐚ₜ)을 설명하며, 보상 함수 R: 𝒮 × 𝒜 × 𝒮 → ℝ는 전이를 스칼라 보상으로 매핑합니다. 에이전트는 상태 sₜ에 직접 접근할 수 없으며, 이미지 관측 xₜ ∈ 𝒪를 통해서만 환경을 볼 수 있습니다. 이러한 이미지 관측은 관측 확률 p(𝐱ₜ ∣ 𝐬ₜ)에 따라 관측 함수 O: 𝒮 × 𝒪 → [0, 1]에 의해 방출됩니다. 목표는 관측값을 행동으로 매핑하여 예상 할인 보상 𝔼ₚ[∑ₜ≥0 γ^t rₜ]를 최대화하는 정책 π를 얻는 것입니다. 여기서 γ ∈ [0, 1]는 할인율입니다. 월드 모델(Ha and Schmidhuber, 2018)은 환경의 생성 모델, 즉 p(sₜ₊₁, rₜ ∣ sₜ, aₜ)을 모델링합니다. 이러한 모델은 샘플 효율적인 방식으로 RL 에이전트를 훈련하기 위한 시뮬레이션 환경으로 사용될 수 있습니다(Sutton, 1991) (Wu et al., 2023). 이 패러다임에서는 훈련 절차가 보통 다음 세 가지 단계를 반복합니다: 실제 환경에서 RL 에이전트를 통해 데이터를 수집한다; 수집된 모든 데이터로 월드 모델을 훈련한다; 월드 모델 환경에서 RL 에이전트를 훈련한다(일명 "상상 속에서").

 

-----

핵심 포인트는:

1. 에이전트는 전체 상태(𝒮)를 직접 볼 수 없음
2. 대신 카메라 이미지(𝒪)를 통해서만 환경을 관측
3. 관측 함수(O)는 실제 상태가 어떻게 불완전한 이미지로 변환되는지를 설명
4. 전이 함수(T)는 행동의 결과로 상태가 어떻게 변하는지를 확률적으로 설명

이것이 "부분적으로 관측 가능한(Partially Observable)" 이유입니다:

• 자율주행 차량은 카메라로만 세상을 봄
• 카메라 이미지만으로는 다른 차량의 정확한 속도나 의도를 알 수 없음
• 카메라가 볼 수 없는 사각지대가 존재
• 센서 노이즈로 인해 관측이 부정확할 수 있음

-----

 

2.2 점수 기반 확산 모델

확산 모델(Sohl-Dickstein et al., 2015)은 비평형 열역학에서 영감을 받아 노이즈 과정을 역으로 되돌려 샘플을 생성하는 일종의 생성 모델입니다.

 

우리는 연속 시간 변수 τ ∈ [0, 𝒯]로 인덱싱된 확산 과정 {𝐱_τ}를 고려하며, 이에 상응하는 주변 확률 분포 {p_τ}, 그리고 경계 조건 p₀ = p_data와 p_𝒯 = p_prior를 설정합니다. 여기서 p_prior는 가우시안과 같은 다루기 쉬운 비구조적 사전 분포입니다. 확산 과정 시간에는 τ를 사용하고, 환경 시간에는 t를 사용하는 이유는 구분을 명확히 하기 위함입니다.

 

이 확산 과정은 표준 확률 미분 방정식(SDE) (Song et al., 2020)의 해로 기술될 수 있습니다:

d𝐱 = 𝐟(𝐱, τ)dτ + g(τ)d𝐰, (1)

 

여기서 𝐰는 Wiener 과정(브라운 운동)을 나타내며, 𝐟는 드리프트 계수 역할을 하는 벡터 값 함수이고, g는 과정의 확산 계수로 알려진 스칼라 값 함수입니다.

 

노이즈에서 데이터로의 매핑을 수행하는 생성 모델을 얻기 위해 이 과정을 역으로 되돌려야 합니다. 놀랍게도 Anderson(1982)은 역 과정 역시 시간을 거꾸로 진행하는 확산 과정이며 다음 SDE로 기술될 수 있음을 보였습니다:

 

d𝐱 = [𝐟(𝐱, τ) - g(τ)²∇_𝐱 log p_τ(𝐱)]dτ + g(τ)d𝐰̄, (2)

 

여기서 𝐰̄는 역 시간 Wiener 과정이고, ∇_𝐱 log p_τ(𝐱)는 (Stein) 점수 함수로, 로그 주변 분포의 지지에 대한 기울기를 나타냅니다. 따라서, 앞선 노이즈 과정을 역으로 되돌리기 위해서는 함수 f와 g를 정의하고(3.1절에서 다룰 예정), 주변 분포 {p_τ}와 관련된 미지의 점수 함수 ∇_𝐱 log p_τ(𝐱)를 추정해야 합니다. 실질적으로는 시간 종속 점수 모델 𝐒_θ(𝐱, τ)를 사용하여 이러한 점수 함수를 추정할 수 있습니다(Song et al., 2020).

 

특정 시점에서 점수 함수를 추정하는 것은 간단하지 않습니다. 실제 점수 함수에 접근할 수 없기 때문입니다. 다행히도, Hyvärinen(2005)은 점수 매칭 목표를 소개하여, 기본 점수 함수의 지식 없이도 데이터 샘플로부터 점수 모델을 훈련할 수 있게 했습니다. 주변 분포 p_τ로부터 샘플에 접근하기 위해, 우리는 시간 0에서 τ까지의 앞으로의 과정을 시뮬레이션해야 합니다. 이는 일반적으로 비용이 많이 드는 작업이지만, f가 선형이라면, 단일 단계에서 가우시안 섭동 커널 p₀_τ를 깨끗한 데이터 샘플에 적용하여 어떤 시간 τ에서도 도달할 수 있습니다(Song et al., 2020). 커널이 미분 가능하기 때문에 점수 매칭은 denoising 점수 매칭 목표(Vincent, 2011)로 단순화됩니다:

 

ℒ(θ) = 𝔼[‖𝐒_θ(𝐱_τ, τ) - ∇_𝐱_τ log p₀_τ(𝐱_τ ∣ 𝐱₀)‖²], (3)

 

여기서 기대값은 확산 시간 τ와 노이즈 샘플 𝐱_τ ~ p₀_τ(𝐱_τ ∣ 𝐱₀)에 대해 계산되며, 이는 τ 수준의 섭동 커널을 깨끗한 샘플 𝐱₀ ~ p_data(𝐱₀)에 적용하여 얻어집니다. 중요한 점은, 커널 p₀_τ가 알려진 가우시안 분포이기 때문에, 이 목표는 간단한 L2 재구성 손실로 변환됩니다:

 

ℒ(θ) = 𝔼[‖𝐃_θ(𝐱_τ, τ) - 𝐱₀‖²], (4)

여기서 𝐃_θ(𝐱_τ, τ) = 𝐒_θ(𝐱_τ, τ)σ²(τ) + 𝐱_τ로 재매개변수화됩니다. 여기서 σ(τ)는 τ 수준 섭동 커널의 분산입니다.

 

2.3 월드 모델링을 위한 확산

2.2절에서 설명한 점수 기반 확산 모델은 무조건적 생성 모델인 p_data를 제공합니다. 월드 모델로 활용하기 위해서는 환경 동역학의 조건부 생성 모델 p(𝐱ₜ₊₁ ∣ 𝐱_≤ₜ, a_≤ₜ)가 필요합니다. 여기서 우리는 POMDP의 일반적인 경우를 고려하며, 마코프 상태 sₜ는 알려지지 않았고 과거의 관측과 행동을 통해 근사될 수 있습니다. 이 확산 모델은 이 이력을 조건으로 다음 관측값을 직접 추정하고 생성할 수 있습니다(그림 1 참조). 이는 식 4를 다음과 같이 수정합니다:

 

ℒ(θ) = 𝔼[‖𝐃_θ(𝐱ₜ₊₁_τ, τ, 𝐱_≤ₜ₀, a_≤ₜ) - 𝐱ₜ₊₁₀‖²]. (5)

 

훈련하는 동안, 우리는 에이전트의 재생 데이터셋에서 궤적 세그먼트 𝐱_≤ₜ₀, a_≤ₜ, 𝐱ₜ₊₁₀를 샘플링하고, τ 수준 섭동 커널을 적용하여 노이즈가 추가된 다음 관측값 𝐱ₜ₊₁_τ ~ p₀_τ(𝐱ₜ₊₁_τ ∣ 𝐱ₜ₊₁₀)를 얻습니다. 요약하면, 월드 모델링을 위한 이 확산 과정은 2.2절에서 설명한 표준 확산 과정과 유사하며, 점수 모델이 과거 관측과 행동을 조건으로 하는 것이 특징입니다.

 

다음 관측을 샘플링하기 위해서는 식 2에 있는 역 SDE를 반복적으로 풀어야 합니다. 그림 1에서 설명한 것처럼 원칙적으로는 어떤 ODE 또는 SDE 솔버라도 사용할 수 있지만, 샘플링 품질과 함수 평가 수(NFE) 사이에는 고유한 트레이드오프가 존재하며, 이는 확산 월드 모델의 추론 비용을 직접 결정합니다(자세한 내용은 부록 A를 참조하십시오).

 

-----

Forward Process (식 1) x₀ (Clean) f(x,τ)dτ + g(τ)dw xₜ (Noisy) Reverse Process (식 2) xₜ (Noisy) f(x,τ) - g(τ)²∇log p_τ(x) x₀ (Clean) World Model (식 5) xₜ (Current) aₜ (Action) xₜ₊₁ (Next) Ground Truth

주요 이점들:

1. 데이터 효율성

• 하나의 샘플로 여러 노이즈 레벨의 학습 데이터 생성 가능
• 같은 메모리 사용량으로 더 풍부한 학습 신호 획득

1. 장기 예측 성능

# 일반 월드 모델의 오차 누적
error = initial_error
for t in range(prediction_steps):
    error *= error_multiplication_factor  # 오차가 기하급수적으로 증가

# 확산 모델의 오차 제어
error = initial_error
for t in range(prediction_steps):
    error = error_correction_via_score_function(error)  # 점수 함수로 오차 보정

-----

 

3. 방법론

3.1 실용적인 확산 패러다임 선택

2장에서 제공된 배경을 바탕으로, 이제 우리는 다이아몬드를 확산 기반 월드 모델의 실용적인 구현으로 소개합니다. 특히, 우리는 2.2장에서 언급된 드리프트 계수(𝐟)와 확산 계수(g)를 정의하며, 이는 특정 확산 패러다임의 선택에 해당합니다. 역사적으로 ddpm(Ho et al., 2020)이 자연스러운 후보로 여겨지지만(부록 B에서 설명), 우리는 대신 Karras et al.(2022)이 제안한 edm 포뮬레이션을 기반으로 구축합니다. 이러한 선택의 실질적인 영향은 5.1장에서 논의합니다. 이제 우리는 edm을 우리의 확산 기반 월드 모델로 어떻게 적응시켰는지 설명합니다.

 

우리는 섭동 커널 p₀_τ(𝐱ₜ₊₁_τ ∣ 𝐱ₜ₊₁₀) = 𝒩(𝐱ₜ₊₁_τ; 𝐱ₜ₊₁₀, σ²(τ)𝐈)를 고려하며, 여기서 σ(τ)는 확산 시간의 실수 값 함수로, 노이즈 스케줄을 나타냅니다. 이는 드리프트 계수 𝐟(𝐱, τ) = 𝟎(선형)과 확산 계수 g(τ) = 2σ˙(τ)σ(τ)로 설정하는 것에 해당합니다.

 

우리는 Karras et al.(2022)이 도입한 네트워크 전처리를 사용하여 식 5에서 𝐃θ를 노이즈 관측값과 신경망 예측값 𝐅θ의 가중합으로 매개변수화합니다.

 

𝐃θ(𝐱ₜ₊₁_τ, yₜ_τ) = c_skip_τ 𝐱ₜ₊₁_τ + c_out_τ 𝐅θ(c_in_τ 𝐱ₜ₊₁_τ, yₜ_τ), (6)

 

여기서 간결하게 하기 위해 yₜ_τ ≔ (c_noise_τ, 𝐱_≤ₜ₀, a_≤ₜ)를 정의하여 모든 조건 변수들을 포함합니다.

 

전처리기 c_in_τ와 c_out_τ는 네트워크의 입력과 출력을 어떤 노이즈 수준 σ(τ)에서도 단위 분산을 유지하도록 선택됩니다. c_noise_τ는 노이즈 수준의 경험적 변환이며, c_skip_τ는 σ(τ)와 데이터 분포의 표준 편차 σ_data를 사용하여 정의됩니다. c_skip_τ = σ_data² / (σ_data² + σ²(τ))로 계산됩니다. 이러한 전처리기들에 대한 자세한 설명은 부록 C에 있습니다.

 

식 5와 6을 결합하면 𝐅θ의 훈련 목표에 대한 통찰을 제공합니다:

 

ℒ(θ) = 𝔼[‖𝐅θ(c_in_τ 𝐱ₜ₊₁_τ, yₜ_τ) - 1/c_out_τ (𝐱ₜ₊₁₀ - c_skip_τ 𝐱ₜ₊₁_τ)‖²]. (7)

 

네트워크 훈련 목표는 노이즈 수준 σ(τ)에 따라 신호와 노이즈를 적응적으로 혼합합니다. σ(τ) ≫ σ_data일 때, c_skip_τ → 0이 되어, 𝐅θ의 훈련 목표는 깨끗한 신호 𝐱ₜ₊₁₀에 의해 지배됩니다. 반대로 노이즈 수준이 낮을 때, σ(τ) → 0이 되어 c_skip_τ → 1이 되며, 목표는 깨끗한 신호와 섭동된 신호 간의 차이, 즉 추가된 가우시안 노이즈가 됩니다. 직관적으로 이는 저노이즈 환경에서 훈련 목표가 자명해지는 것을 방지합니다. 실제로, 이 목표는 노이즈 스케줄의 극단적인 부분에서 높은 분산을 가집니다. 따라서 Karras et al.(2022)은 훈련을 중간 노이즈 영역에 집중하기 위해 노이즈 수준 σ(τ)를 경험적으로 선택한 로그 정규 분포에서 샘플링합니다(부록 C 참조).

 

우리는 벡터 필드 𝐅θ에 대해 표준 U-Net 2D(Ronneberger et al., 2015)를 사용하고, 모델을 조건화하기 위해 과거의 L개의 관측값과 행동들을 버퍼로 유지합니다. 이러한 과거 관측값들을 다음 노이즈 관측값과 채널별로 연결하고, U-Net의 잔차 블록(He et al., 2015)에서 적응형 그룹 정규화 계층(Zheng et al., 2020)을 통해 행동을 입력합니다.

 

2.3절 및 부록 A에서 논의된 것처럼, 훈련된 확산 모델에서 다음 관측값을 생성하기 위한 많은 샘플링 방법이 존재합니다. 우리의 코드베이스는 다양한 샘플링 스키마를 지원하지만, 우리는 오일러 방법이 높은 차수의 샘플러가 요구하는 추가적인 함수 평가 수(NFE)나 불필요한 복잡성을 가지는 확률적 샘플링 없이도 효과적이라는 것을 발견했습니다.

 

3.2 상상 속에서의 강화 학습

3.1절의 확산 모델을 바탕으로, 우리는 이제 강화 학습 에이전트를 상상 속에서 훈련하기 위해 필요한 보상 및 종료 모델을 포함하여 우리의 월드 모델을 완성합니다. 보상과 종료를 추정하는 것은 스칼라 예측 문제이기 때문에, 우리는 부분 관측성을 처리하기 위해 표준 CNN(LeCun et al., 1989; He et al., 2015)과 LSTM(Hochreiter and Schmidhuber, 1997; Gers et al., 2000) 계층으로 구성된 별도의 모델 Rψ를 사용합니다. 강화 학습 에이전트는 정책 및 가치 헤드를 가진 공유 CNN-LSTM으로 매개변수화된 액터-크리틱 네트워크를 포함합니다. 정책 πϕ는 가치 기준과 함께 REINFORCE로 훈련되며, 가치 네트워크 Vϕ는 Micheli et al.(2023)과 유사하게 λ-반환을 사용하는 벨먼 오류로 훈련됩니다. 에이전트는 2.1절에서 설명한 대로 상상 속에서 전적으로 훈련됩니다. 에이전트는 데이터 수집을 위해서만 실제 환경과 상호작용합니다. 각 수집 단계 이후, 현재 월드 모델은 지금까지 수집된 모든 데이터를 훈련하여 업데이트됩니다. 그런 다음, 에이전트는 업데이트된 월드 모델 환경에서 강화 학습으로 훈련되며, 이러한 단계들이 반복됩니다. 이 절차는 알고리즘 1에 자세히 설명되어 있으며, Kaiser et al.(2019), Hafner et al.(2020), Micheli et al.(2023)과 유사합니다. 아키텍처 세부사항, 하이퍼파라미터, RL 목표는 각각 부록 D, E, F에 제공합니다.

 

4. 실험

4.1 Atari 100k 벤치마크

다이아몬드(diamond)의 종합적인 평가를 위해, 우리는 26개의 게임으로 구성된 Atari 100k 벤치마크(Kaiser et al., 2019)를 사용했습니다. 이 벤치마크는 다양한 에이전트의 능력을 테스트합니다. 각 게임에서 에이전트는 평가 전에 학습하기 위해 환경에서 10만 번의 행동만 허용되며, 이는 대략 2시간의 인간 게임 플레이와 동일합니다. 참고로, 제한이 없는 Atari 에이전트들은 보통 5천만 단계 동안 훈련되며, 이는 경험 측면에서 500배의 증가를 의미합니다. 우리는 각 게임에서 5개의 무작위 시드로 다이아몬드를 처음부터 훈련시켰습니다. 각 실행은 약 12GB의 VRAM을 사용하였으며, Nvidia RTX 4090 한 대에서 약 2.9일(총 1.03 GPU년)이 소요되었습니다.

 

표 1: Atari 100k 벤치마크의 26개 게임에서 2시간의 실시간 경험 후 얻은 반환값과 인간-정규화된 집계 지표입니다. 볼드체로 표시된 숫자는 가장 성능이 좋은 방법을 나타냅니다. 다이아몬드는 평균 점수 측면에서 다른 월드 모델 기반의 베이스라인보다 뛰어난 성능을 보였습니다.

 

우리는 표 1에서 다이아몬드를 포함해 storm(Zhang et al., 2023), DreamerV3(Hafner et al., 2023), iris(Micheli et al., 2023), twm(Robine et al., 2023), 그리고 SimPle(Kaiser et al., 2019) 등 다른 최근 월드 모델 기반의 방법들과 비교했습니다. bbf(Schwarzer et al., 2023)와 EfficientZero(Ye et al., 2021)를 포함한 모델-프리 및 탐색 기반 방법들과의 더 넓은 비교는 부록 J에 제공됩니다. 이들 방법은 현재 이 벤치마크에서 가장 성능이 좋은 방법들입니다. bbf와 EfficientZero는 우리의 접근 방식과는 직접적으로 비교할 수 없는 독립적인 기술을 사용합니다. 예를 들어 bbf는 하이퍼파라미터 스케줄링과 결합된 주기적인 네트워크 리셋을 사용하며, EfficientZero는 계산 비용이 많이 드는 Monte-Carlo 트리 탐색을 사용합니다. 이러한 추가적인 구성 요소들을 우리의 월드 모델과 결합하는 것은 미래 연구의 흥미로운 방향이 될 것입니다.

 

4.2 Atari 100k 벤치마크의 결과

그림 2: 평균 및 사분위수 평균 인간 정규화 점수(HNS). 다이아몬드(파란색)는 평균 HNS 1.46과 IQM 0.64를 달성했습니다.

 

표 1은 모든 게임에 대한 점수와 인간-정규화 점수(HNS)의 평균 및 사분위수 평균(IQM)을 제공합니다(Wang et al., 2016). Agarwal et al.(2021)의 점 추정 한계에 대한 권고에 따라, 우리는 그림 2에서 평균 및 IQM에 대한 층화된 부트스트랩 신뢰 구간을 제공하며, 부록 H에는 성능 프로파일과 추가적인 메트릭도 제공합니다.

 

우리의 결과는 다이아몬드가 벤치마크 전반에서 강력한 성능을 발휘함을 보여줍니다. 다이아몬드는 11개의 게임에서 인간 플레이어를 능가하였으며, 평균 HNS 1.46의 초인적 성능을 달성하여 월드 모델 내에서 완전히 훈련된 에이전트들 중에서 새로운 최고 성과를 기록했습니다. 다이아몬드는 또한 IQM에서 storm과 대등한 성과를 보였으며, 다른 모든 베이스라인보다 높은 성과를 보였습니다. 다이아몬드는 특히 Asterix, Breakout, Road Runner와 같이 작은 세부 사항을 포착하는 것이 중요한 환경에서 뛰어난 성능을 보였습니다. 우리는 월드 모델의 시각적 품질에 대한 추가적인 정성적 분석을 5.3장에서 제공합니다.

 

5. 분석

5.1 확산 프레임워크 선택

2장에서 설명한 바와 같이, 우리는 원칙적으로 월드 모델에서 어떤 확산 모델 변형이든 사용할 수 있습니다. 다이아몬드는 3장에서 설명한 대로 edm(Karras et al., 2022)을 사용하고 있지만, ddpm(Ho et al., 2020)도 많은 이미지 생성 응용에서 사용된 자연스러운 후보입니다(Rombach et al., 2022; Nichol and Dhariwal, 2021). 이 섹션에서는 이러한 설계 결정을 정당화합니다.

 

ddpm과 우리의 edm 구현을 공정하게 비교하기 위해, 우리는 두 변형을 동일한 네트워크 아키텍처로, 게임 Breakout에서 전문가 정책을 사용해 수집한 100k 프레임의 공유 정적 데이터셋으로 훈련했습니다. 2.3장에서 논의한 바와 같이, 디노이징 단계 수는 월드 모델의 추론 비용과 직접적으로 관련되며, 단계 수가 적으면 상상된 궤적에서 에이전트를 훈련하는 비용이 줄어듭니다. Ho et al.(2020)은 천 번의 디노이징 단계를 사용하며, Rombach et al.(2022)은 Stable Diffusion에서 수백 단계를 사용합니다. 그러나 우리의 월드 모델이 다른 월드 모델 베이스라인과 계산적으로 비교 가능하기 위해서는(예를 들어 iris는 각 시간 단계에서 16개의 NFE가 필요합니다) 최대 수십 개의 디노이징 단계가 필요하며, 가능하다면 더 적은 단계가 바람직합니다. 불행히도, 디노이징 단계 수를 너무 낮게 설정하면 시각적 품질이 저하되어 누적 오류를 초래합니다.

 

확산 변형들의 안정성을 조사하기 위해, 우리는 디노이징 단계 수 n ≤ 10에 대해 최대 t = 1000 타임스텝까지 자귀회귀적으로 생성된 상상된 궤적을 그림 3에 표시했습니다. ddpm을 사용하는 경우(그림 3(a)), 이와 같은 설정에서는 심각한 누적 오류가 발생하여 월드 모델이 빠르게 분포 밖으로 벗어나는 것을 확인할 수 있습니다. 반면, edm 기반의 확산 월드 모델(그림 3(b))은 디노이징 단계를 하나만 사용해도 장기적인 시간 수평선에서 훨씬 더 안정적으로 보입니다. 이러한 누적 오류에 대한 정량적 분석은 부록 K에 제공됩니다.

 

(a) ddpm 기반 월드 모델 궤적.

 

(b) edm 기반 월드 모델 궤적.

그림 3: ddpm(왼쪽)과 edm(오른쪽)을 기반으로 하는 확산 월드 모델의 상상된 궤적. t = 0에서의 초기 관측은 동일하며, 각 행은 디노이징 단계 수 n이 감소하는 것을 나타냅니다. ddpm 기반 생성에서는 누적 오류가 발생하며, 디노이징 단계 수가 적을수록 오류가 더 빨리 누적됩니다. 반면, 우리의 edm 기반 월드 모델은 n = 1인 경우에도 훨씬 더 안정적인 모습을 보입니다.

 

이 놀라운 결과는 식 7에서 설명된 개선된 훈련 목표 덕분으로, ddpm에서 사용된 더 단순한 노이즈 예측 목표와 비교됩니다. 노이즈 예측은 중간 정도의 노이즈 수준에서는 잘 작동하지만, 이 목표는 노이즈가 지배적일 때 모델이 항등 함수를 학습하게 만듭니다(σ_noise ≫ σ_data ⟹ ξ_θ(𝐱ₜ₊₁_τ, yₜ_τ) → 𝐱ₜ₊₁_τ), 여기서 ξ_θ는 ddpm의 노이즈 예측 네트워크입니다. 이는 샘플링 절차의 시작 부분에서 점수 함수에 대한 나쁜 추정을 제공하며, 생성 품질을 저하시켜 누적 오류로 이어집니다.

 

반대로, 3.1장에서 설명된 edm의 신호와 노이즈의 적응적 혼합은 모델이 노이즈가 지배적일 때 깨끗한 이미지를 예측하도록 훈련되도록 만듭니다(σ_noise ≫ σ_data ⟹ 𝐅_θ(𝐱ₜ₊₁_τ, yₜ_τ) → 𝐱ₜ₊₁₀). 이는 신호가 없는 경우에도 점수 함수에 대한 더 나은 추정을 제공하므로, 모델은 더 적은 디노이징 단계로도 높은 품질의 생성을 할 수 있게 됩니다. 이는 그림 3(b)에서 볼 수 있듯이 명확하게 나타납니다.

 

5.2 디노이징 단계 수 선택

우리의 edm 기반 월드 모델이 단일 디노이징 단계에서도 매우 안정적이라는 것을 발견했지만(그림 3(b) 마지막 행의 Breakout에서 보여준 바와 같이), 이러한 선택이 일부 경우에서 모델의 시각적 품질을 제한할 수 있음을 여기서 논의합니다. 이에 대한 정량적 분석은 부록 L에 제공합니다.

 

2.2장에서 논의된 것처럼, 우리의 점수 모델은 L2 재구성 손실을 사용해 훈련된 디노이징 오토인코더(Vincent et al., 2008)와 동일합니다. 따라서 최적의 단일 단계 예측은 주어진 노이즈 입력에 대한 가능한 재구성에 대한 기대값이며, 이 후방 분포가 다중 모드일 경우 분포 밖의 결과일 수 있습니다. Breakout과 같은 일부 게임은 단일 디노이징 단계로 정확하게 모델링될 수 있는 결정적 전이를 가지지만(그림 3(b) 참조), 다른 게임에서는 부분적인 관찰 가능성으로 인해 다중 모드 관측 분포가 발생할 수 있습니다. 이러한 경우에는 특정 모드를 향해 샘플링 절차를 유도하기 위해 반복적인 솔버가 필요하며, 이는 그림 4의 Boxing 게임에서 보여집니다. 결과적으로, 우리는 모든 실험에서 n = 3으로 설정하였습니다.

 

그림 4: Boxing에서 단일 단계(위쪽 행)와 다단계(아래쪽 행) 샘플링. 검은색 플레이어의 움직임은 예측할 수 없으므로 단일 단계 디노이징은 가능한 결과 사이에서 보간하여 흐릿한 예측을 만듭니다. 반면, 다단계 샘플링은 특정 모드를 향해 생성 과정을 유도하여 선명한 이미지를 생성합니다. 흥미롭게도, 정책은 흰색 플레이어를 제어하므로 그의 행동은 월드 모델에 알려져 있습니다. 이러한 정보는 모호성을 제거하며, 따라서 단일 단계와 다단계 샘플링 모두 흰색 플레이어의 위치를 정확히 예측함을 관찰할 수 있습니다.

 

5.3 iris와의 시각적 질적 비교

이제 우리는 iris(Micheli et al., 2023)와 비교합니다. iris는 이미지들을 이산 토큰으로 변환하기 위해 이산 오토인코더(Van Den Oord et al., 2017)를 사용하고, 이 토큰들을 자귀회귀 변환기(Radford et al., 2019)로 시간에 따라 구성하는 잘 알려진 월드 모델입니다. 공정한 비교를 위해, 우리는 두 월드 모델을 전문가 정책으로 수집한 100k 프레임의 동일한 정적 데이터셋에서 훈련했습니다. 이 비교는 아래의 그림 2에 표시되어 있습니다.

 

(a) iris

(b) diamond

그림 5: iris(왼쪽)와 diamond(오른쪽)으로 상상된 연속 프레임. 하얀 박스는 프레임 간의 불일치를 강조하며, 이는 iris로 생성된 궤적에서만 발생하는 것을 볼 수 있습니다. Asterix(위쪽 행)에서는 적(오렌지색)이 두 번째 프레임에서 보상(빨간색)으로 바뀌었다가, 세 번째 프레임에서는 다시 적으로, 네 번째 프레임에서는 다시 보상으로 바뀝니다. Breakout(가운데 행)에서는 벽돌과 점수가 프레임 간에 불일치합니다. Road Runner(아래쪽 행)에서는 보상(도로 위의 작은 파란색 점들)이 프레임 간에 일관되게 렌더링되지 않습니다. 이러한 불일치는 diamond에서는 전혀 발생하지 않으며, Breakout에서는 빨간 벽돌이 부서질 때 점수가 +7로 신뢰성 있게 업데이트됩니다.

 

그림 2에서 볼 수 있듯이, diamond가 상상한 궤적은 일반적으로 iris가 상상한 궤적보다 시각적 품질이 더 높고 실제 환경에 더 충실합니다. 특히, iris가 생성한 궤적은 프레임 간의 시각적 불일치(하얀 박스로 강조됨)를 포함하고 있으며, 적이 보상으로 표시되거나 그 반대의 경우가 있습니다. 이러한 불일치는 생성된 이미지에서 몇 개의 픽셀만 차지할 수 있지만, 강화 학습에 중요한 영향을 미칠 수 있습니다. 예를 들어, 에이전트는 일반적으로 보상을 목표로 하고 적을 피해야 하므로, 이러한 작은 시각적 불일치는 최적의 정책을 학습하는 것을 더 어렵게 만들 수 있습니다.

 

시각적 세부 사항의 이러한 일관성 개선은 일반적으로 이러한 게임들에서 더 높은 에이전트 성능으로 나타납니다(표 1 참조). 이러한 방법들의 에이전트 구성 요소가 유사하기 때문에, 이 개선은 월드 모델 덕분이라고 볼 수 있습니다.

 

마지막으로, 이 개선이 단순히 계산의 증가로 인한 것이 아니라는 점을 강조하고자 합니다. 두 월드 모델 모두 동일한 해상도(64×64)에서 프레임을 렌더링하며, diamond는 프레임당 3개의 NFE만 필요로 하지만 iris는 프레임당 16개의 NFE가 필요합니다. 이 점은 diamond가 iris보다 훨씬 적은 파라미터를 가지고 있고 훈련하는 데도 시간이 적게 걸린다는 사실로도 확인할 수 있으며, 이에 대한 자세한 내용은 부록 H에 제공합니다.

 

6. Counter-Strike: Global Offensive로 확장된 확산 월드 모델

이 섹션은 NeurIPS 채택 이후, 이후 CS:GO 실험에 대한 커뮤니티의 관심에 따라 추가되었습니다.

 

다이아몬드의 확산 월드 모델이 더 복잡한 3D 환경을 모델링할 수 있는 능력을 조사하기 위해, 우리는 이 월드 모델을 인기 있는 비디오 게임 Counter-Strike: Global Offensive (CS:GO)의 정적 데이터를 사용하여 독립적으로 훈련했습니다. 우리는 Pearce와 Zhu(2022)가 맵 Dust II에서 16Hz로 캡처한 온라인 인간 게임플레이의 5.5M 프레임(95시간)으로 구성된 온라인 데이터셋을 사용했습니다. 이 중 0.5M 프레임(500개의 에피소드, 8시간에 해당)을 테스트용으로 무작위로 보류하고, 나머지 5M 프레임(87시간)을 훈련에 사용했습니다. 이 실험에는 강화 학습 에이전트나 온라인 데이터 수집이 포함되지 않았습니다.

 

계산 비용을 줄이기 위해, 월드 모델링을 위한 해상도를 (280×150)에서 (56×30)으로 줄였습니다. 이후, 생성된 이미지를 원래 해상도에서 개선하기 위해 업샘플러로 작은 두 번째 확산 모델을 도입했습니다(Saharia et al., 2022b). 우리는 U-Net의 채널을 확장하여 Atari 모델의 4M 파라미터에서 CS:GO 모델의 381M 파라미터(이 중 업샘플러는 51M)를 증가시켰습니다. 이 결합된 모델은 RTX 4090에서 12일 동안 훈련되었습니다.

 

마지막으로, 우리는 확률적 샘플링을 도입하고 업샘플러의 디노이징 단계 수를 10으로 늘렸습니다. 이는 생성물의 시각적 품질을 개선하는 데 도움이 되었으며, 동역학 모델은 그대로 유지했습니다(특히 여전히 3개의 디노이징 단계만 사용). 이를 통해 시각적 품질과 추론 비용 간의 합리적인 균형을 달성할 수 있었으며, 모델은 RTX 3090에서 10Hz로 작동합니다. 모델의 대표적인 생성물은 아래의 그림 6에 제공되어 있습니다.

 

그림 6: 키보드와 마우스를 사용하여 다이아몬드의 확산 월드 모델 내에서 플레이하는 사람들로부터 캡처된 이미지입니다. 이 모델은 87시간의 정적 Counter-Strike: Global Offensive (CS:GO) 게임플레이(Pearce and Zhu, 2022)에서 훈련되었으며, 인기 있는 맵인 Dust II에 대해 상호작용이 가능한 신경망 게임 엔진을 생성했습니다. 비디오로 보려면 https://diamond-wm.github.io에서 확인하십시오.

 

우리는 모델이 수백 타임스텝 동안 안정적인 궤적을 생성할 수 있음을 확인했지만, 맵에서 자주 방문되지 않는 영역에서는 분포에서 벗어날 가능성이 더 높아진다는 것도 발견했습니다. 모델의 제한된 메모리로 인해 벽에 접근하거나 가시성을 잃게 되면 모델이 현재 상태를 잊고 대신 새로운 무기나 맵의 다른 영역을 생성할 수 있습니다. 흥미롭게도, 우리는 모델이 점프의 효과를 장면의 지형에 일반화하면서 연속적인 점프를 잘못 허용한다는 것을 발견했습니다. 이는 훈련된 게임플레이에서 여러 번의 점프가 충분히 자주 등장하지 않아, 모델이 공중에서의 점프를 무시해야 한다는 것을 학습하지 못했기 때문입니다. 우리는 모델과 데이터의 확장이 이러한 한계 중 다수를 해결할 수 있을 것으로 예상하지만, 모델의 메모리 문제는 여전히 남아 있을 것입니다. CS:GO 월드 모델의 능력에 대한 정량적 측정 및 이러한 한계를 해결하려는 시도는 향후 연구에 남겨둡니다.

 

7. 관련 연구 월드 모델
신경망 월드 모델의 상상 속에서 강화 학습(RL)을 수행하는 아이디어는 Ha와 Schmidhuber(2018)에 의해 처음 도입되었습니다. SimPLe(Kaiser et al., 2019)은 월드 모델을 Atari에 적용하고, 샘플 효율성에 초점을 맞춘 Atari 100k 벤치마크를 소개했습니다. Dreamer(Hafner et al., 2020)는 반복적 상태 공간 모델(RSSM)의 잠재 공간에서 강화 학습을 도입했으며, DreamerV2(Hafner et al., 2021)는 이산 잠재 변수 사용이 누적 오류를 줄이는 데 도움이 될 수 있음을 보여주었습니다. DreamerV3(Hafner et al., 2023)는 고정된 하이퍼파라미터로 다양한 도메인에서 인간 수준의 성능을 달성했습니다. TWM(Robine et al., 2023)은 DreamerV2의 RSSM을 변형 아키텍처에 사용하도록 조정했으며, STORM(Zhang et al., 2023)은 DreamerV3를 비슷한 방식으로 조정하지만 다른 토큰화 방식을 사용합니다. 대안적으로, IRIS(Micheli et al., 2023)는 이산 오토인코더를 사용해 이미지 토큰의 언어를 구축하고, 이를 자귀회귀 변환기를 통해 시간에 따라 구성합니다.

 

생성 비전 모델
이러한 월드 모델과 이미지 생성 모델 사이에는 유사점이 있어, 생성 비전 모델의 발전이 월드 모델링에 이점을 제공할 수 있음을 시사합니다. 자연어 처리에서 변형기의 부상에 따라(Vaswani et al., 2017; Devlin et al., 2018; Radford et al., 2019), VQGAN(Esser et al., 2021)과 DALL·E(Ramesh et al., 2021)는 이산 오토인코더(Van Den Oord et al., 2017)를 사용해 이미지를 이산 토큰으로 변환하고, 자귀회귀 변형기의 시퀀스 모델링 능력을 활용해 강력한 텍스트-이미지 생성 모델을 구축했습니다. 동시에, 확산 모델(Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020)이 인기를 얻으며(Dhariwal and Nichol, 2021; Rombach et al., 2022), 고해상도 이미지 생성의 주요 패러다임이 되었습니다(Saharia et al., 2022a; Ramesh et al., 2022; Podell et al., 2023).

 

비디오 생성 방법의 최근 개발에서도 동일한 트렌드가 있었습니다. VideoGPT(Yan et al., 2021)는 이산 오토인코더와 자귀회귀 변형기를 결합해 최소한의 비디오 생성 아키텍처를 제공합니다. Godiva(Wu et al., 2021)는 텍스트 조건화를 가능하게 하여 유망한 일반화를 달성합니다. Phenaki(Villegas et al., 2023)는 순차적 프롬프트 조건화를 통해 임의 길이의 비디오 생성을 허용합니다. TECO(Yan et al., 2023)는 MaskGit(Chang et al., 2022)을 사용해 자귀회귀 모델링을 개선하고 입력 시퀀스 임베딩을 압축하여 더 긴 시간 종속성을 가능하게 합니다. 또한, 확산 모델은 고품질이지만 짧은 비디오 생성을 위해 3D U-Net을 사용해 비디오 생성에서 부활을 이루었습니다(Singer et al., 2023; Bar-Tal et al., 2024). 최근에는 DiT(Peebles and Xie, 2023)과 Sora(Brooks et al., 2024)와 같은 변형기 기반 확산 모델이 각각 이미지 및 비디오 생성에서 향상된 확장성을 보여주었습니다.

 

강화 학습을 위한 확산
확산 모델을 강화 학습과 결합하는 것에 대한 많은 관심이 있었습니다. 여기에는 확산 모델의 유연성을 정책으로 활용하는 것(Wang et al., 2022; Ajay et al., 2022; Pearce et al., 2023), 계획자로 사용하는 것(Janner et al., 2022; Liang et al., 2023), 보상 모델로 사용하는 것(Nuti et al., 2023), 오프라인 RL에서 데이터 증강을 위한 궤적 모델링으로 사용하는 것(Lu et al., 2023; Ding et al., 2024; Jackson et al., 2024)이 포함됩니다. 다이아몬드는 상상 속에서 온라인 학습을 위해 확산 모델을 월드 모델로 사용하는 최초의 사례입니다.

 

생성형 게임 엔진
완전히 신경망에서 실행되는 플레이 가능한 게임은 최근 그 범위가 확장되고 있습니다. GameGAN(Kim et al., 2020)은 GAN(Goodfellow et al., 2014)을 사용해 게임의 생성 모델을 학습하며, Bamford와 Lucas(2020)는 Neural GPU(Kaiser and Sutskever, 2015)를 사용합니다. 동시에 진행된 연구에는 이미지 프롬프트로부터 플레이 가능한 플랫폼 환경을 생성하는 Genie(Bruce et al., 2024)와, 유사하게 확산 모델을 활용하여 대규모로 게임 DOOM의 고해상도 시뮬레이터를 얻는 GameNGen(Valevski et al., 2024)이 있습니다.

 

8. 한계점
우리 연구의 한계를 미래 연구를 위해 세 가지로 정리했습니다. 첫째, 우리의 주요 평가는 이산 제어 환경에 집중되어 있으며, 다이아몬드를 연속적인 도메인에 적용하면 추가적인 통찰을 제공할 수 있습니다. 둘째, 조건화를 위한 프레임 스태킹은 과거 관측의 기억을 제공하는 최소한의 메커니즘입니다. Peebles와 Xie(2023)와 같은 접근 방식을 사용해 환경 시간에 대해 자귀회귀 변형기를 통합하면 더 장기적인 기억과 더 나은 확장성을 가능하게 할 수 있습니다. 부록 M에서 잠재적인 교차 주의 아키텍처에 대한 초기 조사를 포함했지만, 초기 실험에서는 프레임 스태킹이 더 효과적이었습니다. 셋째, 보상/종료 예측의 확산 모델로의 통합 가능성은 향후 연구에 남겨둡니다. 이러한 목표들을 결합하고 확산 모델에서 표현을 추출하는 것은 간단하지 않으며(Luo et al., 2023; Xu et al., 2023), 우리의 월드 모델을 불필요하게 복잡하게 만들 수 있기 때문입니다.

 

9. 결론 및 광범위한 영향
우리는 확산 월드 모델에서 훈련된 강화 학습 에이전트인 다이아몬드를 소개했습니다. 우리는 확산을 월드 모델링에 적응시키기 위해 선택한 주요 설계 결정과, 적은 디노이징 단계로도 월드 모델을 긴 시간 수평선에서 안정적으로 유지하기 위한 방법을 설명했습니다. 다이아몬드는 잘 알려진 Atari 100k 벤치마크에서 평균 인간 정규화 점수 1.46을 달성하며, 월드 모델 내에서 완전히 훈련된 에이전트 중 새로운 최고 기록을 세웠습니다. 우리는 일부 게임에서 개선된 성능을 분석했고, 이는 중요한 시각적 세부 사항을 더 잘 모델링한 결과임을 확인했습니다. 우리는 또한 다이아몬드의 확산 월드 모델이 3D 환경을 성공적으로 모델링하고 정적 Counter-Strike: Global Offensive 게임플레이를 훈련하여 실시간 신경망 게임 엔진 역할을 할 수 있음을 증명했습니다.

 

월드 모델은 실제 세계에서 에이전트를 훈련할 때 샘플 효율성과 안전 문제를 해결할 수 있는 유망한 방향을 제시합니다. 그러나 월드 모델의 불완전성은 최적이 아닌 에이전트 행동이나 예기치 않은 행동을 초래할 수 있습니다. 우리는 보다 충실하고 상호작용 가능한 월드 모델의 개발이 이러한 위험을 줄이는 더 광범위한 노력에 기여할 수 있기를 바랍니다.

 

감사의 말 및 자금 공개
Andrew Foong, Bálint Máté, Clément Vignac, Maxim Peter, Pedro Sanchez, Rich Turner, Stéphane Nguyen, Tom Lee, Trevor McInroe, Weipu Zhang에게 유익한 논의와 의견을 주신 것에 감사드립니다. Adam과 Eloi는 Microsoft Research Cambridge의 인턴십에서 만나 이 프로젝트의 초기 단계에서 지원해준 Game Intelligence 팀, 즉 Anssi Kanervisto, Dave Bignell, Gunshi Gupta, Katja Hofmann, Lukas Schäfer, Raluca Georgescu, Sam Devlin, Sergio Valcarcel Macua, Shanzheng Tan, Tabish Rashid, Tarun Gupta, Tim Pearce, Yuhan Cao에게 감사하며, 멋진 여름을 보냈음을 전합니다.

 

2405.12399v2.pdf

4.15MB