π₀-FAST

π₀-FAST는 범용 로봇 제어를 위한 Vision-Language-Action 모델로, 오토레그레시브 다음 토큰 예측을 사용하여 연속적인 로봇 동작을 모델링합니다.

모델 개요

π₀-FAST은 Vision-Language Models의 강점을 FAST(Frequency-space Action Sequence Tokenization)라는 새로운 액션 토큰화 기법과 결합합니다. 이를 통해 표준 바인닝 기반 이산화로는 불가능한 고도의 민첩한 작업을 수행하는 오토레그레시브 VLA를 학습할 수 있으며, π₀와 같은 확산 기반 접근법보다 최대 5배 빠르게 학습할 수 있습니다.

왜 FAST인가요?

로봇 액션 토큰화를 위한 표준 접근법은 차원별, 타임스텝별 단순한 버닝(bin) 스키움을 사용합니다. 단순한 동작에는 어느 정도 적합하지만, 정밀도와 고주파 제어의 정확성이 요구되는 복잡하고 민첩한 스킬에서는 이 방식이 빠르게 한계에 부딪힙니다.

FAST는 신호 처리 기법을 사용하여 액션 시퀀스를 압축하여 이를 해결합니다. 그 결과, 언어 토큰과 마찬가지로 자기회귀적으로 예측할 수 있는 밀집된 액션 토큰 시퀀스가 생성됩니다.

FAST 토큰화가 작동하는 방식

FAST 토크나이저는 다음 단계를 통해 액션 시퀀스를 압축합니다:

1. Normalize: H가 horizon, D가 action dimension인 shape (H, D)의 연속 action chunk를 가져옵니다. 지원되는 정규화 방법 중 하나(외래치 처리를 위해 Quantiles 권장)를 사용하여 정규화합니다.

2. 이산 코사인 변환(DCT): 각 동작 차원에 대해 개별적으로 DCT를 적용합니다(scipy 사용). DCT는 이미지 및 오디오 코덱(JPEG, MP3)에서 일반적으로 사용되는 압축 알고리즘입니다.

3. Quantization: 각 액션 차원에 대해 중요하지 않은 계수를 반올림하고 제거하여 희소 주파수 행렬을 생성합니다.

4. Flatten: 행렬을 1차원 벡터로 평탄화하며, 저주파 성분을 앞에 배치합니다.

5. Byte Pair Encoding (BPE): DCT 계수를 밀집된 액션 토큰으로 압축하기 위해 BPE 토크나이저를 학습하며, 일반적으로 기존 토크나이제이션 방식 대비 10배 압축을 달성합니다.

이 접근법은 이러한 FAST 토큰을 예측하도록 학습함으로써 기존의 어떤 VLM이든 VLA로 변환할 수 있습니다.

설치 요구 사항

1. LeRobot은 설치 가이드를 참고하여 설치하세요.

2. 다음 명령을 실행하여 π₀-FAST의 종속성을 설치합니다:

   복사됨

   pip install -e ".[pi]"

커스텀 FAST 토크나이저 학습

FAST 토크나이저에는 두 가지 옵션이 있습니다:

1. 사전 학습된 토크나이저 사용: lerobot/fast-action-tokenizer 토크나이저는 100만 개 이상의 실제 로봇 액션 시퀀스로 학습되어 범용 토크나이저로 작동합니다.

2. 직접 토크나이저 학습하기: 특정 데이터셋에서 최대 성능을 얻으려면 자체 데이터로 토크나이저를 파인튜닝할 수 있습니다.

직접 토크나이저 학습하기

lerobot-train-tokenizer \
    --repo_id "user/my-lerobot-dataset" \
    --action_horizon 10 \
    --encoded_dims "0:6" \
    --vocab_size 1024 \
    --scale 10.0 \
    --normalization_mode QUANTILES \
    --output_dir "./my_fast_tokenizer" \
    --push_to_hub true \
    --hub_repo_id "username/my-action-tokenizer"

주요 Tokenizer 매개변수

사용

LeRobot에서 π₀-FAST를 사용하려면 정책 유형을 다음과 같이 지정하세요:

policy.type=pi0_fast

학습

π₀-FAST를 학습하려면 LeRobot 학습 스크립트를 사용할 수 있습니다:

lerobot-train \
    --dataset.repo_id=your_dataset \
    --policy.type=pi0_fast \
    --output_dir=./outputs/pi0fast_training \
    --job_name=pi0fast_training \
    --policy.pretrained_path=lerobot/pi0fast-base \
    --policy.dtype=bfloat16 \
    --policy.gradient_checkpointing=true \
    --policy.chunk_size=10 \
    --policy.n_action_steps=10 \
    --policy.max_action_tokens=256 \
    --steps=100000 \
    --batch_size=4 \
    --policy.device=cuda

주요 학습 매개변수

추론

빠른 추론을 위한 KV-Caching

π₀-FAST는 LLM 추론에서 널리 사용되는 최적화 기법인 KV-caching을 지원합니다. 이 기법은 어텐션 메커니즘의 키-밸류 쌍을 캐시에 저장하여 오토레그레시브 디코딩 과정에서의 중복 계산을 방지합니다.

# KV-caching is enabled by default
policy.use_kv_cache=true

추론 예시

from lerobot.policies.pi0_fast import PI0FastPolicy, PI0FastConfig

# Load the policy
policy = PI0FastPolicy.from_pretrained("your-model-path")

# During inference
actions = policy.predict_action_chunk(batch)

모델 아키텍처

π₀-FAST는 PaliGemma 기반 아키텍처를 사용합니다:

• Vision Encoder: 이미지 이해를 위한 SigLIP 비전 타워
• Language Model: 언어 지시사항을 처리하고 액션 토큰을 예측하는 Gemma 2B

모델은 이미지, 텍스트 명령, 로봇 상태를 입력으로 받아 이산 FAST 토큰을 출력하고, 이 토큰은 다시 연속 동작으로 디코딩됩니다.

구성 옵션

π₀와의 비교

π₀Fast 결과 재현

LeRobot 구현을 사용하여 LIBERO 벤치마크에서 π₀Fast의 결과를 재현합니다. LeRobot PiFast 기반 모델 lerobot/pi0fast-base를 가져와 HuggingFace LIBERO dataset을 사용하여 bfloat16으로, 배치 크기 256, 8개의 H100 GPU에서 추가로 40k 스텝 동안 파인튜닝합니다.

여기에서 fine-tuned된 모델을 찾을 수 있습니다:

• π₀Fast LIBERO: lerobot/pi0fast-libero

다음 학습 명령을 사용하여

lerobot-train \
  --dataset.repo_id=lerobot/libero \
  --output_dir=outputs/libero_pi0fast \
  --job_name=libero_pi0fast \
  --policy.path=lerobot/pi0fast-base \
  --policy.dtype=bfloat16 \
  --steps=100000 \
  --save_freq=20000 \
  --batch_size=4 \
  --policy.device=cuda \
  --policy.scheduler_warmup_steps=4000 \
  --policy.scheduler_decay_steps=100000 \
  --policy.scheduler_decay_lr=1e-5 \
  --policy.gradient_checkpointing=true \
  --policy.chunk_size=10 \
  --policy.n_action_steps=10 \
  --policy.max_action_tokens=256 \
  --policy.empty_cameras=1

이후 다음 명령을 실행하여 LeRobot LIBERO 구현으로 fine-tuned 모델을 평가합니다:

tasks="libero_object,libero_spatial,libero_goal,libero_10"
lerobot-eval \
  --policy.path=lerobot/pi0fast-libero \
  --policy.max_action_tokens=256 \
  --env.type=libero \
  --policy.gradient_checkpointing=false \
  --env.task=${tasks} \
  --eval.batch_size=1 \
  --eval.n_episodes=1 \
  --rename_map='{"observation.images.image":"observation.images.base_0_rgb","observation.images.image2":"observation.images.left_wrist_0_rgb"}'

참고: 원본 OpenPI 구현과 유사하게 n_action_steps=10으로 설정했습니다.

결과

LIBERO 벤치마크에서 다음 결과를 얻습니다:

비디오를 포함한 전체 평가 출력 폴더는 여기에서 확인할 수 있습니다.

라이선스

이 모델은 원래의 OpenPI 저장소와 일치하는 Apache 2.0 라이선스를 따릅니다.

참고

• FAST: 효율적인 로봇 액션 토큰화 - Physical Intelligence 블로그
• OpenPI 저장소 - 원본 구현
• Hugging Face의 FAST Tokenizer - 사전 학습된 토크나이저