Scaling Proprioceptive-Visual Learning with Heterogeneous Pre-trained Transformers

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.

한 줄 요약

HPT는 이질적인 로봇 데이터를 기반으로 정책을 사전 학습하여 20% 이상의 미见过 태스크 성능 향상을 달성한다.

핵심 기여도

이질적인 로봇 데이터(52개 데이터셋)를 활용한 HPT 사전 학습 프레임워크 제안.
HPT는 proprioception 및 vision 입력을 공유 토큰 시퀀스로 매핑하여 태스크 및 이질성에 무관한 정책 표현 학습.
52개 데이터셋을 기반으로 확장 실험을 수행, 10억 파라미터 이상의 모델 크기와 50억 토큰 규모의 학습 수행.
미见过 태스크에서 20% 이상의 성능 향상을 기록하며 기존 기법을 초과.

핵심 아이디어

로봇 정책 학습에서 이질성(다양한 로봇 하드웨어, 환경, 센서)은 일반화를 어렵게 만드는 주요 장애물이다. 기존 연구는 특정 로봇과 태스크에 맞춘 데이터를 수집하여 학습하므로 비용이 높고 과적합되기 쉽다. 본 연구는 이질적인 로봇 데이터를 기반으로 **Heterogeneous Pre-trained Transformers (HPT)** 를 제안하여, **공유 가능한 정책 표현**(shared policy representation)을 학습한다. HPT는 **proprioception** 및 **vision 입력**을 각 로봇의 특성에 맞는 **"stem" 토크나이저**로 변환한 후, 공유 **"trunk" 트랜스포머**를 통해 **공통 잠재 토큰 시퀀스**로 매핑한다. 이는 로봇의 이질성을 **공통 언어**(shared language)로 통합하는 방식이다. 학습된 표현은 **"head" 액션 디코더**를 통해 다양한 로봇에 맞는 액션을 생성한다. 이 구조는 인간의 척수 신경 회로와 유사한 피드백 루프를 모방한다.

기술적 접근법

**HPT 아키텍처**:
**Stem**: 로봇의 proprioception 및 vision 입력을 토큰화.
**Trunk**: 공유 트랜스포머, 52개 데이터셋(실제 로봇, 시뮬레이션, 인터넷 영상)을 기반으로 사전 학습.
**Head**: 태스크별 액션 디코더.
**데이터**:
52개 데이터셋(27개 로봇 텔레오퍼레이션, 시뮬레이션, 인터넷 영상 등).
최대 200,000개의 트레이젝토리, 50억 토큰 규모.
**모델 규모**:
HPT-Small: 3.17M 파라미터.
최대 10억 파라미터, 31GFlops 연산.
**학습 설정**:
80,000 iteration, batch size 256, 5B 토큰(시각/proprioception).
ResNet18을 사용한 이미지 인코딩, 정규화 및 필터링 최소한 수행.

주요 결과

**미见过 태스크 성능 향상**:
시뮬레이션 및 실제 로봇 환경에서 기존 기법 대비 **20% 이상의 성능 향상**.
**데이터셋 규모**:
52개 데이터셋, 200,000개 트레이젝토리, 50억 토큰.
**모델 확장성**:
10억 파라미터 이상, 31GFlops 연산.
**사전 학습 효과**:
사전 학습 성능과 미见过 태스크 성능 간 **상관관계** 확인.

의의 및 한계

HPT는 이질적인 로봇 데이터를 기반으로 정책을 사전 학습함으로써, **데이터 수집 비용을 줄이고 일반화 능력을 향상**시킨다. 특히, **proprioception 및 vision 정보를 통합**하여 로봇의 이질성을 극복하는 데 기여한다. 그러나, **데이터 품질**(필터링, 정제) 및 **학습 목표**(supervised learning에 제한)는 아직 개선이 필요하다. 또한, **학습 수렴 속도**가 느리고, **장기적 태스크**(long-horizon)에서는 성능이 제한적이다. 평가 환경이 **단기 조작 태스크**에 제한되어 있어, 더 복잡한 태스크에서의 일반화 능력은 추가 연구가 필요하다.

실용적 활용

HPT는 다양한 로봇 하드웨어(예: 다관절 팔, 이동 로봇) 및 환경(실제/시뮬레이션)에서 **새로운 태스크에 빠르게 적응**할 수 있는 정책을 학습하는 데 활용 가능하다. 특히, **데이터 수집 비용이 높은 산업 로봇** 또는 **다양한 로봇 플랫폼을 운영하는 연구소**에서 유용하다. 또한, **인간 영상 데이터**를 활용한 정책 학습도 가능하여, **저비용 정책 개발**에 기여할 수 있다.