tinyworldformer：在 Raspberry Pi Pico 2 W 上用世界模型运行游戏

使用 2 个小型神经网络，在 Raspberry Pi Pico 2 W 上同时生成并游玩这款游戏。第一个网络创建游戏：根据当前帧和动作，预测下一帧和奖励。第二个网络游玩该游戏：利用生成的帧并选择动作。

上面的每一帧都由世界模型网络采样，动作则由行动者网络执行。

为什么

我开始对循环式 transformer 感兴趣：它不是堆叠不同的层，而是将同一个 transformer 块使用相同权重应用多次。复用权重很节省内存，所以我想看看循环式世界模型能否小到足以在微控制器上运行。我选择 Pico 2 W，因为它价格低廉且带有 FPU；选择小游戏，是因为我可以直接用 numpy 编写它们的真实版本，并将生成的每一帧与之核对。后来我加入了一个流式 RL 智能体，这样游戏板也能学会在自己的梦境中游玩。

目录：这两个网络 · 这些游戏 · 世界模型 · 智能体 · 量化 · 芯片上运行 · 烧录 · 浏览器 · 复现 · 代码布局

世界模型与行动者

大小取决于游戏，有些游戏需要更多参数才能被学会并变得可玩，我们尽量让参数数量保持足够小。

在芯片上，两个网络都不会使用真实游戏：行动者只能看到世界模型提供的帧，世界模型也只能从行动者那里获得动作。二者都运行在 RP2350 的核心 1 上，核心 0 则负责 USB 和显示屏。

预训练的行动者可以替换为一个全新的 stream-x 智能体，让它在板上从零开始学习。按下任意按钮都会将控制权交给你，因此你也可以亲自游玩这个梦境。

这些游戏

每个游戏都是 src/tinyworldformer/games/ 中数百行 numpy 代码。它会渲染成由 token id 组成的小网格，而每一帧始终包含完整的游戏状态，因此隐藏状态必须被绘制到帧中：球的朝向是其 token 的一部分，Doom-lite 还有一张小地图。

“完全正确的帧”是指在留出的转换中，预测帧的每个单元格都与真实游戏相匹配的比例。随机事件之后紧接着的帧会被排除在外，因为这些帧只能被采样，无法被预测。Doom-lite 的 3D 视图包含 768 个单元格，只要有一个阴影墙像素出现偏差，整个帧就算作错误，因此它的数字看起来比实际表现更差；但游戏状态几乎总是正确的（见 检查梦境）。

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

世界模型如何运作

帧被拆分为词元

帧被切成方形图块：Breakout 中为 2×2 个单元格，Doom-lite、Flappy Bird 和 Pong 中为 4×4，Pac-Man 中为 6×6。因此每个游戏会产生 36 到 80 个词元。图块的嵌入向量是其中每个单元格对应的一个学习向量之和；同时为图块内的每个位置设置单独的表，以免丢失单元格的顺序。每个词元还会加入一个位置向量和一个动作向量。

循环式 Transformer

模型是一个单一的 Transformer 块（注意力加 ReLU MLP），共用一组权重，重复应用 1 到 4 次，每次循环前都会将输入重新加回去。每次循环后的输出都以下一帧为训练目标，因此在运行时，循环次数用于在速度和准确率之间取舍。在 Pico 上，为了保持不错的帧率，我们使用两次循环。

动画展示了一个留出的 Doom-lite 步骤。经过一次循环后，模型会画出一个并不存在的敌人；每多循环一次，就会移除其中一部分，错误单元格数依次为 92、35、12，最后为 0。右侧的柱状条表示每种循环次数能在留出帧中取得完全正确的比例。

下面是一个由梦境生成的 Doom-lite 步骤，在 1、2、3 和 4 次循环后解码，红框标出了与真实规则不同的单元格：

按循环次数统计的完全正确留出帧：

解码

模型会为每个单元格和每个词元输出一个分数。事实证明，将这些分数转换成帧的重要性几乎不亚于模型本身，而这一过程最终采用三条规则（如上文 Pong 中所示）：

1. 每个背景单元格（墙、砖块、颗粒、地面）都独立选取最有可能的词元。
2. 单例（球、玩家和每个幽灵）通过在所有（单元格，标记）组合中进行一次选择来放置，因此每种始终恰好有一个。随机事件发生后，会在 logits 与最佳值相差不超过 2.2 的候选中进行采样，因此世界模型会采样新的发球或重生，而不会将它们平均处理。
3. Pac-Man 和幽灵都是 2×2 的精灵图，但每个精灵图中只有一个单元格承载状态（嘴巴记录朝向；幽灵的左下角单元格记录下方是否有小丸）。另外三个单元格由该单元格复制而来，而不是预测得到。

我们实现了规则 2，因为第一个 Breakout 模型同时画出了两个球；采用一次全局选择后，完全正确的帧占比从 91.6% 提高到了 99.7%。

检查世界模型

仅凭记录数据上的准确率还不够。智能体会将梦境推进到数据中从未出现过的状态，而那里的任何错误都可能成为它可以利用的漏洞。因此，主要检查会在智能体游玩时进行：从想象帧（games/<game>.py: read）中读回游戏状态，将真实规则应用于该状态和智能体的动作，再将结果与下一幅想象帧及其奖励进行比较。

每局游戏 5,000 个想象步，进行 2 轮循环，使用随附的智能体（training/evalwm.py --agent）：

训练

首先，智能体游玩真实游戏，每一步都记录为一个样本（帧、动作、下一帧、奖励），每个游戏有 1–3 million 个样本。有些回合从异常状态开始，例如带有缺口的砖块或随机位置，因此数据覆盖了游戏中更多的情况。

接着，世界模型学习以逐单元交叉熵损失预测下一帧：共 30k 步，前 3–4k 步以浮点数运行，随后启用 int8 舍入（QAT）。

int8 权重被写入 Rust 源代码，其中每个浮点数都以其精确比特形式存储。我们使用 cargo test 对其进行检查，以逐比特复现 Python 整数规范，随后固件在 Pico 上运行。

在世界模型中训练的智能体会到达记录数据中从未出现过的状态，因此用该智能体游玩并记录真实游戏，再使用所有数据微调世界模型。

智能体在不同游戏中开始利用这些漏洞，因此以下是针对它的一些修复：

• 按“这个单元格发生了变化”对损失加权，使模型偏向于预测变化，因此它画出了两个球并打碎了两块砖。改为按输入帧中事物可能发生的位置加权，再加上全局单例选择，就解决了这个问题。
• 一个 Breakout 智能体在梦境中拿到 154 分，却还剩 26 块砖，因为砖块被击中时，世界模型没有移除角落中的砖块。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_*，在 tinyworldformer-rs/src/agent.rs 中以 Rust 镜像实现）。例如，在 Breakout 中是球的行、列和朝向；在 Pac-Man 中则是各个方向到最近豆子的最短路径距离，以及每个鬼魂距离当前位置有多近。

下面展示每个游戏的一次训练运行（真实游戏，500k 步）在六个阶段的情况：未训练、训练结束，以及得分达到全程的 25、50、75 和 95% 时的快照。每个面板都从相同起点按贪心策略游玩，并在每帧跳过若干步。"try" 统计迄今已进行的回合数，因此可以看到崩溃情况。

Breakout

Doom-lite

Flappy Bird

Pac-Man lite

Pong

同一次运行的各快照在真实游戏中的贪心回报，取 10 个回合的平均值：

没有奖励塑形时，Doom-lite 智能体会蹲守在起始角落并开火，因为敌人总会送上门来。对于最近 20 步中被访问超过 6 次的格子，智能体每在其中停留一步，--idle-penalty 0.1 就会扣除 0.1。位置是从小地图读取的，因此该惩罚在梦境内也有效。有了它，智能体会巡逻，并获得两倍的击杀数。

tinyworldformer-rs/src/learner.rs 是在 no_std Rust 中实现的同一算法，反向传播由手写实现。其结果与 200 次 PyTorch 更新的误差不超过 3.5e-6。在主机上，一个新智能体在 Rust int8 Flappy Bird 梦境中经过 25k 步后，每局平均通过 6 根管道；100k 步后为 121 根；250k 步后则达到 300 多根。

量化与 Pico

世界模型中的每个矩阵均为 int8，每个输出行都有一个缩放系数（scale = max|w| / 127）。激活值也都是 int8，每个张量和每个循环索引各有一个缩放系数，因为不同循环看到的数值范围不同。矩阵乘法将 int8 × int8 的乘积累加到 int32 中（在 RP2350 上，每条 SMLAD 指令执行两次乘加运算），然后用两个缩放系数做一次浮点乘法。Softmax 使用一个包含 1024 个条目的 exp 表，概率以 uint8 存储。范数、残差和嵌入保持为 float32，因为 RP2350 带有 FPU。

训练采用相同的舍入方式（使用直通估计器的 QAT），因此这个 int8 模型就是实际训练得到的模型。rustexport/intref.py 是用 numpy 实现的 int8 模型，所有浮点运算均按固定顺序执行；Rust 引擎逐位复现它的结果，而 cargo test 会为每个导出的环境检查黄金向量。

闪存和 RAM 占用根据固件 ELF（size -A）测得。该固件使用 Display Pack 驱动程序构建，因此这些占用包含代码、USB、显示驱动程序和设备端学习器。Pico 2 W 配备 4 MB 闪存和 520 KB SRAM，其中 512 KiB 位于主内存映射中。float32 嵌入表占用的空间最大，也可以使用 int8。

src/tinyworldformer/rustexport/ 只导入 numpy 和 torch（有测试确保这一点），因此其他项目也可以复用这一导出结果：基于 rustexport.qat 构建模型（QLinear、LoopedBlock），经过浮点预热后使用 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。运行时不分配任何内存：权重是闪存中的 static 数组，每个缓冲区都是 RAM 中固定大小的 static。

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

1. 行动者从当前帧（agent::observe）中读取特征，并运行其 MLP 来选择按键。如果启用设备端学习，则改由学习智能体选择；如果你正在按键，则使用你的按键操作。
2. int8 transformer 对帧和按键执行 2 次循环（Engine::forward），解码则将其分数转换为下一帧：背景逐格取 argmax，每个单例只放置一次（随机事件使用 xorshift RNG 采样），然后再贴上精灵图案。奖励为奖励头的 argmax；生命值和回合结束次数则是在此基础上维护的计数器。
3. 如果启用学习，学习器会根据这一次状态转换进行更新，随后将其遗忘。
4. 该帧通过双槽通道发送到核心 0。如果核心 0 仍在处理上一帧，这一帧就会被丢弃，因此核心 1 绝不会等待屏幕。

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

Core 0 运行 embassy 异步任务来处理其余所有工作。通过 USB 串口，每帧都会以带颜色的 ANSI 文本输出，并附带一行状态信息（模型耗时、行动器耗时、fps、环境名称和智能体名称），按键输入也会回传。可选的 Display Pack（--features display）配备一块通过 SPI 连接、最高频率达 62.5 MHz 的 320×240 ST7789 显示屏，并且仅重绘自上一帧以来发生变化的单元格。其按键每 20 ms 轮询一次。两个内核仅共享帧通道和几个原子变量：当前由谁行动、按下的按键、学习开关以及循环次数。

内存占用去向：

每个游戏都能装入 4 MB 闪存和 512 KiB SRAM。Breakout 使用 361 KB 闪存和 299 KB SRAM；体积最大的 Pac-Man 使用 600 KB 和 460 KB。

大部分时间花在 int8 点积上（dot_i8 在 src/lib.rs 中）。一次 32 位加载可取回 4 个权重，SXTB16 将字节 0/2 和 1/3 解包成 16 位半字，SMLAD 用一条指令执行两次乘加。注意力运算将 uint8 概率与 int8 值相乘，因此改用 UXTB16。归一化参数、exp 表查找以及最终的逐输出缩放仍由 FPU 以 float32 处理。在 2 次循环下，一帧需要 2–4M 次乘加（Pac-Man 为 8.2M 次）；在桌面 CPU 上，同样的 Rust 代码推演一帧仅需 0.7 ms（Flappy Bird）至 1.5 ms（Breakout）。我尚未在 Pico 本机上测量帧率。

开发板启动时，一个全新智能体开始在设备端学习。通过 USB 按下 l，或在 Display Pack 上按住 B 键 1 秒，即可切换到预训练的行动器；再次执行即可切回。

烧录游戏和智能体

每次导出都并列放在 crate 中：世界模型位于 tinyworldformer-rs/models/envs/<env>/，智能体位于 models/agents/<game>/<agent>/。构建时需要选择两者。

cd tinyworldformer-rs
./flash.sh --list                                            # what is exported
./flash.sh --env doomlite --agent doom_agent_roam_p0.3       # build and flash (hold BOOTSEL while plugging in)
./flash.sh --env breakout --display                          # with the Display Pack
./flash.sh --env pong --host                                 # play it in this terminal
./flash.sh --env flappy --default                            # make it the default for plain cargo
TWF_ENV=doomlite TWF_AGENT=doom_agent_roam_p0.1 cargo test --release   # golden tests for that pair

如果没有 --agent，且 models/default 中列有该环境所对应游戏的智能体，就使用该智能体；否则使用名称排在最前的智能体。为其他游戏准备的智能体会被拒绝，并列出有效选项。屏幕上的状态行会显示当前运行的环境和智能体。

USB 按键：Breakout 为 a/d/s，Doom-lite 为 w/a/d/space/s，Flappy 为 space，Pac-Man 为 w/a/s/d，Pong 为 w/s。随后，g 将控制权交还智能体，1–4 设置循环次数，l 切换设备端学习，r 重新启动学习。

在浏览器中

同一个 no_std int8 引擎也可编译为 WebAssembly（tinyworldformer-rs/web/）。该模块没有导入，也不使用 wasm-bindgen，并会运行固件的循环：推演、智能体，以及即时进行 stream-x 学习。一个静态页面会在画布上绘制这些帧。

rustup target add wasm32-unknown-unknown    # once
cd tinyworldformer-rs/web
./build.sh                                   # one module per env into www/ (or ./build.sh pong flappy)
python3 -m http.server -d www                # open http://localhost:8000, ?env=pacman picks a game

使用相同种子时，wasm 构建会与原生构建逐帧推演出相同的回合，具体速度取决于游戏，每帧为 1–5 ms。每个模块都自带其游戏的默认智能体；TWF_AGENT=... ./build.sh <env> 可选择另一个。按键与 Pico 上相同，另外还有方向键和屏幕按钮；f 会以最快速度运行推演，这在学习过程中很有帮助（l）。

复现

uv sync
G=pong   # breakout, doomlite, flappy, pacman, pong
uv run python -m tinyworldformer.training.data --game $G --transitions 2000000 --varied 0.3   # -> data/$G
uv run python -m tinyworldformer.training.train --data data/$G --d 48 --heads 3 --ff 96 --out models/${G}_wm.pt
uv run python -m tinyworldformer.training.streamx --env real --game $G --out models/${G}_agent.pt
uv run python -m tinyworldformer.training.evalwm --ckpt models/${G}_wm.pt --agent models/${G}_agent.pt   # check the dream
uv run python -m tinyworldformer.training.streamx --env dream --wm models/${G}_wm.pt --out models/${G}_dream_agent.pt
uv run python -m tinyworldformer.training.export --ckpt models/${G}_wm.pt
uv run python -m tinyworldformer.training.streamx --out models/${G}_agent.pt --export tinyworldformer-rs
cd tinyworldformer-rs && TWF_ENV=$G cargo test --release && TWF_ENV=$G cargo run --release

各游戏不同的设置：

代码布局