데이터 효율적 강화 학습을 위한 트랜스포머 월드 모델 개선
Improving Transformer World Models for Data-Efficient RL
February 3, 2025
저자: Antoine Dedieu, Joseph Ortiz, Xinghua Lou, Carter Wendelken, Wolfgang Lehrach, J Swaroop Guntupalli, Miguel Lazaro-Gredilla, Kevin Patrick Murphy
cs.AI
초록
우리는 Craftax-classic 벤치마크에서 새로운 최고 수준의 성능을 달성하는 모델 기반 강화 학습 접근 방식을 제시합니다. Craftax-classic은 요구되는 요소들이 다양한 일반적인 능력을 나타내는 에이전트들에게 필요한 오픈 월드 2D 서바이벌 게임으로, 강력한 일반화, 심층 탐색, 장기적 추론과 같은 능력을 요구합니다. 샘플 효율성을 향상시키기 위한 신중한 설계 선택들의 연속으로, 우리의 MBRL 알고리즘은 환경 단계 1백만 단계 후에 67.4%의 보상을 달성하여 DreamerV3의 53.2%를 크게 능가하며, 처음으로 65.0%의 인간 수준 성능을 초과합니다. 우리의 방법은 SOTA 모델 무관 기준선을 구축하는 것으로 시작하며, CNN과 RNN을 결합한 새로운 정책 아키텍처를 사용합니다. 그런 다음 표준 MBRL 설정에 세 가지 개선을 추가합니다: (a) "워마업이 포함된 Dyna", 실제 및 가상 데이터로 정책을 훈련하는 방법, (b) 이미지 패치에 대한 "최근 이웃 토크나이저"로, 트랜스포머 세계 모델 (TWM) 입력을 생성하는 방법을 개선하고, (c) "블록 티처 포싱", 다음 타임스텝의 미래 토큰에 대해 TWM이 공동으로 추론할 수 있도록 합니다.
English
We present an approach to model-based RL that achieves a new state of the art
performance on the challenging Craftax-classic benchmark, an open-world 2D
survival game that requires agents to exhibit a wide range of general abilities
-- such as strong generalization, deep exploration, and long-term reasoning.
With a series of careful design choices aimed at improving sample efficiency,
our MBRL algorithm achieves a reward of 67.4% after only 1M environment steps,
significantly outperforming DreamerV3, which achieves 53.2%, and, for the first
time, exceeds human performance of 65.0%. Our method starts by constructing a
SOTA model-free baseline, using a novel policy architecture that combines CNNs
and RNNs. We then add three improvements to the standard MBRL setup: (a) "Dyna
with warmup", which trains the policy on real and imaginary data, (b) "nearest
neighbor tokenizer" on image patches, which improves the scheme to create the
transformer world model (TWM) inputs, and (c) "block teacher forcing", which
allows the TWM to reason jointly about the future tokens of the next timestep.Summary
AI-Generated Summary