[논문 리뷰] Facing Off World Model Backbones: RNNs, Transformers, and S4

Fei Deng, Junyeong Park외 1명|arXiv (Cornell University)|2023. 07. 05.

Reinforcement Learning in Robotics인용 수 4

한 줄 요약

이 논문은 S4와 S5와 같은 병렬화 가능한 구조적 상태공간모델(PSSM)과 호환되는 세계모델 프레임워크인 S4WM을 소개한다. 이는 효율적인 장기적 시각적 시퀀스 모델링을 가능하게 한다. S4WM은 장기적 상상력, 맥락 기반 기억 회상, 보상 예측, 기억 기반 추론에서 RNN 및 Transformer 기반 세계모델을 능가하며, 학습 속도가 빠르고 시퀀스 생성도 더 효율적이다.

ABSTRACT

World models are a fundamental component in model-based reinforcement learning (MBRL). To perform temporally extended and consistent simulations of the future in partially observable environments, world models need to possess long-term memory. However, state-of-the-art MBRL agents, such as Dreamer, predominantly employ recurrent neural networks (RNNs) as their world model backbone, which have limited memory capacity. In this paper, we seek to explore alternative world model backbones for improving long-term memory. In particular, we investigate the effectiveness of Transformers and Structured State Space Sequence (S4) models, motivated by their remarkable ability to capture long-range dependencies in low-dimensional sequences and their complementary strengths. We propose S4WM, the first world model compatible with parallelizable SSMs including S4 and its variants. By incorporating latent variable modeling, S4WM can efficiently generate high-dimensional image sequences through latent imagination. Furthermore, we extensively compare RNN-, Transformer-, and S4-based world models across four sets of environments, which we have tailored to assess crucial memory capabilities of world models, including long-term imagination, context-dependent recall, reward prediction, and memory-based reasoning. Our findings demonstrate that S4WM outperforms Transformer-based world models in terms of long-term memory, while exhibiting greater efficiency during training and imagination. These results pave the way for the development of stronger MBRL agents.

연구 동기 및 목표

• 고차원적 시각적 시퀀스에 대해 병렬화 가능한 상태공간모델(PSSM)인 S4와 S5와 호환되는 일반적인 세계모델 프레임워크를 개발하는 것.
• 장기적 메모리 집약적 과제에서 RNN, Transformer, S4를 시각적 세계모델의 백본으로 사용했을 때의 효과성을 조사하는 것.
• 장기적 상상력, 맥락 기반 기억 회상, 보상 예측, 기억 기반 추론 등 핵심 메모리 기능에서 다양한 백본 아키텍처의 성능을 평가하고 비교하는 것.
• S4WM이 RNN 및 Transformer 기반 베이스라인에 비해 장기적 메모리와 추론 효율성 측면에서 뛰어나다는 것을 입증하는 것.

제안 방법

• 이미지 시퀀스 생성을 위한 S4를 백본으로 사용하는 변분 추론 기반의 확률적 잠재변수 모델인 S4WM을 제안한다.
• 매개변수화된 상태 전이를 갖는 구조적 상태공간모델(SSM)을 사용하여 2차 이하의 복잡도를 확보하고 병렬 학습을 가능하게 한다.
• 고차원 관측치를 압축하기 위해 잠재공간 블로킹을 도입하여, 잠재적 상상력을 통한 효율적 이미지 시퀀스 생성을 가능하게 한다.
• 장기적 수평, 맥락 기반 기억, 보상 예측, 다중 도어 키 환경의 네 가지 특화된 환경을 설계하여 메모리 기능을 평가한다.
• 재구성 및 생성 MSE, 상상 정확도, 보상 예측 성능을 평가 지표로 사용한다.
• 오프라인 프로빙 및 모델 예측 제어(MPC) 실험을 통해 S4WM의 다운스트림 정책 학습 잠재력을 검증한다.

실험 결과

연구 질문

• RQ1 S4 기반 세계모델은 고차원 시각적 시퀀스에서 장거리 의존성을 효과적으로 모델링할 수 있는가?
• RQ2 S4WM은 장기적 상상력과 메모리 유지 측면에서 RNN 및 Transformer 기반 세계모델과 비교해 어떻게 성능을 내는가?
• RQ3 S4WM은 시각적 세계모델링에서 RNN과 Transformer보다 더 높은 학습 효율성과 추론 속도를 달성하는가?
• RQ4 S4WM은 동적이고 맥락 기반의 추론 과제에서 정확한 메모리 업데이트를 유지할 수 있는가?
• RQ5 S4WM은 부분 관측 환경에서의 다단계 계획과 같은 복잡한 메모리 집약적 추론 과제로 일반화되는가?

주요 결과

• Long Horizon 과제에서 S4WM은 가장 높은 상상 정확도를 기록했으며, RNN과 Transformer는 장기간에 걸쳐 일관된 상태 표현을 유지하지 못해 성능이 떨어졌다.
• Contextual Recall 과제에서 S4WM은 모든 맥락 길이에서 높은 정확도를 유지했고, TSSM-XL과 RSSM-TBTT는 100단계를 초과해 성능이著しく 떨어졌다.
• S4WM은 유일하게 상상 내에서 보상을 정확히 예측했으며, TSSM-XL은 장기간 시퀀스에서 실패했고, RSSM-TBTT는 랜덤 추측 수준에 가까운 성능을 보였다.
• Multi Doors Keys 환경에서 S4WM은 가장 복잡한 설정(7개의 키, 654개의 쿼리 단계)에서 가장 낮은 생성 MSE(0.10)를 기록했으며, TSSM-XL(9.24)과 RSSM-TBTT(6.28)를 압도적으로 앞섰다.
• S4WM은 RNN 및 Transformer 기반 베이스라인보다 학습 속도가 빠르며, RNN보다 더 높은 상상 처리량을 보여, 뛰어난 계산 효율성을 입증했다.
• 오프라인 프로빙 및 MPC 실험 결과, S4WM은 RSSM보다 높은 작업 성공률을 기록하여, MBRL에서 정책 학습 잠재력이 뛰어나다는 것을 확인했다.