[논문 리뷰] Efficient Transformers in Reinforcement Learning using Actor-Learner Distillation

Emilio Parisotto, Ruslan Salakhutdinov|arXiv (Cornell University)|2021. 04. 04.

Reinforcement Learning in Robotics인용 수 9

한 줄 요약

이 논문은 연속적인 지식 정련 방법인 액터-러닝 정련(Actor-Learner Distillation, ALD)을 제안한다. ALD는 추론 지연 제약 조건이 있는 강화학습 환경에서 대용량 고성능 트랜스포머 러닝 모델의 지식을 소형 빠른 LSTM 액터 모델로 이전한다. 트랜스포머의 샘플 효율성을 활용하면서도 LSTMs의 추론 속도를 유지함으로써, ALD는 더 작은 모델과 유사한 학습 시간 동안 뛰어난 샘플 효율성을 달성한다.

ABSTRACT

Many real-world applications such as robotics provide hard constraints on power and compute that limit the viable model complexity of Reinforcement Learning (RL) agents. Similarly, in many distributed RL settings, acting is done on un-accelerated hardware such as CPUs, which likewise restricts model size to prevent intractable experiment run times. These "actor-latency" constrained settings present a major obstruction to the scaling up of model complexity that has recently been extremely successful in supervised learning. To be able to utilize large model capacity while still operating within the limits imposed by the system during acting, we develop an "Actor-Learner Distillation" (ALD) procedure that leverages a continual form of distillation that transfers learning progress from a large capacity learner model to a small capacity actor model. As a case study, we develop this procedure in the context of partially-observable environments, where transformer models have had large improvements over LSTMs recently, at the cost of significantly higher computational complexity. With transformer models as the learner and LSTMs as the actor, we demonstrate in several challenging memory environments that using Actor-Learner Distillation recovers the clear sample-efficiency gains of the transformer learner model while maintaining the fast inference and reduced total training time of the LSTM actor model.

연구 동기 및 목표

• 추론 속도가 엄격히 제한되는 액터-지연 제약 조건이 있는 강화학습 환경에서 대용량 모델을 배포하는 데 도전하는 데 목적을 두며.
• 로봇과 같은 실세계 응용 분야에서 계산 및 전력 제약으로 인해 추론 시 모델 크기가 제한되는 상황에서 샘플 효율성이 높은 대규모 모델(예: 트랜스포머)을 사용할 수 있도록 하는 데 목적을 두며.
• 고성능 러닝 모델에서 저성능 액터 모델로 지식을 정련하는 연속적인 정련 절차를 개발하여 샘플 효율성을 손상시키지 않도록 하는 데 목적을 두며.
• 작고 빠른 액터 모델을 사용함으로써 분산형 온폴리시 강화학습에서 실험 실행 시간을 최소화하면서도 대규모 러닝 모델의 학습 진전을 그대로 활용하는 데 목적을 두며.
• 특히 부분 관찰 가능한 환경에서 서로 다른 아키텍처 간(특히 트랜스포머와 LSTMs)의 지식 정련이 인덕티브 바이어스와 성능 향상을 효과적으로 전이할 수 있는지 입증하는 데 목적을 두며.

제안 방법

• ALD 프레임워크는 대규모 트랜스포머 모델을 사용하는 중심 러닝 모델을 활용하여 상호작용을 통해 학습하고 샘플 효율성을 향상시킨다.
• 별도로 경량의 LSTM 기반 액터 모델이 추론 속도 제약 조건을 충족시키기 위해 작동하여 저지연 결정을 보장한다.
• 지식은 트랜스포머 러닝 모델에서 LSTM 액터 모델로 연속 정련을 통해 이전된다. 이 과정에서 액터 모델은 러닝 모델의 행동을 모방하도록 훈련된다.
• 정련은 학습 중 실시간으로 적용되며, 러닝 모델의 정책 출력을 소프트 타겟으로 사용해 액터의 정책 업데이트를 지도한다.
• 이 방법은 분산형 온폴리시 강화학습 환경 내에서 작동하며, 액터들이 트레이젝터리를 수집하고 이를 중심 러닝 모델에 전송하여 정책 업데이트를 수행한다.
• 이 접근법은 LSTM 액터의 계산 효율성을 유지하면서도 메모리 집약적인 환경에서 트랜스포머 러닝 모델의 뛰어난 샘플 효율성을 그대로 반영한다.

실험 결과

연구 질문

• RQ1 큰 샘플 효율성의 트랜스포머 모델을 사용해도 추론 속도를 손상시키지 않고 소형 빠른 LSTM 액터 모델을 훈련시킬 수 있는가?
• RQ2 고성능 러닝 모델에서 저성능 액터 모델로의 연속 정련이 부분 관찰 가능한 환경에서 러닝 모델의 샘플 효율성 향상 효과를 유지하는가?
• RQ3 ALD는 작은 모델과 유사한 학습 실행 시간을 유지하면서도 큰 모델의 성능을 그대로 구현할 수 있는가?
• RQ4 액터-지연 제약 조건이 있는 강화학습 환경에서 서로 다른 아키텍처 간(특히 트랜스포머와 LSTMs)의 지식 전이 효과는 어떠한가?
• RQ5 트랜스포머에서 LSTM으로의 정련이 메모리 집약적인 강화학습 과제에서 일반화 능력과 강건성을 향상시키는가?

주요 결과

• ALD는 도전적인 메모리 환경에서 LSTM 액터 모델이 트랜스포머 러닝 모델과 유사한 샘플 효율성을 달성함으로써, 초기 학습에서의 성능을 크게 뛰어넘는다.
• 대규모 트랜스포머를 사용함에도 불구하고, ALD는 LSTM 전용 기준과 유사한 실험 실행 시간을 유지한다.
• 부분 관찰 가능한 환경에서 ALD는 트랜스포머의 샘플 효율성 이점을 회복하면서도 LSTMs의 저지연 추론 성능을 그대로 유지한다.
• 연속 정련 과정은 트랜스포머에서 LSTM으로의 인덕티브 바이어스를 효과적으로 전이하여 학습 안정성과 성능 향상을 이룬다.
• ALD는 서로 다른 아키텍처 간의 지식 정련이 액터-지연 제약 조건이 있는 강화학습에서 매우 효과적일 수 있음을 입증하며, 특히 학생과 교사의 인덕티브 바이어스가 다를 경우에도 효과를 발휘한다.
• 이 방법은 아키텍처 효율성 향상과 수직적(orthogonal)이므로, 효율적인 트랜스포머 변형과 함께 사용해 추가 성능 향상을 이룰 수 있다.