隐状态与潜空间

一句话定位

问题:世界模型的核心是隐状态 ,但 是不可直接观测的。如何从未知的隐状态中推断出最可能的表征?如何确保推断得到的后验与先验转移保持一致?

定位:第二章第二节,承接”观测压缩与状态抽象”,深入探讨后验推断的数学机制与潜空间的概率结构。


前置依赖


核心思想

隐状态是世界的”真实”但不可见的表征——它对应于物理对象的位置、速度、形状等根本属性,而不仅仅是像素级的观测。

潜空间是隐状态所居住的数学空间,其结构由以下因素决定:

  • 先验分布 :智能体对状态空间的先验假设(如高斯分布)
  • 后验分布 :给定观测时对状态的推断
  • 转移分布 :状态随动作的演化规律

关键洞察:我们永远无法直接观测 ,但可以通过变分推断(variational inference)近似它——学习一个识别网络 来逼近真实后验


模型结构图

                隐状态推断与潜空间结构
                ======================

观测序列 o_1, o_2, ..., o_t
    │
    ├── o_1 ──▶ encoder q_φ(s_1|o_1)
    ├── o_2 ──▶ encoder q_φ(s_2|o_≤2)
    │           ...
    └── o_t ──▶ encoder q_φ(s_t|o_≤t)

    联合后验推断(使用 RNN 或 Transformer 聚合历史信息):
    ┌────────────────────────────────────────────────────────┐
    │  q_φ(s_t | o_≤t) = Encoder_RNN( o_1, o_2, ..., o_t )  │
    │                                                         │
    │  形式: h_t = RNN(h_{t-1}, o_t)                         │
    │        s_t ~ q_φ(·|h_t) = N(μ_φ(h_t), σ_φ^2(h_t))      │
    └────────────────────────────────────────────────────────┘

隐状态序列 s_1, s_2, ..., s_t ∈ ℝ^{d_s}
    │
    ├── s_t 的先验: p_θ(s_t | s_{t-1}, a_{t-1}) [来自转移模型]
    │
    └── s_t 的后验: q_φ(s_t | o_≤t) [来自编码器]

    KL 正则项: D_KL(q_φ(s_t|o_≤t) || p_θ(s_t|s_{t-1},a_{t-1}))
    │
    ▼
确保先验-后验一致,使得想象rollout可靠

先验-后验结构图

    真实环境                                    世界模型

  物理状态 s_t*  ←────→  观测 o_t               隐状态 s_t
       │                                        │
       │                                        │
       ▼                                        ▼
  真实后验 p(s_t|o_≤t)              近似后验 q_φ(s_t|o_≤t)
  (不可计算)                                (可计算)
       │                                        │
       │                                        │
       ▼                                        ▼
       └───────────── KL 正则 ──────────────────┘
                    D_KL(q_φ || p_θ)

       先验 p_θ(s_t|s_{t-1},a_{t-1})
       (来自转移模型)

数学推导

1. 隐状态的后验推断

时刻的隐状态,它是导致观测 的”真实”原因,但不可直接观测。

真实后验(贝叶斯公式):

其中:

  • 似然(解码器)
  • 先验(来自上一时刻的转移预测)
  • 证据

问题是真实后验中的先验 不可直接计算,因为它需要知道历史观测的全部信息。

变分近似:引入识别网络 来近似真实后验:

2. 递归后验推断(RSSM)

Dreamer 使用**递归状态空间模型(RSSM)**进行后验推断:

这展开为:

  1. 预测步(Prior)

    • 给定上一时刻后验 和动作
    • 预测当前时刻先验:
  2. 更新步(Posterior)

    • 给定当前观测 和预测先验
    • 推断后验:

数学形式

其中 是解码器。

3. 变分下界与 KL 正则

ELBO 包含后验与先验的 KL 项:

KL 项的物理含义

  • 时,KL = 0,先验与后验完全一致
  • KL > 0 时,编码器推断的后验偏离转移模型的预测

训练目标:最小化 KL 等价于让转移模型学会”预测”编码器看到的东西。

4. 重参数化技巧(Reparameterization Trick)

为了反向传播梯度,需要对随机变量 进行重参数化:

其中 是神经网络输出的均值和标准差。

重参数化的梯度

这样梯度可以通过 反向传播。

5. 潜空间的数学结构

假设隐状态空间 ,且先验

潜空间的结构性质

性质数学描述物理含义
紧致性状态空间是有界的
连通性状态间可以插值
分离性相似状态在空间中接近
语义对齐隐空间梯度对应感知梯度

6. 确定性后验 vs 随机后验

确定性后验(World Models, Ha et al., 2018):

其中 是确定性向量,无随机性。

随机后验(Dreamer, Hafner et al., 2019):

对比

特性确定性后验随机后验
表达能力有限(单点)强(分布)
不确定性建模
训练稳定性中(KL 权重关键)
记忆容量
适用场景简单环境复杂/随机环境

训练细节

数据构造

隐状态推断的训练数据是观测序列(无需直接标注隐状态):

模型从观测序列中自己学习隐状态的表征,无需人工标注。

输入输出(IO)

变量维度说明
历史观测序列
历史动作序列
当前隐状态(分布)
后验均值与方差

损失计算

后验推断的损失函数

对于两个高斯分布,KL 有闭式解:

优化器配置

  • 学习率(Adam)
  • KL 权重 -VAE 或 DreamerV3 的
  • 重参数化:训练时使用,推理时可简化为
  • 梯度裁剪:1.0

训练流程

# 隐状态后验推断训练流程
for batch in dataloader(trajectories):
    # 1. 编码历史观测
    h = zeros(batch_size, hidden_dim
    for t in range(T):
        # 更新 RNN
        h = rnn(h, a_{t-1}, o_t)
        # 输出后验参数
        mu, logvar = encoder_post(h)  # q_φ(s_t|o_≤t)
        std = exp(0.5 * logvar)
 
        # 2. 采样(重参数化)
        eps = torch.randn_like(std)
        s_t = mu + eps * std
 
        # 3. 转移模型预测先验
        prior_mu, prior_logvar = transition_prior(s_{t-1}, a_{t-1})
 
        # 4. 计算 KL 损失
        kl_loss = kl_divergence(
            normal(mu, std),
            normal(prior_mu, exp(0.5 * prior_logvar))
        )
 
        # 5. 重构观测
        o_hat = decoder(s_t)
        recon_loss = F.binary_cross_entropy(o_hat, o_t)
 
    # 6. 总损失
    loss = recon_loss + beta * kl_loss
    optimizer.step()

训练稳定性

  1. KL 权重调度:前期 (允许后验偏离先验),后期 (强制对齐)
  2. Free Bits:DreamerV3 允许每个隐变量有最低 bit 的 KL 预算,防止 posterior collapse
  3. 目标网络:转移模型的参数使用 EMA(指数移动平均)更新,稳定学习

推理/想象 rollout

后验 vs 先验

推理时(真实环境交互)使用后验

想象时(imagination rollout)使用先验

这就是为什么 KL 正则项如此重要:它确保先验与后验足够接近,使得在推理时用先验替代后验不会引入过大的误差。

想象 rollout 中的状态采样

def imagine_rollout(s_0, policy, world_model, H):
    s = s_0
    for k in range(H):
        # 策略基于当前隐状态
        a = policy(s)
 
        # 想象下一步(使用先验)
        # Dreamer 使用随机性(reparameterized sampling)
        s_next = world_model.transition(s, a)  # s_next ~ p(s'|s,a)
 
        s = s_next
    return trajectory

优点与局限

优点

  1. 不确定性建模:随机后验能够建模环境的不确定性(aleatoric uncertainty)
  2. 多模态推断:给定同一观测,后验可以是多峰的(如两个同样合理的物体位置)
  3. 记忆容量:随机性增加了隐空间的容量,允许记忆更多历史信息
  4. 鲁棒性:对噪声观测更加鲁棒,因为后验分布会对冲不确定性

失败模式

  1. Posterior Collapse:后验 退化为先验 ,失去推断能力
  2. 先验偏离:转移模型预测的先验与编码器后验差距过大,导致想象 rollout 失效
  3. 模态坍缩:后验多峰退化为单峰,失去多模态表达能力

与前后内容的衔接

  • 继承:从”观测压缩”的确定性编码发展为随机后验推断
  • 引向:为”离散状态与连续状态”的讨论奠定基础(随机性是多模态的基础)

可复现实现要点

最小模块

class LatentStateInferencer(nn.Module):
    def __init__(self, obs_dim, action_dim, hidden_dim=256, latent_dim=32):
        # RNN 用于聚合历史信息
        self.rnn = nn.GRUCell(obs_dim + action_dim, hidden_dim)
        # 后验网络(给定 RNN hidden state)
        self.posterior_net = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, latent_dim * 2)  # mu + logvar
        )
        # 先验网络(转移模型的一部分)
        self.prior_net = nn.Sequential(
            nn.Linear(latent_dim + action_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, latent_dim * 2)
        )
 
    def forward(self, h_prev, a_prev, o_t, training=True):
        # 更新 RNN
        h = self.rnn(torch.cat([h_prev, a_prev], dim=-1), o_t)
        # 后验推断(给定观测)
        mu, logvar = self.posterior_net(h).chunk(2, dim=-1)
        std = torch.exp(0.5 * logvar)
 
        if training:
            eps = torch.randn_like(std)
            s_t = mu + eps * std
        else:
            s_t = mu
 
        # 先验(给定上一状态和动作)
        prior_mu, prior_logvar = self.prior_net(torch.cat([s_t, a_prev], dim=-1)).chunk(2, dim=-1)
 
        return s_t, h, mu, std, prior_mu, prior_logvar

常见bug

  1. 历史信息丢失:RNN hidden state 未正确传递,导致后验忽略历史上下文
  2. KL 权重为 0:训练时忘记加入 KL 正则,导致先验-后验不一致
  3. 重参数化缺失:推理时使用确定性输出,但训练时需要随机采样

章节摘要

  1. 隐状态 是世界的”真实”表征,但不可直接观测,需要通过后验推断近似。
  2. 变分推断通过识别网络 近似真实后验
  3. ELBO 中的 KL 项 强制先验转移与后验推断保持一致。
  4. 重参数化技巧使得随机后验的梯度可以通过神经网络反向传播。
  5. 确定性后验(RNN hidden)简单但表达能力有限;随机后验(VAE-style)强大但训练更复杂。
  6. 训练时使用后验 编码真实观测;推理时使用先验 进行想象 rollout。
  7. KL 正则的一致性是想象 rollout 可靠性的核心保证。
  8. Free Bits机制防止 posterior collapse,确保每个隐变量至少贡献一定信息量。

关键词

隐状态、后验推断、变分推断、潜空间、ELBO、KL 正则、重参数化技巧、确定性后验、随机后验、RSSM、先验分布、后验分布、转移模型、编码器、不确定性建模、多模态推断、记忆容量、posterior collapse、Free Bits、EMA 更新