TinyWorldFormer ejecuta un modelo de mundo en bucle y un actor de RL en una Raspberry Pi Pico 2 W

tinyworldformer es un proyecto de código abierto de SheepSeb que ejecuta dos pequeñas redes neuronales en una Raspberry Pi Pico 2 W: un modelo de mundo que genera fotogramas de juego y un actor preentrenado que juega con ellos. Ninguna de las redes ve nunca el juego real en el chip: el actor solo lee los fotogramas producidos por el modelo de mundo, y el modelo de mundo solo recibe acciones del actor.

Arquitectura

• Modelo de mundo: transformer en bucle con cuantización int8. 26k–131k parámetros, 2–8M multiplicaciones-sumas por fotograma. La entrada es el fotograma actual + la acción; la salida es el siguiente fotograma + la recompensa.
• Actor: MLP con dos capas ocultas, float32. 9k–16k parámetros, 9k–16k multiplicaciones-sumas por fotograma. La entrada es el fotograma; la salida es la acción.
• Ambos se ejecutan en el núcleo 1 del RP2350; el núcleo 0 gestiona el USB y la pantalla.
• Se puede cambiar un actor preentrenado por un agente de RL stream-x nuevo que aprende desde cero en la placa. Al pulsar un botón, el control pasa a un jugador humano.

Juegos y resultados de benchmark

Cada juego son unos cientos de líneas de NumPy en src/tinyworldformer/games/, que renderizan a una pequeña cuadrícula de IDs de token. Un fotograma contiene todo el estado del juego, por lo que el estado oculto se dibuja en el fotograma (la dirección de la bola forma parte de su token; Doom-lite tiene un minimapa).

• Breakout: fotograma de 20×16, 26k parámetros, 4,3M MACs, 99,99% de fotogramas exactos.
• Doom-lite: vista raycast de 32×32 de un laberinto de 8×8, 55k parámetros, 4,2M MACs, 83,6% de fotogramas exactos (un solo píxel sombreado de pared hace que todo el fotograma cuente como incorrecto).
• Flappy Bird: fotograma de 24×24, 46k parámetros, 2,0M MACs, 100% de fotogramas exactos.
• Pong: fotograma de 20×32, 46k parámetros, 2,3M MACs, 99,45% de fotogramas exactos.
• Pac-Man lite: fotograma de 48×48, laberinto de 23×21, 249 bolitas, 131k parámetros, 8,2M MACs, 86,4% de fotogramas exactos.

«Fotogramas exactos» es la proporción de transiciones de validación en las que cada celda del fotograma predicho coincide con el juego real (se excluyen los fotogramas inmediatamente posteriores a eventos aleatorios).

Cómo funciona el modelo de mundo

Los fotogramas se dividen en parches cuadrados (celdas de 2×2 en Breakout, 4×4 en Doom-lite, Flappy Bird y Pong, 6×6 en Pac-Man), lo que da entre 36 y 80 tokens por juego. Cada embedding de parche suma un vector aprendido por celda con una tabla separada por posición dentro del parche, más vectores de posición y de acción. El modelo es un único bloque transformer (atención + MLP ReLU) con un conjunto de pesos aplicado varias veces —un transformer en bucle—, lo que mantiene la memoria lo bastante baja para un microcontrolador.

Ejecutar los juegos reales

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

Sustituye por breakout, doomlite, flappy o pong. El repositorio también incluye tinyworldformer-rs, documentación, pruebas y una demo para el navegador.

Para los desarrolladores que trabajan en inferencia en el borde o en modelos de mundo pequeños, la conclusión práctica es el equilibrio entre memoria y cómputo: reutilizar los pesos de un único bloque transformer permite que un modelo de 26k parámetros se ejecute en unos pocos millones de MACs por fotograma sin superar los límites de la Pico 2 W.