一句话定位
S4WM 是首个兼容”可并行状态空间模型”（Parallelizable SSM，PSSM，涵盖 S4 及其变体 S5 等）的世界模型骨架；论文以此为支点，在专门设计的、序列长度最长到 2000 步的四类记忆任务上，对 RNN（RSSM-TBTT）、Transformer（TSSM-XL，即 Transformer-XL 版 transdreamer）与 S4（S4WM）三大骨架做了首次系统对比，发现 S4WM 在长期想象、上下文相关召回、奖励预测、记忆推理四项能力上全面超过另外两者，训练速度也最快（NeurIPS 2023；作者 Fei Deng 时属 Rutgers University，Junyeong Park 与通讯作者 Sungjin Ahn 属 KAIST）。
背景与定位
Model-Based RL（MBRL）的世界模型骨架长期由 RNN 主导——以 dreamer-v3 一脉沿用的 RSSM（Recurrent State-Space Model，GRU 为核心）为代表——但 RNN 存在梯度消失、长程记忆容量有限的问题。Transformer 凭借自注意力直接访问任意历史步、缓解梯度消失，transdreamer（论文中称为 TSSM）率先把 Transformer 引入世界模型骨架，但二次复杂度限制了其可处理的序列长度（原始 TSSM 只在 ~100 步序列上评测，同类工作 [41, 54] 的想象视野普遍只有 ~20 步）。与此同时，S4（Structured State Space Sequence model，Gu et al. 2022）在 Long Range Arena 这个低维序列基准上首次解出全部任务，展现出并行训练（像 Transformer）+ 循环生成（像 RNN）+ 亚二次复杂度的组合优势，且已知与 Transformer 存在互补性（Transformer 更擅长局部/上下文相关操作，S4 更擅长长程依赖）。
本文的问题意识是双重的：(1) S4 此前只在低维序列（音频、语言、LRA 合成任务）上验证过，尚不清楚它能否胜任高维图像序列的世界模型骨架，以及如何做到；(2) RNN/Transformer/S4 三条技术路线从未在同一套可控的记忆压力测试下正面对比过。为此作者构建了基于 3D Memory Maze（Pašukonis et al. 2023）与 2D MiniGrid 的专用环境集，刻意让上下文阶段完全揭示环境信息、环境本身确定性且视觉复杂度适中，从而把”世界模型的记忆能力”与”策略学习”和”环境随机性”解耦，用生成 MSE / 奖励预测准确率作为纯净的记忆能力探针。这与同期的 TECO（Yan et al. 2023）路线不同：TECO 关注视觉更复杂但序列上限 300 步的场景，本文反其道而行——用视觉更简单的环境把序列长度推到 2000 步来压测长程记忆的极限。
模型架构
PSSM 接口：论文把 S4、S5 等抽象为满足统一接口的 Parallelizable SSM（PSSM）：训练时全序列并行 y_{1:T}, s_T = PSSM(u_{1:T}, s_0)（离散卷积或并行关联扫描），推理/想象时逐步展开 y_k, s_k = PSSM(u_k, s_{k-1})（线性递推）。S4 的核心是把连续时间 SSM s'(t)=As(t)+Bu(t), y(t)=Cs(t)+Du(t) 离散化后，将 A 参数化为 Diagonal Plus Low-Rank（DPLR）矩阵 A=Λ-PP*（Λ 对角、P 为秩-1 列向量），并用 HiPPO 矩阵初始化 A 以匹配连续时间记忆的理论最优形式；向量值输入输出通过为每个维度复制一份 SSM + 逐位置线性混合层实现。
S4WM 生成式框架：仿照 RSSM 的思路，把动力学建模从原始观测空间搬到紧凑隐空间，定义生成过程
p(x_{1:T} | x_0, a_{1:T}) = ∫ p(z_0|x_0) Π_t p(x_t|z_{≤t},a_{≤t}) p(z_t|z_{<t},a_{≤t}) dz_{0:T}
用一叠 S4 blocks 把历史 (z_{<t}, a_{≤t}) 编码为嵌入 h_t：h_{1:T}, s_T = S4Blocks(g_{1:T}, s_0)（训练并行）/ h_t, s_t = S4Blocks(g_t, s_{t-1})（想象单步），其中 g_t = MLP(concat[z_{t-1}, a_t])。先验 p(z_t|z_{<t},a_{≤t}) = MLP(h_t)，解码 x̂_t = Decoder(concat[h_t, z_t])，似然 p(x_t|z_{≤t},a_{≤t}) = N(x̂_t, 1)。近似后验采用逐帧独立的 CNN 编码（q(z_{0:T}|x_{0:T},a_{1:T}) = Π_t q(z_t|x_t)），使所有后验采样在训练时完全并行，充分发挥 PSSM 的并行优势（附录消融了一个条件于完整历史的 S4WM-Full-Posterior 变体，在 Four Rooms 与之相近，在更难的 Ten Rooms 更优，但训练更慢）。
S4 block 细节（Fig. 21）：一个 S4 block = 一个 S4 层（LayerNorm → S4 → GELU → Dropout → Linear → GLU → Dropout → 残差相加）+ 一个位置逐点 MLP 子块（LayerNorm → Linear → GELU → Dropout → Linear → Dropout → 残差相加）。作者发现在 S4 层后加这个 MLP（原始 S4/LRA 论文没有，但音频生成、语言建模中常见）能显著提升生成质量——去掉它的变体 S4WM-No-MLP 是速度最快但效果明显更差（消融证明 MLP 对长期想象和上下文召回都很关键）。
隐变量：z_{0:T} 为 32×32 的类别（categorical）离散隐变量，直通梯度（straight-through）优化，参数化为 1% 均匀分布 + 99% 网络输出的混合（沿用 DreamerV3 配方），KL 项用 KL balancing（α=0.8，更强调让先验向后验靠拢）。
配置数字（Table 10/11，3D 环境 vs 2D 环境；三个模型公平对比，共享同一套 CNN 编解码器，kernel size 4、stride 2、LayerNorm、SiLU）：
| S4WM | TSSM-XL | RSSM-TBTT | |
|---|---|---|---|
| 3D：blocks / d_model / d_ff（或 cache 长度 m） | 6 / 512 / 2048 | 12 / 512 / 512（m∈{128,256,512}，主实验用 128） | 1（GRU）/ 2048（隐状态）/ – |
| 3D：CNN 层数 / 通道倍数 / MLP 隐层 / 总参数量 | 4 / 48 / 1000 / 41.1M | 4 / 48 / 1000 / 43.8M | 4 / 48 / 1000 / 54.1M |
| 2D：blocks / d_model / d_ff（或 m） | 6 / 512 / 2048 | 10 / 512 / 512（m=128） | 1 / 2048 / – |
| 2D：CNN 层数 / 通道倍数 / MLP 隐层 / 总参数量 | 3 / 32 / 512 / 28.1M | 3 / 32 / 512 / 27.1M | 3 / 32 / 512 / 31.2M |
TSSM-XL 即用 Transformer-XL 做骨架的 TSSM/TransDreamer：把长序列切成块，用缓存长度 m 的中间隐状态作为扩展上下文——m 越大记忆容量潜力越大，但训练更慢、显存更高；论文先专门做了速度/显存对比（Appendix I）来选出与 S4WM 训练开销最接近的 m=128 作为主实验设置的公平对照。
S5WM 变体（Appendix G/D）：把 S4WM 框架里的 S4 层替换为 S5 层（Smith et al. 2023），验证 PSSM 框架的通用性——S5WM 的长期想象质量优于 S4WM（见评测表 9），且因为 S5 的隐状态维度与 h_t 相近（而 S4 的隐状态 s_t 维度通常是 h_t 的 64 倍，直接用于探针不实际），S5WM 被用于 Memory Maze 离线探针实验，取得优于 RSSM 的表现。
数据
纯离线、由脚本策略（scripted policy）采集，无真实世界数据，也不涉及人类示范：
- 3D 环境（Two Rooms / Four Rooms / Ten Rooms，基于 3D Memory Maze）：每个环境采集 30K 条轨迹（28K 训练 / 1K 验证 / 1K 测试），图像观测 64×64×3。
- 2D 环境（Distracting Memory / Multi Doors Keys，基于 MiniGrid）：每个环境采集 10K 条轨迹（8K / 1K / 1K），图像观测 40×40×3。
- 所有环境每个 episode 重新生成一次；数据采集策略被设计为”上下文阶段（context phase，完整遍历环境）+ 查询阶段（query phase，重访部分区域）“两段式，让评测可以聚焦于世界模型对上下文阶段之后的想象/召回质量。
- 具体的 (上下文步数 | 查询步数) 配置：Two/Four/Ten Rooms 长期想象为 (301|200) / (501|500) / (1101|900)；对应的”传送门”（teleport）版本用于上下文相关召回，为 (151|100) / (301|200) / (601|400)；Distracting Memory 三种宽度 100/200/400 对应 (199|51) / (399|101) / (799|201)；Multi Doors Keys 三/五/七把钥匙对应 (76|174) / (170|380) / (296|654)。
- 环境刻意保持确定性且视觉复杂度中等（作者明确说明的简化选择，见”创新点与影响”中的作者自陈局限），从而可以用重建/生成 MSE 作为记忆能力的纯净评测指标，把记忆能力与随机动力学、策略学习解耦。
- 另有一份用于与 TECO 对比的数据：DMLab，遵循 TECO 的评测协议（Appendix F），TECO 与 CW-VAE/Latent-FDM 基线用预训练 VQGAN 编解码器，而 S4WM 用联合训练的简单 CNN 编解码器。
训练方法
- 目标函数：最大化 ELBO（式 9）：
E_q[Σ_t log p(x_t|z_{≤t},a_{≤t}) - KL(q(z_t|x_t), p(z_t|z_{<t},a_{≤t}))]；KL 项用 KL balancing（LKL = α·KL[stopgrad(q)‖p] + (1-α)·KL[q‖stopgrad(p)]，α=0.8）。 - 优化器：AdamW，weight decay 1e-2；S4WM 与 TSSM-XL 梯度裁剪阈值 1000，RSSM-TBTT 为 200。学习率：S4WM 用
1e-3（1000 步线性 warmup + 余弦退火）；TSSM-XL 恒定1e-4；RSSM-TBTT 恒定3e-4（均沿用各自原论文的调度）。 - Batch size / epochs：S4WM、TSSM-XL 训练 batch size = 8；RSSM-TBTT 因用截断反向传播（TBTT）batch size = 32，3D 环境截断窗口 48 步、2D 环境 50 步。3D 环境训练 57 epochs，2D 环境训练 100 epochs。
- 无策略学习：本文只做世界模型的想象质量对比，不训练 RL 策略；Appendix E 额外做了一个基于预定义技能（拾取指定位置的钥匙 / 前往指定位置尝试开锁）的 skill-level MPC agent，用离线训练好的世界模型做两阶段规划（枚举长度为 2 的技能序列 → 用世界模型预测最可能解锁目标门的序列 → 执行第一个技能后重新规划）。
- 实现基础：代码基于公开的 S4 官方实现与 DreamerV3 代码库；三个模型共享同一套 CNN 编解码器以保证公平对比。
Infra（训练 / 推理工程）
- 硬件：速度/显存对比实验（Appendix I，Figure 3）在单张 NVIDIA RTX A6000 上完成；训练侧统一 batch size=8、序列长度=1000、图像 64×64，取 100 个 batch 的均值；想象侧 batch size=64、上下文长度=500、生成长度=500，取 8 个 batch 的均值。
- 速度/显存结论（定性，图中未给出具体每秒数值，仅相对排名）：S4WM 与 TSSM-XL 因可并行训练，训练速度都明显快于 RSSM-TBTT；但 RSSM-TBTT 显存效率显著更高，且在想象阶段吞吐量比 S4WM/TSSM-XL 高约 10×——S4WM 虽然想象时也用递归展开，但多层递归结构中夹杂的 MLP 拖慢了单步速度；解码器在所有模型中占用显存最多，导致三者峰值显存相近。
- 训练总时长 / GPU-小时 / 分布式并行：论文未披露。
- 推理 FPS / control-Hz / 边缘设备延迟：本工作是离线世界模型对比，不涉及在线部署或实时控制，论文未披露。
评测 benchmark
长期想象（Table 1，重建 MSE / 生成 MSE，越低越好）：
| Two Rooms (301|200) Recon/Gen | Four Rooms (501|500) Recon/Gen | Ten Rooms (1101|900) Recon/Gen | |
|---|---|---|---|
| RSSM-TBTT | 1.7 / 62.2 | 1.5 / 219.4 | 1.5 / 323.1 |
| TSSM-XL | 2.5 / 62.9 | 2.4 / 224.4 | 2.6 / 360.4 |
| S4WM | 1.8 / 27.3 | 1.7 / 44.0 | 1.8 / 224.4 |
三模型重建质量都不错，但 S4WM 在 Two/Four Rooms 的长期生成上明显更优；在最难的 Ten Rooms（上下文 1101 步、查询 900 步）三者都退化，作者认为这已触及 S4 模型的建模极限。
上下文相关召回（“传送门”环境，附加最多 20 帧传送后观测作为额外上下文）：TSSM-XL 在上下文较短的 Two Rooms 表现最好（无需额外观测即可召回）；上下文更长的 Four/Ten Rooms 中 S4WM 表现最好。
奖励预测（Table 2，Distracting Memory，推理准确率 = 看到完整序列 / 想象准确率 = 仅看上下文+动作序列）：
| Width | RSSM-TBTT 推理/想象 | TSSM-XL 推理/想象 | S4WM 推理/想象 |
|---|---|---|---|
| 100 | 47.9% / 49.6% | 100.0% / 51.2% | 100.0% / 100.0% |
| 200 | 48.7% / 48.4% | 99.9% / 51.3% | 100.0% / 100.0% |
| 400 | 50.4% / 52.2% | 50.4% / 51.0% | 100.0% / 100.0% |
只有 S4WM 在想象中也能保持 100% 准确率；TSSM-XL 只在能看到完整序列时表现好，一旦切到纯想象就退化到接近随机；RSSM-TBTT 全程接近随机猜测。
记忆推理（Table 3，Multi Doors Keys，Recon/Gen MSE）：
| Three Keys (76|174) | Five Keys (170|380) | Seven Keys (296|654) | |
|---|---|---|---|
| RSSM-TBTT | 0.05 / 5.16 | 0.04 / 6.36 | 0.03 / 6.28 |
| TSSM-XL | 0.05 / 1.27 | 0.03 / 6.05 | 0.02 / 9.24 |
| S4WM | 0.01 / 0.04 | 0.02 / 0.27 | 0.02 / 0.10 |
S4WM 在需要不断更新记忆（钥匙被拾取/消耗）的场景全面占优。
Memory Maze 离线探针（Table 6，Pašukonis et al. 基准，用 concat[h_t 或 s_t, z_t] 训练 4 层×1024 隐层 MLP 探针；实现基于 DreamerV3 复现，RSSM-TBTT 结果略优于原始 Memory Maze 论文报告值）：
| 9×9 Walls Acc↑ | 15×15 Walls Acc↑ | 9×9 Objects MSE↓ | 15×15 Objects MSE↓ | |
|---|---|---|---|---|
| Constant Baseline | 80.8% | 78.3% | 23.9 | 64.8 |
| RSSM-TBTT | 95.0% | 81.7% | 5.4 | 32.6 |
| TSSM-XL | 86.4% | 79.8% | 10.2 | 35.1 |
| S4WM | 88.4% | 80.2% | 8.9 | 34.8 |
| S5WM | 98.3% | 81.8% | 1.8 | 25.3 |
TSSM-XL 与 S4WM 在此基准上都不敌 RSSM-TBTT——作者分析这是因为 RSSM-TBTT 的 h_t（GRU 隐状态）承载更多”全局”信息，而 TSSM-XL、S4WM 的 h_t 更偏”局部”。S4 的内部递归隐状态 s_t 本可承担类似 RSSM h_t 的全局记忆角色，但其维度通常是 h_t 的 64 倍，直接拼接用于探针不现实；换成隐状态维度与 h_t 相近的 S5（即 S5WM）后，在 9×9 迷宫上以 98.3% 大幅超过 RSSM-TBTT 的 95.0%，物体 MSE 也降到全场最低的 1.8；但更大的 15×15 迷宫对所有模型仍具挑战性（S5WM 81.8% 仅略高于 RSSM-TBTT 的 81.7%）。
Skill-level MPC 成功率（Table 7，Multi Doors Keys）：
| Three Keys | Five Keys | Seven Keys | |
|---|---|---|---|
| RSSM-TBTT | 0% | 60% | 42.86% |
| TSSM-XL | 100% | 40% | 57.14% |
| S4WM | 100% | 100% | 100% |
S4WM 世界模型驱动的 MPC agent 在三个环境中都能 100% 解锁目标门，另外两者仅约一半成功率。
DMLab（TECO 评测协议对比，Table 8）：
| 方法 | PSNR↑ | SSIM↑ | LPIPS↓ | 参数量 |
|---|---|---|---|---|
| CW-VAE | 12.6 | 0.372 | 0.465 | 111M |
| Latent FDM（扩散模型） | 17.8 | 0.588 | 0.222 | 31M |
| S4WM | 20.6 | 0.667 | 0.196 | 41M |
| TECO | 21.9 | 0.703 | 0.157 | 169M |
S4WM 大幅超过 CW-VAE 与基于扩散的 Latent FDM，且参数量只有 TECO 的约四分之一时逼近 TECO 的效果。
S4WM vs S5WM 长期想象（Table 9）：S5WM 在 Two/Four/Ten Rooms 三个环境的生成 MSE 上全面优于 S4WM（例如 Ten Rooms 生成 MSE 从 224.4 降到 159.2），说明换用更先进的 PSSM 变体能进一步提升该框架的想象质量。
创新点与影响
- 贡献：(1) 提出 S4WM——首个把 S4 及其变体从低维序列（音频、语言、LRA 合成任务）扩展到高维图像序列世界模型的通用 PSSM 框架，通过隐变量建模 + 紧凑隐空间循环解决了”S4 直接建模像素序列难以规模化”的问题；(2) 首次对 RNN（RSSM-TBTT）/Transformer（TSSM-XL）/S4（S4WM）三大世界模型骨架，在专门设计的、覆盖长期想象/上下文召回/奖励预测/记忆推理四类能力、序列长度最长到 2000 步的环境集上做系统对比。
- 改变了什么：证明了 S4 系骨架不仅在训练速度上快于 RSSM 与 TSSM-XL，其长期记忆能力也全面超过二者；揭示了 Transformer-XL 式的缓存机制虽然能在观测完整序列时捕获长程依赖（奖励推理准确率高），却在纯想象场景中失效，说明”看得到历史”和”能在想象中维持记忆”是两种不同能力；证明了该框架对 PSSM 变体的通用性——换成隐状态维度与 h_t 相近的 S5（S5WM）后解决了 S4 隐状态过大难以用于探针的问题，并在离线探针任务上取得显著提升，为后续工作探索更先进的状态空间模型作为世界模型骨架铺路。
- 作者自陈局限：(1) 主要聚焦视觉简单、确定性的环境以控制计算成本、简化评测流程，未在复杂随机环境上验证；(2) 只评测了想象质量，未把 S4WM 接入实际的策略学习 pipeline 做完整 MBRL agent 测试（仅在附录做了离线探针和 skill-level MPC 的初步验证）；(3) 所有模型在最难的 Ten Rooms 环境（上下文 1101 步）上都表现不佳，暗示已触及当前 S4 建模能力的上限，留给未来工作探索更强的架构。
原始链接
- 论文（arXiv，v1 2023-07-05 / v2 2023-11-09，NeurIPS 2023）：https://arxiv.org/abs/2307.02064
- 论文 PDF：https://arxiv.org/pdf/2307.02064
- 项目主页（含可视化结果，代码”Coming Soon”未发布）：https://fdeng18.github.io/s4wm/
一手源存档（sources/）
- s4wm-world-model-backbones—project-page — 项目主页存档（sources/world-model/2023/s4wm-world-model-backbones—project-page.md）
- 论文全文按约定不入库（arXiv 原文 PDF，不入 git），仅引用上方 arXiv URL；本页数字均取自抓取的 arXiv v2 PDF 正文与附录（Table 1/2/3/6/7/8/9、Appendix I 速度显存实验、Appendix J 超参数表 10/11）。