VAE在基于模型的强化学习中的应用

变分自编码器（VAE）提供了学习数据紧凑且有组织表示的有效方法。这项能力显著扩展到基于模型的强化学习 (reinforcement learning)（MBRL）范围，在此类学习中，理解和预测环境行为对智能体的成功是核心要素。MBRL中的智能体旨在学习环境的动态模型，随后可运用该模型进行规划或优化其策略。当处理高维状态空间时，例如机器人任务中来自摄像头的数据图像，直接在观察空间中学习此模型可能极其困难且需要大量数据。

VAE通过首先学习这些高维观测的低维潜在表示，提供一种巧妙的解决方案。核心思路是训练VAE将观测 $s_t$ 编码为潜在向量 (vector) $z_t = q_\phi(z_t | s_t)$，并将其解码回观测的重构 $\hat{s}_t = p_\theta(s_t | z_t)$。VAE通过最大化证据下界（ELBO）进行训练，这有助于实现良好的重构，同时对潜在空间进行正则化 (regularization)，通常使其接近标准高斯分布 $p(z)$。

一旦VAE完成训练（或正在训练中），可以学习一个独立的模型，以完全在此压缩的潜在空间内预测动态。此潜在动态模型旨在根据当前潜在状态 $z_t$ 和智能体执行的动作 $a_t$ 预测下一个潜在状态 $z_{t+1}$：

$$\hat{z}_{t+1} \sim p_{dyn}(z_{t+1} | z_t, a_t)$$

此动态模型 $p_{dyn}$ 可以是神经网络 (neural network)，例如MLP或RNN（如果时间依赖性重要）。在潜在空间 $z$ 而非高维观测空间 $s$ 中学习动态带来多方面益处：

• 效率: 潜在空间维度低得多，使动态模型更简单，训练和预测速度更快。
• 突出特征: VAE通过其信息瓶颈，倾向于在 $z$ 中捕捉观测最显著的特征，可能滤除噪声或不相关细节。这有助于动态的学习。
• 预测方式: 模型预测状态间的转换，这比预测像素级别的变化更直接。

一种将VAE纳入MBRL智能体的典型架构包含多个彼此关联的组成部分：

VAE驱动的基于模型的强化学习智能体各环节。来自环境的观测被编码成潜在空间。动态模型根据当前潜在状态和动作预测未来潜在状态。策略或规划器借助这些潜在预测来选择动作。VAE的解码器主要用于训练编码器，以生成有用的潜在表示。

在此设定下:

1. 智能体从环境中接收观测 $s_t$。
2. VAE的编码器将 $s_t$ 映射到潜在状态 $z_t$。
3. 策略 $\pi(a_t | z_t)$（或使用已学模型的规划器）选择一个动作 $a_t$。
4. 动作 $a_t$ 在环境中执行，产生新的观测 $s_{t+1}$ 和奖励 $r_t$。
5. 转换 $(s_t, a_t, r_t, s_{t+1})$ 被存储在回放缓冲区中。
6. VAE被训练来重构观测：$s_t \approx p_\theta(s_t | q_\phi(z_t | s_t))$。
7. 潜在动态模型 $p_{dyn}(z_{t+1} | z_t, a_t)$ 被训练来预测下一个状态的编码，$z_{t+1} = q_\phi(z_{t+1} | s_{t+1})$，给定 $z_t = q_\phi(z_t | s_t)$ 和 $a_t$。这通常通过最小化像均方误差这样的损失来完成： \mathcal{L}_{dyn} = || \text{目标_}z_{t+1} - p_{dyn}(z_t, a_t) ||^2 其中 \text{目标_}z_{t+1} 是通过编码 $s_{t+1}$ 获得的“真实”下一个潜在状态（通常梯度会停止）。
8. 可选地，奖励预测器 $p_{rew}(r_t | z_t, a_t)$ 也可以进行训练。

学习到的潜在动态模型 $p_{dyn}$ 和奖励预测器 $p_{rew}$ 随后可用于规划。例如，智能体可以在潜在空间中执行“想象推演”，通过对一系列动作重复应用动态模型：$z_t \rightarrow \hat{z}_{t+1} \rightarrow \hat{z}_{t+2} \dots$。这些想象出的轨迹可以与模型预测控制（MPC）或蒙特卡洛树搜索（MCTS）等技术结合使用，以选择最佳动作。VAE的解码器还可以通过将这些想象出的潜在轨迹转换回观测空间来使其可见，从而提供一种“查看”模型预测内容的方式。

突出架构，例如“Models”（Ha & Schmidhuber, 2018）和“Dreamer”智能体家族（Hafner et al., 2019, 2020, 2023）是这种方法的典型实例。

• Models 明确区分了VAE（V组件）、潜在动态模型（M组件，通常是RNN）和控制器（C组件）。控制器仅使用VAE的潜在表示和M的预测进行操作。
• Dreamer 及其后续版本完全从学习到的模型在紧凑潜在空间中生成的想象轨迹中学习行为。该模型，包括表示学习器、动态模型和奖励预测器，是联合训练的。策略学习随后利用这些想象轨迹进行。

训练考量与目标

整体训练通常涉及同时或迭代优化多个目标：

1. VAE重构损失: 确保 $z_t$ 捕获信息以重构 $s_t$。这是ELBO的一部分： $$\mathcal{L}_{recon} = \mathbb{E}_{q_\phi(z|s)}[\log p_\theta(s|z)]$$
2. VAE KL正则化 (regularization): 对潜在空间进行正则化，通常是 $D_{KL}(q_\phi(z|s) || p(z))$，其中 $p(z)$ 是标准正态先验。
3. 潜在动态损失: 最小化预测下一个潜在状态的误差，如前所示。 $$\mathcal{L}_{dyn} = \mathbb{E}_{z_t, a_t, z_{t+1}} [ \text{距离}(z_{t+1}, p_{dyn}(z_t, a_t)) ]$$
4. 奖励预测损失 (如果适用): $$\mathcal{L}_{rew} = \mathbb{E}_{z_t, a_t, r_t} [ \text{距离}(r_t, p_{rew}(z_t, a_t)) ]$$
5. 策略损失: 策略被训练以最大化预期奖励，通常使用从真实环境或学习模型中生成的轨迹上的Actor-Critic方法。

需要仔细权衡。如果VAE过于侧重像素级完美重构，$z_t$ 可能会保留与动态预测无关的高频细节。反之，如果 $z_t$ 压缩过度或正则化过度，它可能失去对长期预测或控制非常重要的信息。一些方法会修改VAE的目标或架构，以更好地适应控制需求，例如，通过确保未来结果可区分的状态在潜在空间中也能被区分。

优势与挑战

将VAE集成到MBRL中带来多项益处：

• 样本效率: 通过学习模型，智能体可以生成额外经验，减少对成本高昂或缓慢交互的依赖。
• 处理高维数据: VAE有效压缩复杂感知输入（如图像）为可管理的潜在向量 (vector)。
• 规划能力: 学习到的潜在动态模型使得复杂的规划算法能够被应用。

然而，也存在一些挑战：

• 模型准确性: 所学模型是近似的。潜在动态模型中的误差可能在较长的预测周期内累积，导致“模型不匹配”，即想象轨迹与现实显著偏离。
• 表示权衡: 潜在空间必须既利于重构（用于训练VAE），又利于预测动态和奖励。这些目标并非总是完美一致。在 $z_t$ 中实现解耦可能有所帮助，但如果所得表示能很好地服务于动态学习，则不总是主要目标。
• 计算成本: 同时训练VAE、动态模型和策略可能需要大量计算资源。

尽管存在这些挑战，VAE已成为许多先进基于模型的强化学习 (reinforcement learning)智能体的重要组成部分，尤其是那些设计用于处理丰富感知输入的智能体。它们学习有组织潜在表示和在压缩空间中辅助动态建模的能力，持续推动着更智能、数据高效智能体的发展。