lucidrains/pi-zero-pytorch

해결하는 문제

이 프로젝트는 $\pi_0$ (pi-zero) 로봇 파운데이션 모델 아키텍처의 PyTorch 구현을 제공합니다. 시각, 언어 명령 및 로봇 관절의 현재 상태를 기반으로 동작을 생성할 수 있는 범용 로봇 정책을 만드는 것을 목표로 합니다.

작동 원리

이 모델은 사전 학습된 시각-언어 모델(PaliGemma 2B)을 기반으로 구축되었으며, 로봇 동작을 생성하기 위해 여러 아키텍처 선택을 통합합니다:

• Flow Matching: 정책 생성을 위해 전통적인 diffusion 대신 flow matching을 사용하여 종종 더 효율적입니다.
• Joint Attention: 멀티모달 입력을 처리하기 위해 Joint Attention(mmDIT에서 유래)을 통한 파라미터 분리를 채택합니다.
• Hybrid Attention: Flex Attention을 사용하여 자기회귀(autoregressive) 및 양방향(bidirectional) 어텐션 패턴을 혼합합니다.
• Online Learning: EFPO 클래스를 통해 모델이 환경으로부터 경험을 수집하고 해당 메모리로부터 온라인 방식으로 학습할 수 있습니다.

대상 사용자

• 최신 최첨단 로봇 파운데이션 모델을 구현하거나 실험하고자 하는 로봇 공학 연구자 및 박사 과정 학생.
• 로봇 조작을 위한 시각-언어-동작(VLA) 모델을 구축하는 개발자.

주요 특징

• Physical Intelligence에서 제안한 $\pi_0$ 아키텍처 구현.
• 동작 생성을 위한 flow matching 지원.
• 기본 시각-언어 모델로서 PaliGemma 2B와의 통합.
• EFPO 클래스를 통한 온라인 학습 내장 지원.