아타리에서의 세계 모델링을 위한 확산: 시각적 세부 사항의 중요성

Diffusion for World Modeling: Visual Details Matter in Atari

May 20, 2024

저자: Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, François Fleuret

cs.AI

초록

월드 모델(World models)은 안전하고 샘플 효율적인 방식으로 강화 학습 에이전트를 훈련시키기 위한 유망한 접근법입니다. 최근의 월드 모델은 주로 환경 역학을 모델링하기 위해 이산 잠재 변수(discrete latent variables)의 시퀀스에서 작동합니다. 그러나 이러한 간결한 이산 표현으로의 압축은 강화 학습에 중요한 시각적 세부 사항을 무시할 수 있습니다. 동시에, 확산 모델(diffusion models)은 이미지 생성에서 주도적인 접근법으로 자리 잡으며, 이산 잠재 변수를 모델링하는 기존의 잘 정립된 방법들에 도전하고 있습니다. 이러한 패러다임 전환에 영감을 받아, 우리는 DIAMOND(DIffusion As a Model Of eNvironment Dreams)를 소개합니다. DIAMOND는 확산 월드 모델 내에서 훈련된 강화 학습 에이전트입니다. 우리는 확산 모델이 월드 모델링에 적합하도록 만들기 위해 필요한 주요 설계 선택을 분석하고, 개선된 시각적 세부 사항이 어떻게 에이전트 성능 향상으로 이어질 수 있는지 보여줍니다. DIAMOND는 경쟁적인 Atari 100k 벤치마크에서 평균 인간 정규화 점수 1.46을 달성하며, 월드 모델 내에서 완전히 훈련된 에이전트로서 새로운 최고 기록을 세웠습니다. 확산 모델을 이용한 월드 모델링에 대한 미래 연구를 촉진하기 위해, 우리는 코드, 에이전트 및 플레이 가능한 월드 모델을 https://github.com/eloialonso/diamond에서 공개합니다.

English

World models constitute a promising approach for training reinforcement learning agents in a safe and sample-efficient manner. Recent world models predominantly operate on sequences of discrete latent variables to model environment dynamics. However, this compression into a compact discrete representation may ignore visual details that are important for reinforcement learning. Concurrently, diffusion models have become a dominant approach for image generation, challenging well-established methods modeling discrete latents. Motivated by this paradigm shift, we introduce DIAMOND (DIffusion As a Model Of eNvironment Dreams), a reinforcement learning agent trained in a diffusion world model. We analyze the key design choices that are required to make diffusion suitable for world modeling, and demonstrate how improved visual details can lead to improved agent performance. DIAMOND achieves a mean human normalized score of 1.46 on the competitive Atari 100k benchmark; a new best for agents trained entirely within a world model. To foster future research on diffusion for world modeling, we release our code, agents and playable world models at https://github.com/eloialonso/diamond.