世界模型跑上微控制器: 512KB 内存里塞进 5 款游戏

世界模型跑上微控制器: 512KB 内存里塞进 5 款游戏

端侧AI世界模型强化学习Raspberry Pi嵌入式

数据源:tinyworldformer 项目文档 + GitHub 仓库

一块巴掌大小的微控制器开发板,脱离了游戏引擎与外部服务器。它在芯片内部自主演算出物理世界的因果链。在开源项目 tinyworldformer 中,Raspberry Pi Pico 2 W 承担了预测与决策回环。这颗仅有 520KB SRAM 的芯片,同时运行着生成画面的循环 Transformer 和负责动作的决策网络。

世界模型以往总是绑定高能耗云端集群,动辄耗费数十亿参数。tinyworldformer 另辟轻量路径。它将网络压缩到 26k 至 131k 参数,单帧计算量降至 2M 到 8M 次乘加。5 个经典游戏的环境被塞进 4MB Flash 与 520KB SRAM 中。

强化学习训练全程在网络预测帧中展开。智能体从未接触过原生游戏。端到端神经仿真摆脱了云端开销,将世界模型与强化学习拉入微控制器。方案代价体现在画质与 83.6% 至 100% 的精确率,算力开销则被直接省去。

两个网络互为镜像,微控制器上不跑真游戏

Pico 2 W 搭载的 RP2350 具备两个 Cortex-M33 核心。系统将显示刷新与因果推理在物理核心上划清界限。Core 0 专职管理 USB 通信与屏幕输出。Core 1 独立运行两个互为镜像的神经网络,脱离了真实游戏引擎。

世界模型负责环境状态的向前推演。网络输入当前帧与动作,输出预测的下一帧与奖励。其架构由量化为 int8 的循环 Transformer 构成,拥有 26k 到 131k 参数。单帧计算产生 2M 到 8M 次乘加运算(Pac-Man lite 达到 8.2M MACs),保证片上推演跟得上控制循环。

执行器负责感知画面并做出决策。网络采用两层 MLP,保持 float32 精度,参数量仅 9k 至 16k。它每帧仅需 9k 到 16k 次乘加,算力负担极轻。执行器只读世界模型的预测特征,世界模型只接收执行器的离散动作。

这套双网络系统提供了灵活的人机接管机制。板上的预训练执行器可以实时替换为从零起步学习的 stream-x 智能体。开发人员一旦按下板载按键,控制权便立刻交予人类玩家。玩家操作按键时,也是在神经网络实时预测出的虚构画面中游玩。

tinyworldformer 架构与在单片机上运行双网络的概览 图:tinyworldformer 架构与在单片机上运行双网络的概览。来源:SheepSeb/tinyworldformer。

单层权重循环复用,拿计算延迟换取存储空间

传统视觉 Transformer 依赖堆叠多层 Block 扩充容量。但芯片 SRAM 无法容纳多组权重。tinyworldformer 采用循环 Transformer 架构,全局仅保留单个 Transformer Block。模型将同一套权重复用 1 到 4 次,每次复用前把初始输入残差加回。

架构训练时将每一轮循环的中间输出,都对着真实下一帧进行监督回归。运行时的循环次数转变为调节计算延迟换取推演精度的物理旋钮。在 Pico 2 W 硬件上,为了维持可交互的帧率,系统默认固定采用 2 次循环。这种设计把网络深度的选择权交给了边缘执行环境。

以 Doom-lite 的独立测试用例为例,循环次数对画面消除幻觉具有直观效果。单次循环推演时,模型凭空生成了一个不存在的敌人,全图产生 92 个错误网格。将循环增加到 2 次后错误网格降至 35 个。3 次循环进一步降至 12 个,4 次循环时错误全部消除。

项目测试了 5 款游戏在 2 次循环下的表现。Breakout 画面为 20×16 网格,以 26k 参数和 4.3M MACs 实现了 99.99% 的逐帧精确率。Flappy Bird 为 24×24 网格,以 46k 参数和 2.0M MACs 达到 100% 的精确率。Pong 为 20×32 网格,以 46k 参数和 2.3M MACs 录得 99.45% 的精确率,确定性平面力学得到了高质量拟合。

涉及深层状态追踪与伪 3D 渲染的环境,则触碰到了模型容量的边界。Pac-Man lite 在 48×48 网格中塞入复杂迷宫与两只幽灵,参数升至 131k,计算量达 8.2M MACs,精确率为 86.4%。Doom-lite 采用 32×32 第一人称视角和 8×8 迷宫,55k 参数产生 4.2M MACs,精确率为 83.6%。面对多实体隐藏状态,低参数模型必须接受局部网格跳变的折中代价。

世界模型自主生成的 Pac-Man 连续预测帧 图:世界模型在无真机引擎介入下自主生成的 Pac-Man 连续预测帧。来源:SheepSeb/tinyworldformer。

启发式解码压制幻觉,给概率分布立规矩

如果让轻量神经网络直接根据 Softmax 分布无约束采样,生成的物理世界将在几步内走向坍塌。吃豆人容易在地图上分裂出孪生体。小球也会在反弹时凭空蒸发。作者把解码单独列成一节,原因是概率分布不会自己变成合法的游戏画面。

第一条规则面向静态背景。墙壁、砖块、豆子与地板等格子,各自独立选取最大概率的 token。背景元素独立决策,省去了全局图搜索开销。稳定的 argmax 运算锁死地图轮廓,防止迷宫墙体发生蠕动。

第二条规则用于管理单例实体。小球与玩家等单例对象,在全图网格与类别组合中只选一次。物理上的全局唯一由此被锁死。遇到发球或重生等随机事件,解码器在最优解附近 2.2 logits 邻域内采样,既保留了动态重生的多样感,又消除了均值模糊产生的残差。

第三条规则针对占用 2×2 复合网格的移动精灵。吃豆人与幽灵等精灵只有主网格承载状态。它负责记录嘴部朝向,或标记幽灵是否遮挡豆子。其余三格依据预设规则盖章填充,跳过逐点预测,排除了角色撕裂的渲染瑕疵。

全流程压进 int8,一条 SIMD 指令吃下两次乘加

将 Transformer 移植到微控制器,首要障碍在于浮点矩阵乘法的能耗。tinyworldformer 将世界模型的所有权重矩阵压缩至 int8 整数。网络按矩阵输出行独立计算缩放系数(scale = max|w| / 127)。激活张量同样量化为 int8,并根据张量归属与循环轮次单独分配 scale。

RP2350 搭载的 Cortex-M33 架构为整型计算提供了底层加速指令。一次 32 位读取取回 4 个权重,SXTB16 把字节拆成两组 16 位,SMLAD 一条指令完成两次乘加并累加进 32 位寄存器(注意力里的 uint8 概率 × int8 数值走 UXTB16)。矩阵计算管道在 int32 累加完成后,仅需执行一次双 scale 浮点乘积还原。配合 1024 项预计算指数表实现的 uint8 查表 Softmax,整机摆脱了高昂的连续浮点开销。

RP2350 配备的硬件浮点单元并未被闲置。LayerNorm 归一化、残差连接累加以及输入 Embedding 依然保留在 float32 精度。这些非矩阵乘法步骤在整体算力中占比极低。维持浮点格式能够有效规避深层累积误差,兼顾了定点加速与数值稳定。

量化方案的高保真度源于训练与部署的严格吻合。模型在训练阶段就通过感知量化训练与直通估计器引入同样的整数舍入,训练所得即为最终部署的定点网络。导出模块 rustexport/intref.py 严格固化了 NumPy 浮点运算顺序。Rust 推理引擎在微控制器上逐位复现参考实现,消除了跨平台部署的数值漂移。

模型矩阵与激活值的 int8 量化方案 图:模型矩阵与激活值的逐行与按循环轮次 int8 量化方案。来源:SheepSeb/tinyworldformer。

零堆分配锁死 520KB 内存,让端侧承担仿真代价

嵌入式对内存溢出零容忍。用 size -A 分析 ELF,Breakout 占 361KB Flash 与 299KB RAM,Flappy Bird 占 409KB 与 227KB,Pong 占 414KB 与 244KB。Doom-lite 占 471KB 与 357KB,Pac-Man lite 录得 600KB 与 460KB。全套固件被牢牢锁在 Pico 2 W 的 4MB Flash 与 520KB SRAM 之内。

内存剖析揭示了一个反直觉的事实。整个世界模型中,量化后的 int8 权重仅占 19.7KB 至 33.9KB。相比之下,float32 格式的 Embedding、归一化参数与 scale 数组构成了内存消耗大头,占据 31.6KB 到 91.9KB。存储开销主体集中在边缘适配层,紧凑的计算核心只占很小一部分。

为了保障系统的长期稳定性,固件采用 no_std Rust 与 Embassy 异步框架开发。系统达成了全生命周期的零动态堆内存分配。所有模型权重作为 static 数组直接常驻于 Flash 存储器,中间计算激活区与显存缓冲区全部分配为 RAM 里的定长静态数组。消除动态内存管理排除了内存碎片与分配崩溃风险。

板载 stream-x 学习器所需的权重与资格迹仅占 49KB 到 106KB,同样在静态内存中安全运转。双核心的异步流水线维系着稳定的交互响应。Core 1 循环推进特征读取、动作选择、2 次循环推演、启发式解码与奖励预测。计算数据经由两槽深度通道投递给 Core 0,Core 0 繁忙时主动丢帧,确保 Core 1 推演绝不挂起等待。

在桌面 CPU 上该 Rust 引擎单帧耗时仅为 0.7ms(Flappy Bird)至 1.5ms(Breakout)。尽管原作者未实测 Pico 2 W 上的具体毫秒数,人类介入时将帧率放慢至 66ms 到 150ms 即可提供流畅的游玩体验。5 个环境已在 SheepSeb/tinyworldformer 仓库开源,支持 uv 终端运行,也提供 WebAssembly 网页演示版。

世界模型不必依附于云端集群,tinyworldformer 确立了新路径。依靠权重时序复用与定点 SIMD 调度,一块微控制器成功跑起环境仿真。这套方案用画质上限与局部跳变,换取了摆脱外部服务器的端侧自治。几百千字节的微小空间,足以承载智能体专属的物理沙箱。

参考链接:

  • SheepSeb 个人博客:tinyworldformer 技术解析与实现细节
  • GitHub 开源仓库:SheepSeb/tinyworldformer