基于联合嵌入范式的潜在动态预测表征模型

doi:10.19328/j.cnki.2096-8655.2026.01.018

基于联合嵌入范式的潜在动态预测表征模型

doi: 10.19328/j.cnki.2096‐8655.2026.01.018

宋飞 ¹ ，高超 ² ，田园 ³ ，邴琪 ¹ ，步士超 ⁴ ，刘艳阳 ⁴ ，王天成 ⁴ ，陈筠力 ⁵ ，邵晓巍 ¹

1. 上海交通大学航空航天学院，上海 200240

2. 北京跟踪与通信技术研究所，北京 100094

3. 上海交通大学自动化与感知学院，上海 200240

4. 上海卫星工程研究所，上海 201109

5. 上海航天技术研究院，上海 201109

基金项目: 上海航天创新基金资助项目（USCAST2023-11）

详细信息

通信作者

陈筠力,E-mail：xw.shao@sjtu.edu.cn。

中图分类号: TP181；V448

文献标志码: A

Latent Dynamic Prediction Representation ModelBased on Joint Embedding paradigm

SONG Fei ¹ ， GAO Chao ² ， TIAN Yuan ³ ， BING Qi ¹ ， BU Shichao ⁴ ， LIU Yanyang ⁴ ， WANG Tiancheng ⁴ ， CHEN Junli ⁵ ， SHAO Xiaowei ¹

1. School of Aeronautics and Astronautics, Shanghai Jiao Tong University, Shanghai 200240, China

2. Beijing Institute of Tracking and Telecommunications Technology, Beijing 100094, China

3. School of Automation and Intelligent Sensing, Shanghai Jiao Tong University, Shanghai 200240, China

4. Shanghai Satellite Engineering Research Institute, Shanghai 201109, China

5. ShanghaiAcademy of Spaceflight Technology, Shanghai 201109, China

摘要

针对序列决策环境中交互数据稀缺导致的动态建模与泛化难题，提出一种基于联合嵌入预测架构的潜在动态预测表征模型。该模型引入“动作查询式”注意力机制作为核心创新：以动作为查询（Query）、历史状态序列为键（Key）和值（Value），在潜在空间中直接、高效地学习和表征动作——状态的转移关系，从而规避像素级重建的计算负担与信息冗余。在 Atari 学习环境上的实验评估表明：所提模型在训练过的环境中能够准确地进行 15 步开环预测，并在未见过的未知环境中实现了约 3 步的有效外推。研究结果证实：该方法能够在有限交互数据下学习到具备一定泛化能力的世界模型，为实现通用的序列决策提供了有效支撑。

关键词

表征模型 / 联合嵌入 / 动作条件预测 / 多步开环预测 / 泛化

Abstract

To address the challenges of dynamic modeling and generalization in sequential decision-making environments with scarce interaction data，a latent dynamic prediction representation model （PRM）based on the joint embedding predictive architecture （JEPA）is proposed.A core innovation of the model is the "action-query" attention mechanism：using actions as queries and historical state sequences as keys and values. This allows for the direct and efficient learning and representation of action-state transition relationships in latent space，bypassing the computational burden and information redundancy of pixel-level reconstruction. The experimental evaluations in the Atari Learning Environment （ALE）demonstrate that the proposed model can accurately perform 15 steps of open-loop prediction in trained environments and achieve effective extrapolation for about 3 steps in unseen，unknown environments. The results confirm that this method can learn a world model with certain generalization capability under limited interaction data，providing effective support for general sequential decision-making.

Keywords

representation model / joint embedding / action-conditional prediction / open-looped prediction / generalization

0 引言 1 潜在动态预测表征模型 1.1 模型基本架构 1.2 图像编码器与动作编码器 1.3 动作条件预测器 1.4 模型训练与损失函数 2 实验方案与结果 2.1 实验平台与数据来源 2.2 多步开环预测 2.3 消融实验及模型稳健性验证 3 实验结果分析 4 结束语

0 引言

在自主在轨服务、深空探测与未来空间对抗等高动态、高不确定性任务中，智能体对环境动态进行精确建模与长时程预测，是实现前瞻性自主决策的关键能力。在基于模型的强化学习框架下，核心在于学习一个高保真的“世界模型”^[1]：理解环境动态特性，在智能体采取动作后，能够准确预测环境状态的转移结果^[2-3]。如何构建预测表征的世界模型是当前研究的前沿热点。

现有研究根据模型架构的不同可分为以下 3 类路线：基于循环神经网络（Recurrent Neural Network，RNN）的架构，基于 Transformer 网络^[4]的架构及基于扩散模型的架构。

RNN 是早期实现世界模型的经典路径。这类方法学习一个以 RNN 为核心的环境动态模型，能够在给定动作的条件下，预测下一步环境状态，并由此自主生成虚拟的训练数据，用于后续决策模型训练。HA 等^[3]于 2018 年提出基于 RNN 的世界模型，以图像为输入，在同任务的大量数据上训练，可在给定当前观测与历史信息时预测下一时刻的图像；2019 年 HAFNER 等^[5]提出循环状态空间模型（Recurrent State-Space Model，RSSM），利用动作—观测序列学习环境的潜在动态，进而在给定动作下预测未来状态及对应奖励。随后，HAFNER 在Dreamer 系列（如 Dreamer-v1^[6]、Dreamer-v2^[7]）中改进了 RSSM 的表征与损失设计，实现对复杂离散动作任务的时序建模。该类模型训练完成后，可作为“可控模拟器”，给定当前状态与动作序列，便能生成相应的未来状态，从而为决策智能体在“想象”中进行学习与规划提供足量数据，并在 Atari等环境上表现出优于传统无模型强化学习的性能。此外DeepMind 的 MuZero 通过学习隐式动态模型来执行搜索，不依赖环境真实规则信息，在 Atari 游戏和棋类中达到超人水准^[8]。这些成果证明了潜在世界模型结合规划优化的巨大潜力。然而，RNN 架构也存在局限：首先，训练往往依赖海量环境交互数据，且受限于 RNN 网络结构的顺序依赖，序列数据难以并行训练^[9-11]；其次，RNN 捕获长距离关系能力有限，对长序列数据前后关系的建模存在不足^[4，12-14]。

为弥补 RNN 在长时序建模与并行计算上的短板，Transformer 被引入世界模型研究。MICHELI等^[15]提出基于 Transformer 的世界模型 IRIS，其利用离散自动编码器+自回归 Transformer 来建模环境，将动力学学习转化为序列建模问题，显著提高模型的样本效率。IRIS 在等价于仅约 2 h 交互数据的情况下即可学得有效世界模型，并在多个视觉控制任务上取得优异成绩，IRIS 变体 Δ-IRIS 以“离散潜变量编码+自回归 Transformer”建模“相邻步增量”，在 Atari-100k 与 Crafter 等长时任务上兼具数据效率与生成精度^[16]。BRUCE 等^[17] 证明了Transformer 世界模型在数据高效性上的优势，仅通过十万步左右的环境交互，基于 Transformer 的模型即可达到与 RSSM 等结构相当的性能。此外以TrajectoryTransformer^[18] 与 DecisionTransformer^[19]为代表的研究方法，直接将序列决策问题定式为序列生成问题，利用世界模型直接进行决策，充分利用 Transformer 架构对长时序数据前后文关系的学习能力：TrajectoryTransformer 将状态、动作与奖励离散为令牌，按时间顺序输入到 Transformer 中，采用类似大语言模型的自回归训练，期望模型依次输出下一时刻的状态与动作；DecisionTransformer 在相同的数据组织方式上引入“目标回报”条件，直接生成实现目标回报的策略。此类方法在特定序列决策任务上取得了领先结果，但通常需要大规模且同分布的数据支撑，对分布外情形的适配多依赖任务特定的设计。为提高世界模型的泛化性及分布外环境的适应性，ZHANG 等^[20]提出的 WHALE 框架通过行为条件和回溯展开 2 项技术，显著提高世界模型对分布外情景的适应性，并能评估预测轨迹的不确定度。基于此，ZHANG 等^[20-21]构建具备时空 Transformer 架构的 WHALE-ST 模型，在模拟控制任务中取得了更准确的价值评估和更高保真的视频生成效果；模型在大规模数据上训练，在真实机械臂操作任务中，以最小演示数据实现了对复杂技能的泛化，体现出良好的可扩展性。此外，基于Transformer 架构变种（Vision transformer，ViT）的联合嵌入预测架构（Joint Embedding Predictive Architecture，JEPA）构建的世界模型，通过自监督学习在潜在空间理解环境。如基于图像的联合嵌入预测架构（Joint Embedding Predictive Architecture for Images，I-JEPA）可在图像中完成未知结构补全^[22]，视频联合嵌入预测架构（Video Joint Embedding Prediction Architecture，V-JEPA）将该能力扩展至视频^[23-24]。同样，有工作以自监督形式训练基于自注意力的预测编码网络，能够仅凭视觉轨迹判断主体在虚拟世界中的位置^[25]。此类模型能够更好地捕捉序列视觉中的环境信息，而非仅拟合像素，但由于未显式接入动作信号，缺乏对“智能体—环境”交互的建模，难以直接用于决策与规划。

扩散模型作为新型生成架构也被引入世界模型与决策规划领域。JANNER 等^[26]提出 Diffuser 算法，将生成式扩散模型用于规划：其通过迭代去噪生成整段轨迹，以替代逐步预测，从而有效缓解长期滚动预测的误差累积。Diffuser 在长序决策和测试时的灵活适应性方面显示出独特优势。后续研究如 DiffuserLite 进一步针对扩散规划的实时性进行优化，设计渐进细化生成的轻量级扩散架构，利用粗到细的轨迹规划过程，大幅减少冗余信息建模，将决策频率提升一个数量级^[27]。但这类模型需要大量数据进行扩散模型训练，且对环境特性依赖程度高，在交互数据有限或环境改变时，模型缺乏鲁棒性和适应性。

因此，当前研究仍面临一个核心挑战：在交互数据稀缺的条件下，如何有效学习动作与环境状态转移之间的因果关系。为此，本文提出一种基于改进 Transformer 架构的潜在动态表征模型，聚焦于学习“动作—状态转移”的因果规律，获得可迁移的“经验”，以支撑泛化的序列决策。

本文的主要创新与贡献如下：

1）提出基于 JEPA 架构的潜在动态预测表征模型（Latent Dynamic Predictive Representation Model，PRM），模型学习与交互强相关的紧凑环境表征，能够对未来交互结果做出准确预测，并在未知任务的交互中体现出一定的泛化能力，能够对新环境中的智能体—环境互动进行有效表征。

2）设计“动作查询式”注意力机制，以“动作”向量作为查询 Q，以历史状态序列为键 K 与值 V，显式建模“在特定历史条件下，执行某动作对未来的影响”。该机制在多步开环预测任务中表现稳定，模型在已训练环境中可准确预测未来 15 步环境状态，在未知环境中可实现约 3 步的环境状态有效外推。

1 潜在动态预测表征模型

PRM 基于联合嵌入预测架构 JEPA 进行构建，旨在潜在空间中直接学习环境的动作条件转移函数。传统世界模型常依赖像素级的图像重建作为训练目标，此范式不仅计算成本高昂，且会迫使模型关注与任务无关的背景细节，从而影响学习效率与泛化能力。潜在动态预测表征模型的设计遵循JEPA 的核心思想：在紧凑的潜在向量空间中直接进行预测，而非重建高维观测。该策略旨在实现2 个核心目标：

1）学习任务相关的紧凑表征。通过在潜在空间定义损失函数，驱动模型剥离观测信息中的无关变量，专注于捕捉对预测和规划至关重要的动态信息。

2）提升学习效率与泛化能力。规避像素级重建带来的计算负担与不稳定性，将模型的计算资源集中于学习状态转移的内在规律，从而提升训练效率和对未知环境的泛化能力。

1.1 模型基本架构

PRM 模型主要由 3 部分组成：图像编码器 f、动作编码器 h，以及动作条件预测器 g。图像编码器f 将高维的观测图像输入I_t映射到一个紧凑、抽象的潜在状态表征x_t。动作编码器 h 将离散的动作输入M_t映射为一个潜在动作表征a_t。动作条件预测器g 在潜在空间中运作，根据当前潜在状态 x_t 和采取的动作 a_t，来推断下一时刻的潜在状态 x_t'_{+ 1}。模型在训练过程中，希望模型的预测结果与下一时刻潜在状态的真实值相同。模型架构如图 1 所示。

1.2 图像编码器与动作编码器

图像编码器负责将图像输入高效地映射为信息密集的潜在状态表征

x_{t} \in R^{256}

。其主体结构为深度残差网络（Residual Network，ResNet）^[28]，由 1 个初始特征提取层、4 个级联的下采样模块及 1 个最终的特征投影头构成。值得注意的是，每个下采样模块内部集成了注意力模块（Squeeze-and-Excitation，SE）^[29]，以显式建模通道间依赖关系并自适应重标定通道响应，从而使网络聚焦于判别性更强的特征通道。最终，特征通过投影头压缩至 256 维潜在空间，并施加 Tanh 激活函数以约束表示的数值范围、稳定后续注意力计算。

动作编码器面向离散动作输入。首先将动作索引通过嵌入层映射为高维稠密向量

a_{t} \in R^{256}

，随后经过 1 个小型多层感知机，将其投影到与图像潜在特征一致的 256 维空间。该维度对齐确保动作表征与状态表征处于同一潜在空间，便于后续注意力机制对两者进行可比性计算；动作分支同样使用 Tanh激活以维持数值稳定性与与状态分支的一致性。

图 1 PRM 模型架构

Fig. 1 Architectureof the PRM

1.3 动作条件预测器

动作条件预测器是整体架构的核心，旨在学习由历史状态与当前动作共同决定的环境条件转移函数，即在给定 x_{{ 1：t }}与 a_t 的情况下，对未来潜在状态 x_{t +1}（及其多步外推）进行建模。动作条件预测器基于 Transformer 架构，采用“动作查询式”注意力机制，以高效建模动作与状态转移之间的因果关系。

该机制是对标准自注意力机制的改进，采用非对称的交叉注意力设计：其中查询 Q由潜在动作表征a_t 生成，键 K 和值 V 均由潜在状态表征x_t 生成，即：

Q = W_{q} A^{T}

（1）

K = W_{k} X^{T}

（2）

V = W_{v} X^{T}

（3）

A = [a_{t}], t = 1, 2, \dots, t

（4）

X = [x_{t}], t = 1, 2, \dots, t

（5）

式中：W_q、W_k 及 W_v 分别为 Q、K、V 的投影矩阵，由模型训练得到。

为保证模型训练过程稳定，同时减小批次大小对训练的影响，引导模型更加关注于单条序列数据内部各向量间的关系，模型采用组归一化层进行数据归一化方法^[30]。

1.4 模型训练与损失函数

为充分利用 Transformer 的并行建模能力，训练过程在序列维度上进行批处理。每个训练步，从数据集中随机抽取一个批次的转换元组{ D_i}_{i = 1~N}，其中 N 为批处理大小，每条数据含有 k+1 张序列观测图像及k 个动作。

D_{i} = {(I_{t}, m_{t}, I_{t + 1}, \dots, I_{t + k}, m_{t + k}, I_{t + k + 1})}_{i}

（6）

将 k+ 1 张序列观测图像输入图像编码器得到k+ 1 个潜在状态表征(x_t，x_{t + 1}，···，x_{t + k}，x_{t + k + 1})，其中前 k个潜在状态表征(x_t，x_{t + 1}，···，x_{t + k})与对应的潜在动作表征向量(a_t，a_{t + 1}，···，a_{t + k})输入条件预测器，以生成后 k个潜在状态表征的预测值(x_t'_{+ 1}，···，x_t'_{+ k}，x_t'_{+ k + 1})。

损失计算采用负余弦相似度方法，对每个样本、每个时间步独立计算预测向量与目标向量之间的余弦相似度，再对批次与时间求平均，得到最终用于反向传播的标量损失：

L_{vec} = 1 - \frac{1}{N k} (\sum_{N} \sum_{k} \frac{x_{t + 1}^{'} \cdot x_{t + 1}}{‖ x_{t + 1}^{'} ‖ 2 {‖ x_{t + 1} ‖}_{2}})

（7）

该目标通过最大化预测与目标表征的方向一致性来驱动学习，仅关注向量方向而不受范数影响，有助于模型聚焦于状态变化的语义信息，从而提升潜在表征的稳健性与跨场景泛化能力。

在模型训练过程中，阻断作为预测目标值的梯度反向传播以防止模型坍缩。具体而言，每条数据经过图像编码器后得到的目标图像向量(x_{t + 1}，···，x_{t + k}，x_{t + k + 1})梯度不参与计算。此策略能够阻止梯度从损失函数经由预测目标回传至编码器参数，避免模型因预测值与目标值同步更新产生而学到一个平凡解^[31]。但实际效果证明，在多步开环预测中，仅梯度阻断难以阻止模型坍缩，具体情况在第 3 节中有详细分析。

2 实验方案与结果

2.1 实验平台与数据来源

实验基于 Atari 学习环境（Atari Learning Environment，ALE）^[32]开展，用以评估 PRM 模型的学习能力与分布外泛化能力。训练阶段在包含40 种不同 Atari 游戏环境的已知数据集上进行。验证阶段设置 2 类数据集：

1）已知环境验证集（以下简称已知集）：从参与训练的环境中选取 4 个游戏，采集独立的验证数据；

2）未知环境验证集（以下简称未知集）：从未参与训练的环境中选取 4 个游戏，采集验证数据。

需要说明的是，已知集与未知集中的数据均不参与训练。2 类验证数据均通过随机动作策略在对应环境下采集，形成定长的时序轨迹，每条样本的序列步长固定为 16。

2.2 多步开环预测

多步开环预测旨在评估模型，在没有真实环境反馈（即“开环”）的情况下，仅凭初始状态和动作序列推演未来动态的能力。

多步开环预测流程如下：模型训练完成后，在不进行任何参数微调或梯度更新的前提下，模型需依据给定的图像I₁ 及动作序列(a₁，a₂，···，a_m)，分别在验证集上进行多步动作后的潜在状态向量预测，即分别预测在经过 n(1 ≤ n ≤ s)个动作后的潜在状态向量 x_n。开环预测最大步数 s 为 32，预测过程中仅根据编码得到的初始潜在状态 x₁ 与给定动作序列进行前向推断，不引入真实未来观测或任何形式的在线校正。

评估指标采用均方差损失（Mean-square Error Loss，MSE Loss），在每条序列的各预测步长上独立计算预测向量与真实（编码后）向量的误差，取时间平均与样本平均后，再对同一验证类别内的 4 个游戏结果做算术平均，得到该类别的总体损失。选择RSSM 模型作为基线进行对比，在已知集上的多步开环预测性能对比如图 2 所示。结果显示，PRM 模型的预测损失曲线与 RSSM 基线模型基本重合，表明在分布内的环境（已训练过的游戏）中，该模型学习到的动态预测能力与 RSSM 相当。2 个模型的损失均在约 15 步后趋于平稳，这与训练序列长度（16 步）高度相关，在第 3 节中将进一步分析。

图 2 PRM 模型与 RSSM 模型在已知集中多步开环预测验证损失对比

Fig. 2 Multi-steps open-loop prediction loss of the PRM and RSSM on known set

在未知集上的泛化性能评估如图 3 所示。需要特别说明的是，RSSM 模型未参与此项对比。这是因为 RSSM 模型通常采用针对单一任务的“pertask”训练范式，其学习到的动态模型高度特异化于训练时所处的特定环境。因此，一个在环境 A 上训练的 RSSM 模型，其学到的表征和动态函数无法直接迁移应用于未知的环境 B。

图 3 未知集中 PRM 多步开环预测验证损失

Fig. 3 Multi-steps open loop prediction validation loss of the PRMon unknown set

相比之下，PRM 模型的核心优势之一在于采用了多环境混合训练（在 40 种不同游戏上训练），旨在学习跨环境的通用动态表征。图 3 的结果展示了PRM 模型在这一设置下的泛化能力：在分布外的未知环境中，模型仍可在约 3 步内进行有效的外推，之后损失趋于平稳，这证实了模型所学的动态表征具备一定程度的可迁移性。

2.3 消融实验及模型稳健性验证

为进一步验证本文所提模型中关键组件的有效性，并探究其在不同数据分布下的稳健性表现，本节设计针对性的消融实验与补充分析。本节主要包含以下 3 部分：

1）“ 动作查询式 ”注意力是本文区别于标准JEPA 和 Transformer 世界模型的核心机制创新；

2）梯度阻断是保证 JEPA 这类非对比学习范式免于表征坍缩、稳定训练的必要条件。验证这2 点的有效性，对证明本文方法的创新性和完备性至关重要；

3）模型稳健性分析：验证单一模型在不同数据分布（环境）下的性能一致性。

首先，研究验证“动作查询式”注意力机制的必要性。设置对照组，将其替换为标准的自注意力机制。在该对照组中，动作表征a_t与状态表征x_t 被拼接后，共同作为 Q、K、V 的来源。由于这种拼接后的潜变量已无法保证与真实的下一时刻状态 x_{t + 1}对齐，因此无法直接在潜层空间计算损失。为公平对比，“对照组”（标准自注意力）和“实验组”（本文的“动作查询式”注意力）均额外训练一个轻量级图像解码器，转而对比两者在 5 步开环预测后重建图像的质量，L1、L2 分别为预测结果的平均值误差和均方误差。采用峰值信噪比（Peak Signal-to-noise Ratio，PSNR）和结构相似性（Structural Similarity，SSIM）作为衡量指标。结果如图 4 所示，第 1 行为预测结果，第 2 行为真实场景图像。图 4 前 3 列为自注意力机制模型在进行 5 步开环预测后得到的重建图像，第 4 列是“动作查询式”注意力机制模型进行 5 步开环预测后的重建图像。结果表明：在多步预测中，标准自注意力机制无法有效解耦并学习动作与状态转移的因果关系，模型在多步预测后完全无法重建有效的图像；而本文提出的“动作查询式”机制在 5 步开环预测后依然能高度保真地重建图像。

图 4 自注意力与“动作查询式”注意力机制模型 5 步开环预测图像重建结果对比

Fig. 4 Comparison of 5-step open-loop prediction image reconstruction results between the self-attention mechanism model and the “action-query” mechanism model

其次，研究验证梯度阻断策略的有效性。如 1.4节所述，模型在目标分支施加梯度阻断以防止模型坍缩。设计移除该阻断的对比实验，结果如图 5 所示：缺乏梯度阻断的对照组模型，其训练损失依然收敛，但验证损失在短暂下降后迅速爬升并发散。训练及验证损失结果表明模型出现表征坍缩现象。该现象在多步开环预测中得到进一步证实，结果如图 6 所示，对照组的预测损失（MSE Loss）在所有已知集上均近乎为零且方差极小，说明模型在图像编码器部分即发生了数据坍缩。这证明梯度阻断是保证 JEPA 范式稳定训练、避免模型学到平凡解的必要设计。

图 5 缺乏梯度阻断下的 PRM 训练损失及验证损失情况

Fig. 5 Training loss and validation loss of the PRM under the absence of gradient blocking

图 6 缺乏梯度阻断下的 PRM 多步开环预测损失

Fig. 6 Multi-step open-loop prediction loss of PRM under the absence of gradient blocking

最后，为分析模型在不同数据分布下的表现，测试使用同一套固定参数的 PRM 模型在多个不同已知集上的性能表现，结果见表 1。

表 1 模型在不同数据集下的预测结果

Tab. 1 Prediction results of the model on different datasets

由表 1 可知，尽管 Alien-v5、Jamesbond-v5 等环境在视觉风格、动态特性上差异显著，但本文的单一模型在不同分布的数据集上均表现出一致，且呈现稳健的低预测误差。

3 实验结果分析

由图 2、图 3 可知：已知集上的结果体现 PRM 在已知环境下的稳定多步预测能力，未知集上的“3 步有效外推”结果验证模型具备的未知环境泛化能力。传统的“per-task”模型因其训练范式的局限性，无法直接迁移到未知环境 B 进行有效预测。相比之下，PRM 通过多环境混合训练，能够在未见过的环境中依然推演约 3 步的动态，证实模型所学习到的动态先验具备一定的可迁移性。

在已知集上，随预测步长增加，预测损失呈单调上升并在约 15 步后趋于平稳；在未知集上，分步损失在前 3 步同样呈现单调上升趋势。该现象与训练与评估的设定一致。模型每次以相同的起点 I₁ 出发，对每一未来步 x_t 并行产生条件预测，由于Transformer 的并行注意力可在单次前向中对多个时间位点进行信息聚合^[4]，即对任意一对潜在状态向量(x_i，x_j) 的关系匹配复杂度为 O(1)。因此，模型的多步预测总损失函数可设计为所有潜在状态向量 x_j(1 ≤ j ≤ k)到潜在状态向量 x_{k + 1} 的独立预测误差之和：

\begin{matrix} L_{multi} & = \frac{1}{N k} \sum_{i = 1}^{N} \sum_{j = 1}^{k} l ({\hat{x}}_{j + 1}^{(i)}, x_{j + 1}^{(i)}), \\ l & = 1 - \cos (∙, ∙) \end{matrix}

（8）

式中：x^{( i)}_{j + 1} 为第 i 条数据中第 j+1 步的潜在状态向量，

{\hat{x}}_{j + 1}^{(i)}

为 x^{( i)}_{j + 1} 的预测值。

从而在“短步”范围内常见到随步长近线性上升的损失曲线，之后进入平台期。

此外，模型在已知集和未知集上呈现不同的性能表现：在已知集上可预测 15 步；在未知集上，预测步长缩短，仅能进行 3 步的有效外推。分析认为，这种性能退化和平台期的出现，主要由 2 个因素共同导致。首先是训练窗口的局限性：已知集上的 15 步拐点与训练时采用的 16 步序列长度（见 1.4 节）高度吻合。这表明模型主要学习此窗口内的状态转移规律，当开环预测步长超出训练窗口（>15 步）时，外推缺乏有效的监督信号，预测误差不再显著增长，表现为平台期。其次是开环预测中的误差累积与表征坍缩：这一现象在未知集上的 3 步平台期表现得尤为突出。为从实证上验证表征坍缩的发生，额外训练 1 个轻量级图像解码器，对平台期（即“已知集 15 步之后”与“未知集 3 步之后”）的预测潜在状态向量进行可视化重建，随机选择 4 个样本进行重建，重建结果如图 7 所示，其中上一行为真值，下行为样本重建结果。重建结果清晰地揭示表征坍缩的实质：潜在表征丢失了关键的动态信息（如物体的运动），仅保留环境的静态背景信息。这种动态信息的缺失，使得模型无法继续推演有效的状态转移，预测器 g 的输出收敛到背景信息中，在损失曲线上便表现为平台期。

研究进一步从理论上分析误差累积的过程。在开环预测中，第 t + n 步的预测

{\bar{x}}_{t + n}

是基于前一步的预测

{\bar{x}}_{t + n - 1}

和动作 a_{t + n} 生成的，即：

{\bar{x}}_{t + n} = g ({\bar{x}}_{t + n - 1}, a_{t + n})

（9）

而真实动态为

x_{t + n} = g ({\bar{x}}_{t + n - 1}, a_{t + n})

（10）

图 7 潜在状态向量重建结果

Fig. 7 Reconstruction results of Latent state vectors

假设模型的潜在动态函数 g在表征空间中局部满足利普希茨条件^[33]，即：

‖ g (x^{'}, a) - g (x, a) ‖ ⩽ L ‖ x^{'} - x ‖

（11）

式中：L 为利普希茨常数。

且模型单步预测的固有误差为 ε。

第 n 步的累积误差

Δ_{n} = ‖ {\hat{x}}_{t + n} - x_{t + n} ‖

可被界定为

\begin{matrix} Δ_{n} = & ‖ g ({\hat{x}}_{t + n - 1}, a_{t + n - 1}) - \\ g (x_{t + n - 1}, a_{t + n - 1}) ‖ + ε \end{matrix}

（12）

Δ_{n} ⩽ L ‖ {\hat{x}}_{t + n - 1} - x_{t + n - 1} ‖ + ϵ = L Δ_{n - 1} + ε

（13）

式中：

{\hat{x}}_{t + n - 1}

为 x_{t + n - 1} 的预测值。

利用该递推关系，类似于莱布尼茨的误差分析，总误差 Δ_n可表示为

Δ_{n} ⩽ L^{n} Δ_{0} + ε \sum_{i = 0}^{n - 1} L^{i}

（其中 Δ₀是初始状态编码误差，可视为 0）。

如果 L> 1，累积误差将呈指数级增长，即：

Δ_{n} \approx \frac{ε (L^{n} - 1)}{L - 1}

（14）

在未知集中，模型学习到的动态 g 泛化能力有限，由于环境动态不匹配，模型单步预测的固有误差ε显著大于已知集。同时在未见过的状态区域，模型动态可能变得不稳定，导致利普希茨常数 L> 1。在 ε 及 L 的共同作用，导致累积误差 Δ_n 在分布外情况下比分布内情况下更早地呈指数级爆发。当累积误差 Δ_n迅速增大并大于仅含背景信息的图像向量误差时，预测器 g 的输出将进入“无效”区域，此时，模型无法预测有意义的动态，其输出可能收敛到某个或某几个固定的“吸引子”，即无论后续动作如何，预测的表征都变得高度相似或恒定^[31，33-35]。这解释了图 3 中的现象：在前 3 步，误差 Δ_n尚在可控范围，损失随步长增加；约 3 步后，误差累积爆发，导致表征坍缩，预测值

{\bar{x}}_{t + n}

失去时序区分度，其与真实值 x_{t + n} 之间的损失不再随步长 n 变化，而是在较高的误差水平趋于平稳。

4 结束语

本文针对交互数据稀缺条件下的序列决策环境动态建模难题，提出一种基于 JEPA 的潜在动态预测表征模型。该模型摒弃了传统模型对像素级重建的依赖，在潜在空间中实现了对动作—状态转移关系的直接、高效学习。本文的消融实验进一步验证模型关键设计的有效性：核心的“动作查询式”注意力机制被证实在建模动作—状态因果关系上优于标准注意力；梯度阻断策略则是防止表征坍缩、保证 JEPA 范式稳定训练的必要条件。

在 ALE 环境上的实验验证表明，所提模型在分布内的已知环境中展现出与 RSSM 基线模型性能相当的稳定多步预测能力（约 15 步），并在分布外的未知环境中表现出一定的泛化性（约 3 步有效外推）。本文进一步的分析指出，分布外预测性能的瓶颈主要源于开环预测中的误差累积和随之而来的表征坍缩，这解释了长时域外推性能的退化原因。

综上所述，本研究证实了在潜在空间中直接建模动作条件转移是一种在数据稀缺条件下学习泛化世界模型的有效途径。该模型所学的紧凑动态表征，可为后续的规划算法或策略学习提供先验知识支持。未来的工作将聚焦于改进模型架构与训练范式，以缓解表征坍缩，进一步提升模型在未知环境中的长时序预测能力和泛化鲁棒性。

图 1 PRM 模型架构

Fig. 1 Architectureof the PRM

图 2 PRM 模型与 RSSM 模型在已知集中多步开环预测验证损失对比

Fig. 2 Multi-steps open-loop prediction loss of the PRM and RSSM on known set

图 3 未知集中 PRM 多步开环预测验证损失

Fig. 3 Multi-steps open loop prediction validation loss of the PRMon unknown set

图 4 自注意力与“动作查询式”注意力机制模型 5 步开环预测图像重建结果对比

Fig. 4 Comparison of 5-step open-loop prediction image reconstruction results between the self-attention mechanism model and the “action-query” mechanism model

图 5 缺乏梯度阻断下的 PRM 训练损失及验证损失情况

Fig. 5 Training loss and validation loss of the PRM under the absence of gradient blocking

图 6 缺乏梯度阻断下的 PRM 多步开环预测损失

Fig. 6 Multi-step open-loop prediction loss of PRM under the absence of gradient blocking

图 7 潜在状态向量重建结果

Fig. 7 Reconstruction results of Latent state vectors

表 1 模型在不同数据集下的预测结果

Tab. 1 Prediction results of the model on different datasets

图 1 PRM 模型架构

Fig. 1 Architectureof the PRM

图 2 PRM 模型与 RSSM 模型在已知集中多步开环预测验证损失对比

Fig. 2 Multi-steps open-loop prediction loss of the PRM and RSSM on known set

图 3 未知集中 PRM 多步开环预测验证损失

Fig. 3 Multi-steps open loop prediction validation loss of the PRMon unknown set

图 4 自注意力与“动作查询式”注意力机制模型 5 步开环预测图像重建结果对比

Fig. 4 Comparison of 5-step open-loop prediction image reconstruction results between the self-attention mechanism model and the “action-query” mechanism model

图 5 缺乏梯度阻断下的 PRM 训练损失及验证损失情况

Fig. 5 Training loss and validation loss of the PRM under the absence of gradient blocking

图 6 缺乏梯度阻断下的 PRM 多步开环预测损失

Fig. 6 Multi-step open-loop prediction loss of PRM under the absence of gradient blocking

图 7 潜在状态向量重建结果

Fig. 7 Reconstruction results of Latent state vectors

表 1 模型在不同数据集下的预测结果

Tab. 1 Prediction results of the model on different datasets

图(7) / 表(1)

引用本文

宋飞，高超，田园，等 . 基于联合嵌入范式的潜在动态预测表征模型[J]. 上海航天（中英文），2026，43（1）：189-197，210.

复制

计量

图 1 PRM 模型架构

Fig. 1 Architectureof the PRM

图 2 PRM 模型与 RSSM 模型在已知集中多步开环预测验证损失对比

Fig. 2 Multi-steps open-loop prediction loss of the PRM and RSSM on known set

图 3 未知集中 PRM 多步开环预测验证损失

Fig. 3 Multi-steps open loop prediction validation loss of the PRMon unknown set

图 4 自注意力与“动作查询式”注意力机制模型 5 步开环预测图像重建结果对比

Fig. 4 Comparison of 5-step open-loop prediction image reconstruction results between the self-attention mechanism model and the “action-query” mechanism model

图 5 缺乏梯度阻断下的 PRM 训练损失及验证损失情况

Fig. 5 Training loss and validation loss of the PRM under the absence of gradient blocking

图 6 缺乏梯度阻断下的 PRM 多步开环预测损失

Fig. 6 Multi-step open-loop prediction loss of PRM under the absence of gradient blocking

图 7 潜在状态向量重建结果

Fig. 7 Reconstruction results of Latent state vectors

表 1 模型在不同数据集下的预测结果

Tab. 1 Prediction results of the model on different datasets

SUTTON R S, BARTO A G. Reinforcement learning: an introduction [M]. Second edition. Cambridge, Massachusetts: The MIT Press, 2018: 56-287.

HAFNER D, PASUKONIS J, BA J, et al. Mastering diverse control tasks through world models[J]. Nature, 2025, 640(8059): 647-653.

HA D, SCHMIDHUBER J. Recurrent world models facilitate policy evolution [EB/OL]. (2018-09-04) [2025-05-21]. http://arxiv.org/abs/1809.01999.

VASWANI A, SHAZEER N, PARMAR N, et al. Attention is all you need[EB/OL]. (2023-08-02)[2025-07-08]. http://arxiv.org/abs/1706.03762.

HAFNER D, LILLICRAP T, FISCHER I, et al. Learning latent dynamics for planning from pixels[C]//Proceedings of the 36th International Conference on Machine Learning. Long Beach: PMLR, 2019: 2555-2565.

HAFNER D, LILLICRAP T, BA J, et al. Dream to control: learning behaviors by latent imagination[EB/OL]. (2020-03-17) [2025-08-27]. http://arxiv.org/abs/1912.01603.

HAFNER D, LILLICRAP T, NOROUZI M, et al. Mastering atari with discrete world models[EB/OL]. (2022-02-12) [2025-08-27]. http://arxiv.org/abs/2010.02193.

SCHRITTWIESER J, ANTONOGLOU I, HUBERT T, et al. Mastering atari, go, chess and shogi by planning with a learned model[J]. Nature, 2020, 588(7839): 604-609.

GONZALEZ X, WARRINGTON A, SMITH J T H, et al. Towards scalable and stable parallelization of nonlinear RNNs[J]. Advances in Neural Information Processing Systems, 2024, 37: 5817-5849.

LIPTON Z C, BERKOWITZ J, ELKAN C. A Critical review of recurrent neural networks for sequence learning[EB/OL]. (2015-10-17)[2025-10-13]. http://arxiv.org/abs/1506.00019.

PASCANU R, MIKOLOV T, BENGIO Y. On the difficulty of training recurrent neural networks[C]//International Conference on Machine Learning. Atlanta: PMLR, 2013: 1310-1318.

KARITA S, CHEN N, HAYASHI T, et al. A Comparative study on transformer vs RNN in speech applications[C]//2019 IEEE Automatic Speech Recognition and Understanding Workshop (ASRU). Singapore: IEEE, 2019: 449-456.

SOO W, GOUDAR V, WANG X J. Training biologically plausible recurrent neural networks on cognitive tasks with long-term dependencies [J]. Advances in Neural Information Processing Systems, 2023, 36: 32061-32074.

LI X. Comparative analysis and prospect of RNN and transformer[J]. Applied and Computational Engineering, 2024, 75(1): 178-184.

MICHELI V, ALONSO E, FLEURET F. Transformers are sample-efficient world models[EB/OL]. (2023-03-01) [2025-10-13]. http://arxiv.org/abs/2209.00588.

MICHELI V, ALONSO E, FLEURET F. Efficient world models with context-aware tokenization[EB/OL]. (2024-06-27)[2025-10-13]. http://arxiv.org/abs/2406.19320.

BRUCE J, DENNIS M, EDWARDS A, et al. Genie: generative interactive environments[EB/OL]. (2024-02-23)[2025-06-25]. http://arxiv.org/abs/2402.15391.

JANNER M, LI Q, LEVINE S. Offline reinforcement learning as one big sequence modeling problem[J]. Advances in Neural Information Processing Systems, 2021, 34: 1273-1286.

CHEN L, LU K, RAJESWARAN A, et al. Decision transformer: reinforcement learning via sequence modeling[J]. Advances in Neural Information Processing Systems, 2021, 34: 15084-15097.

ZHANG Z, CHEN R, YE J, et al. WHALE: towards generalizable and scalable world models for embodied decision-making[EB/OL]. (2024-11-08)[2025-10-13]. http://arxiv.org/abs/2411.05619.

MAZZAGLIA P, VERBELEN T, DHOEDT B, et al. GenRL: multimodal-foundation world models for generalization in embodied agents[J]. Advances in Neural Information Processing Systems, 2024, 37: 27529-27555.

ASSRAN M, DUVAL Q, MISRA I, et al. Selfsupervised learning from images with a joint-embedding predictive architecture[EB/OL]. (2023-04-13)[2025-06-25]. http://arxiv.org/abs/2301.08243.

BARDES A, GARRIDO Q, PONCE J, et al. V-JEPA: revisiting feature prediction for learning visual representations from video [EB/OL]. (2024-02-15) [2025-06-26]. http://arxiv.org/abs/2404.08471.

ASSRAN M, BARDES A, FAN D, et al. V-JEPA 2: Self-supervised video models enable understanding, prediction and planning[EB/OL]. (2025-06-11)[2025 06-26]. http://arxiv.org/abs/2506.09985.

GORNET J, THOMSON M. Automated construction of cognitive maps with visual predictive coding[J]. Nature Machine Intelligence, 2024, 6(7): 820-833.

JANNER M, DU Y, TENENBAUM J B, et al. Planning with diffusion for flexible behavior synthesis [EB/OL]. (2022-12-21)[2025-10-13]. http://arxiv.org/abs/2205.09991.

DONG Z, HAO J, YUAN Y, et al. Diffuserlite: towards real-time diffusion planning[J]. Advances in Neural Information Processing Systems, 2024, 37: 122556-122583.

HE K, ZHANG X, REN S, et al. Deep residual learning for image recognition[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. Las Vegas: IEEE, 2016: 770-778.

HU J, SHEN L, ALBANIE S, et al. Squeeze-and excitation networks[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. Salt Lake: IEEE, 2018: 7132-7141.

WU Y, HE K. Group Normalization[EB/OL]. (2018 06-11)[2025-10-13]. http://arxiv.org/abs/1803.08494.

BARDES A, PONCE J, LECUN Y. VICReg: variance invariance-covariance regularization for self-supervised learning[EB/OL]. (2022-01-28)[2025-10-11]. http://arxiv.org/abs/2105.04906.

BELLEMARE M G, NADDAF Y, VENESS J, et al. The arcade learning environment: an evaluation platform for general agents[J]. Journal of Artificial Intelligence Research, 2013, 47: 253-279.

MOALLA S, MIELE A, PYATKO D, et al. No representation, no trust: connecting representation, collapse, and trust issues in ppo[J]. Advances in Neural Information Processing Systems, 2024, 37: 69652-69699.

SANSONE E, LEBAILLY T, TUYTELAARS T. Collapse-proof non-contrastive self-supervised learning [EB/OL]. (2025-07-06)[2025-08-28]. http://arxiv.org/abs/2410.04959.

MO S, TONG S. Connecting joint-embedding predictive architecture with contrastive self-supervised learning[J]. Advances in Neural Information Processing Systems, 2024, 37: 2348-2377.

首页

期刊介绍

编委会

期刊指南

下载专区

期刊订阅

联系我们

English

0 引言

1 潜在动态预测表征模型

2 实验方案与结果

3 实验结果分析

4 结束语