π0 & π0-FAST: 범용 로봇 제어를 위한 Vision-Language-Action 모델

π0 & π0-FAST: 범용 로봇 제어를 위한 Vision-Language-Action 모델

이 글은 Hugging Face 블로그의 π0 and π0-FAST: Vision-Language-Action Models for General Robot Control를 한국어로 번역한 글입니다.


우리는 Hugging Face LeRobot에 최초의 로봇 파운데이션 모델을 공개했습니다! Physical Intelligence에서 개발한 π0π0-FAST 모델이 이제 LeRobot repository를 통해 제공되며, 이를 통해 범용 로봇 지능(Generalist Robotic Intelligence)이 Hugging Face 생태계에 새롭게 도입되었습니다. 만약 Vision-Language-Action(VLA) 모델이 Vision-Language Model(VLM)과 어떻게 다른지, 그리고 행동(action)이 어떤 방식으로 표현되는지 궁금하다면, 이 블로그 포스트를 끝까지 읽어보세요!

모델 컬렉션과 PyTorch 버전의 모델은 아래에서 확인할 수 있습니다: Huggingface의 Pi0 모델 컬렉션 | Huggingface의 Pi0+FAST 모델 컬렉션 | LeRobot repo


Introduction

Robert Heinlein은 다재다능한 사람이란 하나의 분야에만 전문화된 사람이 아니라, 지적·육체적으로 폭넓은 과제를 수행할 수 있는 사람이라고 말했습니다. 이 개념을 기계 지능(machine intelligence)에 비유해보면, AI 시스템은 매우 다양한 형태로 존재하지만, 인간 지능은 다재다능하다는 점에서 특히 뛰어납니다. 즉, 인간 지능은 다양한 과제와 환경, 예기치 않은 상황에 유연하게 적응할 수 있습니다. 대형 언어 모델(LLM)과 비전-언어 모델(VLM)은 큰 가능성을 보여주고 있지만, 여전히 물리적 세계와의 상호작용 능력이 부족합니다. 이 격차를 해소하기 위해서는 로봇 데이터를 기반으로 학습된 모델이 필요합니다. 이러한 범용 로봇 모델(generalist robot models)은 다양한 데이터를 활용해 적응력을 높이고, 일반화 성능과 견고성을 향상시킬 수 있습니다. 즉, 개별 작업(task)별로 모델을 따로 학습시키는 대신, LLM과 유사한 방식으로 다양한 로봇 데이터를 사전 학습(pre-training)하면 학습 효율과 성능을 모두 크게 향상시킬 수 있습니다.

범용 로봇 정책(generalist robot policies) 또는 로봇 파운데이션 모델(robot foundation models)을 개발하기 위해서는 다음과 같은 세 가지 주요 과제가 존재합니다:

  1. 대규모 연구의 필요성 — 사전 학습(pre-training)의 이점을 최대한 활용하기 위해서는 방대한 데이터와 실험을 기반으로 한 대규모 연구가 필요합니다.
  2. 모델 아키텍처 설계 — 다양한 데이터 소스를 통합하면서도 복잡한 물리적 상호작용을 포착할 수 있는 아키텍처를 설계해야 합니다. 이와 관련된 중요한 도전 과제는 크로스 임바디먼트 학습(cross-embodiment training)입니다. 이는 로봇의 구조, 제어 공간, 그리고 행동 표현이 서로 다른 다양한 로봇 유형들로부터 모델이 학습해야 하는 과정을 의미합니다. 기존 접근법들은 다음과 같은 방식으로 이를 해결하고자 합니다:
    • 다양한 로봇 플랫폼의 멀티모달 데이터셋을 결합하여 일반화 능력을 향상시키는 방법
    • 공유 표현(shared representations)을 활용하여 단일 팔(single-arm), 양팔(dual-arm), 모바일 매니퓰레이터(mobile manipulator) 등 서로 다른 형태의 로봇 간 격차를 줄이는 방법
  3. 효율적인 학습 레시피 개발 — 최근 NLP와 비전 분야의 발전이 세밀한 사전 학습(pre-training) 및 사후 학습(post-training) 전략에 크게 의존한 것처럼, 로봇 모델에서도 이에 상응하는 학습 절차를 정립하는 것이 중요합니다.

이 글에서는 이러한 도전 과제를 해결하기 위해 Physical Intelligence가 개발한 π0와 π0-FAST — 두 개의 프로토타입 모델 및 학습 프레임워크를 소개합니다.


🔍 π0란 무엇인가?

논문 Jax 코드

π0 (Pi-Zero)Physical Intelligence 팀이 개발한 Vision-Language-Action (VLA) 모델로, 범용 로봇 제어(generalist robot control)를 위해 설계되었습니다.
이 모델은 대규모 사전 학습flow matching 기반의 행동 생성을 결합하여, 다양한 로봇 형태에서 정교한 조작 작업을 수행할 수 있도록 합니다.

π0는 7개의 로봇 플랫폼68개의 고유한 작업 데이터로 학습되었으며, 빨래 개기, 식탁 정리, 식료품 포장, 박스 조립, 물체 회수 등 복잡한 실제 작업에서도 우수한 zero-shotfine-tuning 결과를 보입니다.

기존의 로봇 정책과 달리, π0는 flow matching 기법을 활용하여 50Hz의 주파수로 부드럽고 실시간 행동 궤적(action trajectories)을 생성합니다. 이를 통해 실제 환경에서도 높은 효율성, 정밀성, 적응성을 달성합니다. Flow Matching은 원래 연속 정규화 플로우나 diffusion models의 생성 품질을 개선하기 위해 사용된 기법으로, π0에서도 유사한 원리를 적용합니다. 즉, 무작위 노이즈 상태에서 시작해 점진적으로 의미 있는 모터 동작 시퀀스로 수렴해 나가는 디노이징(denoising) 과정을 거칩니다.

π0를 LeRobot에서 사용하는 방법

먼저, transformers 라이브러리가 새로운 의존성으로 추가되었기 때문에 LeRobot 설치 버전을 최신으로 업데이트해야 합니다. git clone 후 아래 명령어를 실행하세요:

pip install -e ".[pi0]"

π0 모델은 PaliGemma와 마찬가지로 다양한 프레임워크와 환경, 그리고 시각 입력에 유연하게 적용될 수 있도록 설계된 파운데이션 모델입니다. 이 중에서도 기본 모델(Base model)인 π0는 별도의 수정 없이 그대로 사용할 수 있습니다.

π0 Pretrained 모델 실행

python lerobot/scripts/eval.py \
--pretrained_policy.path=/path/to/pretrained/pi0

다만, 이 모델은 JAX에서 PyTorch로 변환된 버전이며 특정 환경에 맞춰 학습된 모델이기 때문에 일부 성능이 저하될 수 있습니다. 따라서 아래와 같이 사용자 환경에 맞게 직접 fine-tuning 하는 것을 권장합니다.

π0 Pretrained 모델 Fine-tuning

openpipi0_base 체크포인트를 사용하여 π0 모델을 fine-tuning하려면 다음 명령어를 실행하세요:

python lerobot/scripts/train.py \
--policy.path=lerobot/pi0 \
--dataset.repo_id=danaaubakirova/koch_test

π0 신경망(neural network)을 PaliGemma와 Expert Gemma와 함께 fine-tuning하려면, 다음 명령어를 실행하세요.
이 두 모델은 π0 fine-tuning 이전에 VLM 기본 파라미터로 사전 학습된 모델입니다.

python lerobot/scripts/train.py \
--policy.type=pi0 \
--dataset.repo_id=danaaubakirova/koch_test

또한, 아래 코드를 사용하면 LeRobot 학습 프레임워크와 독립적으로 사전 학습된 π0 모델을 직접 사용할 수도 있습니다:

policy = Pi0Policy.from_pretrained("lerobot/pi0")

VLM과 VLA의 차이점은 무엇일까?

Vision-Language Models(VLMs)과 Vision-Language-Action Models(VLAs)은 공통적으로 Transformer 아키텍처를 기반으로 합니다. 하지만 두 모델의 핵심적인 차이는 행동의 표현 방식에 있습니다. VLM은 이미지와 텍스트를 기반으로 멀티모달 표현을 학습하고 생성하지만, VLA는 여기에 행동과 관찰 상태 토큰을 추가로 통합합니다. 이러한 토큰이 추가되면, 다음으로 중요한 과제는 어텐션(attention)이 어떻게 계산되는지를 이해하는 것입니다.

로봇 정책에서의 어텐션 메커니즘(Attention Mechanisms)

이제 주요 개념들을 하나씩 살펴보겠습니다.

State Token

  • 로봇의 현재 환경 상태를 나타내는 단일 토큰입니다. (예: 관절 각도, 센서 값, 기타 관련 관찰 정보 등)
  • 마스킹 규칙에 따라, 이 토큰은 prefix 구간의 이미지와 텍스트에 어텐션할 수 있습니다. 즉, state token은 의사결정에 필요한 시각적·언어적 단서를 “볼 수” 있습니다.
  • 또한 이전 상태를 삼각형(triangular) 구조로 참조합니다. 여러 개의 state token이 존재할 경우, 새로운 토큰은 이전 토큰을 볼 수 있지만 그 반대는 불가능합니다.

Action Tokens

  • 로봇의 모터 명령 시퀀스를 나타냅니다.
  • 패딩 영역(padding regions)을 제외하고는 모든 토큰에 대해 완전한 가시성을 가집니다.
    즉, 각 action token은 다음 항목들에 어텐션할 수 있습니다:
    • 비패딩 이미지 토큰 (전체 장면)
    • 비패딩 텍스트 토큰 (명령어 또는 설명)
    • State tokens (현재 및 이전 상태)
    • 다른 Action tokens

Prefix Tokens

  • 전체 장면을 표현하며, PaliGemma와 유사하게 서로 완전히 어텐션할 수 있습니다.

핵심 아이디어

이러한 토큰들은 다음과 같은 요소들을 포함합니다:

  • 로봇이 환경을 이해하는 내부 표현 — 즉, 상태(state)
  • 로봇이 실행하는 명령 또는 제어 신호 — 즉, 행동(action)
  • 시간 또는 단계를 인코딩한 시간 임베딩(time embedding)

이 토큰들은 prefix 구간 — 이미지와 텍스트로 구성된 부분 — 뒤에 추가됩니다. 따라서 prefix는 맥락(context)으로 작용하며, 예를 들어 장면 이미지나 “착한 로봇이 되어라”이나 “큐브를 옮겨라” 같은 언어 지시가 여기에 포함됩니다. 반면, suffix는 로봇의 의사결정 정책(policy)과 관련된 특성을 포착합니다.

전체 토큰 구성 흐름

[Image Tokens] + [Text Tokens] + [State Tokens] + [Action Tokens] + [Time Tokens]

  • Prefix (Image + Text) → 상황과 명령의 맥락(context)
  • Suffix (State + Action) → 로봇의 의사결정(policy)제어(action)
    즉, VLA 모델은 “무엇을 보고(Image), 어떤 지시를 받고(Text), 어떤 상태(State)에서, 어떤 행동(Action)을 할지” 모두 Transformer 시퀀스 하나로 통합하여 학습합니다.

⚡ π0에서 더 빠른 어텐션을 향하여

그러나 π0에서 어텐션을 효율적으로 처리하는 것은 또 다른 도전 과제를 동반합니다.
특유의 어텐션 마스크 형태가 어텐션 계산 방식에 영향을 미치기 때문이죠 — 이제 그 세부 내용을 살펴보겠습니다!

2D 어텐션 마스크 다루기

결과적으로 생성된 2D 인과 마스크(causal mask)는 강한 블록 희소성(block sparsity)을 보입니다. 하지만 각 블록의 경계를 정의하는 일 — 특히 여러 샘플(batch)을 다루는 경우 — 은 꽤 까다롭습니다. 우리는 일반적으로 자기회귀(autoregressive) 모델링에서 삼각형 형태의 causal mask에 익숙하지만, 이번 경우는 그렇지 않습니다.

예를 들어 아래 예시를 보면, 이미지(첫 번째 요소)는 빈 카메라를 나타내는 패딩 토큰(padding tokens)을 포함합니다. 그 다음으로 텍스트 토큰과 state 토큰이 추가됩니다. 이 “prefix” 부분은 PaliGemma에서처럼 완전 비인과(non-causal) 어텐션을 형성합니다. 반면, 그 이후의 “suffix” (즉, state + action/time 토큰) 부분은 인과적(causal) 블록 구조를 가집니다. 이때 단순한 즉시 실행 방식의 구현은 입력 전체에 대해 행렬 곱을 수행하고 softmax를 적용하기 때문에, 매우 비효율적입니다.

VLA Attention Mask

Figure 1: VLA 어텐션 마스크 시각화


FlashAttention2를 사용할 수 있을까?

  • FlashAttention2는 가변 길이(varlen) 인터페이스를 제공하지만, cu_seqlens(누적 시퀀스 길이, cumulative prefix lengths)을 직접 계산해야 합니다. 이 라이브러리는 연속적(또는 엄격히 인과적인) 어텐션 패턴에서, 쿼리의 길이가 동일한 경우에 맞춰 설계되었습니다.
  • 따라서 우리가 필요로 하는 것처럼 불규칙한 블록 마스크나 토큰별로 허용된 위치가 제각각인 경우를 자연스럽게 처리하지 못합니다.
  • 즉, 구현 복잡도를 감수하면 사용할 수는 있지만, 우리는 더 적합한 대안을 선택했습니다.

PyTorch에서 FlexAttention 사용하기

이 문제는 FlexAttention으로 해결할 수 있습니다! FlexAttention은 순수 PyTorch 인터페이스를 제공하며, 두 가지 방법을 시도해볼 수 있습니다:

  1. score_mod를 추가하는 방법은 어텐션이 비활성화된 위치(즉, 주의를 기울이지 않아야 하는 위치)에서 causal mask에 score_mod를 추가하는 방식입니다. 그러나 단순히 스칼라 값을 더하는 것만으로도 FlexAttention의 성능이 크게 저하되었습니다. 그 이유는, 우리의 경우 score_mod가 최적화된 CUDA 커널 외부에서 더해지기 때문입니다.
  2. 정확한 접근법은 인과 마스크(causal mask)를 인덱싱하여, 그 결과로 얻은 시그니처를 사용해 block mask를 생성하는 것입니다. 이렇게 생성된 block mask는 어텐션을 실제로 계산해야 하는 위치와 완전히 건너뛸 수 있는 위치를 효율적으로 구분합니다.
# causal mask를 인덱싱하고 mask_mod를 사용하는 예시
causal_mask = generate_causal_mask(your_samples)  # 형태는 [batch, head, q_len, kv_len] 이어야 함
# 이제 이 마스크를 감싸는 함수를 정의
def precomputed_mask_factory(precomputed_mask: torch.Tensor) -> _mask_mod_signature:
	def mask_mod(b, h, q_idx, kv_idx):
		return precomputed_mask[b][h][q_idx][kv_idx]
	return mask_mod
flex_attention_output = flex_attention(query, key, value, mask_mod=mask_mod)

mask_mod = precomputed_mask_factory(causal_mask)
# 해당 시그니처를 이용해 block mask 생성
block_mask = create_block_mask(
	mask_mod=mask_mod,
	# ...
)

# 이제 FlexAttention을 호출!
attn_output, attention_weights = flex_attention(
	query,
	key,
	value,
	block_mask=block_mask,
)

현재 구현은 정상적으로 작동하지만, torch.compile을 지원하도록 개선하는 작업이 진행 중(WIP)이며, 이를 통해 성능을 최대한 활용할 수 있도록 할 예정입니다!

행동(Action)을 효과적으로 표현하는 방법

이제 행동(action)이 단순히 n차원 벡터로 표현될 수 있고, 이를 토크나이즈할 수 있다는 것을 알게 되었습니다. 그렇다면 이제 Vision-Language-Action (VLA) 모델에서 행동 표현(action representation)이 가지는 도전 과제들을 살펴보겠습니다. 행동이 어떻게 표현되느냐는 효율성, 일반화 능력, 그리고 실행 정확도에 직접적인 영향을 미칩니다.

하나의 접근 방식은 의미 기반 행동 표현(semantic action representation)입니다. 이 방식에서는 행동을 하위 과제(sub-task)나 키포인트(keypoint) 같은 고수준 개념(high-level concept)으로 설명합니다. 이러한 방법은 few-shot이나 zero-shot 학습을 가능하게 하지만, 대개 수작업으로 설계된 저수준 제어기(low-level controller)에 의존하기 때문에 서로 다른 로봇 간의 유연성이 떨어지는 한계가 있습니다. 반대로 저수준 제어 표현(low-level control representation)은 행동을 직접 모터 명령으로 매핑합니다. 이 방식은 정밀한 동작을 가능하게 하지만, 훈련의 안정성이 낮고 확장성이 떨어진다는 단점이 있습니다.

대부분의 기존 VLA 모델은 이산적 행동 토크나이제이션(discrete action tokenization)방식을 사용합니다. 즉, 연속적인 행동을 자가회귀적(autoregressive)으로 생성되는 이산 토큰으로 변환하는 것입니다. 가장 일반적인 방법은 차원별, 시간 단계별 구간화 방식이지만, 이 접근법은 고주파 제어 작업에서 어려움을 겪으며, 결과적으로 손실이 많은 표현과 비효율적인 학습으로 이어집니다. 이에 대한 대안으로 벡터 양자화(Vector Quantization, VQ)나 시계열 압축(time-series compression) 같은 방법들이 제안되었지만, VQ는 하이퍼파라미터에 매우 민감하여, 다양한 로봇 설계 간의 일반성을 확보하기 어렵다는 한계가 있습니다.

이러한 문제를 해결하기 위해, 주파수 공간 행동 시퀀스 토크나이제이션(Frequency-space Action Sequence Tokenization, FAST)이 제안되었습니다. FAST는 이산 코사인 변환(Discrete Cosine Transform, DCT)을 활용한 새로운 시계열 압축 기법으로, 중복을 줄이고, 효율성을 향상시키며, 행동의 충실도를 향상시킵니다.

이를 바탕으로, 우리는 π0-FAST를 소개합니다. π0-FASTπ0의 확장 버전으로, 새로운 토크나이저를 활용해 행동 표현을 개선한 더 빠르고, 자가회귀적(autoregressive) 구조의 모델입니다. 이 모델은 LeRobot repo에서도 사용할 수 있습니다.


🚀 π0-FAST란 무엇인가?

논문 Jax 코드 LeRobot 구현 버전

π0-FAST는 π0의 자가회귀(autoregressive) 버전으로, 새로운 토크나이제이션 방식인 FAST (Frequency-space Action Sequence Tokenization)을 도입하여 모델의 효율성과 성능을 대폭 향상시킨 모델입니다.

π0-FAST의 주요 장점:

  • Diffusion 기반 VLA보다 5배 빠른 학습 속도
  • 행동 시퀀스 내 중복 감소로 인한 더 나은 행동 표현력
  • 보지 못한 환경과 다양한 로봇 형태에 대한 일반화 능력 강화

🔗 π0-FAST 토크나이저는 여기에서 확인할 수 있습니다: FAST Tokenizer
🔗 Pretrained weights는 여기에서 다운로드할 수 있습니다: Pi0+FAST

FAST는 어떻게 작동할까?

FAST는 이산 코사인 변환(Discrete Cosine Transform, DCT)을 사용하여 연속적인 로봇 행동 시퀀스를 이산 토큰(discrete tokens)으로 압축합니다. 그 과정은 아래의 Figure 2에 설명되어 있습니다. 먼저 로봇의 원시 행동(raw robot actions)을 정규화합니다. 각 행동 차원의 1번째 분위수(quantile)와 99번째 분위수를 [-1, 1] 범위에 매핑하여, 서로 다른 로봇 시스템 간의 일관성을 확보하고 이상치에 대한 강건성을 높입니다.

그 다음, 각 행동 차원은 DCT를 통해 독립적으로 변환되어, 시간 도메인(time-domain) 신호가 주파수 도메인(frequency-domain)으로 바뀝니다. 중복을 줄이기 위해, 중요하지 않은 계수들을 scale-and-round 연산을 통해 제거됩니다. 이때 하이퍼파라미터를 통해 압축률과 재구성 정확도 사이의 균형을 조절합니다. 이렇게 생성된 DCT 계수 행렬은 일반적으로 희소(sparse)하며, 이를 1차원 정수 시퀀스로 평탄화합니다. 이때 저주파 성분(low-frequency components)이 우선적으로 배치되어, 핵심 정보가 손실되지 않도록 합니다.

시퀀스를 한층 더 압축하기 위해 Byte Pair Encoding (BPE)이 적용됩니다. 일반적인 방식과 마찬가지로, BPE는 여러 차원에 걸쳐 자주 등장하는 패턴들을 병합하면서도 고정된 크기의 vocabulary를 유지합니다.

image/png

Figure 2: FAST action 토크나이제이션 파이프라인

모든 연산은 가역적(invertible)이기 때문에, 토큰으로부터 행동을 효율적이고 손실 없이 복원할 수 있습니다. FAST의 토크나이제이션 파이프라인은 단 두 개의 하이퍼파라미터만을 가집니다: 반올림 전에 적용되는 스케일링 계수와 BPE vocabulary 크기입니다. 이 두 파라미터는 서로 다른 데이터셋에서도 높은 안정성과 일관성을 유지합니다.

또한, FAST의 범용 버전인 FAST+는 단일 팔(single-arm), 양팔(bimanual), 모바일 매니퓰레이터(mobile manipulation) 로봇 등 다양한 형태의 로봇에서 수집된 100만 개의 행동 시퀀스로 학습되었습니다. 따라서 FAST+는 다양한 로봇 환경에 폭넓게 적용할 수 있습니다. FAST+는 Hugging Face AutoProcessor 형태로 제공되며, 단 몇 줄의 코드만으로 행동 시퀀스를 손쉽게 토큰화할 수 있습니다.

최적의 압축 성능을 얻기 위해서는, 토크나이제이션 전에 입력 행동을 [-1, 1] 범위로 분위수 정규화(quantile-normalization) 하는 것이 좋습니다. 또한, AutoProcessor 모듈을 사용하면 사용자 데이터셋에 맞춰 커스텀 FAST 토크나이저를 직접 학습시킬 수도 있습니다.


FAST 토크나이저 사용 방법

🔗 공식 Repo에서 FAST 토크나이저의 사용법 및 커스텀 action 토크나이저 학습 코드 예시를 확인하세요:
FAST Repo

FAST는 Hugging Face Transformers에 통합되어 있으며, 로봇 행동 시퀀스를 인코딩 및 디코딩하는 데 쉽게 사용할 수 있습니다.

앞으로의 범용 로봇 지능

π0와 π0-FAST는 범용 로봇 지능(Generalist Robot Intelligence)을 향한 중요한 발걸음입니다. 이 두 모델은 확장 가능하고, 효율적이며, 다재다능한 Vision-Language-Action (VLA) 모델을 LeRobot 생태계에 도입했습니다. 특히 FAST 토크나이제이션을 활용함으로써 행동 표현을 더욱 정교하게 만들고, 로봇이 다양한 작업을 더 높은 효율성과 적응력으로 수행할 수 있게 되었습니다. 이러한 발전은 앞으로 다중 로봇 형태(multi-embodiment), 실시간 로봇 정책의 기반을 마련하며, 로봇이 실제 세계에서 수행할 수 있는 일의 한계를 한층 더 확장할 것입니다. 🚀

참고 자료

참고 문헌

@book{heinlein2021time,
  title={Time enough for love},
  author={Heinlein, Robert A},
  year={2021},
  publisher={Penguin}
}

@article{black2024pi_0,
  title={$$\backslash$pi\_0 $: A Vision-Language-Action Flow Model for General Robot Control},
  author={Black, Kevin and Brown, Noah and Driess, Danny and Esmail, Adnan and Equi, Michael and Finn, Chelsea and Fusai, Niccolo and Groom, Lachy and Hausman, Karol and Ichter, Brian and others},
  journal={arXiv preprint arXiv:2410.24164},
  year={2024}
}

@article{pertsch2025fast,
  title={FAST: Efficient Action Tokenization for Vision-Language-Action Models},
  author={Pertsch, Karl and Stachowicz, Kyle and Ichter, Brian and Driess, Danny and Nair, Suraj and Vuong, Quan and Mees, Oier and Finn, Chelsea and Levine, Sergey},
  journal={arXiv preprint arXiv:2501.09747},
  year={2025}
}

hyojung
hyojung Hello, I’m Hyojung Lee! 🤗