[논문 리뷰] Transformers are Meta-Reinforcement Learners

Luckeciano C. Melo|arXiv (Cornell University)|2022. 06. 14.

Reinforcement Learning in Robotics인용 수 8

한 줄 요약

이 논문은 트랜스포머 아키텍처를 사용하여 다중 헤드 자기주의를 통해 에피소딕 메모리를 반복적으로 구축함으로써 빠른 적응과 강건한 일반화를 가능하게 하는 메타강화학습 에이전트 TrMRL을 제안한다. TrMRL은 HalfCheetahVel과 로봇 조작 작업과 같은 연속 제어 환경에서 뛰어난 또는 비교 가능한 샘플 효율성과 분포 외 일반화 성능을 달성한다.

ABSTRACT

The transformer architecture and variants presented remarkable success across many machine learning tasks in recent years. This success is intrinsically related to the capability of handling long sequences and the presence of context-dependent weights from the attention mechanism. We argue that these capabilities suit the central role of a Meta-Reinforcement Learning algorithm. Indeed, a meta-RL agent needs to infer the task from a sequence of trajectories. Furthermore, it requires a fast adaptation strategy to adapt its policy for a new task -- which can be achieved using the self-attention mechanism. In this work, we present TrMRL (Transformers for Meta-Reinforcement Learning), a meta-RL agent that mimics the memory reinstatement mechanism using the transformer architecture. It associates the recent past of working memories to build an episodic memory recursively through the transformer layers. We show that the self-attention computes a consensus representation that minimizes the Bayes Risk at each layer and provides meaningful features to compute the best actions. We conducted experiments in high-dimensional continuous control environments for locomotion and dexterous manipulation. Results show that TrMRL presents comparable or superior asymptotic performance, sample efficiency, and out-of-distribution generalization compared to the baselines in these environments.

연구 동기 및 목표

• 주의 메커니즘이 문맥에 따라 다른 파rameterization을 가지는 특성을 활용하여 메타-RL에서의 빠른 적응 문제를 해결하기 위해.
• 트랜스포머 레이어를 통해 에피소딕 메모리를 반복적으로 정제함으로써 기억 재활성화를 모방하는 메모리 기반 메타-RL 에이전트를 설계하기 위해.
• 안정적인 트랜스포머 기반 아키텍처를 사용하여 연속 제어 환경에서 샘플 효율성과 분포 외 일반화를 향상시키기 위해.
• T-Fixup 가중치 초기화가 트랜스포머 기반 메타-RL 에이전트의 학습 안정성을 높이는 데 효과적인지 조사하기 위해.
• 자기주의가 베이즈 위험을 최소화하는 공통 표현을 계산하여 더 나은 행동 예측을 이끌어내는지 검증하기 위해.

제안 방법

• TrMRL은 인과적 다중 헤드 자기주의를 사용하여 트랜스포머 레이어를 통해 최근 작업 메모리를 임의의 작업 인식 메모리로 연결한다.
• 각 트랜스포머 레이어는 베이즈 위험을 최소화하는 공통 표현을 계산하여 에피소딕 메모리를 정제함으로써 행동 선택을 위한 특징 품질을 향상시킨다.
• 마지재 레이어에서 생성된 최종 에피소딕 메모리는 정책 헤드의 입력으로 사용되어 문맥 의존적 행동 예측을 가능하게 한다.
• T-Fixup 가중치 초기화를 적용하여 학습을 안정화하고 강화학습에서 초기 탐색 단계 동안의 불안정성을 완화한다.
• 에피소딕 메모리는 레이어를 통해 반복적으로 구축되며, 각 레이어는 최적의 행동 예측을 향해 더 나은 공통 표현으로 표현을 정제한다.
• 선형 회귀 모델과 코사인 유사도를 사용하여 각 레이어에서의 표현 품질을 평가하고, 정제 정도와 예측 능력을 측정한다.

실험 결과

연구 질문

• RQ1 트랜스포머의 자기주의 메커니즘이 문맥에 따라 다른 정책을 형성함으로써 메타-RL에서 효과적인 빠른 적응 전략이 될 수 있는가?
• RQ2 트랜스포머 레이어를 통해 에피소딕 메모리를 반복적으로 정제하면, 베이즈 위험을 최소화하는 공통 표현이 도출되어 행동 예측이 향상되는가?
• RQ3 기존의 메타-RL 베이스라인들과 비교할 때 TrMRL의 샘플 효율성과 분포 외 일반화 성능은 어떠한가?
• RQ4 T-Fixup 초기화가 트랜스포머 기반 메타-RL 에이전트의 학습을 얼마나 잘 안정화시키는가?
• RQ5 중간 레이어의 에피소딕 메모리 표현은 전문가 행동을 예측하는 데 의미 있게 사용될 수 있는가? 이는 점진적인 정제를 나타내는가?

주요 결과

• TrMRL는 HalfCheetahVel에서 첫 번째 에피소드부터 높은 성능를 기록하며, 약 20개의 타임스텝 만으로 최적 성능에 도달함으로써 강력한 온라인 적응 능력을 보여준다.
• 목표 속도가 [3.0, 4.0]인 HalfCheetahVel에서 분포 외 일반화 테스트에서 TrMRL은 PEARL과 VariBAD를 포함한 모든 베이스라인을 크게 능가하여 강건한 작업 표현 학습 능력을 보여준다.
• 중간 레이어의 에피소딕 메모리 표현과 최종 레이어 표현 간의 표현 오차가 레이어를 거치며 감소함으로써, 공통 표현으로 향한 점진적인 정제가 확인된다.
• 각 레이어의 에피소딕 메모리 특징을 사용한 선형 회귀 모델은 테스트 세트에서 평균 제곱 오차가 감소함을 보여주며, 깊이가 증가할수록 표현이 전문가 행동 예측에 더 예측 가능해짐을 나타낸다.
• 제거 실험 결과, T-Fixup 초기화는 학습 안정성과 성능에 매우 중요하며, 네트워크 깊이, 시퀀스 길이, 어텐션 헤드 수에 민감함을 확인하였다.
• 에피소딕 메모리 정제 과정은 정리 4.1과 일치하며, 자기주의가 각 레이어에서 베이즈 위험을 최소화하는 공통 표현을 계산함을 입증한다.