TD-MPC2: Temporal Difference MPC for Continuous Control
一句话定位
TD-MPC2是一种用于连续控制的隐式世界模型方法，通过时序差分学习结合模型预测控制，在低样本复杂度下实现高性能的连续控制任务。
前置依赖
- World Models (2019): Haarnoja et al. 提出的分离式世界模型架构
- Dreamer (2020): 基于世界模型的强化学习，使用变分推断
- Soft Actor-Critic (2018): SAC算法的最大熵框架
- MuZero (2020): 隐式世界模型 + MCTS规划
- Temporal Difference Learning: RL中的TD(λ)框架
- Model Predictive Control (MPC): 基于模型的预测控制框架
核心思想
TD-MPC2的核心洞察是：将世界模型学习与价值估计统一到时序差分框架中，无需显式重建原始观察。
传统世界模型方法面临挑战：
- 重建损失过于关注像素细节：无关的背景变化浪费模型容量
- VAE后验坍缩：KL散度可能导致潜在变量失效
- 复合误差累积：多步预测导致误差指数增长
TD-MPC2提出隐式世界模型 + 时序差分MPC：
- 不重建观察，只学习紧凑的状态表征
- 端到端训练价值函数，避免两步走误差
- 使用**想象轨迹(IrTrjectory)**进行规划，而非显式树搜索
模型结构图
输入: 原始观察 o_t
│
▼
┌──────────────────────────────────┐
│ 编码器 (Encoder) │
│ f: O → Z (如 ResNet, CNN) │
│ 将观察编码为潜在状态 z_t ∈ R^D │
└──────────────────────────────────┘
│
▼
┌──────────────────────────────────┐
│ 隐式世界模型 (World Model) │
│ g: (z_t, a_t) → z_{t+1} │
│ 预测下一潜在状态 (无监督) │
└──────────────────────────────────┘
│
├── ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─┤
│ │
▼ ▼
┌───────────────────────┐ ┌───────────────────────┐
│ 立即奖励预测 h(z,a) │ │ 价值估计 V(z) │
│ (reward prediction) │ │ (value estimate) │
└───────────────────────┘ └───────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ 想象轨迹展开 (Imagine Rollout) │
│ z_t ──a_0──> z_{t+1} ──a_1──> z_{t+2} ──...──> z_{t+H} │
│ 累计 r_0 + r_1 + ... + r_{H-1} + γ^H V(z_{t+H}) │
└─────────────────────────────────────────────────────────────────┘
│
▼
┌──────────────────────────────────┐
│ 模型预测控制 (MPC) │
│ 选择最大化累积奖励的动作序列 │
└──────────────────────────────────┘
TD-MPC2架构详解
┌─────────────────────────────────────────┐
│ 智能体 (Agent) │
│ ┌─────────┐ ┌─────────┐ ┌─────────┐ │
│ │ Encoder │ │ World │ │ Reward │ │
│ │ f(o) │ │ Model │ │ Head │ │
│ │ z_t │──▶│ g(z,a) │──▶│ h(z,a) │ │
│ └─────────┘ └─────────┘ └─────────┘ │
│ │ │ │
│ │ │ │
│ ▼ ▼ │
│ ┌─────────┐ │ │
│ │ Value │◀──────────────────────┘ │
│ │ V │ 价值头 (Value Head) │
│ │ Head │ │
│ └─────────┘ │
└─────────────────────────────────────────┘
│
▼
┌────────────────────────────┐
│ 想象轨迹优化 (i.e.规划) │
│ 从当前 z_t 开始，生成H步 │
│ 动作序列，最大化回报估计 │
└────────────────────────────┘
数学推导
1. 编码器
编码器
潜在空间维度通常设为
2. 隐式世界模型
世界模型
关键设计：确定性转移，不预测分布。这简化了规划过程，因为：
- 给定
， 是唯一确定的 - 允许反向传播通过dynamics
3. 奖励预测头
奖励头是一个小型MLP，预测立即奖励。这用于想象轨迹中的奖励累积。
4. 价值估计
价值头输出状态
5. 想象轨迹与回报估计
TD-MPC2的核心规划机制。给定当前状态
for h = 0 to H-1:
a_h = action_head(noise + z_{t+h}) # 从当前状态采样动作
z_{t+h+1} = g(z_{t+h}, a_h)
r_h = h(z_{t+h}, a_h)
G_imagine = Σ_{i=0}^{H-1} γ^i r_i + γ^H V(z_{t+H})
6. 时序差分学习
TD-MPC2使用时序差分(TD)学习训练价值函数。对于n步回报：
TD误差：
价值损失：
7. 端到端训练目标
总损失函数：
7.1 价值损失
其中
7.2 奖励损失
7.3 动作损失（策略学习）
其中
7.4 动力学损失
8. 模型预测控制公式
在每个时间步，执行以下优化：
其中
实际实现中使用**交叉熵方法(CEM)**进行优化：
1. 初始化动作序列分布 (高斯混合)
2. for iteration = 1 to N:
采样M个动作序列
使用世界模型计算每个序列的价值
更新分布：保留top-k样本，重新拟合高斯
3. 返回均值作为规划动作
9. 多步预测与复合误差
TD-MPC2通过以下技术缓解复合误差：
- 短期想象+价值估计：不进行超长视界预测
- 亚稳态表示：编码器学习对任务相关的状态抽象
- TD学习：直接优化价值函数，不依赖完美世界模型
数学上，可证明在某些条件下，即使世界模型有误差，价值估计仍是无偏的：
当
训练细节
训练循环
1. 使用当前策略与环境交互
2. 存储转换 (o_t, a_t, r_t, o_{t+1}) 到replay buffer
3. 定期采样batch进行训练:
a. 用编码器处理观察
b. 计算动力学/奖励/价值损失
c. 更新网络参数
4. 周期性更新目标网络
超参数
- 潜在维度:
- 想象视界:
到 - CEM迭代次数:
- CEM样本数:
- 折扣因子:
- 目标网络更新频率: 每1000步
- Batch size: 256
- Replay buffer: 100万步
分布式实现
TD-MPC2支持分布式数据收集：
- 多个actor并行收集数据
- 中心replay buffer
- 单learner更新网络
与Dreamer/MuZero的比较
| 特性 | Dreamer | MuZero | TD-MPC2 |
|---|---|---|---|
| 世界模型 | 概率(VAE) | 确定性 | 确定性 |
| 规划方法 | CEM | MCTS | CEM |
| 状态空间 | 变分后验 | 抽象状态 | 隐表示 |
| 价值学习 | 重建回报 | n步回报 | TD学习 |
| 适用场景 | 连续控制 | 离散/连续 | 连续控制 |
| 样本效率 | 中等 | 高 | 高 |
推理/rollout过程
def forward(obs, num_steps=10):
z = encoder(obs)
actions = []
for t in range(num_steps):
# CEM优化
best_action = cem_planning(z, horizon=5)
actions.append(best_action)
# 世界模型前向
z = world_model(z, best_action)
return actions
优点与局限
优点
- 高样本效率: 在DMC套件上，TD-MPC2用10倍少的样本达到与Dreamer相当性能
- 训练稳定: 端到端TD学习，避免变分推断的优化困难
- 连续控制友好: CEM优化自然处理连续动作空间
- 隐式世界模型: 避免像素级重建的浪费
- 可扩展性强: 分布式数据收集支持大规模训练
局限
- 计算密集推理: CEM优化需要多次世界模型前向
- 超参数敏感: CEM参数需要仔细调优
- 局部最优: CEM可能陷入局部最优，尤其是高维动作空间
- 表示学习不足: 无重建目标可能导致表示质量下降
- 长视界规划有限: 想象视界受限于复合误差
与前后内容的衔接
前置依赖
TD-MPC2继承了多代世界模型的发展：
- World Models的分层结构 → TD-MPC2的编码器+世界模型
- Dreamer的价值学习 → TD-MPC2的TD框架
- MuZero的隐式模型 → TD-MPC2的确定性dynamics
后置发展
TD-MPC2的思想影响了后续工作：
- DreamerV3 (2023): 结合TD-MPC2的TD学习与Dreamer的变分方法
- Prodigy (2024): 改进TD学习的收敛性
- 掩码世界模型: 探索重建与对比学习的结合
可复现实现要点
编码器实现
class Encoder(nn.Module):
def __init__(self, obs_shape, latent_dim):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(obs_shape[0], 32, 4, stride=2),
nn.ReLU(),
nn.Conv2d(32, 64, 4, stride=2),
nn.ReLU(),
nn.Conv2d(64, 128, 4, stride=2),
nn.ReLU(),
)
self.fc = nn.Linear(...)
self.fc_out = nn.Linear(128, latent_dim)
def forward(self, obs):
x = self.conv(obs)
x = x.reshape(x.size(0), -1)
return self.fc_out(torch.relu(self.fc(x)))世界模型实现
class WorldModel(nn.Module):
def __init__(self, latent_dim, action_dim):
super().__init__()
self.net = nn.Sequential(
nn.Linear(latent_dim + action_dim, 1024),
nn.ReLU(),
nn.Linear(1024, 1024),
nn.ReLU(),
nn.Linear(1024, latent_dim),
)
def forward(self, z, action):
x = torch.cat([z, action], dim=-1)
return self.net(x)CEM规划实现
def cem_planning(world_model, z_init, horizon, num_iters=10, num_samples=100):
mean = torch.zeros(horizon, action_dim)
std = torch.ones(horizon, action_dim)
for _ in range(num_iters):
# 采样
actions = torch.randn(num_samples, horizon, action_dim) * std + mean
actions = actions.to(z_init.device)
# 评估
returns = []
z = z_init
for h in range(horizon):
z = world_model(z, actions[:, h])
r = reward_head(z, actions[:, h])
returns.append(r)
returns = torch.stack(returns)
G = (returns * gamma ** torch.arange(horizon)).sum(dim=0)
# 更新分布
top_k = G.topk(num_samples // 10, dim=0).indices
mean = actions[top_k].mean(dim=0)
std = actions[top_k].std(dim=0)
return mean[0]章节摘要
TD-MPC2是连续控制世界模型的重要里程碑，通过隐式世界模型与时序差分学习的结合，实现了高效样本的连续控制。核心创新在于：确定性dynamics避免概率模型的训练困难，CEM规划自然适配连续动作空间，端到端TD学习直接优化价值函数。与MuZero在离散/混合空间的MCTS规划形成对比，TD-MPC2展示了CEM优化在连续控制中的有效性。
关键词
- TD-MPC2, 时序差分学习, 隐式世界模型, 模型预测控制
- 交叉熵方法(CEM), 想象轨迹, 连续控制
- 样本效率, 价值估计, 确定性dynamics