[논문 리뷰] Transformer-based World Models Are Happy With 100k Interactions

Jan Robine, Marc Höftmann외 2명|arXiv (Cornell University)|2023. 03. 13.

Reinforcement Learning in Robotics인용 수 4

한 줄 요약

이 논문은 단지 100만 번의 환경 상호작용만을 사용하여 Atari 100k 벤치마크에서 최신 기술 성능을 달성하는 Transformer 기반 월드 모델(TWM)을 제안한다. Transformer-XL 아키텍처를 활용해 관측치, 행동, 예측 보상치를 순차적으로 모델링하고, 예측된 보상치를 다시 모델에 피드백함으로써, TWM은 샘플 효율적인 정책 학습을 가능하게 하여 기존의 모델 프리(MF) 및 모델 기반(MB) 방법들을 능가하면서도 추론 시 계산 비용을 낮게 유지한다.

ABSTRACT

Deep neural networks have been successful in many reinforcement learning settings. However, compared to human learners they are overly data hungry. To build a sample-efficient world model, we apply a transformer to real-world episodes in an autoregressive manner: not only the compact latent states and the taken actions but also the experienced or predicted rewards are fed into the transformer, so that it can attend flexibly to all three modalities at different time steps. The transformer allows our world model to access previous states directly, instead of viewing them through a compressed recurrent state. By utilizing the Transformer-XL architecture, it is able to learn long-term dependencies while staying computationally efficient. Our transformer-based world model (TWM) generates meaningful, new experience, which is used to train a policy that outperforms previous model-free and model-based reinforcement learning algorithms on the Atari 100k benchmark.

연구 동기 및 목표

• 일반적으로 인간 학습자보다 훨씬 많은 수십억 번의 환경 상호작용이 필요한 딥 강화학습의 샘플 비효율성 문제를 해결한다.
• 재귀적 월드 모델의 한계를 극복하기 위해 자기주의적 주의(self-attention) 메커니즘을 통해 과거 상태에 대한 직접적이고 장거리의 주의를 가능하게 한다.
• 생성적 월드 모델를 사용해 잠재적 상상 속에서 정책을 학습함으로써 샘플 효율성을 향상시킨다.
• 새로운 손실 함수와 샘플링 전략을 통해 학습을 안정화시켜 일반화 능력과 하이퍼파rameter에 대한 강건성을 향상시킨다.
• Transformer 기반 월드 모델가 실환경 상호작용을 최소한으로(특히 Atari에서 100만 번) 수행하면서도 뛰어난 성능을 달성할 수 있음을 입증한다.

제안 방법

• 고차원 이미지 관측치를 잠재 상태로 압축하기 위해 CNN을 사용하며, 이를 Transformer 입력 시퀀스에 통합한다.
• 잠재 상태, 행동, 예측 보상치를 입력 토큰으로 삼아 Transformer에 입력함으로써, 세 모odal 간에 민첩하게 주의를 기울일 수 있도록 한다.
• 장기적 의존성을 효율적으로 모델링하고 추론 시 계산 비용을 줄이기 위해 상대적 위치 인코딩을 사용한 Transformer-XL 아키텍처를 적용한다.
• Transformer의 은닉 상태에 조건을 두어 다층 퍼셉트론(MLP)을 사용해 다음 잠재 상태, 보상치, 할인 요인을 예측한다.
• 월드 모델의 목적 함수에서 엔트로피와 교차엔트로피 항 간의 트레이드오프를 정밀하게 제어하기 위해 균형 잡힌 KL 발산 손실을 도입한다.
• 정책 엔트로피를 학습 중 안정화시키기 위해 임계값이 적용된 엔트로피 손실을 적용하여 하이퍼파rameter 선택에 대한 민감도를 감소시킨다.
• 최근 경험에 우선순위를 두어 학습 데이터셋의 데이터 분포 이탈을 수정하기 위해 균형 잡힌 샘플링 절차를 구현한다.

실험 결과

연구 질문

• RQ1 100만 번의 환경 상호작용만으로 훈련된 Transformer 기반 월드 모델이 Atari 100k 벤치마크에서 기존의 모델 프리 및 모델 기반 강화학습 방법들을 능가할 수 있는가?
• RQ2 예측된 보상치를 Transformer 입력에 피드백함으로써 월드 모델의 품질과 후속 정책 성능에 어떤 영향을 미치는가?
• RQ3 제안된 균형 잡힌 샘플링 전략은 균일한 샘플링 대비 학습 안정성과 성능 향상에 어느 정도 기여하는가?
• RQ4 임계값이 적용된 엔트로피 손실이 다양한 Atari 게임에서 정책 학습을 안정화시키고 하이퍼파rameter 민감도를 감소시키는가?
• RQ5 Transformer-XL 아키텍처를 사용함으로써 RNN 기반 월드 모델보다 더 나은 장기 의존성 모델링이 가능할까, 동시에 추론 비용은 낮게 유지되는가?

주요 결과

• 제안된 TWM은 Atari 100k 벤치마크에서 인간 정규화 점수 기준 최신 기술 성능을 달성하며, 이전의 모델 프리 및 모델 기반 방법들을 모두 능가한다.
• 예측 보상치를 기반으로 Transformer를 조건화함으로써 다수의 게임에서 성능 향상이 뚜렷하게 관찰되었으며, 이는 추이 분석(ablation studies)를 통해 입증되었다.
• 균형 잡힌 샘플링 절차는 최근 경험에 초점을 맞추어 데이터 분포 이탈을 줄임으로써 정책 성능 향상에 기여한다.
• 임계값이 적용된 엔트로피 손실은 정책 엔트로피를 학습 중 안정화시켜 하이퍼파rameter 튜닝을 단순화하고 게임 간 일반화 능력을 향상시킨다.
• 균형 잡힌 KL 발산 손실은 엔트로피와 교차엔트로피 항 간의 트레이드오프를 정밀하게 제어함으로써 월드 모델 학습을 향상시킨다.
• 실환경 상호작용을 단지 100만 번으로 제한함으로써 높은 성능을 달성하였으며, 정책 추론 시 Transformer를 사용하지 않아 계산 비용이 낮고 샘플 효율성이 뛰어나다는 점을 입증하였다.