[논문 리뷰] StARformer: Transformer with State-Action-Reward Representations for Visual Reinforcement Learning

Jinghuan Shang, Kumara Kahatapitiya외 2명|arXiv (Cornell University)|2021. 10. 12.

Reinforcement Learning in Robotics인용 수 4

한 줄 요약

StARformer는 짧은 시간 윈도우 내에서 상태-행동-보상(StAR) 표현을 명시적으로 모델링함으로써 마르코프 성향의 유도 편향을 주입하는 새로운 트랜스포머 아키텍처를 소개한다. 이는 장기적 시퀀스 모델링을 향상시킨다. ViT 유사한 이미지 패치와 컨볼루션 특징을 조합하고, 번갈아가며 스텝 및 시퀀스 트랜스포머를 통해 특징을 융합함으로써, 오프라인-RL 및 이터레이션 학습 설정에서 Atari 및 DM Control Suite 벤치마크에서 최신 기술을 초월한다.

ABSTRACT

Reinforcement Learning (RL) can be considered as a sequence modeling task: given a sequence of past state-action-reward experiences, an agent predicts a sequence of next actions. In this work, we propose State-Action-Reward Transformer (StARformer) for visual RL, which explicitly models short-term state-action-reward representations (StAR-representations), essentially introducing a Markovian-like inductive bias to improve long-term modeling. Our approach first extracts StAR-representations by self-attending image state patches, action, and reward tokens within a short temporal window. These are then combined with pure image state representations -- extracted as convolutional features, to perform self-attention over the whole sequence. Our experiments show that StARformer outperforms the state-of-the-art Transformer-based method on image-based Atari and DeepMind Control Suite benchmarks, in both offline-RL and imitation learning settings. StARformer is also more compliant with longer sequences of inputs. Our code is available at https://github.com/elicassion/StARformer.

연구 동기 및 목표

• 지역적 상태-행동-보상 전이를 명시적으로 모델링하여 시각적 강화학습에서 장기적 시퀀스 모델링을 향상시키기.
• 국소화된 StAR 표현을 통해 자기주의 주의 메커니즘의 부담을 줄이기 위해 마르코프 성향의 유도 편향을 주입하기.
• 전체 이미지 CNN 특징 대신 ViT 스타일 패치 임베딩을 사용하여 이미지 상태의 세밀한 공간 정보를 유지하기.
• 단기 전이 모델링과 장거리 맥락 학습을 분리함으로써 더 긴 입력 시퀀스로의 일반화 및 확장성 향상시키기.
• 반환-목표, 단계별 보상, 보상 없음(이터레이션) 설정을 포함한 다양한 보상 설정에서 StARformer의 효과성 평가하기.

제안 방법

• 모델은 단일 시간 단계 내에서 이미지 상태 패치, 행동 토큰, 보상 토큰을 자기주의 주의로 처리하여 국소적 StAR 표현을 학습하는 스텝 트랜스포머를 사용한다.
• 이미지 상태는 세밀한 공간 세부 정보를 유지하기 위해 ViT 유사 패치로 인코딩되며, 전역 맥락을 위한 순수한 이미지 특징은 컨볼루션 네트워크를 통해 추출된다.
• 시퀀스 트랜스포머는 전체 궤적에 걸쳐 StAR 표현과 전역 컨볼루션 특징을 결합하여 인과 자기주의 주의를 통해 행동 예측을 생성한다.
• 층별 융합 메커니즘은 각 트랜스포머 층에서 StAR 표현과 전역 이미지 특징을 융합하여 모델 깊이를 늘리지 않고도 표현 능력을 향상시킨다.
• 아키텍처는 반환-목표, 단계별 보상, 보상 없음(이터레이션)을 포함한 다수의 보상 설정을 지원하여 다양한 학습 철학에 대한 분석을 가능하게 한다.
• 다양한 융합 전략(예: P+C, C+P, StAR 융합, StAR 스택)을 사용한 모델 변종을 평가하여 최적의 특징 통합 방식을 도출한다.

실험 결과

연구 질문

• RQ1 지역적 상태-행동-보상 전이의 명시적 모델링이 시각적 강화학습에서 장기적 시퀀스 모델링을 향상시키는가?
• RQ2 ViT 스타일 패치 임베딩과 컨볼루션 특징을 융합하는 것이 단일 유형의 표현을 사용하는 것보다 성능 향상에 기여하는가?
• RQ3 반환-목표, 단계별 보상, 또는 보상 없음의 보상 신호 선택이 모델 성능 및 일반화에 미치는 영향은 무엇인가?
• RQ4 제안된 번갈아가며 스텝-시퀀스 트랜스포머 아키텍처는 스택형 또는 덧셈 융합 변종보다 더 효과적인가?
• RQ5 보상 신호 없이도 StARformer가 이터레이션 학습에서 뛰어난 성능을 내는가? 이는 강력한 시퀀스 모델링 능력을 시사하는가?

주요 결과

• StARformer는 오프라인-RL 및 이터레이션 학습 설정에서 Atari 및 DeepMind Control Suite 벤치마크에서 최신 기술인 Decision Transformer를 초월한다.
• 패치 임베딩을 스텝 트랜스포머에, 컨볼루션 특징을 시퀀스 트랜스포머에 사용하는 P+C 융합 전략에서 뛰어난 성능을 기록하여 세밀한 특징에서 거시적 특징으로 향하는 특징 계층의 유용성을 확인한다.
• 모델은 반환-목표, 단계별 보상, 보상 없음(이터레이션)을 포함한 모든 보상 설정에서 강력한 성능 유지를 보이며, 특히 반환-목표를 사용한 Decision Transformer를 초월하는 성능을 기록한다.
• 층별 융합 메커니즘(원본 StAR 모델)은 요소별 덧셈(StAR 융합) 및 스택 아키텍처(StAR 스택)보다 우수한 성능을 보이며, 다중 수준 추상화의 중요성을 시사한다.
• 어텐션 맵은 행동 토큰이 관련된 공간 영역(예: Pong에서 페달)을 의미 있게 주시하고 있음을 보여주며, 모델이 의미론적으로 기반한 표현을 학습하고 있음을 확인한다.
• 모델은 더 긴 입력 시퀀스에 대한 적응 능력이 향상되어, 이전의 트랜스포머 기반 접근보다 더 나은 확장성을 보인다.