Política DiT Multitarefa

A política Multitask Diffusion Transformer (DiT) é uma evolução da arquitetura original de Diffusion Policy, que utiliza um grande DiT com condicionamento de texto e visão para aprendizado de robôs multitarefa. Esta implementação suporta tanto objetivos de difusão quanto de correspondência de fluxo para a geração de ações, permitindo que robôs executem diversas tarefas de manipulação condicionadas a instruções em linguagem natural.

Visão geral do modelo

O modelo usa:

• CLIP Vision Encoder: Processa imagens RGB de múltiplas visões de câmera
• CLIP Text Encoder: Codifica as instruções de tarefas de linguagem (pesos congelados com projeção aprendível)
• Diffusion Transformer: Prediz sequências de ações condicionadas a observações e linguagem
• Dois objetivos: Suporta tanto difusão (DDPM/DDIM) quanto flow matching para geração de ações

Este modelo é empolgante porque é possível alcançar uma destreza extremamente alta, competitiva com VLAs de vários bilhões de parâmetros, com apenas ~450M de parâmetros e significativamente menos treinamento.

Requisitos de instalação

A política Multitask DiT possui dependências adicionais. Instale-a com:

pip install 'lerobot[multi_task_dit]'

Isso instalará todas as dependências necessárias, incluindo a biblioteca HuggingFace Transformers para modelos CLIP.

Uso

Para usar o Multitask DiT na sua configuração do LeRobot, especifique o tipo de política como:

policy.type=multi_task_dit

Treinamento

Comando Básico de Treinamento

Aqui está um comando completo de treinamento para treinar o Multitask DiT no seu conjunto de dados:

lerobot-train \
  --dataset.repo_id=YOUR_DATASET \
  --output_dir=./outputs/multitask_dit_training \
  --batch_size=32 \
  --steps=5000 \
  --save_freq=500 \
  --log_freq=100 \
  --policy.type=multi_task_dit \
  --policy.device=cuda \
  --policy.repo_id="HF_USER/multitask-dit-your-robot" \
  --trackio.enable=true

Parâmetros hiperbólicos recomendados e detalhes do conjunto de dados (frequência de controle de 30Hz)

Para um desempenho confiável, comece com os seguintes hiperparâmetros padrão sugeridos:

lerobot-train \
  --dataset.repo_id=YOUR_DATASET \
  --output_dir=./outputs/mutitask_dit_training \
  --batch_size=320 \
  --steps=30000 \
  --policy.type=multi_task_dit \
  --policy.device=cuda \
  --policy.horizon=32 \
  --policy.n_action_steps=24 \
  --policy.objective=diffusion \
  --policy.noise_scheduler_type=DDPM \
  --policy.num_train_timesteps=100 \
  --policy.repo_id="HF_USER/multitask-dit-your-robot" \
  --trackio.enable=true

Parâmetros principais:

• Tamanho do lote: 192-320 - Se você tiver acesso a uma GPU que possa suportar isso, obterá as melhores dinâmicas de treinamento
• Horizon: 32 - número de passos de ação a prever, ~1,0 s a 30 Hz
• n_action_steps: 24 - ~0,8 segundos a 30Hz
• Objetivo: diffusion - comece com diffusion e experimente com flow matching se a qualidade da geração for baixa
• Etapas de treinamento: >30k etapas recomendadas para uma única tarefa

Parâmetros de Configuração de Treinamento

Seleção de Objetivo

Escolha entre difusão e correspondência de fluxo:

# Diffusion objective (default)
--policy.objective=diffusion \
--policy.noise_scheduler_type=DDPM \
--policy.num_train_timesteps=100 \
--policy.num_inference_steps=10 \
--policy.beta_schedule=squaredcos_cap_v2 \
--policy.prediction_type=epsilon \
--policy.clip_sample=true \
--policy.clip_sample_range=1.0

# Flow matching objective
--policy.objective=flow_matching \
--policy.timestep_sampling_strategy=beta \
--policy.num_integration_steps=100 \
--policy.integration_method=euler \
--policy.sigma_min=0.0

Arquitetura Transformer

Ajuste a capacidade do modelo com base no tamanho do conjunto de dados:

# Small datasets (< 100 examples)
--policy.num_layers=4 \
--policy.hidden_dim=512 \
--policy.num_heads=8  # should ideally be hidden_dim // 64

# Medium datasets (100-5k examples) - default
--policy.num_layers=6 \
--policy.hidden_dim=512 \
--policy.num_heads=8  # should ideally be hidden_dim // 64

# Large datasets (> 5k examples)
--policy.num_layers=8 \
--policy.hidden_dim=512 \
--policy.num_heads=8   # should ideally be hidden_dim // 64

Opções de Codificação Posicional:

O modelo suporta dois métodos de codificação posicional para sequências de ações:

# Rotary Position Embedding (RoPE) - default, recommended
--policy.use_rope=true \
--policy.rope_base=10000.0  # Base frequency for RoPE

# Absolute positional encoding
--policy.use_positional_encoding=true  # Disables RoPE when true

Outros Parâmetros do Transformer:

--policy.dropout=0.1  # Dropout rate for DiT blocks (0.0-1.0)
--policy.timestep_embed_dim=256  # Timestep embedding dimension

Configuração do Codificador de Visão

# Use different CLIP model for more expressivity at the cost of inference time
# experiment with larger or smaller models depending on the complexity of your tasks and size of dataset
--policy.vision_encoder_name=openai/clip-vit-large-patch14

# Use separate vision encoder per camera
# This may be useful when cameras have significantly different characteristics, but
# be wary of increased VRAM footprint.
--policy.use_separate_rgb_encoder_per_camera=true

# Image preprocessing (you may need to resize your images for inference speed ups)
--policy.image_resize_shape=[XXX,YYY] \
--policy.image_crop_shape=[224,224] \
--policy.image_crop_is_random=true  # Random during training, center at inference

Configuração do Codificador de Texto

# Use different CLIP text encoder model
# same as vision: experiment with larger or smaller models depending on the
# complexity of your tasks and size of dataset
--policy.text_encoder_name=openai/clip-vit-large-patch14

Configuração da Taxa de Aprendizado

O codificador de visão usa um multiplicador de taxa de aprendizado separado, onde 1/10 é sugerido como ponto de partida ideal:

--policy.optimizer_lr=2e-5 \
--policy.vision_encoder_lr_multiplier=0.1  # Vision encoder LR = 0.1 * optimizer_lr

Diretrizes de ajuste para treinamento

1. Flow Matching com Amostragem Beta

A implementação original de difusão aqui é baseada no trabalho descrito em TRI’s LBM paper

Além disso, implementamos um objetivo de correspondência de fluxo, que é descrito em alto nível no post do blog da Boston Dynamics.

Considere testar o objetivo de flow-matching e avaliar as diferenças de desempenho para sua tarefa:

--policy.objective=flow_matching \
--policy.timestep_sampling_strategy=beta \
--policy.timestep_sampling_alpha=1.5 \
--policy.timestep_sampling_beta=1.0 \
--policy.timestep_sampling_s=0.999

Isso não se mostrou uma solução definitiva para todos os casos de uso, mas ocasionalmente resulta em ações mais suaves e consistentes.

2. Número de Camadas do Transformer

Ajuste a capacidade do modelo ao tamanho do seu conjunto de dados:

• Conjuntos de dados pequenos (< 100 exemplos): Reduza para 4 camadas
• Conjuntos de dados grandes (> 5k exemplos): Aumente para 8 camadas

3. Ajuste de horizonte

O modelo pode ser sensível ao horizonte que você escolher. Comece com um horizonte de aproximadamente 1 segundo com base na sua frequência de controle:

• Frequência de 30 Hz: horizon=30
• Frequência de 10 Hz: horizon=10

Depois, experimente aumentando a partir daí. O horizonte determina até que ponto no futuro o modelo prevê ações.

4. Sensibilidade de n_action_steps

O modelo também pode ser muito sensível a n_action_steps. Comece com um valor em torno de 0,8 segundos, com base na sua frequência de controle, e ajuste a partir daí:

• Valores mais baixos: Mais reativo, mas potencialmente menos estável para tarefas de longo prazo
• Valores mais altos: Melhor para execução de longo prazo, mas as falhas em malha aberta têm recuperação limitada

Ajuste de Inferência

Para inferência mais rápida, use DDIM com menos etapas de amostragem:

--policy.noise_scheduler_type=DDIM \
--policy.num_inference_steps=10

Retomando o Treinamento

Para retomar o treinamento a partir de um checkpoint:

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

O diretório de checkpoint deve conter os arquivos model.safetensors e config.json (salvos automaticamente durante o treinamento). Ao retomar, a configuração é carregada do checkpoint, portanto, não é necessário especificar outros parâmetros.

Modos de falha comuns e depuração

O treinamento desses modelos pode ser delicado. Aqui estão os modos de falha comuns e as abordagens de depuração:

Ocioso / Sem Movimento

O modelo pode “colapsar” durante a inferência, resultando em movimento estático ou ausente. Isso pode ocorrer quando:

1. Insuficiência de dados de treinamento: Se você tiver apenas 20 a 50 exemplos, tente aumentar aproximadamente o tamanho do seu conjunto de dados. Se, mesmo com mais de 300 exemplos, você ainda estiver vendo isso, a tarefa pode ser muito complexa.

2. Várias tarefas semelhantes: Quando seu conjunto de dados contém várias tarefas semelhantes (por exemplo, pegar 2 objetos diferentes), o modelo pode depender demais da condicionamento por linguagem, que pode não ser rico o suficiente.

Dicas de depuração:

• Aumente o tamanho do conjunto de dados (dobre até chegar a mais de 300 exemplos)
• Treine por mais tempo, até 100 mil passos, mesmo quando a perda se estabiliza
• Verifique se o modelo está recebendo instruções de idioma adequadas ou aumente a diversidade das instruções

Executando a Tarefa Errada

Às vezes, o robô ignorará completamente sua instrução e executará outra tarefa. Isso geralmente só acontece se você tiver treinado em múltiplas tarefas.

Causas potenciais:

• Ambiguidade na instrução de idioma
• Dados insuficientes de treinamento específicos para a tarefa
• Confusão do modelo entre tarefas semelhantes no conjunto de dados multitarefa

Dicas de depuração:

• Verifique a especificidade da instrução de idioma, especialmente se as descrições forem semelhantes entre várias tarefas
• Verifique a distribuição de tarefas no seu conjunto de dados de treinamento e adicione ponderação à tarefa que está falhando/ignorado
• Considere o ajuste fino específico para a tarefa

Instabilidade no Treinamento

Se a perda de treinamento for instável ou divergente:

• Tente ajustar a taxa de aprendizado entre 1e-5 e 3e-4
• Aumente o tamanho do lote, se possível
• Verifique se a normalização do seu conjunto de dados está correta
• Verifique se o pré-processamento de imagens está funcionando corretamente

Considerações de desempenho

Requisitos de GPU

• Inferência: Recomenda-se pelo menos uma RTX 5070 Ti (ou GPU equivalente) para um desempenho de velocidade razoável
• Treinamento: Uma GPU com VRAM suficiente para carregar tamanhos de lote >64 é ideal, o que variará dependendo do número de observações de imagem, etc.

Recomendações de Tamanho de Lote

• Mínimo: 64 (valores inferiores a isso podem resultar em treinamento instável)
• Recomendado: 256-320 (melhor desempenho, requer GPU maior)

Exemplo: Treinamento em Conjunto de Dados Personalizado

Aqui está um exemplo completo de treinamento em um conjunto de dados personalizado:

lerobot-train \
  --dataset.repo_id=YOUR_DATASET \
  --output_dir=./outputs/mutitask_dit_training \
  --batch_size=320 \
  --steps=30000 \
  --save_freq=1000 \
  --log_freq=100 \
  --env_eval_freq=1000 \
  --policy.type=multi_task_dit \
  --policy.device=cuda \
  --policy.horizon=32 \
  --policy.n_action_steps=24 \
  --policy.objective=diffusion \
  --policy.noise_scheduler_type=DDPM \
  --policy.num_layers=6 \
  --policy.hidden_dim=512 \
  --policy.vision_encoder_name=openai/clip-vit-base-patch16 \
  --policy.image_resize_shape=[320,240] \
  --policy.image_crop_shape=[224,224] \
  --policy.repo_id="HF_USER/multitask-dit-your-robot" \
  --trackio.enable=true \
  --trackio.project=multitask_dit

Resultados do Libero

lerobot-train \
  --dataset.repo_id=HuggingFaceVLA/libero \
  --policy.type=multi_task_dit \
  --policy.push_to_hub=false \
  --output_dir="./outputs/multitask_dit_libero" \
  --job_name="multitask-dit-libero" \
  --trackio.enable=true \
  --trackio.project=multitask_dit_libero \
  --dataset.image_transforms.enable=true \
  --dataset.image_transforms.max_num_transforms=4 \
  --dataset.image_transforms.tfs='{"brightness":{"type":"ColorJitter","kwargs":{"brightness":[0.75,1.25]}},"contrast":{"type":"ColorJitter","kwargs":{"contrast":[0.6,1.4]}},"saturation":{"type":"ColorJitter","kwargs":{"saturation":[0.8,1.2]}},"hue":{"type":"ColorJitter","kwargs":{"hue":[-0.05,0.05]}},"sharpness":{"type":"SharpnessJitter","kwargs":{"sharpness":[0.6,1.4]}},"rotation":{"type":"RandomRotation","kwargs":{"degrees":[-5,5]}},"translation":{"type":"RandomAffine","kwargs":{"degrees":0,"translate":[0.1,0.1]}}}' \
  --dataset.video_backend=torchcodec \
  --policy.use_amp=true \
  --policy.horizon=48 \
  --policy.n_obs_steps=2 \
  --policy.use_rope=true \
  --policy.use_positional_encoding=false \
  --policy.hidden_dim=768 \
  --policy.num_layers=8 \
  --policy.num_heads=12 \
  --policy.dropout=0.1 \
  --policy.timestep_embed_dim=256 \
  --policy.objective=diffusion \
  --policy.optimizer_lr=3e-4 \
  --policy.optimizer_weight_decay=0 \
  --policy.scheduler_warmup_steps=0 \
  --policy.vision_encoder_name=openai/clip-vit-base-patch16 \
  --policy.image_resize_shape=[256,256] \
  --policy.image_crop_is_random=true \
  --policy.text_encoder_name=openai/clip-vit-base-patch16 \
  --policy.vision_encoder_lr_multiplier=0.1 \
  --policy.device=cuda \
  --num_workers=8 \
  --save_freq=4000 \
  --log_freq=100 \
  --steps=100000 \
  --batch_size=320

Resultados:

Referências

Para mais detalhes sobre a implementação técnica e a arquitetura, consulte:

• Uma Análise Cuidadosa de Modelos de Comportamento de Grande Escala para Manipulação Ágil Multitarefa
• Large Behavior Models e Atlas encontram novo rumo
• Dissecando e liberando como código aberto a política do Transformer de Difusão Multitarefa