[논문 리뷰] TransDreamer: Reinforcement Learning with Transformer World Models

Chang Chen, Yifu Wu외 2명|arXiv (Cornell University)|2022. 02. 19.

Reinforcement Learning in Robotics인용 수 8

한 줄 요약

TransDreamer는 장기적 기억과 계획 능력을 향상시키기 위해 꿈꾸는 모델(Dreamer)의 순환 세계 모델을 대체하는 완전히 트랜스포머 기반의 모델 기반 강화학습 에이전트를 소개한다. 이는 새로운 트랜스포머 상태공간 모델(TSSM)을 도입하여 장기적 의존성을 효과적으로 처리한다. 복잡한 2D 및 3D 시각적 제어 과제에서 장기적 추론이 필요한 경우 Dreamer를 능가하며, 짧은 의존성 구조를 가진 단순 과제에서는 성능이 동일하게 유지된다.

ABSTRACT

The Dreamer agent provides various benefits of Model-Based Reinforcement Learning (MBRL) such as sample efficiency, reusable knowledge, and safe planning. However, its world model and policy networks inherit the limitations of recurrent neural networks and thus an important question is how an MBRL framework can benefit from the recent advances of transformers and what the challenges are in doing so. In this paper, we propose a transformer-based MBRL agent, called TransDreamer. We first introduce the Transformer State-Space Model, a world model that leverages a transformer for dynamics predictions. We then share this world model with a transformer-based policy network and obtain stability in training a transformer-based RL agent. In experiments, we apply the proposed model to 2D visual RL and 3D first-person visual RL tasks both requiring long-range memory access for memory-based reasoning. We show that the proposed model outperforms Dreamer in these complex tasks.

연구 동기 및 목표

• 트랜스포머 기반 세계 모델이 모델 기반 강화학습에서 샘플 효율성과 장기적 추론 능력을 향상시킬 수 있는지 조사하기.
• 이전 연구에서 희박한 보상 환경에서 순수 트랜스포머 정책을 훈련하는 데 어려움이 있음이 나타났기에, 트랜스포머 기반 정책의 안정적인 훈련 문제를 해결하기.
• 병렬 훈련과 효과적인 상상 능력을 지원하는 트랜스포머 호환성 있는 확률적, 액션 조건 기반 세계 모델을 개발하기.
• 향상된 세계 모델이 복잡한 장기적 기억 상호작용이 요구되는 과제에서 더 나은 성능을 내는지 평가하기.
• RNN 기반 모델이 성공할 것으로 예상되는 단순 과제에서도 제안된 프레임워크가 강력한 성능을 유지하는지 검증하기.

제안 방법

• 자기주의 어텐션을 사용해 잠재 공간 내 장기적 의존성을 포착하는 확률적, 액션 조건 기반 세계 모델인 트랜스포머 상태공간 모델(TSSM)을 제안한다.
• TSSM를 확률적 동역학 예측을 위한 사전 및 사후 추론 모두를 지원하도록 적응시켜, 상상 내역에서 정확한 후행 추정을 가능하게 한다.
• 정책과 세계 모델 구성 요소 간에 TSSM 세계 모델을 공유함으로써 일관된 잠재 표현을 갖춘 엔드 투 엔드 훈련을 가능하게 한다.
• 변분 추론 목표(ELBO)를 사용해 TSSM를 훈련하며, 확률적 잠재 상태와 결정적 은닉 상태를 다중 헤드 자기주의 어텐션을 통해 업데이트한다.
• TSSM가 생성한 상상된 궤적을 기반으로 트랜스포머 기반 액터-크리틱 정책을 훈련시켜 샘플 효율적인 학습을 실현한다.
• TSSM를 Dreamer 프레임워크에 통합하여 RSSM를 대체하면서도 전체 MBRL 훈련 루프를 유지한다.

실험 결과

연구 질문

• RQ1 장기적 기억 상호작용이 요구되는 장기적 시각 제어 과제에서 트랜스포머 기반 세계 모델이 순환 모델을 능가할 수 있는가?
• RQ2 안정적이고 엔드 투 엔드로 동작하는 트랜스포머 기반 MBRL 에이전트를 구축할 때의 핵심 아키텍처 및 훈련 과제는 무엇인가?
• RQ3 트랜스포머 기반 세계 모델이 상상된 궤적의 품질, 특히 이미지 생성과 보상 예측에서 향상되는가?
• RQ4 짧은 의존성 구조를 가진 과제에서 트랜스포머 기반 MBRL 에이전트의 성능이 순환 기반 기준(Dreamer)과 어떻게 비교되는가?
• RQ5 TSSM 세계 모델은 다양한 시각 환경에 일반화되며 안정적인 훈련과 계산 효율성을 유지할 수 있는가?

주요 결과

• 5-Ball Dense 환경에서 TransDreamer는 장기적 기억과 정확한 공 순서 추론이 필요한 복잡한 2D 시각 제어 과제에서 Dreamer를 능가한다.
• TransDreamer는 더 선명하고 정확한 상상된 프레임을 생성하며, 공 수거 이벤트 시 보상 급증을 정확히 예측하지만, Dreamer는 흐릿한 이미지와 잘못된 보상 예측을 한다.
• 장기적 의존성이 중요한 후속 상상 단계에서 TSSM 세계 모델은 Dreamer의 RSSM보다 더 정확한 미래 프레임 예측과 보상 추정을 생성한다.
• 단기 기억 과제(예: DMC Cheetah Run 및 아케이드 게임)에서는 TransDreamer가 Dreamer와 유사한 최종 성능을 달성하지만 수렴 속도는 느리다.
• TransDreamer는 더 높은 보상 예측 정확도와 더 높은 시각적 상상 정밀도를 보이며, Dreamer의 RNN 기반 대비 더 강력하고 정밀한 세계 모델임을 시사한다.
• TSSM는 완전히 트랜스포머 기반 강화학습 에이전트의 안정적 훈련을 가능하게 하여, 적절한 아키텍처 설계가 적용될 경우 트랜스포머 기반 MBRL가 실현 가능하고 효과적임을 입증한다.