VLA-JEPA

Esta é a portação do LeRobot para o VLA-JEPA, um modelo Vision-Language-Action que combina um backbone de linguagem Qwen3-VL com um modelo de mundo de vídeo auto-supervisionado (V-JEPA2) e uma cabeça de ação flow-matching DiT.

Visão geral da arquitetura

O VLA-JEPA possui três componentes principais:

Fluxo de dados

Treinamento:

1. Um clipe de vídeo de num_video_frames quadros é codificado pelo V-JEPA2 em tokens de patch por quadro.
2. O backbone Qwen3-VL processa imagens multiview + a instrução da tarefa e produz uma sequência de tokens de contexto que inclui tokens de ação especiais (para condicionamento do modelo de mundo) e tokens corporificados.
3. A cabeça de ação recebe esses tokens de contexto como chaves/valores de atenção cruzada e prevê um trecho de ação desruído por meio de correspondência de fluxo.
4. O preditor do modelo de mundo usa os tokens de ação extraídos do Qwen para prever os embeddings de quadros futuros do V-JEPA2; uma perda de regressão sobre essas previsões é adicionada à perda de ação.

Inferência: Apenas o Qwen e a cabeça de ação são usados. O modelo de mundo não é necessário no momento da inferência.

Detalhes da cabeça de ação

Predefinições disponíveis via action_model_type:

O preset define apenas a geometria de atenção, e cada entrada pode ser sobrescrita por action_num_heads / action_attention_head_dim. Duas larguras derivam dele:

• a largura interna do DiT é heads x head_dim (768 para DiT-B), derivada e não configurada;
• a largura da saída do DiT, bem como a largura dos MLPs do decodificador de ação e do codificador de estado, é action_hidden_size (padrão 1024).

Portanto, o DiT-B executa um transformer de largura 768 que projeta para 1024. Os dois são independentes.

Detalhes do modelo de mundo

O preditor de vídeo é um transformer no estilo ViT (ActionConditionedVideoPredictor) que recebe:

• Tokens de quadro: incorporações de patch do V-JEPA2 projetadas para predictor_embed_dim
• Tokens de ação: incorporações de tokens de ação do Qwen projetadas para predictor_embed_dim

Ele usa atenção causal por bloco, de modo que cada passo temporal pode prestar atenção a todos os passos anteriores. A entrada do preditor embed_dim é igual a num_views × video_encoder_hidden_size (por exemplo, 2 views × 1024 = 2048 para os checkpoints pré-treinados).

Pontos de Verificação Pré-treinados

Três checkpoints estão disponíveis diretamente dentro da organização LeRobot aqui: lerobot/VLA-JEPA, convertidos de ginwind/VLA-JEPA:

Todos os checkpoints usam Qwen/Qwen3-VL-2B-Instruct como a base de linguagem.

Configuração

Principais parâmetros em VLAJEPAConfig:

pre_snap_gripper_action e binarize_gripper_action são uma port do loop de avaliação LIBERO do starVLA e estão corretos apenas para a convenção de ação do LIBERO. pre_snap escreve {0, 1} no espaço normalizado, o desnormalizador mapeia esses valores para o ponto médio e o máximo, e binarize então compara esse valor físico com gripper_threshold (0,5). Para um gripper medido em graus, mm ou [0, 100], ambos os valores ficam acima do limiar e o gripper comandado se torna constante. Por esse motivo, eles têm valor padrão False; habilite-os apenas para configurações no estilo LIBERO e defina gripper_threshold nas próprias unidades do gripper, se for o caso. A fábrica de processadores emite um aviso quando as estatísticas do conjunto de dados mostram que a faixa não pode funcionar.

Treinamento

O número de etapas de treinamento pode variar com base no tamanho do conjunto de dados e no orçamento de computação. O artigo original pré-treinou por 50k em ssv2 + droid conjuntamente, seguido de 30k etapas adicionais para LIBERO, mas menos etapas ainda podem resultar em bom desempenho ao fazer fine-tuning a partir dos checkpoints pré-treinados fornecidos.

Treinamento completo do zero

lerobot-train \
  --policy.type=vla_jepa \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=your_org/your_dataset

Ajuste fino a partir de um checkpoint pré-treinado

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=your_org/your_dataset

Se você deseja congelar o backbone do Qwen e treinar apenas a cabeça de ação, defina policy.freeze_qwen=True:

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --policy.freeze_qwen=true \
  --dataset.repo_id=your_org/your_dataset

Fine-tuning em uma encarnação diferente

Quando o robô de destino possui uma dimensionalidade de ação ou estado diferente da do checkpoint pré-treinado, as camadas de projeção de entrada/saída da cabeça de ação terão formatos incompatíveis e não poderão ser carregadas diretamente. reinit_modules permite listar os prefixos de chaves que são permitidos para apresentar incompatibilidade — essas camadas são re-inicializadas aleatoriamente, enquanto todos os outros pesos são reutilizados do checkpoint. Qualquer incompatibilidade de formato fora dos prefixos listados gera um erro.

As camadas que dependem de action_dim e state_dim são:

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --policy.freeze_qwen=true \
  --policy.reinit_modules='["model.action_model.action_encoder", "model.action_model.action_decoder", "model.action_model.state_encoder"]' \
  --dataset.repo_id=your_org/your_dataset

Se o seu robô não possui estado proprioceptivo, omita model.action_model.state_encoder da lista.

Reproduzindo os resultados do LIBERO

Treinamento no LIBERO: inicia o treinamento a partir do checkpoint de Pretrain e treina por 30k passos no conjunto de dados LIBERO. O artigo original menciona treinamento em 8 GPUs com um batch size de 32, o que significa um batch size global de 256.

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=HuggingFaceVLA/libero \
  --steps=30000

Avaliando o checkpoint pré-treinado LIBERO-10:

lerobot-eval \
  --policy.path=lerobot/VLA-JEPA-LIBERO \
  --env.type=libero \
  --env.task=libero_spatial,libero_object,libero_goal,libero_10 \
  --eval.n_episodes=10 \
  --eval.batch_size=5

Para avaliar apenas um subconjunto de tarefas:

lerobot-eval \
  --policy.path=lerobot/VLA-JEPA-LIBERO \
  --env.type=libero \
  --env.task=libero_10 \
  --env.task_ids='[0,1,2]' \
  --eval.n_episodes=10 \
  --eval.batch_size=5

Resultados esperados:

Fine-tuning em conjuntos de dados com um número diferente de câmeras

O preditor do modelo de mundo pré-treinado foi treinado com embed_dim = world_model_num_views × 1024, ou seja, duas vistas de câmera.

Essa contagem de visualizações era lida de jepa_tubelet_size, que também nomeia o tamanho do tubelet temporal do codificador JEPA. world_model_num_views é o campo para isso agora; deixá-lo em None faz com que ele volte para jepa_tubelet_size, de modo que os checkpoints publicados continuem sendo carregados sem alterações.

Comportamento padrão — visualização de preenchimento/trimming (nenhuma ação necessária)

Ao fazer o ajuste fino a partir de VLA-JEPA-Pretrain, o modelo ajusta automaticamente o número de visualizações fornecidas ao modelo de mundo para corresponder a world_model_num_views:

• Conjuntos de dados de visão única (por exemplo, BridgeV2): o latente de visão única é duplicado para produzir uma entrada de modelo de mundo de duas visões, preservando o sinal autorregressivo do JEPA sem qualquer incompatibilidade de pesos.
• >Conjuntos de dados com mais de 2 vistas (por exemplo, DROID com 3 vistas): todas as vistas são passadas para o backbone Qwen (para um contexto mais rico), mas apenas as primeiras world_model_num_views vistas (um pulso + uma terceira pessoa, seguindo a ordem de vista configurada) são usadas para o modelo de mundo.

Opção 1 — Desativar o modelo de mundo

Defina enable_world_model=False para ignorar completamente a perda JEPA. Somente o backbone Qwen e a cabeça de ação são carregados e treinados. Isso é suficiente para um bom desempenho de ação.

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.enable_world_model=false \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=your_org/single_camera_dataset

Opção 2 — Reinitializar a projeção de entrada do predictor

Se você quiser alterar world_model_num_views para um valor diferente de 2, carregue o checkpoint com strict=False e reinitialize model.video_predictor.predictor_embed para o novo embed_dim. Todos os outros pesos dos blocos do preditor (atenção, MLP, norm, projeção de saída) são independentes do número de câmeras e podem ser reutilizados a partir do checkpoint pré-treinado.

Citação

@misc{sun2026vlajepaenhancingvisionlanguageactionmodel,
  title         = {VLA-JEPA: Enhancing Vision-Language-Action Model with Latent World Model},
  author        = {Jingwen Sun and Wenyao Zhang and Zekun Qi and Shaojie Ren and Zezhi Liu and Hanxin Zhu and Guangzhong Sun and Xin Jin and Zhibo Chen},
  year          = {2026},
  eprint        = {2602.10098},
  archivePrefix = {arXiv},
  primaryClass  = {cs.RO},
  url           = {https://arxiv.org/abs/2602.10098},
}

Licença

Os pesos são distribuídos sob os termos de licença do repositório original ginwind/VLA-JEPA (Apache 2.0 License). O código de integração do LeRobot segue a Apache 2.0 License.