离散状态与连续状态
一句话定位
问题:真实世界的状态可以是离散的(棋盘格、离散物体类别)也可以是连续的(位置、速度、加速度)。如何在同一世界模型中同时建模这两种状态?何时选择离散表示,何时选择连续表示?
定位:第二章第三节,承接”隐状态与潜空间”,深入探讨混合状态空间的建模方法,是世界模型理论的核心延伸。
前置依赖
- 1-世界模型的定义:世界模型的基本架构
- 2-理解世界与预测未来:理解与预测的整体框架
- 1-观测压缩与状态抽象:观测到状态的降维映射
- 2-隐状态与潜空间:后验推断与潜空间结构
核心思想
连续状态适合表达物理量(位置、速度、温度)——这些量可以平滑变化,且插值有意义。
离散状态适合表达语义类别(物体 ID、房间类型、任务阶段)——这些量有明确边界,且组合爆炸是主要挑战。
关键洞察:真实世界同时包含两者。世界模型需要能够混合建模:用连续变量描述”怎么动”,用离散变量描述”是什么”。
离散 vs 连续的权衡:
| 维度 | 离散状态 | 连续状态 |
|---|---|---|
| 表示能力 | K 个类别,无法插值 | 无限精度,可插值 |
| 不确定性 | Categorical 分布 | Gaussian 分布 |
| 组合爆炸 | 连续空间平滑 | |
| 梯度传播 | 需要 Gumbel-Softmax | 直接可微 |
| 适用场景 | 语义/概念/分类 | 物理/几何/连续量 |
模型结构图
混合状态空间世界模型
===================
观测 o_t
│
▼
┌─────────────────────────────────────────────────────────┐
│ 编码器 Encoder │
│ │
│ o_t ──▶ Conv Layers ──▶ Feature Map │
│ │ │
│ ▼ │
│ ┌──────────┴──────────┐ │
│ │ │ │
│ ▼ ▼ │
│ 连续后验 q_s(s_t|o) 离散后验 q_z(z_t|o) │
│ N(μ_s, σ_s²) Categorical(π) │
│ │
└─────────────────────────────────────────────────────────┘
│ │
▼ ▼
连续状态 s_t 离散状态 z_t
∈ ℝ^{d_s} ∈ {1,...,K}
│ │
└─────────┬───────────┘
│
▼
┌─────────────────┐
│ 联合隐状态 │
│ h_t = (s_t, z_t)│
└─────────────────┘
│
▼
┌─────────────────────────────────────────────────────────┐
│ 转移模型 Transition │
│ │
│ 输入: (s_t, z_t), a_t │
│ 输出: (s_{t+1}, z_{t+1}) 先验 │
│ │
│ p(s_{t+1}, z_{t+1}|s_t, z_t, a_t) │
│ = p(s_{t+1}|s_t, z_t, a_t) · p(z_{t+1}|s_t, z_t, a_t)│
│ │
└─────────────────────────────────────────────────────────┘
Gumbel-Softmax 图解
Gumbel-Softmax 重参数化
=======================
离散分布参数: π = [π_1, π_2, ..., π_K] (K个类别的概率)
Gumbel 噪声: g = [-log(-log(u_1)), ..., -log(-log(u_K))], u_i ~ Uniform(0,1)
温度参数: τ > 0 (τ→0 趋近硬分布, τ→∞ 趋近均匀)
logits = log(π) + g
exp(logits_i / τ)
π_i' = ─────────────────
Σ_j exp(logits_j / τ)
当 τ → 0 时,π_i' → one_hot(argmax_i(logits_i))
当 τ → ∞ 时,π_i' → 1/K(均匀)
数学推导
1. 连续状态的数学建模
高斯分布后验:
重参数化:
先验转移(连续状态的高斯转移):
2. 离散状态的数学建模
Categorical 分布后验:
其中
先验转移(离散状态的 Categorical 转移):
这可以建模状态机的转移逻辑,如游戏中的阶段切换(开始 → 游戏中 → 结束)。
3. Gumbel-Softmax(离散重参数化)
直接对离散变量采样无法反向传播。Gumbel-Softmax 提供了可微的近似:
Gumbel-Softmax 分布:
其中:
是第 k 个类别的概率(来自 softmax) 是 Gumbel 噪声 是温度参数
温度参数的影响:
| 用途 | ||
|---|---|---|
| 推理时使用硬分布 | ||
| 标准 softmax,接近均匀 | 训练初期探索 | |
| 更平滑,接近均匀 | 训练时增加探索 |
直通估计器(Straight-Through Estimator):
在反向传播时,用
4. 混合状态空间的联合建模
联合隐状态:
联合 ELBO:
分解假设(简化计算):
5. 离散状态的优势:组合爆炸
连续空间的组合问题:如果世界有 10 种物体,每种物体有 100 种可能位置,用连续变量建模需要
离散状态的组合效率:
- 每个离散变量有 K 个取值,N 个变量产生
个组合 - 但只有少数组合在物理上是合理的
- 离散转移矩阵可以建模稀疏的合法转移
例子:Atari 游戏中的状态(如游戏阶段、物体类别)更适合用离散变量建模:
6. 连续与离散的状态转移对比
连续转移(高斯):
- 物理意义:状态在连续空间平滑移动
- 梯度:高斯分布的均值和方差可直接微分
- 不确定性:协方差矩阵捕获状态不确定性
离散转移(Categorical):
- 物理意义:状态在离散类别间跳转
- 梯度:需要 Gumbel-Softmax 或直通估计器
- 不确定性:概率向量捕获跳转的不确定性
训练细节
数据构造
混合状态空间的训练数据与标准世界模型相同:
模型自己学习哪些状态是离散的,哪些是连续的。
输入输出(IO)
| 变量 | 维度 | 说明 |
|---|---|---|
| 连续隐状态 | ||
| 离散隐状态(one-hot 或 Gumbel-Softmax) | ||
| 联合隐状态 |
损失计算
连续部分的 KL:
离散部分的 KL:
优化器配置
- 温度调度:初始
,逐渐降至 (更硬的分布) - KL 权重:连续和离散的 KL 项可以使用不同的权重
和 - 直通估计器:在训练后期使用硬分布(减少梯度方差)
训练流程
# 混合状态空间训练流程
for batch in dataloader(trajectories):
# 1. 编码(连续 + 离散)
mu_s, logvar_s = encoder_continuous(encoder_output)
logit_z = encoder_discrete(encoder_output) # K 维 logit
# 2. 连续状态采样(重参数化)
std_s = exp(0.5 * logvar_s)
eps_s = torch.randn_like(std_s)
s_t = mu_s + eps_s * std_s
# 3. 离散状态采样(Gumbel-Softmax)
gumbel_noise = -log(-log(torch.rand_like(logit_z)))
tau = schedule(step) # 温度调度
z_t = F.gumbel_softmax(logit_z + gumbel_noise, tau=tau, hard=False)
# 4. 先验预测
prior_mu_s, prior_logvar_s, prior_logit_z = transition_prior(s_t, z_t, a_t)
# 5. 计算 KL
kl_s = kl_normal(mu_s, std_s, prior_mu_s, exp(0.5 * prior_logvar_s))
kl_z = kl_categorical(F.softmax(logit_z, dim=-1), F.softmax(prior_logit_z, dim=-1))
# 6. 重构
o_hat = decoder(s_t, z_t)
# 7. 总损失
loss = recon_loss + beta_s * kl_s + beta_z * kl_z
optimizer.step()推理/想象 rollout
离散状态的想象 rollout
def imagine_rollout_mixed(s_0, z_0, policy, world_model, H, tau):
s, z = s_0, z_0
for k in range(H):
# 策略输入(连续 + 离散)
h_k = torch.cat([s, F.gumbel_softmax(z, tau=tau, hard=True)], dim=-1)
a = policy(h_k)
# 先验预测
prior_mu_s, prior_logvar_s, prior_logits_z = world_model.transition_prior(s, z, a)
# 采样下一步(连续 + 离散)
s = prior_mu_s + torch.randn_like(prior_mu_s) * exp(0.5 * prior_logvar_s)
z = F.gumbel_softmax(prior_logits_z, tau=tau, hard=True)
return trajectory优点与局限
连续状态的优点
- 高精度:可以表示任意精度的连续值
- 可微性:梯度传播自然,无需特殊处理
- 插值友好:两个连续状态之间的插值有意义
连续状态的局限
- 组合爆炸:多个连续变量组合产生无限空间
- 不必要精度:对离散语义(如物体类别)不需要高精度
离散状态的优点
- 语义表达:直接表达概念、类别、阶段等语义信息
- 组合效率:
组合但转移矩阵稀疏 - 可解释性:每个类别有明确含义
离散状态的局限
- 梯度传播:需要 Gumbel-Softmax 或直通估计器
- 硬决策:推理时需要 argmax,无法软决策
- 温度调度:需要精心设计温度退火策略
混合状态的优点
- 表达能力最大化:连续描述”怎么动”,离散描述”是什么”
- 任务适配:根据任务特性选择合适的状态类型
与前后内容的衔接
- 继承:从纯连续隐状态(Dreamer V1/V2)扩展到混合状态空间(DreamerV3)
- 引向:为”世界模型主线论文”中的 DreamerV3 实现奠定理论基础
可复现实现要点
最小模块
class MixedLatentSpace(nn.Module):
def __init__(self, obs_dim, action_dim, continuous_dim=32, n_classes=10):
# 连续状态编码器
self.encoder_cont = nn.Sequential(
nn.Linear(obs_dim, 256),
nn.ReLU(),
nn.Linear(256, continuous_dim * 2) # mu + logvar
)
# 离散状态编码器
self.encoder_disc = nn.Sequential(
nn.Linear(obs_dim, 256),
nn.ReLU(),
nn.Linear(256, n_classes)
)
def forward(self, obs, tau=1.0):
# 连续状态
mu_s, logvar_s = self.encoder_cont(obs).chunk(2, dim=-1)
std_s = torch.exp(0.5 * logvar_s)
eps = torch.randn_like(std_s)
s = mu_s + eps * std_s
# 离散状态
logits_z = self.encoder_disc(obs)
z = F.gumbel_softmax(logits_z, tau=tau, hard=False)
return s, z, mu_s, std_s, logits_z常见bug
- 温度参数:
过大导致梯度方差大, 过小导致训练不稳定 - KL 权重不平衡:连续和离散的 KL 项需要分别加权(
vs ) - Gumbel-Softmax 硬模式:推理时忘记使用 hard=True,导致连续近似
章节摘要
- 连续状态用高斯分布建模,适合表达物理量(位置、速度),可直接微分的优点。
- 离散状态用 Categorical 分布建模,适合表达语义类别(物体 ID、阶段),需要 Gumbel-Softmax 处理梯度。
- Gumbel-Softmax通过温度参数
在硬分布( )和软分布( )之间插值,是离散重参数化的核心。 - 混合状态空间
同时包含连续和离散变量,结合两者的表达优势。 - 联合 ELBO需要分别计算连续和离散部分的 KL 项,并使用不同的权重(
)。 - 温度调度是训练离散状态模型的关键:前期高温度(探索),后期低温度(利用)。
- 直通估计器允许在反向传播时使用硬分布的梯度,绕过 Gumbel-Softmax 的采样。
- DreamerV3成功将混合状态空间用于 Atari 等离散-连续混合环境,验证了该方法的有效性。
- 离散状态适合建模组合爆炸的语义空间,但需要稀疏的转移矩阵避免过拟合。
- 连续状态适合建模平滑变化的物理量,但不需要过高精度时可适当离散化。
关键词
离散状态、连续状态、混合状态空间、Gumbel-Softmax、Categorical分布、Gumbel噪声、温度参数、直通估计器、one-hot编码、组合爆炸、连续状态转移、离散状态转移、KL散度分解、温度调度、语义表达、物理连续量、表示学习、变分推断、梯度传播