FLUX 3 Action

FLUX 3 Action prevê ações de robôs a partir de observações de câmera, estado e uma instrução de tarefa. Um Video VAE congelado e um codificador de texto alimentam um transformador de difusão que prevê latentes de vídeo e blocos de ações; apenas as ações são retornadas na inferência. LeRobot suporta inferência e PEFT: LoRA adapta o tronco, enquanto as cabeças de ação e condicionamento do robô são treinadas por completo.

Enviamos flux-3-action-base como modelo de ação base, juntamente com finetunings para os robôs SO-101 flux-3-action-so101 e Franka Panda flux-3-action-droid (finetuned no dataset DROID). O exemplo central aqui é uma adaptação de tarefa via LoRA baseada no modelo flux-3-action-so101.

Para treinar uma nova incorporação de robô a partir do modelo base ou fazer um ajuste fino completo, use nosso repositório independente flux-action.

Instalação

O Video VAE requer NATTEN 0.21.6 ou mais recente, como uma wheel compilada para suas versões exatas de torch e CUDA, a partir de https://whl.natten.org. O NATTEN publica wheels para torch 2.10 a 2.13, mas ainda não para o torch 2.15 que o LeRobot instala por padrão, portanto, instale primeiro um torch suportado. Por exemplo, a partir do checkout do LeRobot, torch 2.13 com CUDA 13.0 junto com as dependências de treinamento, policy, PEFT e EMA. Como no guia de instalação, o TorchCodec vem do índice de CPU:

pip install --index-url https://download.pytorch.org/whl/cu130 \
    'torch==2.13.0+cu130' 'torchvision==0.28.0+cu130' 'numpy>=2.0,<2.3'
pip install --index-url https://download.pytorch.org/whl/cpu 'torchcodec==0.13.0+cpu'
pip install -e ".[training,flux3,peft,diffusion]" \
    'torch==2.13.0+cu130' 'torchvision==0.28.0+cu130' 'torchcodec==0.13.0+cpu'
pip install 'natten==0.21.7+torch2130cu130' --find-links https://whl.natten.org

Para um driver NVIDIA anterior ao 580.65, use a build do CUDA 12.6 em vez disso: o índice cu126, as tags +cu126 para torch e torchvision, e natten==0.21.7+torch2130cu126. Para outra versão do torch, use a wheel do NATTEN cuja tag corresponda a ela, como natten==0.21.6+torch2110cu130 para torch 2.11 com CUDA 13.0.

Adaptação de tarefas do SO-101

Use um dataset LeRobot v3 a 30 Hz, com seis canais de ação comandada absoluta e seis canais de estado medido em ordem e unidades de juntas correspondentes. O garra é o último. Forneça valores brutos, calibração consistente do robô e instruções em task.

Câmeras

Os nomes, a quantidade e o layout das câmeras são configuráveis. O pacote SO-101 nomeia suas duas câmeras de acordo com o que elas mostram:

Você pode usar --rename_map para mapear os nomes das câmeras para as chaves esperadas, por exemplo:

--rename_map='{"observation.images.top":"observation.images.scene","observation.images.gripper_cam":"observation.images.wrist"}'

Use o mesmo mapeamento para o rollout; chaves correspondentes não precisam de mapa. Veja mapas de renomeação.

Carregar e treinar

Carregue a política SO-101 de black-forest-labs/flux-3-action-so101. A configuração salva carrega automaticamente o VAE e o codificador de texto compartilhados de uma revisão fixada de black-forest-labs/flux-3-action-base. Se o acesso for restrito, autentique-se com hf auth login usando uma conta com acesso. Execute a partir do checkout do LeRobot:

lerobot-train \
  --config_path=examples/flux3/lora.json \
  --policy.path=black-forest-labs/flux-3-action-so101 --policy.device=cuda \
  --dataset.repo_id=YOUR_ORG/YOUR_SO101_DATASET \
  --output_dir=outputs/so101_lora

O cache do Hub retém o modelo baixado e os codificadores compartilhados para retomar e para inferência de adaptadores.

O checkpoint fixado SO-101-finetuned usa as chaves de câmera scene / wrist, oito passos de tempo de observação, dois instantâneos visuais codificados independentemente e condicionamento por ação/estado passado. Ele prevê 42 ações, executa as primeiras 32 a 30 Hz e, em seguida, replaneja. Os processadores salvos carregam as estatísticas de normalização correspondentes da base. Essas configurações do checkpoint são preservadas durante a adaptação da tarefa. O arquivo compartilhado examples/flux3/lora.json fornece apenas as configurações de treinamento/execução e o orçamento de treinamento; o mesmo arquivo é incluído como lora.json em cada pacote do Hub. O layout da câmera, o condicionamento, a representação da ação e as configurações de perda vêm do checkpoint. Esses comandos mantêm o comprimento de execução salvo de 32, o guidance 3/3 e a execução eager. Para guidance 4/1 e compilação, adicione --policy.guidance_scale=4 --policy.guidance_scale_action=1 --policy.compile_model=true.

Experimentos de condicionamento: O LeRobot expõe n_obs_steps, history_snapshots e condition_on_past_actions. Alterá-los requer configurações compatíveis de modelo/processor e fine-tuning; mantenha os valores salvos no checkpoint para adaptação de tarefas comuns.

O preset habilita LoRA com rank/alpha 32 e EMA com decay 0.999; nada é enviado automaticamente. Habilite novos adaptadores por meio das configurações de nível superior peft, e não --policy.use_peft=true, que identifica um adaptador salvo.

Em nossos experimentos, o EMA proporcionou benefícios substanciais no DROID, enquanto os resultados do SO-101 têm sido mistos até o momento. Seu benefício depende dos dados e da configuração de treinamento. Recomendamos manter o EMA habilitado e avaliar tanto os checkpoints brutos quanto os checkpoints de EMA sob o mesmo protocolo de rollout antes de escolher qual implantar. A perda de validação offline periódica usa pesos brutos; ela não compara automaticamente as duas variantes.

Comece com o orçamento de GPU única para iteração rápida. Para um orçamento de treinamento maior, use a mesma configuração com estas substituições:

accelerate launch --num_processes=4 --multi_gpu -m lerobot.scripts.lerobot_train \
  --config_path=examples/flux3/lora.json \
  --policy.path=black-forest-labs/flux-3-action-so101 --policy.device=cuda \
  --dataset.repo_id=YOUR_ORG/YOUR_SO101_DATASET \
  --batch_size=8 --steps=14000 \
  --output_dir=outputs/so101_lora_large

O lote efetivo é GPUs × lote por GPU × acúmulo. O LeRobot conta steps, save_freq e eval_steps em micropassos; divida pelo acúmulo para obter as atualizações do otimizador. O preset salva e avalia a cada 500 micropassos (125 atualizações). Mantenha as três configurações divisíveis pelo acúmulo ao alterá-lo. Os comprimentos de congelamento, aquecimento, decaimento e resfriamento do agendador também usam micropassos. Multiplique o número desejado de atualizações do otimizador pelo fator de acúmulo ao definir esses comprimentos de fase. Ambos os orçamentos usam as mesmas configurações de LoRA, otimizador e EMA, e preservam a condicionamento, as dimensões de ação e a normalização da base carregada. Os checkpoints incluem tanto os adaptadores brutos quanto os de EMA. Escolha o orçamento e o checkpoint usando avaliação em conjunto de retenção e rollouts; um lote maior não garante melhores resultados.

O preset reserva os últimos 20% dos episódios por tarefa para avaliação. Ajuste o tamanho do lote e a acumulação para o seu hardware; consulte treinamento multi-GPU e treinamento PEFT.

Adaptação de tarefas DROID

Experimental: examples/flux3/lora.json reutiliza as configurações do otimizador PEFT do SO-101, LoRA, dropout, EMA e orçamento de treinamento. O carregamento da configuração foi verificado, mas o treinamento PEFT DROID de ponta a ponta e o desempenho da tarefa ainda não foram validados.

Carregue a política DROID diretamente pelo ID do repositório; seus codificadores compartilhados são resolvidos automaticamente:

lerobot-train \
  --config_path=examples/flux3/lora.json \
  --policy.path=black-forest-labs/flux-3-action-droid \
  --dataset.repo_id=YOUR_ORG/YOUR_DROID_TASK_DATASET \
  --output_dir=outputs/droid_lora

Isso usa batch efetivo 8 e 2.500 atualizações do otimizador, com EMA habilitado e checkpoints raw/EMA salvos a cada 125 atualizações. Os uploads automáticos para o Hub estão desabilitados. Para o orçamento maior, inicie com quatro GPUs, como no exemplo SO-101, e adicione --batch_size=8 --steps=14000.

O preset carrega o layout de câmeras, a condicionante, a representação de ação, a taxa de controle e os processadores do checkpoint do DROID, sobrescrevendo apenas as configurações de treinamento/execução listadas. Seu conjunto de dados deve corresponder ao contrato do checkpoint: o pacote atual usa três câmeras, oito canais de estado/ação e comandos de posição absoluta das juntas a 15 Hz. Mapeie os nomes das câmeras com --rename_map, se necessário. Ambos os robôs usam o mesmo exemplo de treinamento; nenhuma receita específica de robô é necessária para carregamento ou inferência.

Padrões e normalização

Novas configurações usam esses padrões do PEFT. Carregar --policy.path usa as configurações salvas no checkpoint; sobrescritas explícitas de configuração/CLI têm precedência. Checkpoints do DROID mantêm sua própria configuração.

Novas configurações usam n_obs_steps=1, history_snapshots=1 e condition_on_past_actions=false. Ao carregar um checkpoint, seus valores explícitos são preservados: o pacote SO-101 fixado usa oito passos de histórico, dois instantâneos visuais e condicionamento em ações/estado passados. conditioning="history" seleciona a codificação independente de instantâneos; por si só, não determina o número de passos. O treinamento também lê quadros e comandos futuros como alvos, além do comando anterior para calcular deltas de ação.

Os nomes/ordem das câmeras, as dimensões de estado/ação e os canais de ação absoluta continuam sendo específicos do robô. O checkpoint SO-101 mantém o último canal absoluto (delta_absolute_dims=[-1]) e usa pesos de canal [1, 1, 1, 1, 1, 2]. Os padrões da classe não usam exceções de canal absoluto e utilizam pesos iguais.

Os processadores salvos da base possuem estatísticas de normalização. Forneça estados/ações brutos e mantenha as mesmas estatísticas para treinamento, retomada e inferência. Os quantis do pacote SO-101 assumem sua ordem de juntas, unidades e calibração; eles não são valores universais de calibração para o SO-101. Para usar estatísticas diferentes, prepare uma base separada com processadores correspondentes antes do treinamento. Estatísticas de conjuntos de dados de ação absoluta não podem substituir estatísticas de delta de comando.

O treinamento salva adaptadores e cabeças treinadas com uma referência para a base completa. Mantenha essa base e seus arquivos de codificador congelados acessíveis ao recarregar; os adaptadores não são mesclados automaticamente nela.

Retomar

lerobot-train \
  --config_path=outputs/so101_lora/checkpoints/last/pretrained_model/train_config.json \
  --resume=true

Retome a partir do checkpoint bruto pretrained_model, que restaura o estado do otimizador e do EMA. Para inferência, você pode selecionar seu irmão pretrained_model_ema quando o EMA estiver habilitado. Mantenha a base original e os processadores de cada adaptador; uma base baixada recentemente não é uma substituição para aquela usada por um adaptador mais antigo.

Inferência

Aponte o --policy.path do runner do LeRobot para o checkpoint do adaptador treinado; ele carrega a base referenciada e os processadores salvos. Para chamadas diretas com uma base completa, carregue os processadores junto com ela:

from lerobot.policies.factory import make_pre_post_processors
from lerobot.policies.flux3 import Flux3Policy

checkpoint = "black-forest-labs/flux-3-action-so101"
policy = Flux3Policy.from_pretrained(checkpoint).eval()
preprocessor, postprocessor = make_pre_post_processors(policy.config, pretrained_path=checkpoint)

# Run on every control tick, including while actions are queued:
prepared = preprocessor(observation)
normalized_action = policy.select_action(prepared)
absolute_command = postprocessor(normalized_action)

O pré-processador faz o buffer do histórico de observações/comandos. O pós-processador desfaz a normalização, integra os deltas e atualiza o histórico de comandos compartilhado a cada tick. O primeiro comando é ancorado ao estado medido; os comandos subsequentes integram a partir do comando anterior. As saídas diretas da política são normalizadas e devem ser pós-processadas. Reinicie a política e ambos os processadores nas fronteiras de episódio ou após uma intervenção; o runner síncrono do LeRobot faz isso. Inferência RTC/assíncrona não é suportada.

Para a política padrão de timestep único, as imagens atuais e o estado medido são suficientes. Para previsão offline com uma âncora de delta explícita, forneça observation.command_history com um comando anterior; isso ancora a integração da saída sem adicionar um comando ao condicionamento do modelo. Checkpoints configurados explicitamente para múltiplos timesteps de observação exigem policy.config.n_obs_steps imagens/estados e histórico de comandos. Passe as entradas pelo pré-processador, chame predict_action_chunk e, em seguida, faça o pós-processamento. --policy.compile_model=true habilita a compilação preguiçosa dos codificadores congelados e do DiT de inferência; ele não compila o DiT de treinamento.

Para reduzir a VRAM de inferência, mova o Qwen para a CPU após o carregamento e a chamada final policy.to(...):

policy.frozen.text_encoder.cpu()

Isso move cerca de 8,3 GiB de pesos para a RAM do host; instruções não em cache são codificadas mais lentamente, enquanto os embeddings ainda são transferidos para o dispositivo da política e em cache. Isso não reduz a memória de pico durante o carregamento, e um policy.to(...) posterior move o Qwen de volta junto com a política.

Treinando outras encarnações

Para as implementações em que ainda não oferecemos um checkpoint finetuned, comece pelo modelo base em flux-action. Esse fluxo de trabalho exige a definição dos cabeçalhos de ação e condicionamento do robô, o layout das câmeras, a representação e a normalização da ação. Dimensões iguais, por si só, não tornam os cabeçalhos de outro robô reutilizáveis.

Para inferência ou adaptação de tarefas subsequentes no LeRobot, empacote o modelo treinado como uma política LeRobot completa, incluindo sua configuração e os processadores correspondentes. Um arquivo de pesos bruto, por si só, não é um pacote carregável para --policy.path. Este guia se concentra em adaptar políticas existentes de SO-101 e DROID a novas tarefas.

Todas as políticas compilam apenas os streams de conteúdo video e video_cond, além dos streams de ação do robô. Pesos de imagem/áudio não utilizados são filtrados antes do carregamento, economizando cerca de 28% dos parâmetros do DiT sem alterar as previsões de ação. As listas de streams salvas não ativam os ramos de imagem/áudio; o carregamento e a salvagem de uma base completa registram o layout reduzido. Adaptadores que visam streams removidos exigem sua base original e uma revisão de código anterior.

Políticas completas são carregadas de forma estrita por padrão. Após filtrar streams conhecidos como não utilizados, pesos ausentes, inesperados ou com formato incorreto geram um erro. A especificação explícita de strict=False permite apenas a inicialização de cabeças de embodiment ausentes ou o redimensionamento de cabeças com formato incompatível; os pesos do tronco permanecem obrigatórios e verificados quanto ao formato.

Validação

Os testes de CPU abrangem o empacotamento, os processadores, as atualizações de PEFT e o checkpoint/retomada. Verificações anteriores com pesos reais em H200 antecederam o refatoramento atual do processador. Esta revisão exata ainda precisa de validação em GPU e robô; as verificações existentes não estabelecem a capacidade de quatro GPUs nem a qualidade do PEFT em uma nova encarnação.