π₀-FAST (Pi0-FAST)

π₀-FAST es un modelo de visión-lenguaje-acción para el control general de robots que utiliza la predicción autoregresiva del siguiente token para modelar acciones continuas de robots.

Resumen del modelo

π₀-FAST combina el poder de los modelos de visión-lenguaje con un enfoque novedoso de tokenización de acciones llamado FAST (Frequency-space Action Sequence Tokenization). Esto permite entrenar VLAs autoregresivos en tareas altamente diestras que son imposibles con la discretización estándar basada en binning, mientras que el entrenamiento es hasta 5 veces más rápido que los enfoques basados en difusión como π₀.

¿Por qué FAST?

Los enfoques estándar para la tokenización de acciones de robots utilizan esquemas de binning simples por dimensión y por paso de tiempo. Si bien son aceptables para comportamientos simples, este enfoque se desmorona rápidamente ante habilidades complejas y diestras que requieren precisión y control de alta frecuencia.

FAST resuelve esto comprimiendo las secuencias de acciones mediante técnicas de procesamiento de señales, lo que da como resultado una secuencia densa de tokens de acción que se pueden predecir de forma autoregresiva, al igual que los tokens de lenguaje.

Cómo funciona la tokenización FAST

El tokenizador FAST comprime las secuencias de acciones mediante los siguientes pasos:

1. Normalizar: Toma un trozo de acciones continuas de forma (H, D) donde H es el horizonte y D es la dimensión de la acción. Normaliza usando uno de los métodos de normalización admitidos (se recomienda Quantiles para manejar valores atípicos).

2. Transformada Discreta del Coseno (DCT): Aplica DCT (mediante scipy) a cada dimensión de acción por separado. DCT es un algoritmo de compresión comúnmente utilizado en códecs de imagen y audio (JPEG, MP3).

3. Cuantización: Redondear y eliminar coeficientes insignificantes para cada dimensión de acción, produciendo una matriz de frecuencia dispersa.

4. Flatten: Aplana la matriz en un vector 1D, con los componentes de baja frecuencia primero.

5. Codificación de pares de bytes (BPE): Entrena un tokenizador BPE para comprimir los coeficientes DCT en tokens de acción densos, logrando típicamente una compresión de 10x sobre los enfoques de tokenización anteriores.

Este enfoque puede transformar cualquier VLM existente en un VLA entrenándolo para predecir estos tokens FAST.

Requisitos de instalación

1. Instala LeRobot siguiendo nuestra Guía de instalación.

2. Instale las dependencias de π₀-FAST ejecutando:

   Copiado

   pip install -e ".[pi]"

Entrenamiento de un tokenizador FAST personalizado

Tienes dos opciones para el tokenizador FAST:

1. Usar el tokenizador preentrenado: El tokenizador lerobot/fast-action-tokenizer se entrenó con más de 1M de secuencias de acciones de robots reales y funciona como un tokenizador de uso general.

2. Entrena tu propio tokenizador: Para obtener el máximo rendimiento en tu conjunto de datos específico, puedes ajustar el tokenizador con tus propios datos.

Entrenar su propio tokenizador

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"

Parámetros clave del Tokenizer

Uso

Para usar π₀-FAST en LeRobot, especifique el tipo de política como:

policy.type=pi0_fast

Entrenamiento

Para entrenar π₀-FAST, puedes usar el script de entrenamiento de 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

Parámetros clave de entrenamiento

Inferencia

KV-Caching para una inferencia rápida

π₀-FAST admite el almacenamiento en caché de pares clave-valor (KV-caching), una optimización ampliamente utilizada en la inferencia de LLM. Esto almacena en caché los pares clave-valor del mecanismo de atención, evitando cálculos redundantes durante la decodificación autoregresiva.

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

Ejemplo de inferencia

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)

Arquitectura del modelo

π₀-FAST utiliza una arquitectura basada en PaliGemma:

• Codificador de visión: torre de visión SigLIP para la comprensión de imágenes
• Modelo de lenguaje: Gemma 2B para procesar instrucciones de lenguaje y predecir tokens de acción

El modelo toma como entrada imágenes, instrucciones de texto y el estado del robot, y genera tokens FAST discretos que se decodifican de nuevo a acciones continuas.

Opciones de configuración

Comparación con π₀

Reproduciendo los resultados de π₀Fast

Reproducimos los resultados de π₀Fast en el benchmark LIBERO utilizando la implementación de LeRobot. Tomamos el modelo base LeRobot PiFast lerobot/pi0fast-base y lo afinamos durante 40k pasos adicionales en bfloat16, con un tamaño de lote de 256 en 8 GPUs H100, utilizando el conjunto de datos HuggingFace LIBERO.

El modelo afinado se puede encontrar aquí:

• π₀Fast LIBERO: lerobot/pi0fast-libero

Con el siguiente comando de entrenamiento:

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

A continuación, evaluamos el modelo afinado utilizando la implementación de LeRobot LIBERO, ejecutando el siguiente comando:

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"}'

Nota: Establecemos n_action_steps=10, de forma similar a la implementación original de OpenPI.

Resultados

Obtenemos los siguientes resultados en el benchmark LIBERO:

La carpeta de salida de la evaluación completa, incluidos los videos, está disponible aquí

Licencia

Este modelo sigue la Apache 2.0 License, en línea con el repositorio original OpenPI.

Referencias

• FAST: Tokenización eficiente de acciones de robots - Blog de Physical Intelligence
• Repositorio OpenPI - Implementación original
• Tokenizer FAST en Hugging Face - Tokenizador preentrenado