[논문 리뷰] Scaling Proprioceptive-Visual Learning with Heterogeneous Pre-trained Transformers

Lirui Wang, Xinlei Chen외 2명|arXiv (Cornell University)|2024. 09. 30.

Image Retrieval and Classification Techniques인용 수 5

한 줄 요약

이종형 사전 학습 변환기(HPT)는 구현체별 stems와 작업별 heads를 활용해 공유 정책 트렁크를 다양한 로봇 구현에 걸쳐 사전 학습하고, 새로운 구현체 및 작업으로의 전이를 통해 성능과 확장성을 52 datasets와 over 1B parameters에서 개선합니다.

ABSTRACT

One of the roadblocks for training generalist robotic models today is heterogeneity. Previous robot learning methods often collect data to train with one specific embodiment for one task, which is expensive and prone to overfitting. This work studies the problem of learning policy representations through heterogeneous pre-training on robot data across different embodiments and tasks at scale. We propose Heterogeneous Pre-trained Transformers (HPT), which pre-train a large, shareable trunk of a policy neural network to learn a task and embodiment agnostic shared representation. This general architecture aligns the specific proprioception and vision inputs from distinct embodiments to a short sequence of tokens and then processes such tokens to map to control robots for different tasks. Leveraging the recent large-scale multi-embodiment real-world robotic datasets as well as simulation, deployed robots, and human video datasets, we investigate pre-training policies across heterogeneity. We conduct experiments to investigate the scaling behaviors of training objectives, to the extent of 52 datasets. HPTs outperform several baselines and enhance the fine-tuned policy performance by over 20% on unseen tasks in multiple simulator benchmarks and real-world settings. See the project website (https://liruiw.github.io/hpt/) for code and videos.

연구 동기 및 목표

• 다양한 구현체와 작업에 걸친 확장 가능하고 일반화 가능한 로봇 정책 학습의 필요성을 제시한다.
• 다른 로봇의 proprioception과 vision을 공유 표현으로 정렬하기 위한 모듈형 아키텍처( stems, trunk, heads )를 제안한다.
• 다양한 실제, 시뮬레이션, 인간 비디오 데이터세트에서 데이터, 모델 크기, 계산량에 따른 확장 동작을 입증한다.
• 감독된 사전 학습과 미세 조정을 통해 unseen 구현체, 작업 및 실제 세계 시나리오로의 전이 성능을 보여준다.

제안 방법

• 이질적 입력을 각 모달리티당 고정 토큰 세트(예: 16개)로 매핑하는 proprioception tokenizer와 vision tokenizer인 stems를 도입한다.
• 결합된 토큰을 처리하여 공동 잠재 표현으로 변환하는 공유 트랜스포머 trunk를 사용한다.
• 트렁크의 출력을 각 구현체-작업 쌍의 행동으로 매핑하는 작업별 heads를 활용한다.
• K개의 이질적 데이터셋에서 정규화된 행동에 대해 허버 로스(Huber loss)를 사용하는 행동 복제(bahavior cloning) 목표로 학습하며, 데이터셋별로 stems/heads를 업데이트하고 trunk를 공동으로 업데이트한다.
• 최대 52개의 데이터셋에서 1B 파라미터 이상으로 사전 학습하여 새로운 구현체에 stem/head를 재초기화하고 trunk를 고정시켜 전이를 가능하게 한다.

실험 결과

연구 질문

• RQ1 실제 로봇, 시뮬레이션, 인간 비디오에 걸친 데이터 양과 다양성에 따라 이질적 사전 학습이 어떻게 확장되는가?
• RQ2 다양한 구현체에서 학습된 단일 trunk가 최소한의 적응으로 unseen 구현체와 작업으로의 효과적인 전이를 가능하게 하는가?
• RQ3 모델 크기와 배치 규모가 사전 학습 수렴 및 다운스트림 전이 성능에 미치는 영향은 무엇인가?
• RQ4 사전 학습된 HPT 표현이 시뮬레이션 벤치마크와 실제 로봇 작업 모두에 얼마나 잘 전이되는가?
• RQ5 이질적 데이터를 포함하는 것이 구현체와 환경 간의 강건성과 일반화를 개선하는가?

주요 결과

• HPT는 더 큰 모델, 더 많은 데이터, 더 높은 컴퓨트에서 이점을 확장 시키며 데이터/모드가 증가할수록 검증 손실이 개선된다.
• 사전 학습에 더 많은 구현체를 포함시키면 트렁크 일반화와 작업 간 전이 성능이 향상된다.
• 1B 파라미터(HPT-Huge)까지의 사전 학습은 큰 배치 크기로도 개선이 지속되며 깊이 대비 너비 확장의 이점은 최소하다.
• 합성 시뮬레이션 데이터와 인터넷 인간 비디오를 포함한 사전 학습은 보완적인 구현체 데이터를 제공하고 전이 이점을 유지한다.
• 시뮬레이션 벤치마크로의 전이에서 HPT는 처음부터 학습하는 것 또는 트렁크가 없는 경우보다 작업 성공률을 개선하며, finetuned HPT 변형은 baselines보다 더 높은 성능(HPT-XL 등)에 도달한다.
• 실세계 테스트에서 사전 학습된 정책은 시야 구성 및 물체 다양성에 대해 베이스라인에 비해 강건성과 일반화가 향상된다.