原文链接：https://ieeexplore.ieee.org/abstract/document/10522953

1. 世界模型的发展

A. 世界模型的结构基础

在这里插入图片描述

世界模型包含4个关键组件，以模拟人类连贯的思考和决策过程。

a）感知模块使用如变分自编码器（VAE）、掩膜自编码器（MAE）或离散自编码器（DAE）等编码模块，将环境输入（如图像、视频、文本、控制指令等）处理并压缩为更易于管理的格式。

b）记忆模块用于记录和管理过去、现在和预测的世界状态及其相关代价或奖励。

c）控制/行为模块评估当前状态和世界模型的预测，以决定实现目标（如最小化代价或最大化奖励）的最优行为序列。该模块与世界模型模块区分开来，使用小参数集独立训练。

d）世界模型模块估计当前世界状态丢失的信息并预测未来的环境状态。

模型预测的确定性和真实世界现象的不确定性之间的协调是一个挑战，因此，可以使用温度变量引入不确定性，或使用结构化框架（如循环状态空间模型（RSSM）与联合嵌入预测结构（JEPA）），以微调预测精度与灵活性的平衡。此外，使用Top-K采样，以及从CNN模型转换到Transformer模型（如Transformer状态空间模型（TSSM）和时空Patchwise Transformer（STPT））通过将世界模型的预测与真实世界的可能发展对齐，能更好地近似真实世界的复杂程度和不确定性。

世界模型中最常用的核心结构是RSSM与JEPA。

1）循环状态空间模型（RSSM） 在隐空间中实现纯正向预测。
在这里插入图片描述
上图c展示了潜在动态模型的3次步进。随机变量（圈）和确定性变量（方块）在模型结构内交互——实线表示生成过程，虚线表示推断过程。

与无法捕捉多种潜在未来的、确定性初始状态的RNN（图a）和完全随机、无信息保留能力的状态空间模型（SSM，图b）不同，RSSM将状态分解为随机性状态和确定性状态，有效利用确定性元素的预测稳定性和随机性元素的自适应潜力。

记观测和行为序列为 $(x_0,a_1,x_1,a_2,x_2,\cdots,a_T,x_T)$ 。观测到 $x_t$ 后，智能体进行行为 $a_{t+1}$ ，并接收下一步观测 $x_{t+1}$ 。RSSM将观测和状态转移建模为下列生成过程：
$p(x_{0:T}|a_{1:T})=\int\prod_{t=0}^{T}p(x_t|z_{\leq t},a_{\leq t})p(z_t|z_{<t},a_{\leq t})dz_{0:T}$

其中 $z_{0:T}$ 为随机隐状态。近似后验定义为：
$q(z_{0:T}|x_{0:T},a_{1:T})=\prod_{t=0}^Tq(z_t|z_{<t},a_{\leq t},x_t)$

RSSM使用共享的GRU来将 $z_{<t}$ 和 $a_{\leq t}$ 压缩为确定性编码 $h_t$ ：
$h_t=GRU(h_{t-1},MLP(concat[z_{t-1},a_t]))$

随后，用于计算先验、似然和后验的充分统计量：
$p(z_t|z_{<t},a_{\leq t})=MLP(h_t)\\ p(x_t|z_{\leq t},a_{\leq t})=\mathcal N(\hat x_t,1)\\ \hat x_t=Decoder(concat[h_t,z_t])\\ q(z_t|z_{<t},a_{\leq t},x_t)=MLP(concat[h_t,e_t])\\ e_t=Encoder(x_t)$

训练目标是最大化证据下界（ELBO）：
$\log p(x_{0:T}|a_{1:T})\geq\mathbb E_q[\sum_{t=0}^T\log p(x_t|z_{\leq t},a_{\leq t})-\mathcal L_{KL}(q(z_t|z_{<t},a_{\leq t},x_t),p(z_t|z_{<t},a_{\leq t}))]$

2）联合嵌入预测结构（JEPA） 关注表达空间而非直接且具体的预测。
在这里插入图片描述
如图所示，JEPA通过双编码器将输入 $x$ 和目标 $y$ 抽象为表达 $s_x,s_y$ ，并使用隐变量 $z$ 进行预测。其预测过程更简单，且能保证结果的相关性和可靠性。

JEPA的核心是能量函数 $E_w(x,y,z;\theta)$ ，其捕捉了模型的预测误差。这里 $\theta$ 表示模型参数。能量函数定义为
$E_w(x,y,z;\theta)=\|s_y-Pred(s_x,z;\phi)\|_2^2+\lambda\|z\|_2^2$

其中，L2范数平方衡量了预测表达与目标表达 $s_y$ 的欧式距离，即模型预测误差；预测函数 $P re d$ 将输入表达 $s_x$ 和隐变量 $z$ 映射到目标空间，其参数为 $\phi$ 。正则化项 $\lambda\|z\|_2^2$ 则惩罚模型的复杂性，以避免过拟合。

优化的目标是通过寻找 $\theta,\phi,z$ 最小化能量函数 $E_w$ ，可表达为以数据分布为约束的复值拉格朗日优化问题：
$L(\theta,\phi,z;x,y,\alpha)=E_w(x,y,z;\theta)-\alpha(h(x,y,z;\theta,\phi)-c)$

其中 $L(\theta,\phi,z;x,y,\alpha)$ 是拉格朗日方程， $\alpha$ 为拉格朗日乘子，强制满足约束 $h(x,y,z;\theta,\phi)=c$ ； $h(x,y,z;\theta,\phi)$ 为约束函数，其参数为 $\theta$ 和 $\phi$ ， $c$ 为函数 $h$ 的目标值（常数）。

JEPA的训练包含高阶优化方法，考虑二阶导数以保证收敛：
$\theta_{t+1}=\theta_t-\eta\nabla_\theta^2L(\theta_t,\phi_t,z_t;x,y,\alpha_t)$

其中 $\theta_{t+1}$ 为更新后的参数向量， $\eta$ 为学习率； $\eta\nabla_\theta^2L(\theta_t,\phi_t,z_t;x,y,\alpha_t)$ 为拉格朗日函数 $L$ 对 $\theta$ 二阶偏导数的Hessian矩阵。

考虑到 $z$ 的高维特性和多模态分布的可能性，JPEA可采用变分近似方法处理后验 $p(z|x,y;\theta)$ ，得到变分下界：
$\log p(y|x;\theta,\phi)\geq\mathbb E_{q(z|x;\psi)}[\log p(y|x,z;\theta;\phi)]-KL[q(z|x;\psi)||p(z|x;\theta)]$

其中 $\log p(y|x;\theta,\phi)$ 为数据 $y$ 在给定 $x$ 和模型参数 $\theta,\phi$ 时的对数似然， $\mathbb E_{q(z|x;\psi)}[\cdot]$ 为关于分布 $q(z|x;\psi)$ 的期望， $KL[q(z|x;\phi)||p(z|x;\theta)]$ 为变分分布 $q(z|x;\psi)$ 和先验分布 $p(z|x;\theta)$ 的KL散度。