FastWAM

FastWAM é uma política de World Action Model para controle de robôs. A integração com LeRobot expõe o FastWAM por meio da API padrão de políticas, permitindo que ele seja configurado com policy.type=fastwam, treinado com lerobot-train e carregado por meio da interface de políticas pré-treinadas do LeRobot.

Visão geral do modelo

O FastWAM mantém a modelagem de vídeo durante o treinamento, mas usa a predição direta de ações no momento da inferência, em vez de gerar iterativamente observações futuras. Essa política do LeRobot encapsula o modelo de ações do FastWAM, adapta os lotes do LeRobot para amostras de treinamento do FastWAM e fornece o pipeline padrão de processador para normalização e pós-processamento de ações.

A implementação inicializa os componentes do modelo de mundo visual de Wan-AI/Wan2.2-TI2V-5B por padrão e prevê blocos de ação com a forma [batch, action_horizon, action_dim].

O que a Integração do LeRobot Abrange

• Configuração padrão policy.type=fastwam por meio do LeRobot
• Adaptação de lotes para tarefas de imagem, estado, ação e linguagem
• Inferência de chunks de ações por meio de select_action e predict_action_chunk
• Salvamento/carregamento de checkpoints por meio das APIs de política do LeRobot
• Pós-processamento configurável da ação do garra do LIBERO

Requisitos de instalação

Instale o LeRobot a partir do código-fonte e, em seguida, instale as dependências do FastWAM:

pip install -e ".[fastwam]"

Isso instala o extra de política FastWAM a partir de pyproject.toml: transformers, diffusers, ftfy e regex, além das dependências base do LeRobot.

Para a avaliação do LIBERO, instale também as dependências do benchmark:

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

Isso instala os dois extras. Além das dependências do FastWAM acima, o extra libero instala as dependências do dataset LeRobot, hf-libero no Linux e scipy.

O FastWAM usa o backbone Wan2.2 TI2V. O id do modelo padrão é:

policy.model_id=Wan-AI/Wan2.2-TI2V-5B

Requisitos de Dados

O FastWAM espera um dataset LeRobot com:

• uma ou mais observações visuais cujas larguras concatenadas correspondem a policy.image_size[1]
• observation.state quando policy.proprio_dim não é None
• action
• uma instrução de tarefa de linguagem por meio do campo de tarefa do conjunto de dados, ou tensores context e context_mask pré-computados

O setup visual padrão é um recurso de imagem chamado observation.images.image com forma (3, 224, 448). Se o conjunto de dados usar duas câmeras, configure policy.input_features para que suas alturas correspondam a 224 e suas larguras somem 448.

Uso

Crie uma nova política FastWAM com:

lerobot-train \
  --dataset.repo_id=your-org/your-dataset \
  --policy.type=fastwam \
  --policy.action_dim=7 \
  --policy.proprio_dim=8 \
  --policy.action_horizon=32 \
  --policy.n_action_steps=10 \
  --policy.image_size='[224,448]' \
  --output_dir=./outputs/fastwam_training \
  --job_name=fastwam_training \
  --steps=300000 \
  --batch_size=8 \
  --policy.device=cuda

Avalie um checkpoint existente no formato LeRobot no LIBERO-10 com:

lerobot-eval \
  --policy.path=lerobot/fastwam_libero_uncond_2cam224 \
  --policy.device=cuda \
  --policy.dtype=float32 \
  --policy.n_action_steps=10 \
  --env.type=libero \
  --env.task=libero_10 \
  --env.observation_height=224 \
  --env.observation_width=224 \
  --eval.batch_size=1 \
  --eval.n_episodes=50 \
  --seed=0 \
  --env.episode_length=600

Para libero_goal, libero_spatial e libero_object, use --env.episode_length=300.

Este checkpoint usa o pré-processamento VISUAL: IDENTITY. O FastWAM reescala as imagens para [-1, 1] dentro do modelo, portanto, aplicar a normalização de imagem no pré-processador as normalizaria duas vezes.

Para o rollout em robôs reais, use o mesmo caminho de checkpoint:

lerobot-rollout \
  --robot.type=so101_follower \
  --robot.port=/dev/ttyACM0 \
  --policy.path=your-org/fastwam-real-robot

Notas de configuração

Recursos de Imagem

policy.image_size é o tamanho do tensor de imagem FastWAM concatenado como (height, width). Cada recurso de imagem configurado deve ter a forma (3, height, camera_width), e todas as larguras de câmera devem somar a largura configurada.

Chunking de Ações

policy.action_horizon controla o número de ações futuras supervisionadas durante o treinamento e previstas durante a inferência. policy.n_action_steps controla quantas ações são consumidas antes que a política preveja um novo bloco. policy.n_action_steps deve ser menor ou igual a policy.action_horizon.

Ações Relativas

O FastWAM prevê ações absolutas por padrão. Com policy.use_relative_actions=true, ele prevê deslocamentos a partir do observation.state atual, portanto policy.proprio_dim deve ser definido. As juntas nomeadas em policy.relative_exclude_joints (padrão ["gripper"]) permanecem absolutas. Recalcule primeiro as estatísticas do conjunto de dados no espaço relativo, com --operation.chunk_size definido como policy.action_horizon (veja ações relativas do π₀.₅):

lerobot-train \
  --dataset.repo_id=your-org/your-dataset \
  --policy.type=fastwam \
  --policy.use_relative_actions=true \
  --policy.relative_exclude_joints='["gripper"]' \
  ...

Checkpoints salvos antes da existência de ações relativas, como lerobot/fastwam_base, recebem as etapas do processador de ações relativas adicionadas ao fazer fine-tuning a partir deles com esta opção. Isso se aplica apenas no momento do fine-tuning: a avaliação e o rollout usam os processadores salvos com o checkpoint como estão.

Dispositivo do Codificador de Texto

O codificador de texto UMT5-XXL congelado ocupa cerca de 11 GB. policy.text_encoder_device=cpu mantém esse componente fora da GPU: os prompts são codificados na CPU e os embeddings são movidos para policy.device. Durante a inferência, a tarefa codificada é armazenada em cache até policy.reset(), de modo que a codificação mais lenta na CPU é executada uma vez por episódio, em vez de uma vez por chunk. No treinamento, cada lote é codificado, portanto, a alocação na CPU torna cada passo mais lento.

Componentes do Wan

O FastWAM carrega o VAE do Wan, o DiT de vídeo, o codificador de texto e o tokenizer a partir do diretório do modelo Wan configurado ou do id do modelo no Hugging Face Hub. Os checkpoints do FastWAM no formato LeRobot salvos por save_pretrained também copiam os arquivos locais dos componentes do Wan necessários para o from_pretrained.

Backend de Atenção

O DiT do FastWAM usa o scaled_dot_product_attention (SDPA) do PyTorch para toda a atenção. Ele não usa FlashAttention: seu roteamento Mixture-of-Transformers (MoT) precisa de máscaras de atenção booleanas arbitrárias [query, key], que a API varlen do FlashAttention não consegue expressar. Portanto, instalar o pacote flash-attn não tem efeito no caminho do FastWAM. (Observe que o próprio SDPA ainda pode selecionar internamente o kernel flash / memory-efficient / math do PyTorch — isso não tem relação com o pacote flash-attn.)

Inferência de Ação Compilada

O FastWAM pode compilar seus grafos de pré-preenchimento de vídeo e de desruidamento de ação em cache para reduzir a sobrecarga de lançamento de kernels Python e CUDA ao longo das etapas repetidas de inferência:

--policy.compile_action_infer=true

A opção é desativada por padrão. A primeira previsão de ação compila um grafo para as formas de entrada atuais e, portanto, é mais lenta; previsões posteriores com as mesmas formas o reutilizam. Alterações de forma podem acionar a recompilação. O caminho eager continua disponível em configurações de PyTorch ou acelerador não suportadas.

Alternância de Ação LIBERO

Os checkpoints do LIBERO do FastWAM usam policy.toggle_action_dimensions=[-1] por padrão para corresponder à convenção de ação do garfo usada pelo pipeline de avaliação original do FastWAM:

--policy.toggle_action_dimensions='[-1]'

Resultados

Avaliado no LIBERO com lerobot/fastwam_libero_uncond_2cam224:

Reproduzir: lerobot-eval --policy.path=lerobot/fastwam_libero_uncond_2cam224 --policy.device=cuda --policy.dtype=float32 --policy.n_action_steps=10 --env.type=libero --env.task=libero_spatial --env.observation_height=256 --env.observation_width=256 --eval.batch_size=1 --eval.n_episodes=50 --seed=0 --env.episode_length=300 (1x H20 140 GB).

Referências

• artigo Fast-WAM
• Página do projeto Fast-WAM
• Código do Fast-WAM
• Checkpoints lançados upstream
• Wan2.2 TI2V 5B

Citação

@article{yuan2026fastwam,
  title = {Fast-WAM: Do World Action Models Need Test-time Future Imagination?},
  author = {Tianyuan Yuan and Zibin Dong and Yicheng Liu and Hang Zhao},
  journal = {arXiv preprint arXiv:2603.16666},
  year = {2026},
  url = {https://arxiv.org/abs/2603.16666}
}