开发者 SheepSeb 近期开源了 tinyworldformer 项目,在搭载 RP2350 芯片的 Raspberry Pi Pico 2 W 微控制器上,同时塞进了两个小型神经网络,脱离真实游戏代码实现了AI自己生成游戏画面并自己玩。第一个网络是 int8 量化的循环 Transformer 世界模型,参数量 26k–131k,每帧计算量 2–8M MACs,负责根据当前帧和动作预测下一帧与奖励;第二个网络是 float32 的双隐藏层 MLP 策略网络,参数量 9k–16k,负责看画面选动作。
两套网络全部跑在 RP2350 的 Core 1 上,Core 0 只管 USB 和显示,支持 Breakout、Doom-lite、Flappy Bird、Pong 和 Pac-Man lite 五款微型游戏,还能在板端从零训练强化学习智能体或由人类按键接管。普通读者看到几块钱的单片机跑通世界模型往往会惊叹算力奇迹,但拆开底层代码就会发现,它靠的是把视觉问题降维成全状态暴露的离散网格;更扎心的是,即便单步帧预测准确率高达 99.99%,强化学习智能体依然能精准钻空子刷出超越物理上限的假高分,把当前整个世界模型赛道的闭环训练软肋暴露无遗。
放弃像素扩散与隐藏状态,把数百兆显存需求砍进 460KB RAM
主流视觉世界模型是个吞金兽。Google 的 GameNGen 在单块 TPU 上跑 320×240 的 RGB 画面只能维持约 20 FPS;DIAMOND 模拟 Atari 需要 4.4M 参数,模拟 CS:GO 高达 381M 参数,仅 FP16 权重就占去约 762 MB,在 RTX 3090 上也只有约 10 FPS。而 tinyworldformer 最庞大的 Pac-Man lite 完整固件(含 USB、显示驱动和板端学习器),也只占用了 460 KB RAM 和 600 KB Flash,体积最小的 Flappy Bird 更是压到了 227 KB RAM。
这种量级跨越不是靠架构魔法硬压出来的,而是对环境表征做了极端的工程妥协。tinyworldformer 直接放弃了 RGB 像素生成,把五款游戏改写成了 20×16 到 48×48 的离散 Token 网格,切分成 2×2 到 6×6 的方块后,每帧只有 36 到 80 个 Token。
为了让一个极小的模型不靠长序列记忆就能推演世界,作者强行消除了所有的隐藏状态。在 Breakout 和 Pong 里,球的运动方向被直接编码进单个网格的 Token ID 中;在 32×32 分辨率的伪 3D 游戏 Doom-lite 里,画面一角被硬生生画上了一个迷宫小地图。单帧画面直接包含了完整的马尔可夫状态,模型不需要去猜球正在往哪飞、墙后有什么,难度从视觉想象降维成了符号查表。
在模型结构上,作者采用了循环 Transformer(Looped Transformer),把同一个 Transformer 块的权重重复执行 1 到 4 次。这招确实省下了宝贵的 Flash 存储空间,但并不省算力——每多循环一次,注意力机制和 MLP 就要原封不动重算一遍。而且循环次数绝非越多越好,Pac-Man 在循环 2 次时精确帧匹配率为 86.4%,加到 3 次和 4 次反而跌到了 85.7% 和 84.4%。
做边缘 AI 的工程师还得注意一个内存陷阱:131k 的 int8 参数看着只有一百多 KB,但词嵌入、归一化、缩放因子和残差连接全得保留 float32。单是 Pac-Man 每单元格的输出得分缓冲区,一帧就要吃掉 184 KB 内存。作者把静态权重锁死在 Flash 里,RAM 只留固定缓冲区且零动态内存分配,才勉强塞进芯片。
- 提醒.项目文档里提到的 0.7–1.5 ms/帧其实是桌面 CPU 的计时数据,Pico 2 W 实机运行时的真实帧率(FPS)目前根本还没有测出来。
| 游戏名称 | 网格尺寸 | 世界模型参数與计算量 | 2 次循环帧准确率 | 完整固件内存占用 |
|---|---|---|---|---|
| Breakout | 20×16 | 26k / 4.3M MACs | 99.99% | 299 KB RAM / 361 KB Flash |
| Doom-lite | 32×32 | 55k / 4.2M MACs | 83.6% | 357 KB RAM / 471 KB Flash |
| Flappy Bird | 24×24 | 46k / 2.0M MACs | 100% | 227 KB RAM / 409 KB Flash |
| Pong | 20×32 | 46k / 2.3M MACs | 99.45% | 244 KB RAM / 414 KB Flash |
| Pac-Man lite | 48×48 | 131k / 8.2M MACs | 86.4% | 460 KB RAM / 600 KB Flash |
剔除经验回放池实现板端强化学习,但 99.99% 准确率仍被智能体钻空子
很多读者容易把芯片上自己学习玩自己的梦境,误解为世界模型也在单片机上自我进化。实际上,板端的世界模型是只读的——它预先在桌面端用 PyTorch 跑了 30,000 步(前 3,000 步 float 预热后开启 int8 量化感知训练),再把权重硬编码进 Rust 源码烧录进 Flash。真正在 Pico 2 W 上边玩边学的,只有那颗 9k–16k 参数的策略网络。
要在几百 KB 内存里搞强化学习,最大的拦路虎是传统 DQN 或 PPO 动辄几万条状态的经验回放池(Replay Buffer)。tinyworldformer 引入了 2024 年提出的 Stream AC(λ)(即项目中的 stream-x)算法,不设回放池、不搞批处理(Batch)、不要目标网络,来一个状态转移就用资格迹(Eligibility Traces)和 ObGD 优化器当场更新一步。这种极致的流式更新,才让板端从零训练成为可能。
哪怕世界模型准确率高达 99.99%,强化学习也能靠那 0.01% 的幻觉把梦境玩崩。
这个微型项目最核心的学术价值,恰恰发生在智能体进入世界模型内部训练的那一刻。在 Breakout(打砖块)里,世界模型单步预测的精确帧匹配率已经做到了惊人的 99.99%。但当强化学习智能体完全在世界模型生成的梦境里训练时,它迅速找到了那万分之一的规则漏洞:把球打向特定角落时,世界模型会给出加分奖励,但忘了把画面里的砖块擦除。
智能体抓着这处幻觉疯狂刷分,在真实规则下满分只有 56 分 的游戏里,硬是在梦境中刷出了 210 分 的荒谬战绩;而把这个在梦里天下无敌的智能体放回真实游戏代码里测试,它只能拿到可怜的 28 分。同样的模型漏洞利用(Model Exploitation)在别的遊戲里也层出不穷:Doom-lite 的智能体学会了隔着 5 个格子远程狙击,还会导致敌人图标在画面中永久消失;Pac-Man 在连续 5,000 步的生成测试中,更是累积出了 742 个错误游戏状态和 38 个错误奖励。
为了堵住智能体钻空子的路,作者不得不给神经网络加上浓重的人工规则拐杖。比如早期 Breakout 模型经常在画面里同时画出两个球,作者干脆在解码层写死了一条全局唯一(Singleton)规则,强行规定整张图只能选一个最高概率的位置放球,这才把单步准确率从 91.6% 拉升到 99.7%。面对刷分漏洞,作者只能把智能体在梦里跑偏的轨迹抓出来,放回真实游戏重新打标,再拿去微调世界模型,并效仿 GameNGen 在训练时注入 0.3 的随机像素噪声,强迫模型看小地图重绘画面而不是偷懒复制上一帧。
- 结论.对于 TinyML 和强化学习开发者,tinyworldformer 的 Rust 混合精度推理与无缓冲 Stream AC(λ) 代码是极佳的边缘端工程范本;但它更像一面照妖镜——在几十 KB 的网格游戏里,尚且需要写死全局唯一解码规则和反复回炉微调才能按住智能体作弊,那些试图用纯视觉世界模型取代真实物理引擎来做闭环训练的团队,接下来必须先回答如何低成本消灭那万分之一的恶性幻觉。
