TWM：把世界模型换成 Transformer，10 万步玩转 Atari
🎯 为什么重要：把「记忆」从「压缩」改成「直查」
Atari 100k 是个「穷日子」benchmark：每款游戏只能玩 10 万次交互（≈2 小时真人游戏时长，比标准 200M 帧训练少 500 倍）。怎么用这点数据学到能打的策略？
Transformer 的逻辑完全不同——不压缩，直接查。每一步都能 attend 到任意过去时刻的潜变量、动作、奖励。
论文 Fig.1 那张架构图就是整篇文章的核心：三个模态（latent $z$ / action $a$ / reward $r$）的序列，全摊开喂进 Transformer-XL，输出一个 hidden $h_t$，再用 MLP 预测奖励、终止、下一个潜变量。
💡 核心思想：三模态全摊开 + Transformer-XL 做骨架
世界模型的形式化定义很简洁（论文 §2.1）：
$$h_t = f_\psi(z_{t-\ell:t},\; a_{t-\ell:t},\; r_{t-\ell:t-1})$$
$$\hat r_t,\;\hat\gamma_t,\;\hat z_{t+1} = \text{MLP}(h_t)$$
三件东西一起进 Transformer：潜变量 $z$（32 个 32 类的分类分布，照搬 DreamerV2）、动作 $a$、奖励 $r$。三个模态各走各的线性 embedding，论文照搬 Decision Transformer 的做法。然后用 Transformer-XL 跑因果自注意，输出对应「动作位置」的 token 当 $h_t$（其他两个模态的输出丢弃）。
🛠️ 五个工程关键：让 Transformer 真的在 100k 数据上赢
光换 RNN 为 Transformer 还赢不了——论文 §2 给了五个配套设计，每个都对应一个具体的「坑」：
| 设计 | 解决什么 | 做法 |
|---|---|---|
| ① Transformer-XL 骨架 | vanilla Transformer 推理慢（每步重算 KV） | XL 的相对位置编码+片段递归——训练并行、推理缓存。还去掉了对「绝对时刻」的依赖 |
| ② 反馈奖励进 transformer | 模型不知道「之前发的奖励」是什么 | 把预测的 $\hat r$ 作为下一步的输入 token——模型能对自己的奖励做出反应（论文消融确认增益显著） |
| ③ 解耦的 balanced cross-entropy | DreamerV2 的 balanced KL 把 entropy / cross-entropy 绑死 | 显式拆成 $\alpha_1 H + \alpha_2 H_{ce}$ 两项，分别调权——一致性 + 熵正则可控 |
| ④ 阈值化熵损失 | 每款游戏最优熵不同，调参累 | 归一化到 $[0,1]$，只在熵低于阈值 $\Gamma$ 时惩罚（hinge loss）——跨游戏用同一超参 |
| ⑤ Balanced dataset sampling | 数据集慢慢长大，均匀采样偏向旧数据 → 过拟合早期 | 用访问计数 $v_i$ 算 softmax $\propto \exp(-v_i/\tau)$，新数据被过采——τ 控制强度 |
世界模型派的痛点是「推理时还要跑世界模型」——Dreamer 系列、IRIS 都得在部署时跑 RNN/Transformer，慢且贵。TWM 把 policy 训练成只吃 latent $z_t$（再加 frame stack 提供短时记忆），部署时完全不需要 Transformer。模型规模再大也不影响推理速度。这是论文卖点之一。
📊 结果：Atari 100k 上四个聚合指标全 SOTA
论文用 Agarwal et al. 2021 的统计方法（median / IQM / mean / optimality gap + 95% 置信区间），跑 5 runs/game，对比 5 个强基线：
| 方法 | Normalized Mean ↑ | Normalized Median ↑ | 类别 |
|---|---|---|---|
| SimPLe | 0.332 | 0.134 | model-based（视频预测） |
| DER | 0.350 | 0.189 | model-free（Rainbow 变体） |
| CURL | 0.261 | 0.092 | model-free + 对比学习 |
| DrQ(ε) | 0.465 | 0.313 | model-free + 数据增广 |
| SPR | 0.616 | 0.396 | model-free + 自监督一致 |
| TWM（本文） | 0.956 | 0.505 | model-based（transformer 世界模型） |
消融研究（论文 §3.3）确认了设计 ②（奖励反馈）和 ⑤（balanced sampling）各自带来的提升——尤其 balanced sampling，把训练时间从「一半花在前 19K 数据」纠正到「大致均衡」，对低数据 regime 至关重要。
🌐 在领域中的位置
TWM 是「把 Transformer 引入想象式世界模型」的早期典范，与 IRIS（Micheli et al. 2022）同期。它确立了「训练时用重型 Transformer，推理时丢弃」的范式，被后续机器人/视频世界模型广泛继承。详见 DreamerV3、IRIS。
❓ 常见误区
TWM 是 model-based RL 吗？
是，但很特别：它学一个世界模型，但 policy 在latent 想象空间里用 model-free 的 actor-critic（A2C + GAE）训练。所以是「model-based 学习，model-free 推理」。
100k 交互很多吗？
反直觉地少。100k 交互 = 400k 帧（4 帧跳帧）≈ 2 小时游戏时间。主流 DQN/PPO 通常用 200M 帧（500 倍），所以这玩意儿测的是「数据效率」而非绝对性能。
它为什么不在 Atari 50M（满血版）上比？
论文没做。TWM 的核心卖点是「低数据 regime」——10 万步下大幅领先。在 5000 万步下，能否超过 DreamerV2/MuZero 是另一个问题（后续 IRIS 主要仍聚焦 100k；Transformer 路线在长训练 regime 的验证待核）。
「Transformer」是真不可或缺，还是其他 trick 才是关键？
论文消融没有单独拆 Transformer 本身，但拆了奖励反馈 + balanced sampling——两者各自都带来显著提升。可以理解为「Transformer 让历史信息到位，sampling 让训练数据到位，两个一起才把 100k 喂饱」。