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

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

上方的每一帧都由世界模型网络采样生成，动作则由演员网络给出。

为什么

我对循环式 transformer 产生了兴趣，即同一个 transformer 块以相同权重被多次应用，而不是堆叠不同的层。复用权重在内存上开销很低，所以我想看看循环式世界模型能否小到在微控制器上运行。我选择 Pico 2 W 是因为它便宜且带有 FPU，选择小游戏是因为我可以直接把真实版本写成 numpy 代码，并逐帧对照检查每个生成的帧。后来我加入了一个流式 RL 智能体，让开发板也能在自己的梦境中学习游玩。

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

世界模型与演员

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

在芯片上，两个网络都从不使用真实游戏：演员只看到世界模型生成的帧，而世界模型只从演员获取动作。两者都在 RP2350 的 core 1 上运行，core 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）。

世界模型如何工作

帧被拆分为 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 上为了保持可接受的帧率，我们使用两次循环。

动画展示了一个留出的 Doom-lite 步骤。一次传递后，模型画出了一个并不存在的敌人，之后每次传递都会消除它的一部分：92、35、12，然后是 0 个错误单元格。右侧的条形图是每个循环次数完全正确的留出帧所占比例。

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

按循环次数统计的精确留出帧：

解码

模型为每个单元格和每个 token 输出一个分数。将这些分数转换为帧的重要性几乎与模型本身相当，最终归结为三条规则（上方以 Pong 展示）：

1. 背景单元格（墙壁、砖块、豆子、地板）各自独立地取最可能的 token。
2. 单例（球、玩家、每个幽灵）在所有（单元格，token）对上进行一次选择来放置，因此每种总是恰好有一个。随机事件之后，选择会在与最佳分数相差 2.2 logits 以内的候选中采样，因此世界模型会采样新的发球或重生，而不是对它们取平均。
3. Pac-Man 和幽灵是 2×2 精灵，但每个只有一个单元格携带状态（嘴巴持有朝向；幽灵的左下单元格持有其下方是否有豆子）。其他三个单元格是从它盖章而来，而不是预测的。

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

检查世界模型

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

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

训练

首先，智能体玩真实游戏，每一步都被记录为一个样本（帧、动作、下一帧、奖励），每局 100 万到 300 万个。有些回合从异常状态开始，比如带洞的砖块或随机位置，这样数据能覆盖游戏的更多情况。

然后世界模型学习用逐单元交叉熵损失预测下一帧：3 万步，其中前 3–4 千步用浮点运行，之后开启 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_*，在 Rust 中镜像为 tinyworldformer-rs/src/agent.rs）。例如，Breakout 中球的行、列和方向，或者在 Pac-Man 中，对每个方向，到最近豆子的最短路径距离以及每个幽灵有多近。

下面是每个游戏的一次训练运行（真实游戏，50 万步），分为六个阶段：未训练、结束，以及分数覆盖 25%、50%、75% 和 95% 时的快照。每个面板从相同起点贪婪地玩，并且每帧跳过几步。“try”统计到目前为止玩过的回合数，所以你可以看到崩溃。

Breakout

Doom-lite

Flappy Bird

Pac-Man lite

Pong

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

如果没有塑造机制，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 步后平均每局通过 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 为每个导出的环境检查黄金向量。

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）构建模型，在浮点预热后用 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 运行游戏。其循环的一轮（firmware/src/main.rs 中的 dream_loop）就是一帧：

1. actor 从当前帧读取特征（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 文本形式输出，并附带一行状态信息（模型时间、actor 时间、fps、环境与 agent 名称），同时接收按键输入。可选的 Display Pack（--features display）是一块 320×240 的 ST7789，通过 SPI 以最高 62.5 MHz 运行，并且只重绘自上一帧以来发生变化的单元格。其按键每 20 ms 轮询一次。两个核心只共享帧通道和少量原子变量：谁在播放、按下的键、学习开/关以及循环计数。

内存去向：

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

大部分时间花在 int8 点积上（src/lib.rs 中的 dot_i8）。一次 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 本身的帧率。

开发板上电时会启动一个在设备上学习的新 agent。通过 USB 按 l，或在 Display Pack 上按住 B 键 1 秒，可切换到预训练的 actor，再次操作则切换回来。

烧录游戏与 agent

每个导出都并排放在 crate 中：世界模型在 tinyworldformer-rs/models/envs/<env>/，agent 在 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 中列出的该环境游戏的 agent（如果存在），否则按名称取第一个。为不同游戏提供的 agent 会被拒绝，并列出有效的 agent。屏幕上的状态行显示正在运行的环境和 agent。

USB 按键：Breakout a/d/s，Doom-lite w/a/d/space/s，Flappy space，Pac-Man w/a/s/d，Pong w/s。然后 g 交还给 agent，1–4 设置循环计数，l 切换设备端学习，r 重新启动它。

在浏览器中

同样的 no_std int8 引擎也可编译为 WebAssembly（tinyworldformer-rs/web/）。该模块没有导入，也没有 wasm-bindgen，它运行固件的循环：做梦、agent 以及即时流式学习。一个静态页面在 canvas 上绘制帧。

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，具体取决于游戏。每个模块携带其游戏的默认 agent；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

每个游戏不同的设置：

代码布局