π₀-FAST (Pi0-FAST)

π₀-FAST est un modèle Vision-Language-Action pour le contrôle général de robots qui utilise la prédiction auto-régressive du jeton suivant pour modéliser les actions continues des robots.

Aperçu du modèle

π₀-FAST combine la puissance des Vision-Language Models avec une approche novatrice de tokenisation des actions appelée FAST (Frequency-space Action Sequence Tokenization). Cela permet d'entraîner des VLAs autoregressifs sur des tâches très dextres, impossibles à réaliser avec une discrétisation standard basée sur le binning, tout en étant jusqu'à 5 fois plus rapide que les approches basées sur la diffusion comme π₀.

Pourquoi FAST ?

Les approches standard de tokenisation des actions robotiques utilisent des schémas de binage simples, par dimension et par pas de temps. Bien que suffisants pour des comportements simples, ces schémas s’effondrent rapidement face à des compétences complexes et dextères exigeant précision et contrôle haute fréquence.

FAST résout ce problème en compressant les séquences d’actions à l’aide de techniques de traitement du signal, ce qui aboutit à une séquence dense de jetons d’action pouvant être prédite de manière autorégressive, tout comme les jetons de langage.

Comment fonctionne la tokenisation FAST

Le tokeniseur FAST compresse les séquences d’actions au moyen des étapes suivantes :

1. Normaliser : Prendre un bloc d’actions continu de forme (H, D) où H est l’horizon et D la dimension de l’action. Normaliser en utilisant l’une des méthodes de normalisation prises en charge (Quantiles recommandée pour gérer les valeurs aberrantes).

2. Transformation en cosinus discrète (DCT) : Applique la DCT (via scipy) à chaque dimension d'action séparément. La DCT est un algorithme de compression couramment utilisé dans les codecs d'images et d'audio (JPEG, MP3).

3. Quantification : arrondir et supprimer les coefficients non significatifs pour chaque dimension d'action, produisant une matrice de fréquences creuse.

4. Flatten : Aplati la matrice en un vecteur 1D, en plaçant d'abord les composantes à basse fréquence.

5. Byte Pair Encoding (BPE) : entraîner un tokeniseur BPE afin de compresser les coefficients DCT en jetons d'action denses, atteignant typiquement une compression 10 fois supérieure par rapport aux approches de tokenisation précédentes.

Cette approche peut transformer n'importe quel VLM existant en VLA en l'entraînant à prédire ces jetons FAST.

Prérequis d'installation

1. Installez LeRobot en suivant notre Guide d'installation.

2. Installez les dépendances de π₀-FAST en exécutant :

   Copié

   pip install -e ".[pi]"

Entraîner un tokeniseur FAST personnalisé

Vous disposez de deux options pour le tokeniseur FAST :

1. Utiliser le tokenizer pré-entraîné : le tokenizer lerobot/fast-action-tokenizer a été entraîné sur plus d’1 million de séquences d’actions réelles de robots et fonctionne comme un tokenizer polyvalent.

2. Entraîner votre propre tokenizer : Pour des performances maximales sur votre jeu de données spécifique, vous pouvez affiner le tokenizer sur vos propres données.

Entraîner votre propre tokenizer

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"

Paramètres clés du Tokenizer

Utilisation

Pour utiliser π₀-FAST dans LeRobot, spécifiez le type de politique comme suit :

policy.type=pi0_fast

Entraînement

Pour l’entraînement de π₀-FAST, vous pouvez utiliser le script d’entraînement 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

Paramètres d'entraînement clés

Inférence

Mise en cache des clés et valeurs pour une inférence rapide

π₀-FAST prend en charge le cache KV, une optimisation couramment utilisée lors de l'inférence des LLM. Cette technique met en cache les paires clé-valeur issues du mécanisme d'attention, évitant ainsi les calculs redondants lors du décodage autorégressif.

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

Exemple d'inférence

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)

Architecture du modèle

π₀-FAST utilise une architecture basée sur PaliGemma :

• Encodeur visuel : tour visuelle SigLIP pour la compréhension d'images
• Modèle de langage : Gemma 2B pour le traitement des instructions linguistiques et la prédiction des jetons d'action

Le modèle prend en entrée des images, des instructions textuelles et l'état du robot, et produit des jetons FAST discrets qui sont ensuite décodés en actions continues.

Options de configuration

Comparaison avec π₀

Reproduire les résultats de π₀Fast

Nous reproduisons les résultats de π₀Fast sur le benchmark LIBERO en utilisant l'implémentation LeRobot. Nous prenons le modèle de base LeRobot PiFast lerobot/pi0fast-base et l'ajustons pendant 40k étapes supplémentaires en bfloat16, avec une taille de lot de 256 sur 8 GPU H100 en utilisant le jeu de données LIBERO de HuggingFace.

Le modèle ajusté fin peut être trouvé ici :

• π₀Fast LIBERO : lerobot/pi0fast-libero

Avec la commande d'entraînement suivante :

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

Nous évaluons ensuite le modèle affiné à l'aide de l'implémentation LeRobot LIBERO en exécutant la commande suivante :

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

Remarque : Nous définissons n_action_steps=10, de manière similaire à l'implémentation originale d'OpenPI.

Résultats

Nous obtenons les résultats suivants sur le benchmark LIBERO :

Le dossier de sortie d'évaluation complet, y compris les vidéos, est disponible ici

Licence

Ce modèle est soumis à la licence Apache 2.0, conformément au dépôt OpenPI d'origine.

Références

• FAST : Tokenisation efficace des actions robotiques - Blog de Physical Intelligence
• Dépôt OpenPI - Implémentation originale
• FAST Tokenizer sur Hugging Face - Tokenizer pré-entraîné