Tinyworldformer：在 Raspberry Pi Pico 2 W 上运行循环 Transformer

概述

Tinyworldformer 展示了通过使用循环 Transformer 和 int8 量化，复杂的世界模型和强化学习智能体可以在资源受限的硬件上运行。该系统使 Raspberry Pi Pico 2 W 能够实时生成游戏帧并学习游玩，有效地让微控制器在其自身的模拟环境中“做梦”并采取行动。

架构：循环 Transformer 与智能体

该系统依赖于在 RP2350 芯片上运行的两个不同的神经网络：

• 世界模型： 一个循环 Transformer，根据当前帧和动作预测下一帧和奖励。通过多次重复使用具有相同权重的单个 Transformer 块，该模型实现了显著的内存节省，使其在微控制器上变得可行。根据游戏复杂度的不同，它使用 26k–131k 个参数。
• 智能体 (Actor)： 一个多层感知机 (MLP)，用于观察生成的帧并选择动作。它非常轻量，仅使用 9k–16k 个参数。

这两个网络都在 RP2350 的核心 1 上运行，而核心 0 负责管理 USB 通信和显示输出。该系统在推理过程中避免使用“真实”的游戏逻辑，完全依赖世界模型的预测。

游戏实现与性能

该项目支持五款游戏：Breakout、Doom-lite、Flappy Bird、Pong 和 Pac-Man。每款游戏均由几百行 NumPy 代码实现，渲染为一个小型的 token ID 网格。世界模型的准确性因游戏而异，在 Flappy Bird 等简单游戏中实现了高保真度（100% 精确帧预测），而在 Pac-Man 等复杂游戏中（2 次循环时为 86.4%）也表现良好。

为了确保模型保持稳定，该项目采用了特定的解码规则，例如放置“单例”对象（如球或玩家），以防止模型产生同一对象的多个幻觉实例。世界模型通过在智能体训练期间将其预测与实际游戏规则进行比较来验证。

量化与部署

为了适应 Pico 2 W 的 512 KB SRAM 和 4 MB Flash，该项目使用了量化感知训练 (QAT)。世界模型中的每个矩阵都被转换为 int8 格式。固件使用 no_std Rust 编写，确保了确定性的性能，且无需动态内存分配。

最大游戏 Pac-Man 的性能指标显示，该系统消耗约 600 KB 的 Flash 和 460 KB 的 SRAM，完全在硬件限制范围内。RP2350 的 FPU 的使用允许高效处理范数和残差，而基于整数的点积则使用特定的汇编指令 (SMLAD) 进行了优化。

关于微控制器上 AI 的见解

社区讨论强调了这种方法在未来 AI 应用中的潜力。评论者指出，这类专注于最小化、本地运行模型的项目代表了嵌入式 AI 未来的重要方向。

"我真心认为，这些精巧的项目（尽可能小、本地运行的 LLM）才是 AI 长远发展的真正未来。" — skrellm

其他人则指出，虽然硬件限制令人印象深刻，但其底层逻辑可以集成到传统游戏开发中。通过将这些模型用于程序化生成或自适应 NPC 行为，开发者可能能够解决标准程序化生成技术中常见的局限性。