Long-WAM

Artigo · Projeto e demonstrações · Código

Long-WAM é um modelo de mundo-ação de contexto longo: ele usa uma instrução em linguagem natural, o estado do robô e o histórico visual para prever a dinâmica futura e gerar ações.

Visão geral do modelo

• Aprenda a usar o histórico. O pré-treinamento de vídeo autoregressivo do LongLive2.0-Robot aprende as dependências de histórico para o futuro antes da adaptação às ações do robô.
• Prediga, depois aja. O IDM condiciona a geração de ações no histórico observado e nos latentes futuros previstos, sem decodificar as imagens RGB futuras.
• Escale o contexto temporal. Janelas de observação mais longas capturam movimento e progresso da tarefa; o comprimento do histórico é separado do horizonte de ação.

LeRobot suporta tanto IDM quanto co-denoising de vídeo/ação conjunta (COD), com carregamento nativo de checkpoints, fine-tuning e inferência síncrona de chunks de ação. Os backbones Wan vêm do Diffusers e o codificador de texto UMT5 vem do Transformers.

Resultados

Resultados do lerobot-eval reportados para os checkpoints nativos do LeRobot publicados do IDM e COD:

Veja Avaliação para exemplos de comandos e limites de episódios. Os comandos foram testados em NVIDIA GB200 com BF16 e inferência síncrona.

Instalar

A partir de um checkout de fonte do LeRobot:

pip install -e ".[longwam]"

Checkpoints

Os checkpoints nativos do LeRobot incluem os pesos completos do modelo e os pré-processadores e pós-processadores salvos. O tokenizer é carregado de google/umt5-xxl.

Ambos usam histórico P48 (2,4 s a 20 Hz), duas visões de câmera de 224×224, ações de 7-D e um estado do robô de 8-D. Use o checkpoint correspondente ao modo de desruidamento desejado; os pacotes originais do Long-WAM model.pt não podem ser carregados diretamente.

Requisitos de dados

Use um dataset LeRobot com câmeras RGB, observation.state, action e instruções de tarefa. As chaves de câmera do LIBERO liberado são observation.images.image e observation.images.image2, nessa ordem. Forneça tensores RGB em [0, 1]; o processador redimensiona e aplica center-crop em cada visão, e a política concatena as visões horizontalmente. Mantenha a ordem das câmeras e as convenções de ação/estado alinhadas com o checkpoint.

Inferência

from lerobot.policies import make_pre_post_processors
from lerobot.policies.longwam.modeling_longwam import LongWAMPolicy

checkpoint = "Efficient-Large-Model/Long-WAM-LIBERO-IDM-LeRobot"
policy = LongWAMPolicy.from_pretrained(checkpoint).to("cuda").eval()
preprocess, postprocess = make_pre_post_processors(
    policy.config,
    pretrained_path=checkpoint,
    preprocessor_overrides={"device_processor": {"device": "cuda"}},
)
policy.reset()
# observation contains both current camera images, observation.state, and task.
# Repeat on every control tick, including ticks that consume queued actions.
action = postprocess(policy.select_action(preprocess(observation)))

Chame select_action em cada tick de controle, inclusive enquanto um chunk de ação está sendo executado, para que o histórico permaneça atualizado. predict_action_chunk retorna [batch, prediction_steps, action_dim]; cada chamada registra o frame atual fornecido. reset() e drop_queued_actions() limpam tanto o histórico quanto as ações enfileiradas. O async/RTC genérico ainda não foi validado nesta integração.

Por padrão, inference_seed=1 reinicia os geradores de ruído de vídeo/ação separados com a mesma semente para cada bloco, preservando o comportamento de inferência liberado. Defina inference_seed=None no Python (--policy.inference_seed=null na CLI) para usar o RNG global em vez disso; reset() não redefina a semente desse fluxo. O treinamento sempre amostra ruído novo.

Avaliação

Instale o extra de benchmark e avalie um checkpoint nativo:

pip install -e ".[longwam,libero]"

lerobot-eval --policy.path=Efficient-Large-Model/Long-WAM-LIBERO-IDM-LeRobot \
  --policy.device=cuda --policy.dtype=bfloat16 \
  --env.type=libero --env.task=libero_spatial \
  --env.observation_height=256 --env.observation_width=256 \
  --env.episode_length=300 \
  --eval.batch_size=1 --eval.n_episodes=50 --seed=1

Use libero_object ou libero_goal para as outras suítes curtas. Para o LIBERO-Long, defina --env.task=libero_10 --env.episode_length=600. Substitua IDM por COD no caminho do checkpoint para avaliar o co-denoising.

História e divisão de ações

A duração do histórico é past_obs_steps / control_frequency. O início do histórico é preenchido com a primeira observação. Use um checkpoint treinado para o histórico pretendido; alterar a janela na inferência não é treinamento de escala de contexto.

Ajuste fino

lerobot-train --policy.path=Efficient-Large-Model/Long-WAM-LIBERO-IDM-LeRobot \
  --policy.device=cuda --policy.dtype=bfloat16 \
  --accelerator.mixed_precision=bf16 \
  --dataset.repo_id=YOUR_ORG/YOUR_DATASET \
  --output_dir=outputs/longwam --batch_size=1 --steps=30000 \
  --policy.push_to_hub=false

--policy.dtype seleciona o armazenamento dos parâmetros antes do carregamento; os gradientes e momentos nativos do AdamW usam esse dtype. --accelerator.mixed_precision controla o autocast, não os pesos mestres em FP32. O exemplo usa BF16 para ambos, correspondendo ao dtype original dos parâmetros. Para parâmetros/gradientes/momentos em FP32 com autocast em BF16, use --policy.dtype=float32 --accelerator.mixed_precision=bf16; o VAE congelado e o codificador de texto também usam FP32, aumentando a memória. Escolha isso antes do carregamento: converter pesos carregados em BF16 para FP32 não recupera a precisão perdida.

O pós-processador de garras liberado converte a última dimensão da ação de [0, 1] para comandos com sinal. Use --policy.toggle_action_dimensions=[] se seus dados de ajuste fino já utilizarem a convenção de abrir/fechar com sinal do ambiente. A configuração ativa também se aplica aos processadores salvos; postprocessor_overrides explícitos têm precedência. As estatísticas de normalização salvas são mantidas, a menos que sejam substituídas, portanto, alinhe-as com seus dados de ajuste fino.

O treinamento supervisiona os próximos blocos de vídeo e ação; a inferência recebe apenas as observações passatuais/atuais. Ajuste --batch_size de acordo com a memória da GPU e mantenha a ordem das câmeras do checkpoint, as definições de ação/estado e o modo de desruidamento.

Citação

@misc{huang2026longwamscalingcontextworldaction,
  title={Long-WAM: Scaling the Context of World-Action Models},
  author={Wei Huang and Bohan Zhang and Chenzhi Liu and Isabella Liu and
          Shuai Yang and Weian Mao and Luozhou Wang and Yicheng Xiao and
          Weifeng Lin and Qixin Hu and Bryan Chu and Sifei Liu and
          Linxi Fan and Xiaojuan Qi and Song Han and Yukang Chen},
  year={2026},
  eprint={2610.10528},
  archivePrefix={arXiv},
  primaryClass={cs.RO},
  url={https://arxiv.org/abs/2610.10528},
}