← 返回 2026-01-15

流等变世界模型:在部分观测动态环境中构建结构化记忆 Flow Equivariant World Models: Memory for Partially Observed Dynamic Environments

Hansen Jin Lillemark, Benhao Huang, Fangneng Zhan, Yilun Du, Thomas Anderson Keller 📅 2026-01-03 👍 6 2026-07-13 08:35
世界模型 等变性 视频预测 记忆网络 部分可观测

用流等变记忆提升部分观测动态世界模型的长程一致性预测。

前置知识

流等变(Flow Equivariance)

流等变是群等变在时间参数化变换下的推广:若输入按由向量场 $\nu$ 生成的流 $\psi_t(\nu)$ 运动,模型输出也按同一流变换。本文将内部状态视为多‘速度通道’的地图,每步先用流 $\psi_1(\nu)$ 平移隐藏状态,再叠加观测更新,使得特征在自运动与外部运动下保持空间对齐。

该性质保证未观测区域的记忆不因时间推移而错位,是本文方法能长程预测且少幻觉的根本约束。

部分可观测环境

智能体视野有限,无法一次获得全局状态;未观测区域仍会继续演化。本文场景中,智能体只能看到局部图像,并通过动作改变视角。世界模型需在持续部分观测下维护可迁移的记忆以支持未来观测预测。

记忆结构决定是否能稳定回溯‘离开视野’的物体,否则长程预测会退化或产生幻觉。

世界模型(World Model)

世界模型用于在给定初始状态与动作序列时预测环境未来状态。本文关注视觉观测序列的条件预测,要求模型既建模外部物体动力学,也要对自运动产生一致响应。

作为下游规划与控制的基础,世界模型的长程稳定性和一致性直接决定控制成功率与样本效率。

等变神经网络

等变网络通过结构化权重共享,使网络输出在群变换下遵循 $\phi(g\cdot f)=g\cdot \phi(f)$。在物理模拟等领域,引入平移/旋转等已知对称性可大幅提升数据效率与泛化。

本文将此思想拓展到时间维度的流,构建更高效、更可泛化的动态记忆与预测模型。

研究动机

现有视频扩散或循环世界模型在部分观测动态场景下难以维护一致记忆:自回归扩散依赖滑动窗口,窗口外信息被丢弃,长序列易产生幻觉或遗忘;SSM 记忆则逐步衰减,离开视野的物体信息会消失。作者在 2D/3D 基准中观察到,DFoT 在 150 帧外推时生成‘合理但错误’的物体,MSE 高于‘全黑’基线,DFoT-SSM 物体逐渐褪色。根源是记忆结构与世界动态不对齐。

本文的目标是提出一种将记忆按世界的时间-动力学结构组织的框架:让内部状态随自运动与外部物体运动一致变换,从而在长程预测中保留未观测区域信息,降低幻觉率,并在训练长度外稳定外推;同时提供统一群论语言处理自运动与外部运动,兼顾理论可解释性与工程可扩展性。

与已有工作不同的是,以往记忆方案多以‘静态视角缓存’或‘体素重建’形式存在,缺乏统一的群论语言同时处理自运动与外部运动;且多依赖显式深度投影或精确等变编码器,限制扩展性。本文选择在抽象潜空间引入流等变,不强求编码器严格等变,而是通过时间-流约束诱导学习等变表示,兼顾表达力与可扩展性。

核心方法

FloWM 将世界状态表示为潜空间地图(多速度通道 token 或卷积图),并在每步以流 $\psi_1(\nu)$ 平移地图以模拟外部物体运动,再以动作流 $\psi_1(-a_t)$ 平移以对齐自运动;随后用编码器将当前观测写入视野对应位置并更新记忆。最终通过解码器(MLP 或 ViT 解码器)从地图视野区域预测下一帧。

核心是‘流等变记忆’:隐藏状态按向量场生成的流作时间-群变换,保证未观测区域信息随世界运动而对齐。与以往方案不同,本文不要求编码器严格等变,而是让潜空间在流约束下学习等变表示,从而在 3D 第一人称视角到抽象俯视地图的复杂映射中保持可扩展性。

方法步骤详情

每步先对潜记忆施加外部流 $\psi_1(\nu)$ 与自运动流 $T_{a_t}^{-1}$,使地图与世界坐标对齐;然后用编码器 $E_\theta$ 将当前观测与对应视野记忆输入 ViT,得到更新 token;再用更新操作 $U_\theta$ 将 token 写回地图对应位置;最后从视野记忆解码预测下一帧。2D 实现中记忆为卷积图,速度通道通过 max-pool 聚合;3D 实现中记忆为 token 地图,解码器通过交叉注意力查询视野 token。

技术新颖性

技术新颖性在于:(1) 将内部与外部运动统一为流,并在潜空间严格施加 $h_{t+1}(\nu)=\psi_1(\nu)\cdot U_\theta(h_t(\nu); E_\theta(f_t,h_t))$ 的等变递推;(2) 通过多速度通道排列($\nu-\hat{\nu}$)同时追踪不同运动源;(3) 在 3D 第一人称场景中,用动作变换 $T_{a_t}$ 与流组合实现自运动等变,不依赖显式深度投影。实验显示,这种结构化记忆带来更快收敛与更强的长程外推。

Flow Equivariant World Models maintain structured dynamic memory in partially observed environments.
Figure 1: Flow Equivariant World Models maintain structured dynamic memory in partially observed environments.
Visualization of the Simple Recurrent FloWM on MNIST World.
Figure 3: Visualization of the Simple Recurrent FloWM on MNIST World.
Transformer-Based FloWM: ViT encoder/decoder updates structured token map.
Figure 4: Transformer-Based FloWM: ViT encoder/decoder updates structured token map.

实验结果

在 2D MNIST World 上,FloWM 在训练长度 20 帧与长度外推 150 帧均显著优于基线:MSE 从 DFoT 的 0.1448/0.2111 降至 0.0005/0.0018,PSNR 提升至 32.99/27.56 dB,SSIM 达 0.99/0.98。在 3D Dynamic Block World(70/210 帧)上,MSE 从 DFoT 的 0.0118/0.0217 降至 0.000603/0.001539,PSNR 提升至 32.19/28.13 dB。可视化显示 FloWM 在 210 帧仍保持物体位置与颜色一致,而基线出现幻觉或模糊。消融表明速度通道与自运动等变均必要,缺一时性能退化显著。

Rollout metrics on 2D MNIST World.
Table 1: Rollout metrics on 2D MNIST World.
Rollout metrics on 3D Dynamic Block World.
Table 2: Rollout metrics on 3D Dynamic Block World.
Planning Metrics on Block World.
Table 3: Planning Metrics on Block World.
FloWM generalizes further and trains faster.
Figure 5: FloWM generalizes further and trains faster.
FloWM accurately predicts moving objects in a 3D environment over long time-spans.
Figure 6: FloWM accurately predicts moving objects in a 3D environment over long time-spans.
查看结构化数据
任务指标本文基线提升
2D MNIST World rollout MSE(20/150) 0.0005 / 0.0018 DFoT 0.1448 / 0.2111 约 289.6×(20)、117.3×(150)
2D MNIST World rollout PSNR(20/150) 32.99 / 27.56 DFoT 8.394 / 6.755 +24.60 dB(20)、+20.81 dB(150)
2D MNIST World rollout SSIM(20/150) 0.9900 / 0.9813 DFoT 0.4045 / 0.2434 +0.5855(20)、+0.7379(150)
3D Dynamic Block World rollout MSE(70/210) 0.000603 / 0.001539 DFoT 0.011759 / 0.021684 约 19.5×(70)、14.1×(210)
3D Dynamic Block World rollout PSNR(70/210) 32.19 / 28.13 DFoT 19.30 / 16.64 +12.89 dB(70)、+11.49 dB(210)
3D Dynamic Block World rollout SSIM(70/210) 0.9673 / 0.9525 DFoT 0.9377 / 0.8885 +0.0296(70)、+0.0640(210)

局限与改进

论文在受控几何环境中验证,动作空间为离散平移/旋转,场景为刚体运动与简单墙壁反弹;未涉及复杂纹理、遮挡、多类型物体交互或开放世界语义动作。作者亦承认:为保持抽象潜空间,3D 编码器未严格等变,依赖时间流约束诱导;地图尺寸与视野固定,面对大规模或无限地平线场景需自适应地图与更强解码器。此外,当前评估聚焦预测质量,未系统评估学习效率与计算代价在复杂真实任务中的权衡。

独立分析的弱点

(1)动作空间受限:离散动作与刚体几何难以覆盖语义动作或可变形体;改进方向是扩展为连续 Lie 群与学习动作变换。(2)编码器等变近似:当前仅通过流约束诱导,未保证严格等变;可引入几何先验或自监督等变目标提升稳定性。(3)地图规模固定:视野与地图尺寸固定,难以适应开放地平线;可引入层级/可扩展地图与注意力机制。(4)计算与推理代价:3D ViT 解码需大量 token 交叉注意力;可研究稀疏或局部查询以降低推理时间。

未来方向

作者提出将框架扩展至更复杂数据集、更强大解码器与自适应地图,以支撑自动驾驶/机器人场景;可进一步将 FloWM 与规划目标(如 MPC、强化学习)深度耦合,利用结构化记忆做更长程规划。另一方向是将速度通道从离散逼近扩展为连续群表示,以及将记忆与自监督表示学习(如 JEPA)结合,提升潜空间抽象质量。

复现评估

论文提供代码与项目页面,2D/3D 数据集可自生成,模型配置与训练超参在附录给出;但 3D 实验依赖特定渲染环境与动作分布,复现需搭建 Miniworld 数据管线与 ViT 训练基础设施。总体复现难度中等:关键算法与损失清晰,训练步数与架构细节完整,但需一定工程投入以生成数据、配置解码器与评估流程。