Исследователи из Вашингтонского университета и Sony AI представили Semantic World Models (SWM) — новый подход к моделированию мира для робототехники. Вместо предсказания будущих кадров в пикселях, модель отвечает на вопросы о семантической информации в будущем. Система построена на базе PaliGemma 3B и обучается на данных изображение-действие-текст.
Традиционные модели мира предсказывают будущие кадры пиксель за пикселем, что часто не коррелирует с качеством планирования. SWM переосмысливает эту задачу как visual question answering (VQA) о будущем. Например, вместо генерации изображения модель отвечает на вопросы типа «Приблизилась ли рука к объекту?» или «Опрокинулся ли красный куб?». Это позволяет использовать готовые инструменты от vision-language моделей и наследовать их способности к обобщению.
Архитектура SWM основана на трех ключевых компонентах: трансформерной языковой модели Gemma, энкодере изображений SigLIP и проекционной матрице для действий размером act_dim×2048. Модель принимает текущие наблюдения, последовательность действий и естественно-языковой запрос о будущем, генерируя текстовый ответ. Обучение проводится методом supervised finetuning с оптимизацией стандартной cross-entropy функции потерь.
Для планирования действий авторы используют два подхода: семплирование на основе MPPI и градиентную оптимизацию. Каждая задача определяется набором вопросов, желаемых ответов и весов. Скалярная оценка действий вычисляется как взвешенная сумма вероятностей желаемых ответов модели. Градиентный метод требует меньше семплов и работает быстрее крупных моделей.
Тестирование на средах LangTable и OGBench показало значительное улучшение производительности. В LangTable средняя эффективность выросла с 14.4% до 81.6%, в OGBench — с 45.33% до 76%. SWM превзошла базовые методы IDQL и pixel-based world modeling. Модель также продемонстрировала композиционное обобщение при работе с новыми цветами блоков и изменениях фона, показав улучшение на 20% по сравнению с базовыми политиками.
Анализ карт внимания показал, что модель корректно фокусируется на релевантных объектах в зависимости от языкового запроса. При вопросе «Касается ли красная луна синего куба?» внимание концентрируется на соответствующих участках изображения. Модель успешно обобщается даже на вопросы о трех объектах, хотя обучалась только на двухобъектных примерах.
Источник новости и обложки: arxiv.org