На Raspberry Pi Pico 2 W запустили модель мира, которая сама генерирует игры

Проект tinyworldformer на GitHub показывает, как на микроконтроллере Raspberry Pi Pico 2 W работают сразу две маленькие нейросети. Первая, модель мира: по текущему кадру и действию она предсказывает следующий кадр и награду, то есть сама «рисует» игру. Вторая, актор: она берёт сгенерированный кадр и выбирает действие. На самой плате ни одна из сетей не обращается к настоящей игре: актор видит только кадры модели мира, а модель мира получает действия только от актора. Обе сети работают на втором ядре (core 1) чипа RP2350, первое ядро занято USB и дисплеем. Предобученного актора можно заменить свежим агентом, который учится с нуля прямо на плате, а любое нажатие кнопки передаёт управление человеку, так что в сгенерированную игру можно играть самому.

Мотивация, по словам автора README: его заинтересовали циклические трансформеры (looped transformers), где один блок применяется несколько раз с одними и теми же весами вместо стопки разных слоёв. Повторное использование весов дёшево по памяти, и он захотел проверить, можно ли уместить такую модель мира в микроконтроллер. Pico 2 W выбран, потому что он дешёвый и имеет блок вычислений с плавающей запятой (FPU). Игры маленькие, каждая это несколько сотен строк на numpy, чтобы можно было сравнивать каждый сгенерированный кадр с настоящей реализацией. Это Pac-Man, Breakout, Pong, Flappy Bird и Doom-lite (упрощённая 3D-игра, не сам Doom); в настоящие версии можно играть в терминале.

Устройство модели. Кадр, сетка токенов; его режут на квадратные патчи: 2×2 клетки в Breakout, 4×4 в Doom-lite, Flappy Bird и Pong, 6×6 в Pac-Man, что даёт от 36 до 80 токенов на игру. Модель, один трансформерный блок (внимание и MLP с ReLU) с одним набором весов, применяемый от 1 до 4 раз; выход после каждого прохода обучается предсказывать следующий кадр, поэтому на этапе работы число циклов обменивает скорость на точность. На Pico для приемлемой частоты кадров используют два цикла. Пример из README: на одном отложенном шаге Doom-lite после первого прохода модель рисует несуществующего врага, а затем число неверных клеток падает: 92, 35, 12, затем 0 (всего в 3D-виде 768 клеток).

Превращение выхода модели в кадр, по словам автора, важно почти так же, как сама модель, и сводится к трём правилам. Фоновые клетки (стены, кирпичи, точки, пол) берут самый вероятный токен независимо. Одиночные объекты (мяч, игрок, каждый призрак) ставятся одним общим выбором по всем парам «клетка, токен», поэтому каждый присутствует ровно один раз; после случайного события выбор делается случайно среди кандидатов в пределах 2,2 логита от лучшего. Спрайты Pac-Man и призраков 2×2 хранят состояние только в одной клетке, остальные три штампуются. Правило 2 ввели потому, что первая модель Breakout рисовала два мяча сразу: один общий выбор поднял долю точных кадров с 91,6% до 99,7% (на отложенной выборке).

Проверка. Точности на записанных данных оказалось мало: агент заводит «сон» в состояния, которых не было в данных, и любая ошибка становится лазейкой. Поэтому главная проверка идёт во время игры агента: из сгенерированного кадра считывают состояние, применяют настоящие правила и действие агента и сравнивают результат со следующим сгенерированным кадром и наградой (по 5000 шагов «сна» на игру при двух циклах). Агенты действительно находили лазейки: агент в Breakout набрал 154 очка, когда осталось 26 кирпичей, потому что модель не убирала кирпичи в углу; агенты Doom-lite научились стрелять с 5 клеток. Агенты, обученные только во «сне», набрали там 210 очков (реальный максимум 56), а в настоящей игре 28, как и предсказывали Ха и Шмидхубер в работе World Models. Каждый раунд «собрать данные, дообучить модель» закрывал найденную лазейку. Среди других исправлений: шумовая аугментация (--noise 0,3, как в GameNGen) помогла, когда в Doom-lite модель теряла спрайт врага и больше не рисовала его; для Breakout ширину модели пришлось увеличить с d=32 до d=48, а для 48×48 кадра Pac-Man, до d=64 (текст источника в этом месте искажён).

Агенты. Используется Stream AC(λ) (Эльсайед, Васан и Махмуд, 2024): обучение по одному переходу за раз, без буфера воспроизведения, батчей и целевой сети. Сети политики и оценки состояния имеют по два скрытых слоя из 64 нейронов (32 на Pico) с LayerNorm. Для Doom-lite добавлен штраф за топтание на месте (--idle-penalty 0,1): без него агент сидит в стартовом углу и стреляет, с ним патрулирует и убивает вдвое больше. Алгоритм обучения написан и на Rust (no_std) с ручным обратным проходом и совпадает с 200 обновлениями PyTorch с точностью до 3,5e-6; на хосте свежий агент в int8-«сне» Flappy Bird даёт в среднем 6 труб за эпизод после 25 тыс. шагов, 121 после 100 тыс. и больше 300 после 250 тыс.

Квантизация и прошивка. Все матрицы модели, int8 с одним масштабом на строку выхода, активации тоже int8 (масштаб на тензор и на номер цикла); умножения идут int8×int8 в int32, нормы, остаточные связи и эмбеддинги остаются во float32. Softmax использует таблицу экспонент на 1024 записи. Обучение сразу идёт с тем же округлением (QAT): 30 тыс. шагов, первые 3-4 тыс. во float. Rust-движок воспроизводит эталон на numpy бит в бит, это проверяет cargo test. Прошивка, no_std Rust на embassy без выделения памяти во время работы: веса лежат статическими массивами во флеше. Когда играет человек, кадр занимает 66-150 мс в зависимости от игры. Кадры уходят по USB-последовательному порту как цветной ANSI-текст; есть и необязательный дисплей Display Pack (ST7789, 320×240). Текст README в источнике обрывается на описании дисплея.

Ключевые факты

• Две маленькие нейросети на Raspberry Pi Pico 2 W: циклический трансформер-модель мира генерирует следующий кадр и награду, актор играет в этой сгенерированной игре.
• Один трансформерный блок применяется 1-4 раза с общими весами; на плате используют два цикла. В примере Doom-lite число неверных клеток по циклам: 92, 35, 12, 0.
• Правило «один общий выбор для одиночных объектов» подняло долю точных кадров Breakout с 91,6% до 99,7%.
• Агенты находили лазейки в модели мира: 154 очка при 26 оставшихся кирпичах; обученные только во «сне» набрали 210 очков там и 28 в настоящей игре.
• Веса квантованы в int8 (QAT), прошивка на Rust без выделения памяти, расчёт на плате воспроизводит эталон на numpy бит в бит.

Почему это важно

Проект показывает, что модель мира, то есть нейросеть, которая сама генерирует игровую среду по кадру и действию, можно довести до размеров микроконтроллера: на Pico 2 W её запускают вместе с агентом, играющим внутри этой генерации. Ставка автора, циклический трансформер, который многократно применяет один блок с одними весами и потому экономит память. Работа ценна ещё и подробным разбором того, что ломалось по пути: двойные мячи, потерянные спрайты, лазейки, которые находит агент.

Кому это важно

Тем, кто интересуется моделями мира и обучением с подкреплением, а также разработчикам встраиваемых систем и квантизации. Отдельный модуль rustexport, по словам автора README, зависит только от numpy и torch, так что экспорт квантованной модели в Rust-код могут переиспользовать другие проекты.

Как это применить

В доступной части README описано: запуск настоящих игр в терминале командой uv run python -m tinyworldformer.games.play --game pacman (также breakout, doomlite, flappy, pong); обучение агентов и проверка «сна» скриптами из каталога training; экспорт квантованной модели в Rust через rustexport. Для запуска на плате нужен Pico 2 W, вывод идёт по USB-порту в виде ANSI-текста, дисплей Display Pack необязателен. Разделы про прошивку, браузерную версию и воспроизведение в доступном тексте оборваны.

Можно ли доверять

Источник, README самого проекта, все цифры авторские и независимо не проверялись. Поэтому доверять стоит скорее как описанию демонстрации с открытым кодом, который можно запустить самому, чем как рецензированному результату. Имя автора в тексте README не указано, а доступный текст обрывается в разделе про дисплей.

Риски и подводные камни

Игры маленькие и написаны автором специально; Doom-lite, упрощённая игра, а не Doom. Метрика «точные кадры» исключает кадры сразу после случайных событий, потому что их можно только сэмплировать, а у Doom-lite она выглядит хуже реального качества: достаточно одного неверного пикселя стены. Агенты в «сне» эксплуатируют ошибки модели, а в README этот эффект закрывают раундами дообучения. Данных о частоте кадров при игре самого агента на плате в доступном тексте нет, а крупнейшую часть занимают эмбеддинги во float32, которые, по словам автора, можно тоже квантовать.

«Повторное использование весов дёшево по памяти, поэтому я захотел проверить, можно ли сделать циклическую модель мира достаточно маленькой, чтобы она работала на микроконтроллере.»

— из README проекта tinyworldformer