# π₀-FAST π₀-FAST는 **범용 로봇 제어를 위한 Vision-Language-Action 모델**로, 오토레그레시브 다음 토큰 예측을 사용하여 연속적인 로봇 동작을 모델링합니다. ## 모델 개요 π₀-FAST은 Vision-Language Models의 강점을 **FAST(Frequency-space Action Sequence Tokenization)**라는 새로운 액션 토큰화 기법과 결합합니다. 이를 통해 표준 바인닝 기반 이산화로는 불가능한 고도의 민첩한 작업을 수행하는 오토레그레시브 VLA를 학습할 수 있으며, π₀와 같은 확산 기반 접근법보다 **최대 5배 빠르게** 학습할 수 있습니다. An overview of Pi0-FAST ### 왜 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은 [설치 가이드](./installation)를 참고하여 설치하세요. 2. 다음 명령을 실행하여 π₀-FAST의 종속성을 설치합니다: ```bash pip install -e ".[pi]" ``` ## 커스텀 FAST 토크나이저 학습 FAST 토크나이저에는 두 가지 옵션이 있습니다: 1. **사전 학습된 토크나이저 사용**: `lerobot/fast-action-tokenizer` 토크나이저는 100만 개 이상의 실제 로봇 액션 시퀀스로 학습되어 범용 토크나이저로 작동합니다. 2. **직접 토크나이저 학습하기**: 특정 데이터셋에서 최대 성능을 얻으려면 자체 데이터로 토크나이저를 파인튜닝할 수 있습니다. ### 직접 토크나이저 학습하기 ```bash 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 매개변수 | 매개변수 | 설명 | 기본값 | | ---------------------- | --------------------------------------------------------------------------------- | ------------ | | `--repo_id` | LeRobot 데이터셋 저장소 ID | 필수 | | `--action_horizon` | 각 청크의 미래 동작 수 | `10` | | `--encoded_dims` | 인코딩할 쉼표로 구분된 차원 범위 (예: `"0:6,7:23"`) | `"0:6,7:23"` | | `--vocab_size` | BPE 어휘 크기 | `1024` | | `--scale` | 양자화를 위한 DCT 스케일링 계수 | `10.0` | | `--normalization_mode` | 정규화 모드 (`MEAN_STD`, `MIN_MAX`, `QUANTILES`, `QUANTILE10`, `IDENTITY`) | `QUANTILES` | | `--sample_fraction` | 에피소드당 샘플링할 청크 비율 | `0.1` | ## 사용 LeRobot에서 π₀-FAST를 사용하려면 정책 유형을 다음과 같이 지정하세요: ```python policy.type=pi0_fast ``` ## 학습 π₀-FAST를 학습하려면 LeRobot 학습 스크립트를 사용할 수 있습니다: ```bash 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 ``` ### 주요 학습 매개변수 | 매개변수 | 설명 | 기본값 | | -------------------------------------- | -------------------------------------------------- | ------------------------------- | | `--policy.gradient_checkpointing=true` | 학습 중 메모리 사용량을 크게 줄입니다 | `false` | | `--policy.dtype=bfloat16` | 효율성을 위해 혼합 정밀도 학습 사용 | `float32` | | `--policy.chunk_size` | 예측할 액션 스텝 수(액션 호라이즌) | `50` | | `--policy.n_action_steps` | 실행할 액션 단계 수 | `50` | | `--policy.max_action_tokens` | 액션 청크당 최대 FAST 토큰 수 | `256` | | `--policy.action_tokenizer_name` | 사용할 FAST 토크나이저 | `lerobot/fast-action-tokenizer` | | `--policy.compile_model=true` | torch.compile를 활성화하여 더 빠른 학습 | `false` | | `--policy.optimizer_fused=true` | fused AdamW 사용; 재개 시 무시 | `None` | ## 추론 ### 빠른 추론을 위한 KV-Caching π₀-FAST는 LLM 추론에서 널리 사용되는 최적화 기법인 **KV-caching**을 지원합니다. 이 기법은 어텐션 메커니즘의 키-밸류 쌍을 캐시에 저장하여 오토레그레시브 디코딩 과정에서의 중복 계산을 방지합니다. ```python # KV-caching is enabled by default policy.use_kv_cache=true ``` ### 추론 예시 ```python 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 토큰을 출력하고, 이 토큰은 다시 연속 동작으로 디코딩됩니다. ## 구성 옵션 | 매개변수 | 설명 | 기본값 | | -------------------- | ----------------------------------------------- | ---------- | | `paligemma_variant` | VLM 백본 변형 (`gemma_300m`, `gemma_2b`) | `gemma_2b` | | `max_state_dim` | 최대 상태 벡터 차원(패딩됨) | `32` | | `max_action_dim` | 최대 액션 벡터 차원(패딩됨) | `32` | | `temperature` | 샘플링 온도 (그리디의 경우 0.0) | `0.0` | | `max_decoding_steps` | 최대 디코딩 단계 | `256` | | `use_kv_cache` | 추론 속도를 높이기 위해 KV 캐싱 활성화 | `true` | ## π₀와의 비교 | 기능 | π₀ | π₀-FAST | | --------------------- | ------------------------- | ---------------------------- | | 액션 표현 | Flow Matching (확산) | 자기회귀 토큰(FAST) | | 학습 속도 | 1x | **5배 빠름** | | 민첩성 | 높음 | 높음 | | 추론 방법 | 반복적 노이즈 제거 | 자기 회귀 디코딩 | | KV 캐싱 | N/A | 지원 | ## π₀Fast 결과 재현 LeRobot 구현을 사용하여 LIBERO 벤치마크에서 π₀Fast의 결과를 재현합니다. LeRobot PiFast 기반 모델 [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base)를 가져와 [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero)을 사용하여 bfloat16으로, 배치 크기 256, 8개의 H100 GPU에서 추가로 40k 스텝 동안 파인튜닝합니다. 여기에서 fine-tuned된 모델을 찾을 수 있습니다: - **π₀Fast LIBERO**: [lerobot/pi0fast-libero](https://huggingface.co/lerobot/pi0fast-libero) 다음 학습 명령을 사용하여 ```bash 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 모델을 평가합니다: ```bash 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 벤치마크에서 다음 결과를 얻습니다: | 모델 | LIBERO Spatial | LIBERO Object | LIBERO 목표 | LIBERO 10 | 평균 | | ----------- | -------------- | ------------- | ----------- | --------- | -------- | | **π₀-fast** | 70.0 | 100.0 | 100.0 | 60.0 | **82.5** | 비디오를 포함한 전체 평가 출력 폴더는 [여기](https://drive.google.com/drive/folders/1HXpwPTRm4hx6g1sF2P7OOqGG0TwPU7LQ?usp=sharing)에서 확인할 수 있습니다. ## 라이선스 이 모델은 원래의 [OpenPI 저장소](https://github.com/Physical-Intelligence/openpi)와 일치하는 **Apache 2.0 라이선스**를 따릅니다. ## 참고 - [FAST: 효율적인 로봇 액션 토큰화](https://www.physicalintelligence.company/research/fast) - Physical Intelligence 블로그 - [OpenPI 저장소](https://github.com/Physical-Intelligence/openpi) - 원본 구현 - [Hugging Face의 FAST Tokenizer](https://huggingface.co/physical-intelligence/fast) - 사전 학습된 토크나이저