AI Research 技术调研

Vector Quantized Models for Planning

19分钟阅读

一句话定位

在 muzero 的隐空间 MCTS 框架里，把对手/环境的下一步响应也编码成离散 VQVAE latent code，让搜索树同时在「agent 动作节点」和「环境 latent-code 节点」上展开，从而把 MuZero-式规划从确定性、全可观测环境推广到随机、部分可观测环境——在把对手动作从「可枚举」变成「环境的一部分」的单人视角象棋任务上恢复了双人版 MuZero 的水平，并在高维视觉、部分可观测的 DeepMind Lab 上验证了可扩展性。ICML 2021。

背景与定位

muzero 用价值等价的隐模型 + MCTS 统一了 AlphaZero 式规划与无需环境模拟器的学习，但仍隐含四条假设：环境确定性、agent 拥有完整状态观测、搜索发生在对手/环境动作可枚举的层面、且搜索粒度与单步交互对齐。本文动机是同时松开前两条：作者认为 MuZero 在两人棋类游戏里之所以有效，很大程度依赖于训练与规划时都能拿到”对手视角”的数据；一旦把对手的落子重新定义为环境动力学的一部分（agent 只看到自己的状态和动作），MuZero 的确定性假设就被打破，性能会灾难性下降。

解法是把 muzero 的动作搜索机制搬到离散隐变量上：借鉴 VQVAE（van den Oord et al. 2017; Razavi et al. 2019 的 VQ-VAE-2）用离散自编码器捕捉一个动作在随机环境下的多种可能后果，再用一种”随机化的 MCTS”变体同时对 agent 动作和这些离散隐变量做搜索。文中明确指出最相似的先前工作是 dreamer-v2（Hafner et al. 2020，“Mastering Atari with Discrete World Models”）——同样用离散自编码器学随机转移模型，但该工作只把模型用于生成 Dyna 式的合成数据训练策略，并未把模型直接接入 MCTS 做规划；本文则是把离散隐变量本身当作 MCTS 可以搜索的”环境动作”节点。这也使其成为后续一系列离散 token / 隐空间世界模型（如同属 world-model 类目下用 VQ token 做自回归世界模型的 iris）谱系中，较早把 VQ 离散化与基于搜索的规划直接结合的工作。

模型架构

两阶段流水线：先训练一个条件式 state VQVAE（编码器把状态-动作历史压成离散 latent code，解码器用 code 重建状态），再在冻结的离散 latent 之上训练一个自回归 transition model（预测 policy、value、reward，外加对离散 latent code 的 policy τ）。作者在正文脚注中明确说明：端到端联合训练 VQVAE 和 transition model 是一个有前景的未来方向——即当前方法本身不是端到端可微的单阶段训练。

State VQVAE：编码器 f_enc(s_1:t+1, a_1:t) → k_{t+1}（把到 t+1 为止的状态序列和到 t 为止的动作序列编码成离散 code），解码器 f_dec(s_1:t, a_1:t, k_{t+1}) → ŝ_{t+1}。训练用标准 VQVAE 损失 L = L_r(x̂,x) + β‖z_u − sg(e_c)‖²（重建损失 + commitment 损失，straight-through 梯度估计，codebook 用滑动平均更新)。

Transition model / MCTS：延续 muzero 的隐空间展开，但规划路径变成动作与离散 latent code 交替的序列 s, a⁰, k¹, a¹, …, a^{M-1}, k^M；模型除了 MuZero 原有的 policy 头 π、value 头 v、reward 头 r 之外，新增一个离散 latent code 头 τ。MCTS 搜索树含两类节点：动作节点（沿用 MuZero 的 UCB 式 argmax[Q(s,a)+P(a|s)U(s,a)]）和随机（stochastic）节点（argmax[Q̂(s,k)+P(k|s)U(s,k)]，P(k|s) 由 τ 给出）；Q̂(s,k) 可设为 0（中性/近似 Expectimax）、Q(s,k)（合作）或 −Q(s,k)（对抗），作为对环境性质的先验注入——棋类等零和博弈用对抗设定。

论文提出三种规划路径变体：VQHybrid（离散 latent 与 agent 动作交替，正文主要实验设定）、VQPure（要求一个因子化的 state VQVAE，把环境随机性和 agent 动作分别编码成两套独立 latent，从而让 transition model 完全脱离真实动作空间、纯粹在离散 latent 上展开）、VQJumpy（state VQVAE 直接跳步预测 s_{t+m} 而非 s_{t+1}，引入时间抽象）；VQJumpy 论文中只是提出概念，未做任何实验，VQPure 仅在象棋上有部分结果，两者均被作者称为”largely unexplored, left for future works”。

棋类配置数字（原文与附录存在不一致，两处均如实记录）：正文 5.1.2 节称 state VQVAE 量化层为 2 个 codebook，每个 128 codes × 64 维（拼接成 2-hot 编码）；附录 A.2.3 则称同一 state VQVAE 编码器/解码器为 16 层残差卷积、256 hiddens / 64 bottleneck hiddens，量化层为 2 个 codebook，每个 64 codes × 256 维。Transition model 新增的 g'（消费离散 latent code 的动力学函数）为 30 层残差卷积、1024 hiddens / 256 bottleneck hiddens，与 MuZeroChess 基线的 g/f（20 层残差卷积、1024 hiddens / 256 bottleneck hiddens，所有预测头隐藏维 256）同规格。棋盘用符号表示（每个棋子编码为整数），动作空间沿用 Silver et al. (2017a) 协议的 8×8×73 离散动作。

DeepMind Lab 配置数字：三阶段流水线——(1) frame-level VQVAE（前馈卷积编码/解码器，codebook 512 codes × 64 维）把观测帧压成 frame-level 离散表示；(2) 在 frame-level 表示之上训练 state VQVAE（卷积编码器躯干 + 编码器 LSTM + 量化层 + 解码器 LSTM + 卷积解码头），量化层由每个时间步 32 个独立的 latent variable 组成，每个都有自己的 codebook（512 codes × 64 维），解码头把 transition 离散 latent 映射回 frame-level VQ 空间；(3) transition model 用深层因果 Transformer（Vaswani et al. 2017）在 VQHybrid 规划路径上拟合。生成新样本时：从 prior 网络采样 state latent → state VQVAE 解码器映射回 frame-level 离散空间 → frame-level VQVAE 解码器映射回像素空间。

数据

象棋（离线 RL 设定，固定数据集，无在线交互）：训练集为 Million Base 数据集（250 万局）+ FICS Elo>2000 数据集（96 万局） 的组合；验证集为 2017 年 FICS Elo>2000 的 4.5 万局；训练集玩家 Glicko-2 rating 分布见附录直方图（Figure 9）。行为数据即人类专家对局，动作已内含在棋谱记录里，无需单独动作标注。

DeepMind Lab：用一个 A2C agent（Mnih et al. 2016）在 explore_rat_goal_locations_small 关卡采集数据集，共 101,325,000 个 episode，其中 675,500 个作为测试集留出;每个 episode 128 个时间步，agent 在迷宫式环境中随机出生、第一人称观测，捕获随机放置的苹果后获得奖励并被传送到新的随机位置。由于数据采集从随机初始化的 A2C agent 开始、并随训练推进持续进行，数据集本身混合了不同经验水平的行为轨迹（而非单一专家策略）。全部实验均为离线设定，数据只采集一次、模型训练时不再与环境交互。

训练方法

两阶段（象棋）/三阶段（DeepMind Lab）训练：先训练 state VQVAE（及 DM Lab 的 frame-level VQVAE），冻结后再用**教师强制（teacher forcing）**在行为策略（人类专家棋谱 / A2C 生成轨迹）产生的真实状态、动作、离散 latent code 序列上训练 transition model。

Transition model 损失：对展开长度 M 的规划路径，损失为动作头的交叉熵 CE(a^m, π(h^{2m})) + 离散 latent code 头的交叉熵 CE(k^m, τ(h^{2m-1})) + value 损失（象棋用 Monte Carlo 回报作为 value target）+ reward 损失（真实环境 reward），四项按系数 α、β 加权平均——与 MuZero 损失几乎一致，唯一区别是新增了对离散 latent code τ 的预测项。

含 VQ 的 MCTS：搜索树扩展为动作节点与随机节点交替；随机节点用 Q̂(s,k) 的三种设定（中性 / 合作 / 对抗）编码环境先验知识——消融显示（棋类实验）用对抗先验的 VQHybrid/VQPure 显著优于中性假设的 VQHybrid Neutral/VQPure Neutral，说明方法仍能从环境先验知识中获益，并非完全假设无关。

Quasirandom 采样（附录 A.3）：随机节点扩展/评估离散 latent code 时，若假设环境中性，理论上应从 τ 给出的分布 i.i.d 采样，但作者发现选择 p_i/(N(i)+1) 最大的 code（一种拟随机 Monte Carlo 采样）在实践中效果更好，其经验分布比朴素 i.i.d 采样更接近真实分布（Figure 10 定量对比欧氏距离误差）。

跨智能体评估技巧：VQ agent 与 Q-value / imitation / MuZeroChess 等确定性智能体对弈时，为避免两个确定性策略对局变得完全确定、评估失去意义，前 30 步采用随机化的树策略——保留访问次数 ≥ 最大访问次数 1% 的动作池，按访问次数比例采样。

关键超参：象棋 state VQVAE——batch size 1024，Adam，lr 3e-4，指数衰减率 0.9、衰减步数 100000，训练 120 万步；MuZeroChess 基线与 VQ transition model 共用一套设置——batch size 2048，Adam lr 3e-4，同样的衰减率/步数，梯度裁剪最大值 1；MuZeroChess 训练 20 万步，VQHybrid 训练 20 万步，VQPure 训练 40 万步。MCTS 超参：discount 1.0，UCB 常数 c_base=19652.0、c_init=1.25，acting 阶段的访问计数策略不加温度，动作限制在合法着法内。

Infra（训练 / 推理工程）

GPU/TPU 型号与数量、并行方式、GPU-hours、推理 FPS / 控制频率、边缘硬件时延：全文（含附录 A–C）均未披露——通读正文与附录未出现任何硬件配置或墙钟时间数字（对照可比的 muzero 原文披露过 TPU 数量与训练时长，本文完全没有给出对应工程细节）。已披露的仅是上述训练超参（batch size、优化器、学习率衰减、训练步数）。

评测 benchmark

象棋（图表呈现，正文未给出可提取的具体数值）：

  • Figure 4：单人视角 MuZeroChess（不观测对手动作）对阵 Stockfish 10 的表现明显差于双人视角 MuZeroChess；VQHybrid、VQPure（对抗先验/最坏情况 chance node）恢复到双人版 MuZeroChess 的水平；VQHybrid Neutral、VQPure Neutral（中性先验）表现不及对抗先验版本。
  • Figure 6：中性假设下，增大 MCTS 模拟预算并不能提升表现（作者归因于中性假设下搜索的是训练数据里覆盖的宽泛 Elo 分布对手的期望行为，而非 Stockfish 特定强度的行为，预算越大越暴露这种分布错配）。
  • Figure 5：VQHybrid 与 imitation agent 对局表现同双人版 MuZeroChess 接近；VQHybrid 与 MuZeroChess 直接对局胜率相近。
  • 附录 Figure 11（1/10/100/400/800/1200 次模拟 × 5 档 Stockfish 强度的完整扫描）：把 MuZeroChess 训练计算量增加 3 倍（MuZeroChess@600k）只带来轻微提升；单人视角 MuZeroChess Single@600k 即便计算量增加 3 倍也没有改善，1200 次模拟时其胜率在所有 Stockfish 等级上反而下降——表明单人视角 MuZeroChess 无法通过增加算力弥补对环境随机性建模的缺失。
  • 以上均为图表结果，原文正文/附录文本中未给出可提取的具体百分比数值（图 4–6、11 均为柱状图，无配套数值表）。

DeepMind Lab（Mean Best Reconstruction Error, MBRE）：评测协议——测试集每条轨迹取 T0=16 帧前缀，从模型采样 k=1000 条 rollout，与真实轨迹在目标时间步 T=128 比较累计 MSE 的最优匹配样本，取全测试集平均。基线：确定性 LSTM 下一帧预测（重建 PSNR ≈24dB）、两档 GECO 连续 VAE（目标失真 25dB / 33dB PSNR，分别对标 LSTM 基线与 VQVAE 解码器≈32dB 的重建质量）。结果（Figure 8，图表形式）：LSTM 基线显著差于所有隐变量模型（因确定性模型无法表达多模态随机结果，退化为预测均值）；连续 VAE 基线体现”重建质量 vs 长程预测性能”的权衡；VQM 在两个指标上都取得有竞争力的表现。同样，具体 MBRE 数值原文以图表呈现，文本中无可提取的表格数字。

重要限定：DeepMind Lab 部分只评测了世界模型的 rollout 生成质量（MBRE），并未在该环境上运行 MCTS agent 做实际规划/控制评测——这是对世界模型可扩展性的验证，而非完整的 VQM-MCTS 智能体评测。

创新点与影响

  • 核心贡献：把 muzero 的隐空间 MCTS 从”只搜索 agent 动作”扩展为”同时搜索 agent 动作和环境的离散 VQ latent code”，用一套统一机制（动作节点 + 随机节点，Q̂(s,k) 可调环境先验）把 MuZero 式规划推广到随机、部分可观测环境，且不需要访问真实模拟器或对手视角数据。
  • 改变了什么：证明了 MuZero 在双人游戏上的成功一定程度依赖于对手视角数据的可得性——单人视角 MuZero 性能会灾难性下降,且无法靠增加计算量挽回;而只用 agent 自身视角数据训练/规划的 VQ 方法能恢复到双人版 MuZero 的水平。相比 dreamer-v2 等把离散隐变量模型仅用于生成合成训练数据的路线，本文是较早把离散 VQ latent 直接接入 MCTS 树搜索的工作，也验证了这套机制能扩展到 DeepMind Lab 这类高维视觉、部分可观测的第一人称 3D 环境。
  • 作者自陈局限：(1) VQVAE 与 transition model 目前是两阶段分开训练，端到端联合训练被明确列为未来方向；(2) VQPure 仅在象棋上做了部分实验、VQJumpy 完全未实现评测，两者”largely unexplored, left for future works”；(3) DeepMind Lab 上只验证了世界模型的生成/rollout 质量，未运行 MCTS 规划智能体;(4) 中性环境先验下模型表现不如对抗/合作先验，说明方法仍需要一定环境先验知识而非完全通用;(5) 全部实验限定在离线 RL 设定（固定数据集、无在线交互），未探讨在线场景下的表现。

原始链接

一手源存档（sources/）

  • arXiv 全文见上方链接（arXiv 原文 PDF，不入 git）。
  • 项目页 https://sites.google.com/view/vqmodels/home（论文 Figure 7 标注的补充视频页）多次抓取均因 net::ERR_CONNECTION_CLOSED 失败，未能archive；页面内容据论文引用仅为 rollout 视频演示，不含正文之外的额外方法细节。
  • 未发现独立于 arXiv 论文的官方博客 / GitHub 代码仓 / HuggingFace 模型卡（本文未开源代码，通读全文未见代码仓库链接）。

关系图谱

反向链接