[논문 리뷰] Masked World Models for Visual Control

Younggyo Seo, Danijar Hafner외 5명|arXiv (Cornell University)|2022. 06. 28.

Advanced Vision and Imaging인용 수 10

한 줄 요약

이 논문은 마스크된 월드 모델(Masked World Models, MWM)을 제안한다. MWM은 시각적 표현 학습과 동역학 학습을 분리한 강화학습 프레임워크로, 세밀한 시각적 표현 학습을 위해 시각 트랜스포머 기반의 컨volutional feature-masked autoencoder를 훈련하고, 별도의 잠재 동역학 모델을 통해 계획을 수행한다. 작업에 관련된 세부 정보를 인코딩하기 위해 보조 보상 예측을 도입함으로써, MWM는 50개의 Meta-World 시각적 조작 작업에서 81.7%의 성공률을 기록하여 이전의 SOTA 기준인 67.9%보다 유의하게 뛰어난 성능을 달성한다.

ABSTRACT

Visual model-based reinforcement learning (RL) has the potential to enable sample-efficient robot learning from visual observations. Yet the current approaches typically train a single model end-to-end for learning both visual representations and dynamics, making it difficult to accurately model the interaction between robots and small objects. In this work, we introduce a visual model-based RL framework that decouples visual representation learning and dynamics learning. Specifically, we train an autoencoder with convolutional layers and vision transformers (ViT) to reconstruct pixels given masked convolutional features, and learn a latent dynamics model that operates on the representations from the autoencoder. Moreover, to encode task-relevant information, we introduce an auxiliary reward prediction objective for the autoencoder. We continually update both autoencoder and dynamics model using online samples collected from environment interaction. We demonstrate that our decoupling approach achieves state-of-the-art performance on a variety of visual robotic tasks from Meta-world and RLBench, e.g., we achieve 81.7% success rate on 50 visual robotic manipulation tasks from Meta-world, while the baseline achieves 67.9%. Code is available on the project website: https://sites.google.com/view/mwm-rl.

연구 동기 및 목표

• 로봇과 소형 물체 간의 상호작용을 정확히 모델링하는 데 도전한다.
• 엔드 투 엔드 훈련에서 발생하는 표현 품질과 동역학 정확도 간의 트레이드오프를 해결한다.
• 표현 학습과 동역학 학습을 분리함으로써 시각 기반 강화학습의 표본 효율성과 점진적 성능을 향상시킨다.
• 컨볼루션 특징 마스킹을 통한 자기지도 표현 학습이 픽셀 패치 마스킹보다 세밀한 시각적 세부 정보를 더 잘 포착할 수 있는지 탐색한다.
• 보조 보상 예측이 작업에 관련된 정보를 포함하여 후속 정책 학습을 위한 시각적 표현 품질을 향상시킬 수 있는지 조사한다.

제안 방법

• 이미지 패치가 아닌 마스크된 컨볼루션 특징을 사용하여 시각 관측치를 재구성하는 데 목표로 하는 시각 트랜스포머 기반 오토인코더를 훈련한다.
• 오토인코더가 추출한 표현을 기반으로 작동하는 별도의 잠재 동역학 모델을 사용하여 향후 상태를 예측한다.
• 오토인코더 훈련 중에 보조 보상 예측 헤드를 도입하여 작업에 관련된 정보를 시각적 표현에 통합한다.
• 환경 상호작용 중에 수집된 온라인 경험을 바탕으로 오토인코더와 동역학 모델을 지속적으로 업데이트한다.
• 훈련 과정을 분리한다: 먼저 재구성 및 보상 예측을 통해 오토인코더를 업데이트하고, 고정된 오토인코더 특징을 사용하여 동역학 모델을 업데이트한다.
• 컨볼루션 특징의 마스킹 비율을 점진적으로 증가시키는 커리큘럼 학습을 적용하여 안정성과 특징 품질을 향상시킨다.

실험 결과

연구 질문

• RQ1 시각적 표현 학습과 동역학 학습을 분리함으로써 시각 기반 강화학습의 성능 향상이 가능할까?
• RQ2 이미지 패치가 아닌 컨볼루션 특징을 마스킹함으로써 시각 제어 작업에서 세밀한 시각적 세부 정보를 더 잘 포착할 수 있을까?
• RQ3 표현 학습 중에 보조 보상 예측을 도입하면 후속 정책 학습을 위한 시각적 특징의 품질이 향상될까?
• RQ4 MWM는 엔드 투 엔드 월드 모델과 다른 분리형 베이스라인에 비해 표본 효율성과 최종 성능 측면에서 어떻게 비교되는가?
• RQ5 사전 훈련된 표현이 한 작업에서 학습된 후, 예측할 수 없는 조작 작업으로 일반화되는 정도는 어느 정도인가?

주요 결과

• MWM는 Meta-World의 50개의 시각적 로봇 조작 작업에서 81.7%의 성공률을 기록하여 이전의 SOTA 기준인 67.9%보다 뚜렷이 뛰어나다.
• 절단 실험 결과, 보상 예측 기능이 있는 MWM가 동일한 설정에서 보상 예측 없이 훈련된 경우보다 성능이 뛰어나, 표현 학습에서 작업에 관련된 신호의 유용성을 확인한다.
• Push 작업에서 사전 훈련된 고정된 표현은 Pick Place, Push Back, Drawer Open와 같은 예측할 수 없는 작업으로 잘 일반화되며, 처음부터 MWM를 훈련하는 것과 비교해도 성능이 유사하거나 뛰어나다.
• 상태 회귀 분석 결과, 보상 예측을 함께 훈련한 표현은 프로피오셉티브 상태 예측 오차가 낮아지며, 이는 향상된 특징 품질을 의미한다.
• MWM는 대조적 표현 학습을 사용하는 DreamerV2보다 복잡한 작업인 Pick Place에서 뛰어난 성능을 보이며, 분리된 마스크된 특징 접근 방식의 우수성을 입증한다.
• 이미지넷 분류 성능는 유사하지만, 시각 제어 작업에서는 픽셀 패치 마스킹보다 컨볼루션 특징 마스킹이 더 높은 성능을 보이며, 작업 특화된 이점을 강조한다.