Huajishe Notes

1-MuZero

12 min read

MuZero: Mastering Atari without Knowing the Rules

一句话定位

MuZero通过自监督学习学习环境动力学模型，结合MCTS规划算法，在Atari等复杂任务上实现了超越人类的性能，同时具备AlphaZero的强规划能力与模型无关的通用性。

前置依赖

  • AlphaZero (2018): 使用MCTS进行决策，但需要完整的游戏规则模拟器
  • World Models (2019): V. Mnih等提出的分离式世界模型（VAE + MDN-RNN）
  • Dreamer (2020): 基于世界模型的强化学习，使用crossentropy method进行规划
  • R2D2 (2019): Recurrent Replay Distributed Deep RL，展示RNN对于序列决策的重要性
  • MCTS基础理论: 包含UCB公式、备份机制、展开与模拟

核心思想

MuZero的核心洞察是：规划不需要精确重建真实环境，只需在抽象的潜在空间中进行即可。

传统方法面临两难：

  1. 基于模型的方法 (World Models, Dreamer): 学习 ，但可能浪费容量在无关的视觉细节上
  2. 无模型方法: 浪费样本效率，无法进行前瞻性规划

MuZero提出隐式世界模型(Implicit World Model)：

  • 不直接预测观察序列 ，而是预测抽象状态
  • 状态转换是确定性的（不同于Dreamer的概率模型）：
  • 在抽象空间进行MCTS规划，利用学习到的状态转移函数

模型结构图

输入观察 o_t
    │
    ▼
┌─────────────────────────┐
│   表示函数 (Representation) │  h_0 = f(o_0)
│   将原始观察编码为隐藏状态    │
└─────────────────────────┘
    │
    ▼
┌─────────────────────────┐
│   动态函数 (Dynamics)      │  s_{t+1} = g(s_t, a_t)
│   状态转移 + 立即奖励预测    │  r_t = h(s_t, a_t)
└─────────────────────────┘
    │
    ▼
┌─────────────────────────┐
│   预测函数 (Prediction)    │  π_t, V_t = φ(s_t)
│   从状态预测策略和价值      │
└─────────────────────────┘

三函数协同工作

        MCTS规划循环
             │
             ▼
    ┌────────────────┐
    │  从当前隐藏状态 │
    │     h_t 开始     │
    └───────┬────────┘
            │
            ▼
    ┌────────────────┐
    │  动态函数扩展树  │
    │ s_{t+1}=g(s,a) │
    │ r_t = h(s,a)   │
    └───────┬────────┘
            │
            ▼
    ┌────────────────┐
    │  预测函数评估   │
    │ π, V = φ(s)   │
    └───────┬────────┘
            │
            ▼
    ┌────────────────┐
    │  UCB选择最优动作│
    │ 选择最大UCB的  │
    │ 子节点继续      │
    └────────────────┘

数学推导

1. 表示函数 (Representation)

表示函数 将历史观察序列编码为初始隐藏状态：

其中 是最近 帧的观察序列。

2. 动态函数 (Dynamics)

动态函数 预测下一状态和立即奖励：

这里的关键设计是确定性转移——不同于Dreamer的随机MDN-RNN。这种设计使得：

  • MCTS可以在确定性的搜索树中高效展开
  • 每个节点的状态可以被精确复现

3. 预测函数 (Prediction)

预测函数 从隐藏状态预测策略和价值：

是 softmax 输出的动作概率分布， 是价值估计。

4. MCTS算法细节

MuZero使用坦南特舒尔曼树(Truncated Rollout MCTS)，包含以下组件：

4.1 节点表示

每个节点存储：

  • 访问计数
  • 累计动作价值
  • 隐状态
  • 子节点字典

4.2 UCB公式

对于节点 选择动作 ：

其中：

  • 是从预测函数得到的动作价值
  • 是先验概率（来自预测函数）
  • 是探索常数（通常设为1.25）
  • 是父节点的根号访问总数

4.3 搜索策略

MuZero的MCTS搜索包括**重放缓冲区(replay buffer)和双重量值(Double Value)**技术：

对于每个模拟：

  1. 从根节点开始
  2. 使用UCB公式选择子节点，直到达到叶节点
  3. 叶节点使用动态函数展开(k=1步)：
    • 计算
    • 获取奖励 和下一状态
    • 评估
  4. 备份：沿路径反向更新Q值

4.4 价值备份方程

其中 是从模拟 得到的回报：

5. 端到端训练目标

整个网络通过联合损失函数端到端训练：

5.1 价值损失

使用均方误差：

其中 是n步回报：

5.2 策略损失

使用交叉熵：

5.3 奖励损失

5.4 表示损失

这确保连续状态序列的一致性。

6. 规划与推理

6.1 推理阶段

在推理时，MuZero执行以下流程：

for each inference step:
    1. 获取初始隐藏状态 h_0 = f(o_current)
    2. 运行N次MCTS模拟
    3. 选择访问计数最多的根节点动作
    4. 环境执行动作，获取新观察
    5. 更新根节点（re-use tree）

6.2 树重用(Tree Reuse)

MuZero的一个关键效率优化是在推理时复用部分搜索树：

  • 选定动作后，不是重新开始，而是将所选子节点作为新的根节点
  • 被剪枝分支的信息被保留用于后续搜索

训练细节

训练循环

1. 使用当前策略与环境交互，收集轨迹数据
2. 存储到replay buffer
3. 定期从buffer采样batch进行训练
4. 更新网络参数最小化联合损失
5. 周期性更新target network

超参数

  • Hidden state维度: 或
  • 观察序列长度: 帧
  • MCTS模拟次数: (Atari), (围棋/象棋)
  • 折扣因子: (Atari长视界任务)
  • 探索常数:
  • 重放缓冲区大小: 100万步
  • 训练batch size: 1024

自定义奖励缩放

对于Atari游戏，不同游戏的奖励尺度差异巨大。MuZero使用：

  • 运行统计奖励的均值和标准差
  • 对奖励进行归一化：
  • 价值输出使用tanh压缩到

优点与局限

优点

  1. 样本效率高: 相比无模型方法，Dota2中仅需约8000个expert demonstrations
  2. 通用性强: 单一架构适用于Atari、围棋、象棋、Dota2等多种任务
  3. 端到端学习: 无需手动设计世界模型，自动学习相关抽象
  4. 可解释性: 搜索树提供了决策的可视化解释
  5. 超人类性能: 在多个领域达到或超越人类顶级水平

局限

  1. 计算开销大: 推理时需要大量MCTS模拟，不适合实时控制
  2. 训练不稳定: 端到端训练可能导致循环依赖
  3. 长视界规划困难: 折扣因子导致远期奖励信号衰减
  4. 隐式模型不可解释: 无法直接查看学习到的世界模型
  5. 部分可观测性处理: 虽然使用LSTM处理，但长期记忆仍有限

与前后内容的衔接

前置衔接 (AlphaZero → MuZero)

AlphaZero展示了MCTS+强先验策略的威力，但依赖完美模拟器。MuZero解决了：

  • 如何学习模拟器而非手工设计
  • 如何在不完美模型下进行可靠规划

后置发展

MuZero的隐式世界模型思想影响了后续工作：

  • DreamerV3 (2023): 继承了隐式世界模型+价值估计的框架
  • EfficientZero (2021): 改进MuZero的样本效率
  • Sample-Efficient RL (2022): 进一步降低数据需求

可复现实现要点

关键代码模块

  1. 表示网络: 3-4层CNN + LSTM，输出固定维度隐藏状态
  2. 动态网络: Residual MLPs，输入，输出
  3. 预测网络: 两层MLP，分别输出策略和价值头
  4. MCTS实现: 递归展开 + 并行模拟

环境接口

class MuZeroEnv:
    def __init__(self, game_name):
        self.game = atari_env(game_name)
    
    def reset(self):
        return self.game.reset()
    
    def step(self, action):
        obs, reward, done = self.game.step(action)
        return obs, reward, done

自定义 Atari 环境

MuZero使用自家AlphaZero项目组的Atari环境实现，包含：

  • 动作空间: 18种离散动作
  • 观察空间: (210, 160, 3) RGB图像
  • 奖励归一化处理
  • Episodic life endings (每条命结束后重置，而非游戏结束后)

章节摘要

MuZero是AlphaZero的深层进化版本，核心创新在于隐式世界模型+MCTS规划。通过表示函数将观察编码为抽象状态，动态函数学习状态转移规律，预测函数评估策略价值，三者协同实现端到端规划能力。相比依赖精确规则的AlphaZero，MuZero能够从原始观察中自主学习环境动力学，极大扩展了适用范围。论文首次展示了在Atari等视觉丰富且规则未知的环境中，纯粹通过学习也能达到超人类性能。

关键词

  • MuZero, 隐式世界模型, MCTS规划, 深度强化学习, 坦南特舒尔曼树
  • Representation Function, Dynamics Function, Prediction Function
  • UCB公式, 价值备份, 自监督学习, 端到端规划
  • Atari, AlphaZero, 模型预测控制

Graph View