pi0와 pi0-FAST: 일반 로봇 제어를 위한 Vision-Language-Action 모델

Hugging Face는 π0와 π0-FAST, Physical Intelligence가 개발한 범용 Vision-Language-Action(VLA) 모델을 LeRobot 저장소에 통합했습니다. 이러한 모델은 대규모 사전 훈련과 혁신적인 행동 표현 기술을 활용하여 다양한 하드웨어 구성에서 복잡하고 정교한 조작 작업을 수행할 수 있게 합니다.

π0: Flow Matching을 통한 일반 로봇 제어

π0 (Pi-Zero)는 일반 로봇 제어를 위해 설계된 VLA 모델로, 7개의 로봇 플랫폼과 68개의 고유한 작업 데이터로 훈련되었습니다. 이 모델은 zero-shot 및 fine-tuned 설정 모두에서 옷 개 folding, 식료품 포장, 상자 조립, 물체 검색과 같은 실제 작업을 수행할 수 있습니다.

표준 로봇 정책과 달리, π0는 flow matching을 사용하여 50Hz에서 부드럽고 실시간인 행동 경로를 생성합니다. 이 과정은 랜덤 노이즈에서 시작하여 점차적으로 행동 시퀀스로 수렴하며, 배포 중에 모델이 효율적이고 정확하도록 보장합니다.

π0-FAST: 자기회귀 효율성과 FAST 토큰화

π0-FAST는 훈련 속도와 행동 정확성을 높이기 위해 **Frequency-space Action Sequence Tokenization (FAST)**를 활용하는 π0의 자기회귀 버전입니다.

π0-FAST의 주요 장점

π0-FAST는 diffusion-based VLAs에 비해 다음과 같은 개선점을 제공합니다:

Training Speed: 5배 더 빠른 훈련을 달성합니다.
Action Representation: 행동 시퀀스의 중복을 줄여 표현력을 향상시킵니다.
Generalization: 로봇 형태와 보지 못한 환경 전반에 걸쳐 더 강한 일반화를 보여줍니다.

FAST 토큰화가 작동하는 방식

FAST는 **Discrete Cosine Transform (DCT)**를 사용하여 연속적인 행동 시퀀스를 이산 토큰으로 압축합니다. 파이프라인은 다음 단계로 구성됩니다:

Normalization: 원시 로봇 행동은 서로 다른 시스템 간의 일관성을 보장하기 위해 [-1, 1] 범위로 quantile-정규화됩니다.
DCT Transformation: 각 행동 차원은 시간 영역에서 주파수 영역으로 변환됩니다.
Compression: 규모-반올림 작업을 통해 insignificant 계수들이 제거됩니다.
Flattening: 결과적으로 얻은 희소 DCT 계수 행렬은 저주파 구성 요소를 우선시하여 1D 정수 시퀀스로 평탄화됩니다.
BPE Encoding: Byte Pair Encoding (BPE)는 고정된 크기의 어휘를 유지하기 위해 빈번한 패턴을 병합합니다.

이 작업들은 가역적이므로 행동을 무손실로 재구성할 수 있습니다. **FAST+**라는 보편적인 버전은 모바일, 양팔, 단일 팔 로봇에서 1백만 개의 행동 시퀀스로 훈련되었으며, Hugging Face AutoProcessor로 제공됩니다.

기술 아키텍처: VLA 주의 메커니즘

VLA는 행동 및 관찰 상태 토큰을 통합하여 Vision-Language Models(VLM)를 확장합니다. π0는 이러한 멀티모달 입력을 처리하기 위해 특정 주의 구조를 사용합니다:

Prefix Tokens: 전체 장면(이미지 및 텍스트)을 나타내며, PaliGemma와 유사하게 서로 완전히 주의를 기울입니다.
State Tokens: 로봇의 현재 환경 상태(예: 관절 각도)를 나타냅니다. 이들은 접두사 토큰 및 이전 상태 토큰과 삼각형(인과적) 방식으로 주의를 기울입니다.
Action Tokens: 모터 명령 시퀀스를 나타냅니다. 이러한 토큰은 이미지 토큰, 텍스트 토큰, 상태 토큰 및 기타 행동 토큰에 대한 전체 가시성을 가지며, 패딩은 제외합니다.

FlexAttention을 사용한 주의 최적화

결과된 2D 인과 마스크(강한 블록 희소성을 나타냄)를 처리하기 위해 구현은 PyTorch FlexAttention을 활용합니다. 팀은 불규칙한 블록 마스크에서 어려움을 겪는 naive 행렬 곱셈 또는 FlashAttention2 대신, 인과 마스크를 인덱싱하여 생성된 블록 마스크를 사용합니다. 이 접근 방식은 불필요한 계산을 효율적으로 건너뛰어 성능을 향상시킵니다.

LeRobot에서의 구현

π0와 π0-FAST는 이제 LeRobot 저장소에서 사용할 수 있습니다. 사용자는 사전 훈련된 모델에 대해 추론을 수행하거나 특정 환경에 맞게 미세 조정할 수 있습니다. 미세 조정은 lerobot/scripts/train.py 스크립트를 통해 실행할 수 있으며, 이를 통해 기본 π0 모델을 특정 로봇 작업 및 환경에 맞게 조정할 수 있습니다.