GeoWorld: Geometric World Models¶

会议: CVPR 2026
arXiv: 2602.23058
代码: https://steve-zeyu-zhang.github.io/GeoWorld
领域: 强化学习
关键词: 几何世界模型、双曲空间、JEPA、强化学习、长期规划

一句话总结¶

GeoWorld 将预测式世界模型的潜在表征从欧氏空间映射到双曲流形上，通过 Hyperbolic JEPA 保持几何结构和层级关系，并提出 Geometric Reinforcement Learning 来优化多步规划，在 CrossTask 和 COIN 上实现了约 3% SR（3步）和 2% SR（4步）的提升。

研究背景与动机¶

领域现状：世界模型分为生成式和预测式两大类。生成式世界模型（如 VideoWorld）显式生成像素或视觉 token 来预测下一步，但缺乏对整条轨迹结构和能量景观的全局感知。预测式世界模型（如 JEPA、V-JEPA 2）不生成像素，而是在潜在空间中学习一个能量景观来度量当前状态和目标状态之间的兼容性，支持多步层级规划。

现有痛点：现有预测式世界模型面临两个关键问题：(1) 几何忽视——潜在表征在欧氏空间中学习，无法保持状态间的内在几何结构和层级关系，导致学到的能量景观不能捕捉有意义的测地距离；(2) 多步退化——多步视频数据稀缺且昂贵，模型主要在单步转换上训练，虽然概念上能做长期规划，但性能随 horizon 增长快速下降。

核心矛盾：欧氏空间的"平坦"结构无法自然编码现实世界中状态的层级关系（如"切菜"是"做饭"的子步骤），而强行在欧氏空间中建模长程依赖会导致几何漂移。

本文目标 (1) 如何在潜在空间中保持几何和层级结构？(2) 如何在训练数据有限的情况下提升多步规划的稳定性？

切入角度：双曲空间的树状结构天然适合编码层级关系（距离随层级呈指数增长），且双曲测地线提供了最自然的"最短路径"概念。结合强化学习来优化预测器的能量函数，使轨迹沿测地线前进。

核心 idea：将 JEPA 的潜在动力学从欧氏空间提升到双曲流形上，并用几何强化学习优化多步规划中的测地线一致性。

方法详解¶

整体框架¶

GeoWorld 接收当前视频观测和目标视频，输出到达目标的动作序列。系统由三部分组成：(1) Hyperbolic JEPA (H-JEPA)——冻结编码器将观测映射到欧氏空间后经指数映射投影到 Poincaré 球上，然后用条件预测器在双曲空间中预测下一状态；(2) Geometric Reinforcement Learning (GRL)——将多步规划建模为最小化双曲能量的优化问题，用三角不等式正则化保持轨迹的测地线一致性；(3) 能量规划——使用 Cross-Entropy Method (CEM) 搜索使双曲距离最小化的最优动作序列。

关键设计¶

Hyperbolic JEPA (H-JEPA):
- 功能：在双曲流形上保持几何结构和层级关系的潜在动力学建模
- 核心思路：使用预训练冻结编码器 \(E_\theta\) 将观测 \(x_t\) 编码为欧氏嵌入 \(s_t^x \in \mathbb{R}^n\)，然后将其视为原点切空间中的切向量，通过 Poincaré 球的指数映射投影到双曲空间：\(s_{t,\mathbb{H}}^x = \exp_0(s_t^x) = \tanh(\sqrt{c}\|s_t^x\|) \frac{s_t^x}{\sqrt{c}\|s_t^x\|}\)，其中曲率 \(c\) 是可学习参数。条件预测器 \(P_\phi\) 接收双曲状态序列和动作序列，预测下一状态。训练目标包括 teacher-forcing loss（单步预测精度）和 rollout loss（两步递归预测的一致性），都基于双曲测地距离 \(d_\mathbb{H}\)
- 设计动机：双曲空间的容积随半径呈指数增长，天然适合嵌入树状层级结构。状态间的测地距离比欧氏距离更能反映真实世界中动作之间的语义层级关系
Geometric Reinforcement Learning (GRL):
- 功能：通过能量优化直接改进预测器的多步规划能力
- 核心思路：定义能量代价为预测状态和目标状态之间的双曲距离 \(c_t = d_\mathbb{H}(\hat{s}_{t+1,\mathbb{H}}^x, s_{t+1,\mathbb{H}}^x)\)，奖励为负能量代价。路径价值函数是期望累积奖励 \(V = \mathbb{E}[\sum \gamma^{t-1} r_t]\)，最优价值函数等价于最小化累积双曲距离。关键创新是三角不等式正则化 \(\mathcal{L}_\Delta = \frac{1}{T-2}\sum[d_\mathbb{H}(\hat{s}_t, \hat{s}_{t+2}) - d_\mathbb{H}(\hat{s}_t, \hat{s}_{t+1}) - d_\mathbb{H}(\hat{s}_{t+1}, \hat{s}_{t+2})]_+\)，强制预测轨迹满足测地线性质。总损失 \(\mathcal{L}_{\text{GRL}} = \text{累积测地距离} + \beta \mathcal{L}_\Delta\)
- 设计动机：仅靠监督学习无法充分改善多步规划——训练数据中多步视频有限。RL 的优化目标直接面向最终的规划质量。GRL 的巧妙之处在于不需要额外训练策略网络或奖励模型，直接优化预测器参数
能量规划 (Energy-Based Planning with CEM):
- 功能：在推理阶段搜索最优动作序列
- 核心思路：给定当前观测和目标，编码到双曲空间后定义能量代价函数 \(C = d_\mathbb{H}(P((\hat{a}_t); s_{1,\mathbb{H}}^x), s_{1+T,\mathbb{H}}^x)\)。使用 CEM（N=800 个采样，K=80 个精英，I=10 轮迭代）搜索使能量最小化的动作序列
- 设计动机：CEM 是一种高效的零阶优化方法，适合在高维动作空间中搜索

损失函数 / 训练策略¶

两阶段训练：(1) 监督后训练阶段：\(\mathcal{L}_{\text{SFT}} = \lambda \mathcal{L}_{\text{TF}} + (1-\lambda) \mathcal{L}_{\text{rollout}}\)，AdamW 优化器，warmup → constant → decay 学习率调度，batch size 256，约 94500 iterations。(2) GRL 阶段：更小学习率和更短调度，batch size 128，约 25000 iterations，\(\gamma=0.99\)，\(\beta=0.1\)。预测器是 ~300M 参数 Transformer（24层，16头，1024维）。4 节点 32 卡 H100 训练。

实验关键数据¶

主实验 — Procedural Planning (图像input)¶

方法	CrossTask T=3 SR	CrossTask T=4 SR	COIN T=3 SR	COIN T=4 SR
V-JEPA 2 ViT-g384	45.58	31.36	34.08	23.43
GeoWorld ViT-g384	47.47	31.48	34.85	27.79
SCHEMA (LLM)	38.93	24.50	32.09	22.02
MTID (Generative)	40.45	24.76	30.44	22.74

主实验 — Visual Planning with Videos¶

方法	CrossTask T=3 SR	CrossTask T=4 SR	COIN T=3 SR	COIN T=4 SR
V-JEPA 2 ViT-g384	50.16	35.01	42.74	31.63
GeoWorld ViT-g384	51.71	37.04	45.29	33.29
GPT-5	50.03	30.20	43.84	32.64
VideoWorld	41.59	25.50	34.88	23.74

长期规划 (CrossTask)¶

方法	T=3	T=4	T=5	T=6
V-JEPA 2 ViT-g384	50.16	35.01	23.17	16.88
GeoWorld ViT-g384	51.71	37.04	24.83	18.26

关键发现¶

GeoWorld 在所有模型规模（ViT-L/H/g/g384）上都一致超过 V-JEPA 2，说明改进不依赖于特定规模
随着规划 horizon 增加（T=3→T=6），GeoWorld 相对 V-JEPA 2 的优势逐渐扩大，在 T=6 时 SR 提升 1.38%（18.26 vs 16.88），验证了长期规划稳定性的增强
mIoU 指标提升最显著（如 PP 设置下 CrossTask T=3: 86.55 vs 69.42），说明双曲表征对轨迹重叠度的提升最大
GeoWorld ViT-g384 在视频规划设置下超越了 GPT-5 和 Gemini 2.5 Pro

亮点与洞察¶

将双曲几何引入世界模型是一个非常自然且优雅的设计。日常活动的层级结构（如"煮咖啡"→"磨豆→烧水→冲泡"）本质上就是树状的，双曲空间恰好是嵌入树的最佳空间。这种 geometric inductive bias 比单纯增大模型更高效
GRL 的设计巧妙之处在于直接把测地距离作为 reward，不需要额外的奖励模型。三角不等式正则化确保了"走两步不应比直接到达更远"的几何一致性，这是欧氏空间中不容易施加的约束
两阶段训练（SFT → GRL）类似 LLM 的 SFT → RLHF 流程，但目标函数完全基于几何原则，不依赖人类偏好

局限与展望¶

只在 CrossTask 和 COIN 两个数据集上验证，且这两个数据集的动作空间相对简单（日常活动）。在更复杂的机器人操作或游戏环境中的表现未知
Poincaré 球的曲率 \(c\) 需要学习，但在所有状态间共享同一曲率可能不够灵活，不同层级的子空间可能需要不同曲率
编码器完全冻结，H-JEPA 只训练预测器。如果编码器本身的表征不适合双曲空间，性能天花板受限
SR 绝对值在长 horizon 下仍然较低（T=6 只有 18.26%），说明长期规划仍是核心难题

评分¶

新颖性: ⭐⭐⭐⭐⭐ 双曲空间+世界模型+几何RL的组合在视觉规划领域是首次，理论动机和技术实现都很新颖
实验充分度: ⭐⭐⭐⭐ 两个数据集、多种基线、多种模型规模。但只有两个数据集且缺乏消融细节
写作质量: ⭐⭐⭐⭐ 数学推导严谨，但符号较多，读起来有一定门槛
价值: ⭐⭐⭐⭐ 提出了将几何思维引入世界模型的有前景方向，超越了 GPT-5 级别基线