SheepSeb/tinyworldformer

来源: https://github.com/SheepSeb/tinyworldformer 用两个小型神经网络在 Raspberry Pi Pico 2 W 上同时生成并游玩游戏。第一个网络负责创建游戏：给定当前帧和动作，预测下一帧和奖励。第二个网络负责游玩这个游戏：使用生成的帧选择动作。

上方的每一帧都是由世界模型网络采样生成的，动作由策略网络做出。

为什么我对循环 Transformer 感兴趣

循环 Transformer 指的是将同一个 Transformer 块用相同的权重重复应用多次，而不是堆叠不同的层。复用权重对内存非常友好，所以我想要验证循环世界模型是否能小到在微控制器上运行。我选择了 Pico 2 W，因为它价格低廉且带有 FPU；选择小游戏是因为我可以把真实版本直接写成 numpy，然后将每一帧生成结果与真实版本逐一比对。后来我加入了一个流式 RL 智能体，使得开发板也可以在自己的“梦境“中学习游玩。

目录：两个网络 · 游戏 · 世界模型 · 智能体 · 量化 · 芯片上运行 · 刷写 · 浏览器 · 复现 · 代码结构

世界模型与策略网络

参数量取决于游戏，有些游戏需要更多参数才能学会可玩的策略，我们尽量把参数量控制在较小范围。在芯片上，两个网络都不使用真实游戏：策略网络只看到世界模型生成的帧，世界模型只收到策略网络的动作。两者都运行在 RP2350 的核心 1 上，核心 0 负责 USB 和显示屏。预训练的策略网络可以替换为一个全新的 stream-x 智能体，它在开发板上从零开始学习。此外，任何按键都会将控制权交给你，所以你也可以亲自体验这个梦境。

游戏

每个游戏都是几百行 numpy 代码，位于 src/tinyworldformer/games/ 中。游戏渲染到一个由 token id 组成的小网格上，一帧必须包含完整的游戏状态，因此隐藏状态必须绘制到帧中：球的朝向编码在其 token 中，Doom-lite 还有一个小地图。

“精确帧“是指在留出的转换数据中，预测帧的每个格子都与真实游戏完全一致的比例。紧跟随机事件之后的帧被排除在外，因为这些帧只能被采样而无法被预测。Doom-lite 的 3D 视图有 768 个格子，哪怕只有一个带阴影的墙面像素出错，整帧就算错误，所以它的数字看起来比实际表现差；游戏状态几乎总是正确的（见检查梦境）。

你可以在终端中玩真实的游戏： uv run python -m tinyworldformer.games.play --game pacman（或 breakout、doomlite、flappy、pong）。

世界模型的工作原理

将帧切分为 token（图块）

帧被切成方形图块：Breakout 中为 2×2 格子，Doom-lite、Flappy Bird 和 Pong 中为 4×4，Pac-Man 中为 6×6。每个游戏产生 36 到 80 个 token。一个图块的嵌入是每个格子对应的可学习向量之和，图块内每个位置有独立的表，因此格子顺序不会丢失。每个 token 还会附加位置向量和动作的向量。

循环 Transformer

模型是一个单独的 Transformer 块（注意力 + ReLU MLP），只有一组权重，被应用 1 到 4 次，每次传递前输入会被加回。每次传递后的输出都用下一帧进行训练，因此在运行时，循环次数可以权衡速度与精度。在 Pico 上为了保持较好的帧率，我们使用 2 次循环。

动画展示了一个留出的 Doom-lite 步骤。经过一次传递后，模型画出了一个不存在的敌人，每一次后续传递都去除了它的一部分：92、35、12，最后 0 个错误格子。右侧的柱状图显示每种循环次数在留出帧上精确匹配的比例。

这里是一个 Doom-lite 步骤在 1、2、3、4 次循环后解码的梦境画面，红框标记了与真实规则不同的格子：

解码

模型为每个格子和每个 token 输出一个分数。将这些分数转化为帧的过程，其重要性几乎不亚于模型本身，最终归纳为三条规则（上方 Pong 的演示中可见）：

1. 背景格子（墙、砖块、豆子、地板）各自独立选择最可能的 token。
2. 单例对象（球、玩家、每个幽灵）通过一次全局选择放置在所有（格子, token）对中，确保每种对象恰好只有一个。随机事件发生后，选择从距最优解 2.2 logits 以内的候选者中采样，因此世界模型会采样出新的发球或重生，而不是对它们取平均。
3. Pac-Man 和幽灵是 2×2 精灵，但只有其中一个格子携带状态（嘴巴编码朝向；幽灵左下角的格子编码其下是否有豆子）。其余三个格子由该格子盖印而来，而非独立预测。

我们实现了规则 2，因为最早的 Breakout 模型会同时画两个球；一次全局选择将精确帧从 91.6% 提升到 99.7%。

检查世界模型

在录制数据上的准确率还不够。智能体会把梦境推进到数据中从未出现的状态，那里的任何错误都可能被利用。因此主要的检查在智能体游玩时进行：从梦境帧中读回游戏状态（games/.py: read），对该状态和智能体的动作应用真实规则，然后将结果与下一梦境帧和奖励进行比较。

使用随附的智能体，在 2 次循环下检查每个游戏 5,000 个梦境步骤（training/evalwm.py --agent）：

训练

首先，智能体游玩真实游戏，每一步被记录为一条示例（帧, 动作, 下一帧, 奖励），每个游戏 1–3 百万条。部分 episode 从不常见的状态开始，比如带孔洞的砖块布局或随机初始位置，因此数据覆盖了更多的游戏状态。

然后世界模型学习预测下一帧，使用逐格交叉熵损失：30k 步训练，前 3–4k 步以 float 运行，之后 int8 舍入开启（QAT）。

int8 权重以源码形式写入 Rust，每个 float 存储其精确的二进制位。我们用 cargo test 验证它们逐位复现了 Python 的整数规格，然后固件就可以在 Pico 上运行了。

在世界模型中训练出的智能体会到达录制数据从未覆盖的状态，因此用该智能体录制真实游戏，并在所有数据上微调世界模型。该智能体在不同游戏中开始利用模型漏洞，以下是对应的一些修复：

• 按“该格子发生了变化“加权损失会偏向模型预测变化，导致它画出两个球和碎裂两块砖。改为按输入帧中事物可能发生的位置加权，加上全局单例选择，修复了该问题。
• Dream 中的 Breakout 智能体在剩余 26 块砖的情况下获得 154 分，因为世界模型在砖块被击中后没有将其从角落移除。Doom-lite 智能体学会了在 5 格距离外狙击敌人。
• 只在梦境中训练的智能体在那里得到 210 分（真实最高分是 56），在真实游戏中只得到 28 分，正如 World Models（Ha & Schmidhuber）所预言的那样。每一轮收集-微调都修复了智能体发现的漏洞。
• Doom-lite 的梦境有时会丢失敌人精灵且不再重新绘制：模型直接复制了上一帧的视图，而不是从小地图重新绘制。噪声增强修复了这一点（train --noise 0.3，同 GameNGen）：输入视图会被随机添加像素，并偶尔丢失精灵。
• d=32 无法吸收更宽泛的 Breakout 数据，d=48 可以。Pac-Man 的 48×48 帧需要 d=64。

智能体：stream-x

智能体使用 Stream AC(λ)（Elsayed, Vasan & Mahmood, 2024），每次从单个转换中学习：没有经验回放缓冲区，没有批量，没有目标网络。资格迹将信用度随时间回传，ObGD 会缩小任何会超出 TD 目标的步长。策略网络和价值网络各有两个 64 单位的隐藏层（Pico 上为 32），带 LayerNorm。输入是从帧中读取的特征向量（training/streamx.py: observe_*，在 Rust 中镜像于 tinyworldformer-rs/src/agent.rs）。例如，Breakout 中球的行、列和朝向；或在 Pac-Man 中，每个方向上到最近豆子的最短路径距离以及每个幽灵的距离。

以下是每个游戏的一次训练运行（真实游戏，500k 步）的六个阶段：未训练、结束时、以及得分分别完成 25%、50%、75% 和 95% 的快照。每个面板都从相同的起点贪心游玩，并每帧跳过几步。“try” 表示到目前为止玩了多少 episode，因此你可以看到失败次数。

Breakout Doom-lite Flappy Bird Pac-Man lite Pong

真实游戏中的贪心回报，10 个 episode 的均值，同一运行的不同快照：

如果没有塑形奖励，Doom-lite 的智能体会待在起始角落反复开火，因为敌人总是会走向它们。--idle-penalty 0.1 会对在最近 20 步内被访问超过 6 次的格子中的每一步扣 0.1 分。位置从小地图读取，因此该惩罚在梦境中同样生效。有了这个惩罚，智能体会巡逻，击杀数翻倍。

tinyworldformer-rs/src/learner.rs 是相同算法的 no_std Rust 实现，反向传播是手写的。它与 200 次 PyTorch 更新的误差在 3.5e-6 以内。在宿主机上，一个全新的智能体在 Rust int8 Flappy Bird 梦境中训练 25k 步后平均每 episode 通过 6 个管道，100k 步后 121 个，250k 步后超过 300 个。

量化与 Pico

世界模型中每个矩阵都是 int8，每个输出行有一个缩放因子（scale = max|w| / 127）。激活值也是 int8，每个张量和每个循环索引各有一个缩放因子，因为不同循环看到的数值范围不同。矩阵乘法将 int8 × int8 求和为 int32（RP2350 上每条 SMLAD 指令执行两次乘加），然后进行一次 float 乘法，乘上两个缩放因子。Softmax 使用 1024 项的 exp 查找表，概率存储为 uint8。归一化、残差和嵌入保持 float32，因为 RP2350 有 FPU。

训练时应用相同的舍入（QAT，使用 straight-through 估计器），因此 int8 模型就是被训练出来的模型。rustexport/intref.py 是 int8 模型的 numpy 实现，每个 float 操作都固定顺序；Rust 引擎逐位复现它，cargo test 会检查每个导出环境的黄金向量。

各游戏资源占用：

Flash 和 RAM 是从固件 ELF 文件（size -A）中测量的，使用 Display Pack 驱动构建，因此包含了代码、USB、显示驱动和设备端学习器。Pico 2 W 有 4 MB Flash 和 520 KB SRAM，其中 512 KiB 在主内存映射中。float32 的嵌入表是最大的部分，也可以改为 int8。

src/tinyworldformer/rustexport/ 仅依赖 numpy 和 torch（有测试保证这一点），因此其他项目可以复用这个导出流程：从 rustexport.qat 构建模型（QLinear、LoopedBlock），在 float 预热后用 ActQuant.enabled = True 训练，用 rustexport.convert.looped_model 转换，然后用 int_model_arrays、mlp_arrays 和 codegen.write_module 写出 Rust 源码。rustexport.Golden 生成 Rust 测试检查的向量。

在芯片上运行

固件（tinyworldformer-rs/firmware）是运行在 embassy 上的 no_std Rust。运行时不做任何内存分配：权重是 Flash 中的 static 数组，每个缓冲区都是 RAM 中固定大小的 static。

核心 1 运行游戏。其循环的一次迭代（dream_loop，位于 firmware/src/main.rs）就是一帧：

1. 策略网络从当前帧读取特征（agent::observe），运行其 MLP 选择一个按键。如果设备端学习已开启，则由学习中的智能体选择；如果你正在按键，则使用你的按键。
2. int8 Transformer 对帧和按键运行 2 次循环（Engine::forward），解码将其分数转化为下一帧：背景逐格 argmax，每个单例只放置一次（随机事件使用 xorshift RNG 采样），然后进行精灵盖印。奖励取奖励头的 argmax；命数和 episode 结束在此基础上计数。
3. 如果学习开启，学习器从这单个转换中更新，然后将其遗忘。
4. 帧通过双槽通道发送给核心 0。如果核心 0 仍在处理上一帧，这一帧会被丢弃，因此核心 1 永远不会等待屏幕。

当你游玩时，核心 1 会减速到可玩的速度（每帧 66–150 ms，取决于游戏）。智能体则以芯片允许的最快速度运行。

核心 0 运行 embassy 异步任务处理其余一切。通过 USB 串口，每帧以彩色 ANSI 文本输出，并附带状态行（模型时间、策略时间、fps、环境和智能体名称），按键也从串口传入。

可选的 Display Pack（--features display）是 SPI 上最高 62.5 MHz 的 320×240 ST7789 屏幕，只重绘与上一帧相比发生变化的格子。其按键每 20 ms 轮询一次。

两个核心只共享帧通道和少量原子变量：当前是谁在玩、按下的键、学习开/关、循环次数。