# π₀-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 π₀.

<img
  src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/lerobot-pifast.png"
  alt="An overview of Pi0-FAST"
  width="85%"
/>

### ¿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](./installation).
2. Instale las dependencias de π₀-FAST ejecutando:

   ```bash
   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

```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"
```

### Parámetros clave del Tokenizer

| Parámetro              | Descripción                                                                       | Predeterminado      |
| ---------------------- | --------------------------------------------------------------------------------- | ------------ |
| `--repo_id`            | ID del repositorio de datos de LeRobot                                                     | Obligatorio     |
| `--action_horizon`     | Número de acciones futuras en cada bloque                                            | `10`         |
| `--encoded_dims`       | Rangos de dimensiones separados por comas para codificar (p. ej., `"0:6,7:23"`)                   | `"0:6,7:23"` |
| `--vocab_size`         | Tamaño del vocabulario BPE                                                               | `1024`       |
| `--scale`              | Factor de escala DCT para cuantización                                               | `10.0`       |
| `--normalization_mode` | Modo de normalización (`MEAN_STD`, `MIN_MAX`, `QUANTILES`, `QUANTILE10`, `IDENTITY`) | `QUANTILES`  |
| `--sample_fraction`    | Fracción de fragmentos a muestrear por episodio                                          | `0.1`        |

## Uso

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

```python
policy.type=pi0_fast
```

## Entrenamiento

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

### Parámetros clave de entrenamiento

| Parámetro                              | Descripción                                        | Predeterminado                         |
| -------------------------------------- | -------------------------------------------------- | ------------------------------- |
| `--policy.gradient_checkpointing=true` | Reduce significativamente el uso de memoria durante el entrenamiento | `false`                         |
| `--policy.dtype=bfloat16`              | Use el entrenamiento con precisión mixta para mejorar la eficiencia        | `float32`                       |
| `--policy.chunk_size`                  | Número de pasos de acción a predecir (horizonte de acción) | `50`                            |
| `--policy.n_action_steps`              | Número de pasos de acción a ejecutar                  | `50`                            |
| `--policy.max_action_tokens`           | Número máximo de tokens FAST por trozo de acción     | `256`                           |
| `--policy.action_tokenizer_name`       | Tokenizer FAST a utilizar                              | `lerobot/fast-action-tokenizer` |
| `--policy.compile_model=true`          | Activar torch.compile para un entrenamiento más rápido           | `false`                         |
| `--policy.optimizer_fused=true`        | Use AdamW fusionado; se ignora al reanudar                 | `None`                          |

## 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.

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

### Ejemplo de inferencia

```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)
```

## 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

| Parámetro            | Descripción                                     | Predeterminado    |
| -------------------- | ----------------------------------------------- | ---------- |
| `paligemma_variant`  | Variante del backbone del VLM (`gemma_300m`, `gemma_2b`) | `gemma_2b` |
| `max_state_dim`      | Dimensión máxima del vector de estado (rellenada)         | `32`       |
| `max_action_dim`     | Dimensión máxima del vector de acción (rellenada)        | `32`       |
| `temperature`        | Temperatura de muestreo (0.0 para greedy)           | `0.0`      |
| `max_decoding_steps` | Pasos máximos de decodificación                          | `256`      |
| `use_kv_cache`       | Activar el almacenamiento en caché de KV para una inferencia más rápida          | `true`     |

## Comparación con π₀

| Característica               | π₀                        | π₀-FAST                      |
| --------------------- | ------------------------- | ---------------------------- |
| Representación de la acción | Flow Matching (difusión) | Tokens autoregresivos (FAST) |
| Velocidad de entrenamiento        | 1x                        | **5 veces más rápido**                |
| Dexteridad             | Alto                      | Alto                         |
| Método de inferencia      | Desruido iterativo       | Decodificación autorregresiva      |
| KV-Caching            | N/A                       | Soportado                    |

## 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](https://huggingface.co/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](https://huggingface.co/datasets/HuggingFaceVLA/libero).

El modelo afinado se puede encontrar aquí:

- **π₀Fast LIBERO**: [lerobot/pi0fast-libero](https://huggingface.co/lerobot/pi0fast-libero)

Con el siguiente comando de entrenamiento:

```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
```

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

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

**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:

| Modelo       | LIBERO Espacial | LIBERO Object | Objetivo de LIBERO | LIBERO 10 | Promedio  |
| ----------- | -------------- | ------------- | ----------- | --------- | -------- |
| **π₀-fast** | 70.0           | 100.0         | 100.0       | 60.0      | **82.5** |

La carpeta de salida de la evaluación completa, incluidos los videos, está disponible [aquí](https://drive.google.com/drive/folders/1HXpwPTRm4hx6g1sF2P7OOqGG0TwPU7LQ?usp=sharing)

## Licencia

Este modelo sigue la **Apache 2.0 License**, en línea con el repositorio original [OpenPI](https://github.com/Physical-Intelligence/openpi).

## Referencias

- [FAST: Tokenización eficiente de acciones de robots](https://www.physicalintelligence.company/research/fast) - Blog de Physical Intelligence
- [Repositorio OpenPI](https://github.com/Physical-Intelligence/openpi) - Implementación original
- [Tokenizer FAST en Hugging Face](https://huggingface.co/physical-intelligence/fast) - Tokenizador preentrenado

