FastWAM

FastWAM es una política de modelo de acción mundial para el control de robots. La integración de LeRobot expone FastWAM a través de la API estándar de políticas, de modo que puede configurarse con policy.type=fastwam, entrenarse con lerobot-train y cargarse a través de la interfaz de políticas preentrenadas de LeRobot.

Resumen del modelo

FastWAM mantiene la modelización de video durante el entrenamiento, pero utiliza la predicción directa de acciones en tiempo de inferencia en lugar de generar iterativamente observaciones futuras. Esta política de LeRobot envuelve el modelo de acciones de FastWAM, adapta los lotes de LeRobot a muestras de entrenamiento de FastWAM y proporciona el pipeline estándar de procesador para la normalización y el postprocesamiento de acciones.

La implementación inicializa los componentes del modelo de mundo visual desde Wan-AI/Wan2.2-TI2V-5B de forma predeterminada y predice fragmentos de acciones con forma [batch, action_horizon, action_dim].

Lo que abarca la integración de LeRobot

• Configuración estándar policy.type=fastwam a través de LeRobot
• Adaptación de lotes para tareas de imagen, estado, acción y lenguaje
• Inferencia de trozos de acciones a través de select_action y predict_action_chunk
• Guardado y carga de puntos de control a través de las APIs de políticas de LeRobot
• Postprocesamiento configurable de la acción del agarre de LIBERO

Requisitos de instalación

Instale LeRobot desde el código fuente y luego instale las dependencias de FastWAM:

pip install -e ".[fastwam]"

Esto instala el extra de la política FastWAM desde pyproject.toml: transformers, diffusers, ftfy y regex, además de las dependencias base de LeRobot.

Para la evaluación de LIBERO, instale también las dependencias del benchmark:

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

Esto instala ambos extras. Además de las dependencias de FastWAM anteriores, el extra libero instala las dependencias del conjunto de datos LeRobot, hf-libero en Linux y scipy.

FastWAM utiliza el backbone Wan2.2 TI2V. El id de modelo predeterminado es:

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

Requisitos de datos

FastWAM espera un conjunto de datos de LeRobot con:

• una o más observaciones visuales cuyas anchuras concatenadas suman policy.image_size[1]
• observation.state cuando policy.proprio_dim no es None
• action
• una instrucción de tarea de lenguaje a través del campo de tarea del conjunto de datos, o tensores context y context_mask precalculados

La configuración visual predeterminada es una característica de imagen denominada observation.images.image con forma (3, 224, 448). Si el conjunto de datos utiliza dos cámaras, configure policy.input_features para que sus alturas coincidan con 224 y la suma de sus anchos sea 448.

Uso

Cree una nueva política de FastWAM con:

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

Evalúe un punto de control existente en formato LeRobot en LIBERO-10 con:

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 y libero_object, usa --env.episode_length=300.

Este punto de control utiliza el preprocesamiento VISUAL: IDENTITY. FastWAM reescala las imágenes a [-1, 1] dentro del modelo, por lo que aplicar la normalización de imágenes en el preprocesador las normalizaría dos veces.

Para el despliegue en un robot real, use la misma ruta de punto de control:

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

Notas de configuración

Características de imagen

policy.image_size es el tamaño del tensor de imagen FastWAM concatenado como (height, width). Cada característica de imagen configurada debe tener la forma (3, height, camera_width), y todas las anchuras de cámara deben sumar la anchura configurada.

Chunking de acciones

policy.action_horizon controla el número de acciones futuras supervisadas durante el entrenamiento y predichas durante la inferencia. policy.n_action_steps controla cuántas acciones se consumen antes de que la política prediga un nuevo bloque. policy.n_action_steps debe ser menor o igual que policy.action_horizon.

Acciones relativas

FastWAM predice acciones absolutas de forma predeterminada. Con policy.use_relative_actions=true predice desplazamientos desde el observation.state actual, por lo que policy.proprio_dim debe estar configurado. Las articulaciones especificadas en policy.relative_exclude_joints (valor predeterminado ["gripper"]) permanecen absolutas. Recalcule primero las estadísticas del conjunto de datos en el espacio relativo, con --operation.chunk_size establecido en policy.action_horizon (consulte acciones relativas de π₀.₅):

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

Los puntos de control guardados antes de que existieran las acciones relativas, como lerobot/fastwam_base, reciben los pasos del procesador de acciones relativas al realizar el ajuste fino a partir de ellos con esta opción. Esto solo se aplica durante el ajuste fino: la evaluación y el despliegue utilizan los procesadores guardados con el punto de control tal cual están.

Dispositivo del codificador de texto

El codificador de texto UMT5-XXL congelado ocupa aproximadamente 11 GB. policy.text_encoder_device=cpu lo mantiene fuera de la GPU: los prompts se codifican en la CPU y los embeddings se mueven a policy.device. Durante la inferencia, la tarea codificada se almacena en caché hasta policy.reset(), por lo que la codificación más lenta en la CPU se ejecuta una vez por episodio en lugar de una vez por chunk. El entrenamiento codifica cada batch, por lo que en ese caso la ubicación en la CPU ralentiza cada paso.

Componentes de Wan

FastWAM carga el VAE de Wan, el DiT de video, el codificador de texto y el tokenizador desde el directorio del modelo Wan configurado o desde el id de modelo de Hugging Face Hub. Los puntos de control de FastWAM en formato LeRobot guardados por save_pretrained también copian los archivos locales de los componentes de Wan necesarios para from_pretrained.

Backend de atención

El DiT de FastWAM utiliza scaled_dot_product_attention (SDPA) de PyTorch para toda la atención. No utiliza FlashAttention: su enrutamiento Mixture-of-Transformers (MoT) requiere máscaras de atención booleanas arbitrarias [query, key], que la API varlen de FlashAttention no puede expresar. Por lo tanto, instalar el paquete flash-attn no tiene ningún efecto en la ruta de FastWAM. (Tenga en cuenta que el propio SDPA aún puede seleccionar internamente el kernel flash / memory-efficient / math de PyTorch; esto no tiene relación con el paquete flash-attn.)

Inferencia de acciones compiladas

FastWAM puede compilar sus grafos de prellenado de video y de desruido de acciones en caché para reducir la sobrecarga de Python y de lanzamiento de kernels de CUDA a lo largo de los pasos de inferencia repetidos:

--policy.compile_action_infer=true

La opción está desactivada de forma predeterminada. La primera predicción de acción compila un grafo para las formas de entrada actuales y, por lo tanto, es más lenta; las predicciones posteriores con las mismas formas lo reutilizan. Los cambios de forma pueden provocar una recompilación. La ruta eager sigue estando disponible en configuraciones de PyTorch o de acelerador no compatibles.

LIBERO Alternar acción

Los checkpoints de FastWAM LIBERO usan policy.toggle_action_dimensions=[-1] de forma predeterminada para coincidir con la convención de acción de la pinza utilizada por el pipeline de evaluación original de FastWAM:

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

Resultados

Evaluado en LIBERO con lerobot/fastwam_libero_uncond_2cam224:

Reproducir: 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).

Referencias

• artículo Fast-WAM
• Página del proyecto Fast-WAM
• Código de Fast-WAM
• Puntos de control publicados aguas arriba
• Wan2.2 TI2V 5B

Cita

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