Крошечный WorldFormer запускает зацикленную мировую модель и RL-актор на Raspberry Pi Pico 2 W

tinyworldformer — это открытый проект от SheepSeb, который запускает две небольшие нейронные сети на Raspberry Pi Pico 2 W: модель мира, генерирующую игровые кадры, и предобученного актора, который в них играет. Ни одна из сетей никогда не видит настоящую игру на чипе — актор читает только кадры, созданные моделью мира, а модель мира получает только действия от актора.

Архитектура

• Модель мира: зацикленный трансформер с int8-квантованием. 26–131 тыс. параметров, 2–8M умножений-сложений на кадр. Вход — текущий кадр + действие; выход — следующий кадр + награда.
• Актор: MLP с двумя скрытыми слоями, float32. 9–16 тыс. параметров, 9–16 тыс. умножений-сложений на кадр. Вход — кадр; выход — действие.
• Оба работают на ядре 1 RP2350; ядро 0 обслуживает USB и дисплей.
• Предобученного актора можно заменить на свежего stream-x RL-агента, который учится с нуля прямо на плате. Нажатие кнопок передаёт управление человеку.

Игры и результаты бенчмарков

Каждая игра — это несколько сотен строк NumPy в src/tinyworldformer/games/, рендерящих в небольшую сетку ID токенов. Один кадр содержит всё состояние игры, поэтому скрытое состояние отрисовывается в кадре (направление мяча — часть его токена; в Doom-lite есть мини-карта).

• Breakout — кадр 20×16, 26 тыс. параметров, 4.3M MAC, 99.99% точных кадров.
• Doom-lite — 32×32 луч-кастинг вид лабиринта 8×8, 55 тыс. параметров, 4.2M MAC, 83.6% точных кадров (один затенённый пиксель стены делает весь кадр неверным).
• Flappy Bird — кадр 24×24, 46 тыс. параметров, 2.0M MAC, 100% точных кадров.
• Pong — кадр 20×32, 46 тыс. параметров, 2.3M MAC, 99.45% точных кадров.
• Pac-Man lite — кадр 48×48, лабиринт 23×21, 249 точек, 131 тыс. параметров, 8.2M MAC, 86.4% точных кадров.

«Точные кадры» — это доля отложенных переходов, где каждая ячейка предсказанного кадра совпадает с реальной игрой (кадры сразу после случайных событий исключены).

Как работает модель мира

Кадры разбиваются на квадратные патчи (2×2 ячейки в Breakout, 4×4 в Doom-lite, Flappy Bird и Pong, 6×6 в Pac-Man), что даёт 36–80 токенов на игру. Эмбеддинг каждого патча суммирует один обученный вектор на ячейку с отдельной таблицей для каждой позиции внутри патча, плюс векторы позиции и действия. Модель представляет собой один блок трансформера (внимание + ReLU MLP) с одним набором весов, применяемым несколько раз — зацикленный трансформер — что сохраняет память достаточно низкой для микроконтроллера.

Запуск настоящих игр

uv run python -m tinyworldformer.games.play --game pacman

Замените на breakout, doomlite, flappy или pong. Репозиторий также включает tinyworldformer-rs, документацию, тесты и браузерную демо-версию.

Для разработчиков, работающих над edge-инференсом или небольшими моделями мира, практический вывод — это компромисс между памятью и вычислениями: повторное использование весов одного блока трансформера позволяет модели с 26 тыс. параметров работать за несколько миллионов MAC на кадр, оставаясь в пределах ограничений Pico 2 W.