# 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:

```bash
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:

```bash
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 é:

```python
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:

```bash
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:

```bash
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:

```bash
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 π₀.₅](./pi05#relative-actions)):

```bash
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 nenhum 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:

```bash
--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 garra usada pelo pipeline de avaliação original do FastWAM:

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

## Resultados

Avaliado no LIBERO com [`lerobot/fastwam_libero_uncond_2cam224`](https://huggingface.co/lerobot/fastwam_libero_uncond_2cam224):

| Suite          | Taxa de sucesso | n_episodes |
| -------------- | -----------: | ---------: |
| libero_spatial |        97.6% |        500 |
| libero_object  |        99.0% |        500 |
| libero_goal    |        95.0% |        500 |
| libero_10      |        94.0% |        500 |
| **média**    |    **96.4%** |       2000 |

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](https://arxiv.org/abs/2603.16666)
- [Página do projeto Fast-WAM](https://yuantianyuan01.github.io/FastWAM/)
- [Código do Fast-WAM](https://github.com/yuantianyuan01/FastWAM)
- [Checkpoints lançados upstream](https://huggingface.co/yuanty/fastwam)
- [Wan2.2 TI2V 5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B)

## Citação

```bibtex
@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}
}
```