[논문 리뷰] STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning

Weipu Zhang, Gang Wang외 3명|arXiv (Cornell University)|2023. 10. 14.

Reinforcement Learning in Robotics인용 수 8

한 줄 요약

STORM은 lookahead search 없이 샘플 효율성과 학습 속도를 개선하기 위해 VAE 인코더가 있는 Transformer 기반의 확률적 세계 모델을 도입하여 Atari 100k에서 새로운 평균 인간-정규화 점수를 달성하고 실시간 학습 속도가 빨라졌습니다.

ABSTRACT

Recently, model-based reinforcement learning algorithms have demonstrated remarkable efficacy in visual input environments. These approaches begin by constructing a parameterized simulation world model of the real environment through self-supervised learning. By leveraging the imagination of the world model, the agent's policy is enhanced without the constraints of sampling from the real environment. The performance of these algorithms heavily relies on the sequence modeling and generation capabilities of the world model. However, constructing a perfectly accurate model of a complex unknown environment is nearly impossible. Discrepancies between the model and reality may cause the agent to pursue virtual goals, resulting in subpar performance in the real environment. Introducing random noise into model-based reinforcement learning has been proven beneficial. In this work, we introduce Stochastic Transformer-based wORld Model (STORM), an efficient world model architecture that combines the strong sequence modeling and generation capabilities of Transformers with the stochastic nature of variational autoencoders. STORM achieves a mean human performance of $126.7\%$ on the Atari $100$k benchmark, setting a new record among state-of-the-art methods that do not employ lookahead search techniques. Moreover, training an agent with $1.85$ hours of real-time interaction experience on a single NVIDIA GeForce RTX 3090 graphics card requires only $4.3$ hours, showcasing improved efficiency compared to previous methodologies.

연구 동기 및 목표

• 시각 환경에서 모델 기반 RL의 샘플 효율성 향상을 고무한다.
• 트랜스포머와 확률적 잠재 표현을 활용하는 효율적인 세계 모델을 개발한다.
• Atari 100k에서 성능을 유지하거나 향상시키면서 예측 오차 누적과 학습 시간을 줄인다.

제안 방법

• 관측치를 이산적 VAE 인코더로 매핑하여 확률적 잠재 z_t(32 카테고리 × 32 클래스)로 변환한다.
• z_t와 행동 a_t를 하나의 토큰 e_t로 융합하고 GPT 유사 트랜스포머를 시퀀스 모델로 삼아 h_t를 생성한다.
• h_t로부터 MLP 헤드를 사용해 보상, 연속 여부 플래그, 그리고 다음 잠재 분포를 예측한다.
• 재구성, 보상, 연속성, 다이내믹스(KL), 표현(KL) 항을 결합한 자기지도 손실로 세계 모델을 학습한다(베타 가중치 포함).
• 상상된 경험에서 전적으로 정책을 학습하고 DreamerV3 스타일의 actor-critic 목표를 사용하며 lambda-return 과 KV-캐시 가속 추론을 활용한다.

실험 결과

연구 질문

• RQ1 stochastic Transformer 기반 세계 모델이 Atari 100k에서 RNN 기반 또는 Transformer-XL 기반 모델보다 더 잘 작동할 수 있는가?
• RQ2 이미지당 하나의 확률적 잠재 표현이 정책 학습에 의해 dynamics를 효과적으로 포착하는가?
• RQ3 제안된 손실 설계와 상상 기반 학습이 샘플 효율성과 계산 효율성에 어떤 영향을 미치는가?
• RQ4 encoder 유형, 상태 표현, 트랜스포머 깊이 등의 세계 모델 디자인 선택이 성능에 미치는 영향은 무엇인가?
• RQ5 STORM을 사용하여 실제 환경과의 상호 작용이 제한된 경우에도 높은 성능을 달성할 수 있는가?

주요 결과

• STORM은 Atari 100k에서 평균 인간-정규화 점수 126.7%를 달성하며 lookahead search가 없는 방법 중 새로운 기록을 세웠습니다.
• RTX 3090에서 약 1.85시간의 실제 데이터로 학습하는데 약 4.3시간이 걸려, 이전 방법들보다 효율이 개선되었습니다.
• SimPLe, TWM, IRIS, DreamerV3와 비교하여 STORM은 Transformer 시퀀스 모델링과 확률적 잠재 표현의 이점을 활용해 큰 보상 객체가 있거나 보상 관련 객체가 다수 있는 게임에서 더 나은 성능을 보입니다.
• 시퀀스 모델로 Transformer를 사용하고 단일 확률적 잠재 및 관찰-행동 토큰의 결합이 효과적임이 확인되었으며, Atari 100k에서 더 큰 Transformer 깊이가 반드시 결과를 개선하지는 않습니다.
• 단일 시演 traj를 도입하면 희박보상 게임(Pong 등)에서 탐색을 개선할 수 있지만 밀집 보상 게임(Ms. Pacman)에서는 방해가 될 수 있습니다.