[논문 리뷰] Switch Trajectory Transformer with Distributional Value Approximation for Multi-Task Reinforcement Learning

Qinjie Lin, Han Liu외 1명|arXiv (Cornell University)|2022. 03. 14.

Reinforcement Learning in Robotics인용 수 4

한 줄 요약

SwitchTT는 샘플 효율성 향상과 계산 비용 감소를 위해 희소 전문가 기반 트랜스포머 아키텍처(Switch Transformer)와 분포가치 추정을 사용하는 다중 작업 오프라인 강화학습 방법을 제안한다. 10개의 희소 보상 작업에서 Trajectory Transformer보다 10% 높은 성능을 달성하고, 학습 속도를 최대 90% 빠르게 하여 다중 작업 환경에서 우수한 확장성과 가치 추정 능력을 입증한다.

ABSTRACT

We propose SwitchTT, a multi-task extension to Trajectory Transformer but enhanced with two striking features: (i) exploiting a sparsely activated model to reduce computation cost in multi-task offline model learning and (ii) adopting a distributional trajectory value estimator that improves policy performance, especially in sparse reward settings. These two enhancements make SwitchTT suitable for solving multi-task offline reinforcement learning problems, where model capacity is critical for absorbing the vast quantities of knowledge available in the multi-task dataset. More specifically, SwitchTT exploits switch transformer model architecture for multi-task policy learning, allowing us to improve model capacity without proportional computation cost. Also, SwitchTT approximates the distribution rather than the expectation of trajectory value, mitigating the effects of the Monte-Carlo Value estimator suffering from poor sample complexity, especially in the sparse-reward setting. We evaluate our method using the suite of ten sparse-reward tasks from the gym-mini-grid environment.We show an improvement of 10% over Trajectory Transformer across 10-task learning and obtain up to 90% increase in offline model training speed. Our results also demonstrate the advantage of the switch transformer model for absorbing expert knowledge and the importance of value distribution in evaluating the trajectory.

연구 동기 및 목표

• 다양하고 이질적인 데이터셋을 포함한 다중 작업 오프라인 강화학습에서 높은 계산 비용과 낮은 샘플 효율성의 과제를 해결한다.
• 다양한 작업 간 공유 정책 파rameter가 다중 작업 환경에서 성능 저하를 초래하는 한계를 극복한다.
• 몬테 카를로 추정기의 낮은 샘플 복잡도로 악영향을 받는 희소 보상 환경에서의 가치 추정을 향상시킨다.
• 추론 비용의 비례 증가 없이도 대규모 다중 작업 데이터셋으로부터 고용량 지식 흡수를 가능하게 한다.

제안 방법

• 입력을 작업별 전문가 네트워크로 라우팅하는 게이팅 네트워크를 갖춘 Mixture-of-Experts (MoE) 레이어를 포함한 스위치 트랜스포머 아키텍처를 채택하여 희소 파라미터 활성화를 통해 계산을 감소시킨다.
• 행동 예측, 동역학 모델링, 수익 회수(returns-to-go, RTG) 예측을 위한 세 개의 트랜스포머 헤드를 통합하여 궤적 생성을 위한 통합된 시퀀스 모델링 프레임워크를 형성한다.
• 기본 예상 가치 추정 방식을 대체로 궤적 수익의 범주형 분포를 모델링하는 분포가치 추정기로 교체하여 불확실성 인식 평가를 향상시킨다.
• 궤적 생성 시 분포가치 추정을 기반으로 비드 서치를 수행하여 계획 중 높은 보상 행동 시퀀스를 선택한다.
• 전체 상태, 행동, 수익, 작업 식별자 시퀀스를 입력 토큰으로 사용하여 오프라인 다중 작업 데이터셋에서 모델을 종합적으로 학습시킨다.
• 스위치 트랜스포머의 고용량을 활용하여 다양한 전문가 시나리오를 포함한 다중 작업에서의 지식 흡수를 가능하게 하되, 계산 효율성은 유지한다.

실험 결과

연구 질문

• RQ1 희소하게 활성화된 트랜스포머 아키텍처는 모델 용량이나 성능을 희생시키지 않고 다중 작업 오프라인 강화학습에서 계산 비용을 줄일 수 있는가?
• RQ2 기대 수익뿐 아니라 궤적 수익의 분포를 모델링할 경우, 희소 보상 환경에서 더 견고하고 정확한 가치 추정이 가능해지는가?
• RQ3 SwitchTT는 표준 Trajectory Transformer와 비교해 다중 작업, 희소 보상 환경에서 성능 및 학습 효율성 측면에서 어떻게 다른가?
• RQ4 스위치 트랜스포머 아키텍처는 이질적이고 다양한 작업 간에 효과적으로 표현을 학습하고 공유할 수 있는가?

주요 결과

• SwitchTT는 gym-mini-grid 환경에서 10개의 희소 보상 작업 전반에서 Trajectory Transformer보다 10% 높은 성능을 달성한다.
• 스위치 트랜스포머의 전문가 계산 희소성 덕분에 표준 Trajectory Transformer 대비 오프라인 모델 학습 시간을 최대 90% 감소시킨다.
• 분포가치 추정기는 샘플 수가 적고 보상이 희소한 상황에서 가치 추정 정확도를 크게 향상시켜 몬테 카를로 추정기의 열악한 샘플 복잡도 문제를 완화한다.
• 스위치 트랜스포머 아키텍처는 작업 분포가 다를 경우에도 더 나은 파라미터 공유와 지식 전이를 가능하게 한다.
• 분포가치 추정치를 기반으로 하는 비드 서치는 더 신뢰할 수 있는 궤적 계획과 더 높은 누적 보상을 이끈다.
• 실험 결과 가치를 기대값이 아닌 분포로 모델링할 경우, 다중 작업 오프라인 강화학습에서 정책 성능과 안정성이 향상됨을 확인할 수 있다.