[논문 리뷰] In-context Reinforcement Learning with Algorithm Distillation

Michael Laskin, Luyu Wang외 12명|arXiv (Cornell University)|2022. 10. 25.

Reinforcement Learning in Robotics인용 수 10

한 줄 요약

이 논문은 강화학습(Reinforcement Learning, RL) 에이전트의 다수의 에피소드에 걸친 학습 역사를 기반으로 행동을 예측하는 인과적 트랜스포머를 훈련시켜, 파rameter 업데이트 없이도 맥락 내에서 정책 개선을 가능하게 하는 알고리즘 정련(Algorithm Distillation, AD)을 소개한다. AD는 전체 학습 트레이젝터리에 대해 순차적 예측 문제로 정책 개선을 모델링함으로써, 원본 RL 알고리즘보다 더 데이터 효율적인 RL 알고리즘을 학습한다.

ABSTRACT

We propose Algorithm Distillation (AD), a method for distilling reinforcement learning (RL) algorithms into neural networks by modeling their training histories with a causal sequence model. Algorithm Distillation treats learning to reinforcement learn as an across-episode sequential prediction problem. A dataset of learning histories is generated by a source RL algorithm, and then a causal transformer is trained by autoregressively predicting actions given their preceding learning histories as context. Unlike sequential policy prediction architectures that distill post-learning or expert sequences, AD is able to improve its policy entirely in-context without updating its network parameters. We demonstrate that AD can reinforcement learn in-context in a variety of environments with sparse rewards, combinatorial task structure, and pixel-based observations, and find that AD learns a more data-efficient RL algorithm than the one that generated the source data.

연구 동기 및 목표

• 모델이 파arameter 미세조정 없이도 시행착오를 통해 개선되는 맥락 내 강화학습을 가능하게 하기 위해.
• 장기적인 맥락을 사용하여 RL 알고리즘 학습 과정을 순차적 예측 문제로 모델링하기 위해.
• 고정된 정책뿐만 아니라 RL 학습 역사에서 정책 개선 연산자 전체를 정련하기 위해.
• 오프라인 RL 학습 역사를 기반으로 한 암기 학습을 통해 맥락 내 강화학습이 가능하다는 것을 보여주기 위해.
• 정련된 알고리즘이 원본 RL 알고리즘보다 더 효율적으로 학습할 수 있도록 일반화될 수 있다는 것을 보여주기 위해.

제안 방법

• 다양한 개별 작업들에 걸쳐 원본 RL 알고리즘으로부터의 전체 학습 역사를 포함한 대규모 다중 작업 데이터셋을 수집한다.
• 각 학습 역사를 상태-행동-보상 토큰의 순서로 표현하여 시간적 및 에피소드 구조를 유지한다.
• 전체 이전 학습 역사를 맥락으로 사용하여 행동을 자동적으로 예측하는 인과적 트랜스포머를 훈련시킨다.
• 과거 상태, 행동, 보상만을 조건으로 삼는 시퀀스 모델링 손실을 사용하며, 수익을 조건으로 삼지 않는다.
• 여러 에피소드에 걸친 정책 개선을 포괄할 수 있도록 충분히 긴 맥락 길이를 확보한다.
• 추론 시 훈련된 트랜스포머를 맥락 내에서 활용한다: 이전 에피소드의 역사를 프롬프트로 제공하고, 가중치 업데이트 없이 행동을 예측한다.

실험 결과

연구 질문

• RQ1 시퀀스 모델은 전체 학습 역사를 참조함으로써 맥락 내에서 정책을 개선할 수 있는가?
• RQ2 RL 학습을 순차적 예측 문제로 모델링하면 맥락 내 탐색 및 시간적 책임 부여가 가능해지는가?
• RQ3 정련된 모델은 새로운 작업으로 일반화되어 원본 RL 알고리즘보다 데이터 효율성 면에서 뛰어나게 학습할 수 있는가?
• RQ4 오프라인 RL 학습 역사를 기반으로 한 암기 학습을 통해 RL 알고리즘의 정책 개선 연산자를 단일 신경망으로 정련할 수 있는가?
• RQ5 정련된 알고리즘의 성능은 샘플 효율성 측면에서 원본 알고리즘과 비교해 어떻게 되는가?

주요 결과

• AD는 파arameter 업데이트 없이도 이전 에피소드의 맥락만으로 맥락 내 강화학습을 달성한다.
• AD는 DMLab의 Watermaze와 같은 부분관측 환경에서 맥락 내 탐색 및 시간적 책임 부여를 보여준다.
• AD는 학습 데이터를 생성하는 데 사용된 원본 RL 알고리즘보다 더 데이터 효율적인 RL 알고리즘을 학습한다.
• 희소 보상, 조합적 구조, 픽셀 기반 관측이 있는 작업 간에 일반화된다.
• AD는 적응을 위해 미세조정이나 전문가 시범을 요구하는 기존의 오프라인 정책 정련 방법보다 뛰어나다.
• AD의 성공은 정책 개선을 포괄할 수 있도록 여러 에피소드를 아우르는 충분히 긴 맥락 창문에 의해 결정적으로 좌우된다.