π₀-FAST (Pi0-FAST)

π₀-FAST é um modelo Vision-Language-Action para controle geral de robôs que usa a previsão autoregressiva do próximo token para modelar ações contínuas de robôs.

Visão geral do modelo

π₀-FAST combina o poder dos Vision-Language Models com uma nova abordagem de tokenização de ações chamada FAST (Frequency-space Action Sequence Tokenization). Isso permite treinar VLAs autoregressivos em tarefas altamente ágeis que são impossíveis com a discretização padrão baseada em binning, enquanto o treinamento é até 5x mais rápido do que abordagens baseadas em difusão, como π₀.

Por que FAST?

As abordagens padrão para a tokenização de ações de robôs utilizam esquemas simples de binning por dimensão e por passo de tempo. Embora sejam aceitáveis para comportamentos simples, essa abordagem se torna rapidamente ineficaz para habilidades complexas e ágeis que exigem precisão e controle de alta frequência.

O FAST resolve isso comprimindo sequências de ações usando técnicas de processamento de sinal, resultando em uma sequência densa de tokens de ação que pode ser prevista de forma autoregressiva — assim como tokens de linguagem.

Como o tokenização FAST funciona

O tokenizador FAST comprime sequências de ações por meio das seguintes etapas:

1. Normalizar: Pega um trecho de ação contínuo de forma (H, D), onde H é o horizonte e D é a dimensão da ação. Normaliza usando um dos métodos de normalização suportados (Quantis é recomendado para lidar com valores extremos).

2. Transformada Cosseno Discreta (DCT): Aplique DCT (via scipy) a cada dimensão de ação separadamente. DCT é um algoritmo de compressão comumente usado em codecs de imagem e áudio (JPEG, MP3).

3. Quantização: Arredonde e remova coeficientes insignificantes para cada dimensão de ação, produzindo uma matriz de frequência esparsa.

4. Flatten: Achatar a matriz em um vetor 1D, com os componentes de baixa frequência primeiro.

5. Byte Pair Encoding (BPE): Treine um tokenizador BPE para comprimir os coeficientes DCT em tokens de ação densos, alcançando tipicamente 10x de compressão em relação a abordagens anteriores de tokenização.

Essa abordagem pode transformar qualquer VLM existente em um VLA, treinando-o para prever esses tokens FAST.

Requisitos de instalação

1. Instale o LeRobot seguindo nosso Guia de Instalação.

2. Instale as dependências do π₀-FAST executando:

   Copiado

   pip install -e ".[pi]"

Treinando um Tokenizador FAST Personalizado

Você tem duas opções para o tokenizador FAST:

1. Use o tokenizer pré-treinado: O tokenizer lerobot/fast-action-tokenizer foi treinado em mais de 1 milhão de sequências reais de ações de robôs e funciona como um tokenizer de uso geral.

2. Treine seu próprio tokenizer: Para obter o máximo de desempenho no seu conjunto de dados específico, você pode fazer o ajuste fino do tokenizer nos seus próprios dados.

Treinando o Seu Próprio 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 Principais do Tokenizer

Uso

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

policy.type=pi0_fast

Treinamento

Para treinar o π₀-FAST, você pode usar o script de treinamento do 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 Principais de Treinamento

Inferência

KV-Caching para Inferência Rápida

π₀-FAST suporta cache de KV, uma otimização amplamente utilizada na inferência de LLMs. Isso faz o cache dos pares chave-valor do mecanismo de atenção, evitando cálculos redundantes durante a decodificação autoregressiva.

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

Exemplo de Inferência

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)

Arquitetura do Modelo

π₀-FAST usa uma arquitetura baseada em PaliGemma:

• Vision Encoder: torre de visão SigLIP para compreensão de imagens
• Modelo de linguagem: Gemma 2B para processar instruções de linguagem e prever tokens de ação

O modelo recebe imagens, instruções de texto e o estado do robô como entrada e produz tokens FAST discretos que são decodificados de volta para ações contínuas.

Opções de Configuração

Comparação com π₀

Reproduzindo os resultados do π₀Fast

Reproduzimos os resultados do π₀Fast no benchmark LIBERO usando a implementação do LeRobot. Utilizamos o modelo base LeRobot PiFast lerobot/pi0fast-base e realizamos o finetuning por mais 40k passos em bfloat16, com tamanho de lote de 256 em 8 GPUs H100, usando o conjunto de dados LIBERO do HuggingFace.

O modelo ajustado pode ser encontrado aqui:

• π₀Fast LIBERO: lerobot/pi0fast-libero

Com o seguinte comando de treinamento:

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

Em seguida, avaliamos o modelo finetuned usando a implementação LeRobot LIBERO, executando o seguinte 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"}'

Observação: Definimos n_action_steps=10, semelhante à implementação original do OpenPI.

Resultados

Obtemos os seguintes resultados no benchmark LIBERO:

A pasta de saída completa da avaliação, incluindo os vídeos, está disponível aqui

Licença

Este modelo segue a Apache 2.0 License, em conformidade com o repositório original OpenPI.

Referências

• FAST: Tokenização eficiente de ações de robôs - Blog da Physical Intelligence
• Repositório OpenPI - Implementação original
• Tokenizer FAST no Hugging Face - Tokenizer pré-treinado