NanoDreamer：通过想象学会行动

来源： https://www.benjamin-bai.com/projects/nanodreamer Benjamin Bai，2026年10月 GitHub (https://github.com/benyebai/NanoDreamer)

简介 (https://www.benjamin-bai.com/projects/nanodreamer#nanodreamer-title) 能有多难呢？ (https://www.benjamin-bai.com/projects/nanodreamer#how-hard-could-it-be) 呃，它没跑通 (https://www.benjamin-bai.com/projects/nanodreamer#uh-it-did-not-work) 改变模型的学习方式 (https://www.benjamin-bai.com/projects/nanodreamer#changing-how-the-model-learned) Actor-Critic 检查 (https://www.benjamin-bai.com/projects/nanodreamer#actor-critic-check) 借鉴一点灵感 (https://www.benjamin-bai.com/projects/nanodreamer#taking-a-little-inspiration) 嗯，修好了吗？ (https://www.benjamin-bai.com/projects/nanodreamer#did-it-fix-it) 最终方案 (https://www.benjamin-bai.com/projects/nanodreamer#final-draft) 我们等待已久的时刻 (https://www.benjamin-bai.com/projects/nanodreamer#the-moment-of-truth) 经验总结与下一步 (https://www.benjamin-bai.com/projects/nanodreamer#takeaways-and-next-steps) 补充说明 (https://www.benjamin-bai.com/projects/nanodreamer#additional-notes) 从零构建 (https://www.benjamin-bai.com/projects/nanodreamer#foundations)

最近我开始接触世界模型和具身智能，而且因为我热爱电子游戏，Dreamer 自然吸引了我的注意力。但我有一个问题——我没有 GPU。

那我做了什么呢？自己造一个足够小的 Dreamer，小到可以在我的 MacBook 上训练。而且因为我想理解每一个组件，我从最底层开始，自己实现了自动求导引擎、张量库、神经网络层、优化器和 Transformer。

经过大量的调试，NanoDreamer 诞生了！一个 75 万参数的世界模型和 actor-critic，学会了玩全部 5 个 MinAtar 游戏。actor 完全通过想象力训练，模型小到可以在我的 M5 上运行和训练！它从零开始，仅通过自身游戏经验进行训练，不依赖专家演示或针对特定游戏的策略，所有游戏使用相同的架构和训练设置，只是权重各自独立。

下面是 NanoDreamer 在每款游戏训练 50 万步后的表现（在 M5 Pro 上约需 3 小时）。

NanoDreamer 玩 Space Invaders

NanoDreamer 正在玩 Space Invaders。怎么玩

沿底部移动你的炮台，射击下落的外星人，躲避它们的子弹。被外星人或敌方子弹击中即游戏结束。

屏幕上有什么

• 你的炮台——位于底部一行
• 外星人——颜色表示它们的移动方向
• 你的子弹——向上飞行
• 敌方子弹——向下飞行

训练奖励

每击中一个外星人 +1 分。移动、开火和多存活一步不会获得额外分数。

在这篇文章中，我将展示让 Dreamer 以极低成本运行所经历的 bug、实验和架构改动。如果你是机器学习新手，我还单独写了一篇基础知识指南 (https://www.benjamin-bai.com/projects/nanodreamer/foundations)，包含我自己的实现和心得！

能有多难呢？

我从阅读 DreamerV3 和 DreamerV4 开始，试图搞清楚到底需要构建什么。这个想法看起来很简单：一个世界模型学习游戏的运行方式，一个actor在预测中选择动作，一个critic估计更长远的奖励。以下是我第一次尝试的架构。

第一版 NanoDreamer 架构

帧数据经过编码器和历史信息进入因果 Transformer 及转移头。Actor 提供所选动作。Critic 估计未来回报。训练摘要分别展示重构、预测和行为目标。绿色标记点标识了可运行版本中的改动。

01 / 表征与预测

解码器 MLP · 重构损失 Actor 读取 zt；后期为 [zt, Δzt] 帧 编码器 MLP · 预训练后冻结 隐变量历史 Transformer 因果 · 包含当前动作 下一隐状态 z′ = z + Δz 奖励 标量 · MSE 延续 概率 · 二元交叉熵 预测损失 组合学习信号 Critic 标量 V（actor 输入）

02 / 行为训练

Actor 目标 −log π(a | state) × 优势 回报目标 奖励 + critic · λ-回报 Critic 目标 MSE(V, 想象回报)

点击任一组件以查看其解释。

我从一个 MLP 自编码器开始，将每帧压缩为 64 维隐嵌入。当它能够重构屏幕后，我将其权重冻结。

一个小型因果 Transformer读取隐嵌入，以及动作和位置嵌入。它的三个预测头分别预测下一个隐状态、奖励和游戏是否继续。Actor 和 critic 都是 MLP，输入是当前隐状态与上一帧变化量的拼接。

然后，actor 通过 5 步想象展开进行训练——采样一个动作，将每个预测的隐状态反馈回模型。我们将预测奖励与 critic 对未来回报的估计结合，给 actor 的动作提供反馈。

在我脑海中，这一切在理论上都行得通。现在只需验证它是否真的能工作。

第一版方案内部

架构、损失与想象回报

首先，用 64 个数字描述游戏

一个小型 MLP 自编码器将每个游戏帧压缩为 64 个数字，然后尝试重构它。

大多数格子是空的。因此模型可以通过预测一个几乎空白的屏幕来“答对“许多格子。我使用了加权二元交叉熵，让漏判已占用格子的代价更高。

点开任意公式可查看符号含义。

\mathcal{L}_{\text{reconstruction}} = -\operatorname{mean}_i \left[ w_i x_i \log \hat{x}_i + (1 - x_i) \log(1 - \hat{x}_i) \right]

每个 x_i 是某个通道中某个格子的真实占用状态，取值为 0 或 1。其预测值 \hat{x}_i 是一个概率。带帽子（hat）表示预测值，w_i 根据该格子在其通道中的稀有程度给予已占用格子更高的权重，mean 表示在 batch 上对误差取平均。如果一个格子被占用，其损失为 -w_i \log \hat{x}_i。预测 0.9 约产生 0.105 w_i 的损失，而预测 0.1 约产生 2.303 w_i。自信的错误会受到更大的惩罚。另一项处理空格子。

之后我冻结了自编码器。Transformer 只能利用那 64 个数字保留下来的信息工作。

然后，预测动作的效果

Transformer 读取 8 步的隐状态和动作历史。它有两个因果 block，宽度为 64，4 个注意力头。

三个头读取输出。下一个隐状态头预测对当前隐状态的修正量。奖励头预测一个数值，延续头预测当前回合继续的概率。

\hat{z}_{t+1} = z_t + \Delta z_t, \quad \mathcal{L}_{\text{next latent}} = \operatorname{MSE}(\hat{z}_{t+1}, E(x_{t+1}))

\mathcal{L}_{\text{reward}} = \operatorname{MSE}(\hat{r}_t, r_t), \qquad \mathcal{L}_{\text{continuation}} = \operatorname{BCE}(\hat{c}_t, c_t)

z_t 是当前 64 维隐状态，\Delta z_t 是预测的变化量。E 是冻结的编码器，因此 E(x_{t+1}) 是真实下一帧的编码。MSE 表示均方误差：将每个预测值减去目标值，对差值取平方，然后在隐维度和采样的转移上取平均。

r_t 是真实奖励，\hat{r}_t 是其预测值。MSE 计算平方误差的平均，因此对奖励 1 预测 0.2 会得到 (0.2 - 1)^2 = 0.64。c_t 在回合继续时为 1，结束时为 0。BCE 表示二元交叉熵。它惩罚对延续概率的自信错误预测，正如重构 BCE 惩罚错误的占用概率预测一样。

预测帧损失将预测的隐状态解码后与真实下一帧进行比较。发生变化的格子获得更高的权重，因此一个微小的移动物体更难被忽略。我将四项损失组合如下：

\mathcal{L}_{\text{world}} = \mathcal{L}_{\text{next latent}} + 2\mathcal{L}_{\text{reward}} + 0.2\mathcal{L}_{\text{continuation}} + 0.5\mathcal{L}_{\text{predicted frame}}

这些系数是手动设定的平衡权重。相对于下一个隐状态损失，2 提高了奖励的贡献，而 0.2 和 0.5 降低了延续和预测帧的贡献。这些是原型设定，并非从 Dreamer 借鉴，也未通过受控实验验证为最优。由于各项损失的量级不同，2 并不意味着奖励恰好是两倍重要。

接下来，让 actor 练习

我从回放缓冲区中的真实历史开始。Actor 选择一个动作，冻结的模型预测下一个隐状态、奖励和延续，actor 再选一个动作。如此重复五次想象中的动作。

Critic 估计五步窗口之外的奖励。假设回合继续，简单的回报就是将预测奖励加到该最终估计之上。

G_0 = \hat{r}_0 + \gamma \hat{r}_1 + \gamma^2 \hat{r}_2 + \gamma^3 \hat{r}_3 + \gamma^4 \hat{r}_4 + \gamma^5 V(s_5)

G_0 是起始状态的回报目标。每个 \hat{r} 是世界模型预测的奖励。\gamma 对更远未来的奖励进行折扣。V(s_5) 是 critic 对第五个想象动作之后奖励的估计，它使展开的结束不会被当作游戏结束。这个简化公式假设整个过程中回合都在继续。

实际脚本使用了 λ-回报，将下一个 critic 估计与更长的回报混合。我使用 \gamma = 0.99，\lambda = 0.95。这里 s 是 actor-critic 的输入状态。

G_t^\lambda = \hat{r}_t + \gamma \hat{c}_t \left[ (1 - \lambda) V(s_{t+1}) + \lambda G_{t+1}^\lambda \right], \qquad G_5^\lambda = V(s_5)

带 λ 上标的 G_t 是决策 t 处的回报目标。下一个预测奖励被加到下一个 critic 估计和更长回报的混合值上。\lambda 控制这个混合比例，\gamma 对未来进行折扣，\hat{c}_t 是延续概率。从 G_5 = V(s_5) 开始反向计算。延续概率为零时会去除结束之后的未来奖励。

每个决策都有自己的回报目标。我减去 critic 的预期值来计算优势。正优势使采样到的动作概率增大，负优势使其减小。

A_t = \operatorname{sg}(G_t^\lambda - V(s_t)), \qquad \mathcal{L}_{\text{actor}} = -\mathbb{E}_t[\log \pi(a_t \mid s_t) A_t]

\mathcal{L}_{\text{critic}} = \mathbb{E}_t\left[(V(s_t) - \operatorname{sg}(G_t^\lambda))^2\right]

A_t 是优势，即回报减去 critic 的预期。\pi(a_t \mid s_t) 是 actor 在该状态采样到的动作的概率。将其负对数概率乘以正优势会鼓励该动作，而负优势则抑制它。\mathbb{E} 表示在采样决策上取平均。\operatorname{sg} 表示停止梯度（stop-gradient），即在此更新期间优势作为固定的学习信号。

V(s_t) 是 critic 的估计值，带 λ 上标的 G_t 是回报目标。它们的平方差衡量预测误差。\mathbb{E} 在采样决策上取平均。\operatorname{sg} 表示停止梯度，使 critic 向固定目标学习，而不是同时改变比较的两侧。

这些是核心损失。一个熵奖励鼓励探索。目标被 detach，actor 的梯度通过动作概率流动，而非通过世界模型。

但拟合想象回报并不意味着它们就是正确的。之后我用真实游戏经验的折扣回报验证了 critic。

第一版方案详情到此结束

呃，它没跑通

第一次训练运行了大约 20,000 步。它甚至不开枪……完全是一场灾难。在 30 局游戏中，它平均只有 0.03 分。即使是随机策略也能拿到 4.3 分。

原始 Space Invaders actor 只会左右移动不开火，最后以零分告终。

我的第一个疑问是：模型预测的奖励是否是 actor 需要从中学习的？我在真实游戏和模型中回放了相同的动作序列，然后对比实际发生了什么和模型预测了什么。

以下是我从实验中取的一个例子。

模型的奖励预测偏差很大。模型在第四步预期约 1.19 分，但实际的射击在第五步才得了一分，而此时模型预测几乎什么都没有。它甚至在什么都没得到的时候预测了微小的负奖励。