Light-WAM: Modelos de Acción Mundial Eficientes con Decodificación de Acción mediante Fusión de Estados

Resumen

Los World Action Models (WAMs) extienden el aprendizaje de políticas robóticas al incorporar la predicción futura como un objetivo de entrenamiento adicional, incentivando a la política a codificar la estructura temporal relevante para la tarea en sus representaciones. Los WAMs actuales suelen basarse en arquitecturas generativas a gran escala que incurren en altos costos de entrenamiento y latencia de inferencia, lo que dificulta su despliegue como políticas eficientes de bucle cerrado. Proponemos Light-WAM, un World Action Model ligero para manipulación robótica eficiente. Específicamente, se construye con un backbone de video compacto y realiza supervisión de video futuro en un espacio latente submuestreado, reduciendo el costo del coentrenamiento de video mientras conserva sus beneficios para el aprendizaje de representaciones. Para la predicción de acciones, Light-WAM introduce StateFusionActionExpert, que lee estados adaptados de múltiples capas del backbone, los fusiona mediante agrupación por consultas aprendidas y predice fragmentos de acción directamente en una sola pasada forward. Este diseño proporciona una interfaz eficiente entre las representaciones del backbone de video y las acciones del robot, evitando la necesidad de expertos de acción generativos pesados. Los experimentos demuestran que Light-WAM mantiene un rendimiento sólido en LIBERO y logra un rendimiento multitarea utilizable en RoboTwin 2.0, utilizando solo 0,44 mil millones de parámetros entrenables. También alcanza una latencia de inferencia de 72,03 ms con 4,1 GiB de memoria máxima de GPU y un rendimiento de entrenamiento mejorado.

English

World Action Models (WAMs) extend robot policy learning by incorporating future prediction as an additional training objective, encouraging the policy to encode task-relevant temporal structure in its representations. Current WAMs often rely on large-scale generative architectures that incur high training costs and inference latency, making them difficult to deploy as efficient closed-loop policies. We propose Light-WAM, a lightweight World Action Model for efficient robot manipulation. Specifically, it is built with a compact video backbone and performs future-video supervision in a downsampled latent space, reducing the cost of video co-training while retaining its benefits for representation learning. For action prediction, Light-WAM introduces the StateFusionActionExpert, which reads adapted states from multiple backbone layers, fuses them through learned-query pooling, and directly predicts action chunks in a single forward pass. This design provides an efficient interface between video backbone representations and robot actions, avoiding the need for heavy generative action experts. Experiments demonstrate that Light-WAM maintains strong performance on LIBERO and achieves usable multi-task performance on RoboTwin 2.0, while using only 0.44B trainable parameters. It also achieves 72.03ms inference latency with 4.1GiB peak GPU memory and improved training throughput.