LingBot-VA

LingBot-VA é uma política de modelo de mundo de vídeo-ação autoregressiva construída sobre a pilha de difusão de vídeo Wan2.2. Ela intercala, em uma única sequência autoregressiva, a previsão de latentes de vídeo futuros e ações de robô (“VA” = Vídeo-Ação). A integração com LeRobot conecta o LingBot-VA às interfaces padrão de treinamento, avaliação e processador.

Visão geral do modelo

LingBot-VA é um “mixture-of-transformers” de fluxo duplo: um fluxo de vídeo/latente (patch_embedding_mlp → blocks → proj_out) e um fluxo de ação (action_embedder → blocks → action_proj_out) compartilham os mesmos 30 blocos de transformer e a condicionamento de texto.

Na inferência, a política executa um loop autoregressivo por bloco: ela remove o ruído do fluxo de latentes de vídeo (CFG, ~20 passos) e do fluxo de ações (~50 passos) com dois agendadores de flow-matching independentes, mantendo um cache de KV entre os blocos. Os keyframes observados reais são devolvidos ao cache de KV conforme o bloco é executado (modelagem de mundo em laço fechado).

O que a Integração do LeRobot Abrange

• Configuração padrão policy.type=lingbot_va por meio do LeRobot.
• Checkpoints prontos para uso no formato LeRobot no Hub (convertidos a partir dos originais publicados).
• Inferência dual-stream autoregressiva por trás da interface padrão select_action (avaliação de ambiente único, --eval.batch_size=1).
• Salvamento opcional dos vídeos previstos (imaginados) da política durante a avaliação / o treinamento.
• Avaliação com lerobot-eval em LIBERO e RoboTwin.
• Treinamento / ajuste fino por meio da perda de fluxo de correspondência de fluxo duplo (policy.forward), veja abaixo.

Instalação

1. Instale o LeRobot seguindo o Guia de Instalação.
2. Instale o extra LingBot-VA:

pip install -e ".[lingbot_va]"

Checkpoints

Os checkpoints upstream liberados foram convertidos para o formato LeRobot e enviados para o Hub:

Apenas o transformer treinável de ~5B é armazenado no model.safetensors do LeRobot. O VAE + UMT5 + tokenizer congelados (~20 GB) são obtidos de config.wan_pretrained_path no momento do carregamento (o padrão é o repositório de origem robbyant/*). O codificador de texto UMT5-XXL é executado na CPU por padrão (config.text_encoder_device), de modo que o transformer de 5B + VAE cabem em uma única GPU de 24–32 GB.

Avaliação (LIBERO)

lerobot-eval \
    --policy.path=lerobot/lingbot_va_libero_long \
    --policy.device=cuda \
    --env.type=libero --env.task=libero_10 \
    --env.observation_height=128 --env.observation_width=128 \
    --eval.n_episodes=50 --eval.batch_size=1 \
    --output_dir=outputs/eval/lingbot_va_libero

A inferência em streaming do LingBot-VA (cache de KV + feedback de keyframes observados) é implementada para avaliação em um único ambiente; use --eval.batch_size=1.

Avaliação (RoboTwin)

O RoboTwin 2.0 requer a pilha de simuladores SAPIEN + CuRobo; o estágio robotwin_native de docker/Dockerfile.sim é uma instalação de referência (o LingBot-VA também requer warp-lang==1.3.1 e um pytorch3d fixado, e o CuRobo compilado com a capacidade de computação da GPU em TORCH_CUDA_ARCH_LIST). A configuração do RoboTwin do servidor do simulador usa controle por juntas, portanto o modo de efector final abaixo é executado localmente. O RoboTwin usa controle de pose do efector final, portanto execute com --env.action_mode=ee: a política prevê deltas de xyz+quaternion+gripper por braço (layout latente robotwin_tshape) que são compostos sobre a pose inicial do efector final do episódio e executados via IK do CuRobo.

lerobot-eval \
    --policy.path=lerobot/lingbot_va_robotwin \
    --policy.device=cuda \
    --env.type=robotwin --env.task=beat_block_hammer --env.action_mode=ee \
    --eval.n_episodes=10 --eval.batch_size=1 \
    --output_dir=outputs/eval/lingbot_va_robotwin

Salvando vídeos previstos (imaginados)

Defina --policy.save_predicted_video=true para decodificar adicionalmente os latentes do vídeo previsto com VAE e gravar pred_episode_*.mp4 ao lado dos vídeos eval_episode_*.mp4 renderizados pelo ambiente. A mesma flag funciona para a avaliação periódica durante lerobot-train.

Assistindo ao vídeo imaginado ao vivo

Defina --policy.return_predicted_video=true durante lerobot-rollout --display_data=true para decodificar o vídeo previsto de cada chunk e exibir, conforme cada ação é executada, os quadros de câmera imaginados aos quais ela pertence (prediction.images.<camera>, ao lado de cada câmera real). Isso custa uma decodificação VAE por chunk.

Treinamento / ajuste fino

LingBotVAPolicy.forward(batch) implementa a perda de flow-matching de fluxo duplo (latent_loss + action_loss, ponderada por timestep e mascarada por ação) descrita no artigo: ela codifica os clipes da câmera em latentes de vídeo via VAE, codifica a tarefa via UMT5, adiciona ruído a ambos os fluxos, executa a passagem de treinamento causal por bloco do transformer e retorna (loss, metrics). O preset do otimizador é AdamW com um agendamento de aquecimento linear seguido de constante (em conformidade com o upstream).

Requisitos:

• As máscaras causais de bloco usam o flex-attention do PyTorch, portanto, construa a política com --policy.attn_mode=flex para treinamento (o SDPA torch padrão é apenas para inferência).
• O DiT completo de 5B não cabe em uma única GPU de 24–32 GB com AdamW; faça o fine-tuning com LoRA (--peft.method_type=LORA) e/ou offload do otimizador. get_optim_params retorna apenas os parâmetros treináveis (por exemplo, adaptadores); o VAE + codificador de texto UMT5 permanecem congelados. Instale o extra lerobot[peft] para habilitar o suporte a PEFT (veja o guia de treinamento PEFT).

lerobot-train \
  --policy.path=lerobot/lingbot_va_libero_long --policy.attn_mode=flex \
  --peft.method_type=LORA --peft.r=32 --peft.lora_alpha=32 \
  --peft.target_modules='transformer\.blocks\.\d+\.attn[12]\.(to_q|to_v)' \
  --dataset.repo_id=<your LeRobot-format dataset> \
  --batch_size=1 --steps=... --output_dir=outputs/train/lingbot_va

Diferentemente de SmolVLA / π₀, o LingBot-VA não inclui alvos padrão de LoRA embutidos, portanto, você deve passar --peft.target_modules explicitamente. Apenas self.transformer (o transformer de fluxo duplo Wan) é treinável; o exemplo acima adapta as projeções de query/value de seus blocos de autoatenção (attn1) e de atenção cruzada (attn2), o conjunto padrão de alvos para LoRA. Amplie-o (por exemplo, adicione to_k/to_out ou as camadas ffn) se precisar de um adaptador com maior capacidade. Passar --peft.method_type implica PEFT, portanto, --policy.use_peft=true não é necessário.

O conjunto de dados deve fornecer clipes de câmera (uma janela temporal por câmera, codificada em VAE para frame_chunk_size quadros latentes) e frame_chunk_size * action_per_frame etapas de ação por item.

Formato de dados (canais de ação & ordem das câmeras)

LingBot-VA é uma política de pose do efector final (Cartesiana); ela prevê poses do EEF + garra, não posições de junta. As ações estão em um layout fixo multi-embodiment de 30 dimensões; mapeie as dimensões de ação do seu robô para esses canais e preencha o restante com 0 (used_action_channel_ids seleciona os canais que um determinado checkpoint realmente usa):

• LIBERO usa os canais 0–6: um delta de EEF de 6-DoF (xyz + rotação) + garra (braço único).
• RoboTwin usa os canais [0–6, 28, 7–13, 29]: EEF esquerdo (xyz + quaternião) + garra esquerda + EEF direito + garra direita (16 dimensões). O ambiente converte essas poses em trajetórias articulares via CuRobo IK — as articulações nunca são previstas.

Conjuntos de dados no espaço articulado (ou uma convenção de EEF diferente) devem ser remapeados para este esquema antes do ajuste fino desses checkpoints.

A ordem das câmeras é fixa e sensível à ordem, os latentes por câmera são concatenados espacialmente na ordem de obs_cam_keys, portanto, o mapeamento físico câmera→slot deve corresponder ao do treinamento:

A primeira câmera é a visão externa/de cabeça e as demais são visões de pulso.

Hyperparâmetros de Inferência (LIBERO)

Estes são os padrões de LingBotVAConfig; substitua qualquer um deles via --policy.<name>=....

Observações

• Backend de atenção: a inferência usa o backend SDPA torch (sempre disponível). Os backends flashattn e flex são opcionais; flex é necessário apenas para treinamento.
• Tamanho do modelo: o DiT tem ~5B de parâmetros e o VAE+UMT5 congelados adicionam ~20 GB; a inferência requer aproximadamente 18–24 GB de VRAM.

Licença

LingBot-VA é lançado sob a licença Apache-2.0. Consulte o repositório upstream.