在树莓派Pico 2 W上运行Tiny World Model玩游戏的工程实践

在 Raspberry Pi Pico 2 W 上使用两个小型神经网络同时生成游戏并进行游戏。第一个网络负责创建游戏：给定当前帧和操作，它预测下一帧和奖励。第二个网络负责进行游戏：它使用生成的帧并选择一个操作。上述每一帧均由世界模型网络采样，而操作则由演员网络执行。## 为什么我对循环 Transformer 产生兴趣：即多次应用同一个 Transformer 块（共享权重），而不是堆叠不同的层。复用权重在内存方面非常节省，因此我想看看循环世界模型是否足够小，能够在微控制器上运行。我选择 Pico 2 W 是因为它价格便宜且具备 FPU（浮点运算单元），并且选择小型游戏是因为我可以将真实版本直接写入 numpy，并将每一张生成的帧与它们进行比对检查。后来我添加了一个流式强化学习（RL）智能体，使开发板也能在其自身的“梦境”中学习游玩。内容：两个网络 · 游戏 · 世界模型 · 智能体 · 量化 · 芯片上运行 · 烧录 · 浏览器 · 复现 · 代码布局 ## 世界模型与演员 | 网络 1：世界模型 | 网络 2：演员 任务 | 帧 + 操作 → 下一帧 + 奖励 | 帧 → 操作 架构 | 循环 Transformer，q int8 | 具有两个隐藏层的 MLP，float32 大小 | 26k–131k 参数 | 9k–16k 参数 每帧成本 | 2–8M 次乘加运算 | 9k–16k 次乘加运算 训练方式 | 在记录的真实游戏上（如何） | 使用 stream-x，在真实游戏中或梦境内（如何） 大小取决于游戏，某些游戏需要更多参数才能学会使其可玩，我们试图保持参数量足够小。在芯片上，没有任何网络会使用真实游戏：演员只能看到世界模型生成的帧，而世界模型只能从演员获取操作。两者均在 RP2350 的核心 1 上运行，而核心 0 处理 USB 和显示。预训练的演员可以被一个从头开始在板上学习的最新 stream-x 智能体替换。任何按键按下都会将控制权交给你，这样你就可以亲自游玩这个“梦境”。## 游戏 每个游戏由 src/tinyworldformer/games/ 中几百行 numpy 代码组成。它渲染为一个小型的 token id 网格，且每一帧始终包含完整的游戏状态，因此隐藏状态必须被绘制到帧中：球的朝向是其 token 的一部分，Doom-lite 拥有小地图。game | frame | rules | random event | world model | exact frames, 2 loops Breakout | 20×16 | paddle, ball, 56 bricks, 5 lives | the serve after a miss | 26k params, 4.3M MACs | 99.99% Doom-lite | 32×32 | raycast first-person view of an 8×8 maze; an enemy walks shortest paths to you; you shoot | the enemy's respawn | 55k params, 4.2M MACs | 83.6% (see below) Flappy Bird | 24×24 | gravity, flaps, scrolling pipes | the next gap height | 46k params, 2.0M MACs | 100% Pong | 20×32 | you against a CPU paddle that reacts late, first to 5 | the serve after a point | 46k params, 2.3M MACs | 99.45% Pac-Man lite | 48×48 | 23×21 maze, 249 pellets, two ghosts on shortest paths, 3 lives | the ghosts' respawn | 131k params, 8.2M MACs | 86.4% "Exact frames" is the share of held-out transitions where every cell of the predicted frame matches the real game.随机事件发生后的帧被排除在外，因为这些只能采样，无法预测。Doom-lite 的 3D 视图有 768 个单元格，单个阴影墙壁像素出错即导致整帧被判错，因此其数值看起来比实际表现更差；游戏状态几乎总是正确的（参见 checking the dream）。你可以在终端中玩真实游戏：uv run python -m tinyworldformer.games.play --game pacman（或 breakout, doomlite, flappy, pong）。## How the world model works ## Frames split into tokens The frame is cut into square patches: 2×2 cells in Breakout, 4×4 in Doom-lite, Flappy Bird and Pong, 6×6 in Pac-Man.这使每个游戏的 token 数量在 36 到 80 之间。一个 patch 的嵌入是每个单元格的单个学习向量的总和，每个 patch 内部的位置都有单独的表，以免丢失单元格顺序。每个 token 还会获得位置向量和动作向量。## Looped Transformer The model is a single transformer block (attention plus a ReLU MLP) with one set of weights, applied 1 to 4 times, with the input added back before every pass.每次迭代后的输出都会与下一帧进行训练对比，因此在运行时，循环次数是在速度与精度之间进行权衡。在 Pico 上为了保持合理的帧率，我们使用两次循环。动画展示了一个保留的 Doom-lite 步骤。经过一次迭代后，模型会画出并不存在的敌人，而每一次进一步的迭代都会消除其中的一部分：错误单元格数分别为 92、35、12，然后是 0。右侧的条形图显示了每个循环次数恰好正确预测的保留帧比例。以下是经过 1、2、3 和 4 次循环解码出的一个梦境中的 Doom-lite 步骤，红色方框标记了与真实规则不同的单元格：按循环次数划分的精确保留帧比例：循环次数 | Breakout | Doom-lite | Flappy Bird | Pong | Pac-Man 1 | 99.68% | 59.8% | 99.99% | 98.47% | 79.6% 2 | 99.99% | 83.6% | 100% | 99.45% | 86.4% 3 | 100% | 85.9% | 100% | 99.60% | 85.7% 4 | 99.96% | 84.7% | 100% | 99.53% | 84.4% ## 解码过程 模型为每个单元格和每个 token 输出一个分数。将这些分数转换为帧的过程竟然与模型本身同样重要，最终形成了三条规则（如上所示，以 Pong 为例）：- 背景单元格（墙壁、砖块、豆粒、地板）各自独立地选择最可能的 token。- 单体对象（球、玩家、每个幽灵）在所有 (单元格, token) 对中通过单一选择进行放置，因此每种对象始终恰好存在一个。在随机事件发生后，从距离最佳选项 2.2 logits 范围内的候选项中进行采样，这样世界模型会采样新的发球或重生状态，而不是对它们进行平均。- Pac-Man 和幽灵是 2×2 的精灵图，但每个精灵只有一个单元格携带状态（嘴巴单元格持有方向信息；幽灵的左下角单元格持有下方是否有豆粒的信息）。其余三个单元格是从中冲压而成的，而非预测得出。我们实现了规则2，因为最初的Breakout模型会同时抽取两个球；一次全局选择从91.6%到99.7%的帧数中获取了精确帧。## 检查世界模型在记录数据上的准确性还不够。智能体会将梦境驱动至数据从未出现过的状态，而那里的任何错误都是它可以利用的漏洞。因此主要检查是在智能体游玩时运行：从梦境帧中读取回游戏状态（games/.py: read），对该状态和智能体的动作应用真实规则，并将结果与下一个梦境帧及奖励进行比较。每局游戏5,000个梦境步骤，循环2次，使用发布的智能体（training/evalwm.py --agent）：游戏 | 错误的游戏状态 | 错误的奖励 | 每步错误的单元格 Breakout | 未读取回；0帧包含不可能出现的标记 | 6（在未打破砖块的情况下支付） | – Doom-lite | 1 | 0 | 1.50 Flappy Bird | 0 | 0 | 0.00 Pong | 1 | 0 | 0.00 Pac-Man | 742 | 38 | 0.39 ## 训练首先，智能体玩真实游戏，每一步都被记录为一个示例（帧、动作、下一帧、奖励），每局游戏1–3百万条。一些剧集从不寻常的状态开始，例如有洞的砖块或随机位置，因此数据覆盖了更多的游戏内容。然后，世界模型通过逐单元格交叉熵损失来学习预测下一帧：30k步，前3–4k步以float运行，随后启用int8舍入切换（QAT）。int8权重被写入Rust源代码，每个float都以其精确位存储。我们通过cargo test进行检查，以按位重现Python整数规范，然后固件在Pico上运行。在世界模型中训练的智能体会达到记录数据从未出现过的状态，因此用该智能体记录真实游戏，并在所有内容上微调世界模型。不同游戏中的智能体被发现正在利用这些漏洞，因此对其进行了以下修复：