stable-worldmodel

面向世界模型研究的统一平台，集数据收集、模型训练与模型预测控制评估于一体，支持DeepMind Control Suite、Gymnasium、OGBench等多种标准化环境，内置DINO-WM、LeWM等参考实现和CEM、MPPI等多种规划求解器。亮点在于标准化接口降低研究复现门槛，自带数据格式（LanceDB、HDF5、Video等）转换和CLI工具，并支持通过视觉/物理因子变化（FoV）评测零样本泛化。研究向，非生产级产品。

README

stable-worldmodel

一个用于可复现世界模型（world model）研究和评估的平台。

安装 · 快速开始 · 环境 · 求解器与基线 · 文档 · 论文 · 引用

stable-worldmodel 为世界模型研究的三个阶段——收集数据、训练和利用模型预测控制进行评估——提供了一个统一的接口，覆盖大量标准化环境。它内置了常见基线和规划求解器的参考实现，使研究代码能够聚焦在真正重要的贡献上：模型和目标函数。

安装

通过 PyPI 安装：

pip install stable-worldmodel            # 仅核心库
pip install 'stable-worldmodel[all]'     # + 训练、环境和数据格式

LeRobot 数据集支持是单独的可选扩展（需要 Python 3.12+）：pip install 'stable-worldmodel[lerobot]'

从源码安装（开发模式）：

git clone https://github.com/galilai-group/stable-worldmodel
cd stable-worldmodel
uv venv --python=3.10 && source .venv/bin/activate
uv sync --extra all --group dev

数据集和检查点存储在 $STABLEWM_HOME 下（默认为 ~/.stable_worldmodel/）。可通过设置环境变量来指向你偏好的存储位置。

该库处于活跃开发中。API 可能在次要版本之间发生变化。

快速开始

import stable_worldmodel as swm
from stable_worldmodel.policy import WorldModelPolicy, PlanConfig
from stable_worldmodel.solver import CEMSolver

# 1. 收集数据集
world = swm.World("swm/PushT-v1", num_envs=8)
world.set_policy(your_expert_policy)
world.collect("data/pusht_demo.lance", episodes=100, seed=0)

# 2. 加载数据集并训练你的世界模型（格式自动检测）
dataset = swm.data.load_dataset("data/pusht_demo.lance", num_steps=16)
world_model = ...  # 你的模型

# 3. 通过模型预测控制进行评估
solver = CEMSolver(model=world_model, num_samples=300)
policy = WorldModelPolicy(solver=solver, config=PlanConfig(horizon=10))

world.set_policy(policy)
results = world.evaluate(episodes=50)
print(f"Success Rate: {results['success_rate']:.1f}%")

参考实现位于 scripts/train/ 中：lewm.py 实现了 LeWM，prejepa.py 复现了 DINO-WM。

GPU 利用率对比
在 H200 GPU 上使用 Push-T LanceDB 数据集训练 LeWM 的 GPU 利用率。

数据格式

记录、加载和转换都通过一个简洁的格式注册表进行。选择符合你权衡的后端，或注册你自己的格式。

格式	磁盘布局	最佳用途
`lance`	LanceDB 表（按 episode 连续存储的扁平行）	默认格式——支持追加，快速索引读取
`hdf5`	单个 `.h5` 文件（每列一个数据集）	可移植单文件制品
`folder`	`.npz` 列 + 每步一个 JPEG	检查、部分 key 流式读取
`video`	`.npz` 列 + 每个 episode 一个 MP4（`decord`）	长 episode，紧凑的图像存储
`lerobot`	`lerobot://<repo_id>`（只读适配器）	直接在 LeRobot Hub 数据集上进行训练/评估

world.collect("data/pusht.lance", episodes=100)                  # 默认：lance
world.collect("data/pusht_video", episodes=100, format="video")  # mp4 episodes

ds = swm.data.load_dataset("data/pusht.lance", num_steps=16)     # 自动检测
swm.data.convert("data/pusht.lance", "data/pusht_video",
                 dest_format="video", fps=30)                    # 一次性迁移

每个写入器都接受一个 mode 参数（'append'（默认）、'overwrite'、'error'）；重新运行 world.collect 会扩展已有数据集而不是失败。

吞吐量与存储基准

以下数字由 scripts/benchmark/compare_h5_lance.py 生成，可通过该脚本复现。基准测试使用 LeWorldModel 论文中的 PushT 数据集。

吞吐量

格式	来源	缓存	samples/s	ms/step
HDF5	本地	无缓存	1416.1	45.2
HDF5	本地	已缓存	1474.0	43.4
LanceDB	本地	无缓存	4814.8	13.3
LanceDB	本地	已缓存	4431.3	14.4
Video	本地	-	1330.6	48.1
LanceDB	s3	无缓存	3183.7	20.1
LanceDB	s3	已缓存	3253.2	19.7
HDF5	s3	无缓存	9.1	7032.5
HDF5	s3	已缓存	756.5	84.6

各格式存储大小（本地）

格式	本地大小
HDF5	43.12 GB
LanceDB	13.31 GB
Video	496.29 MB

环境

第一行：默认外观 · 第二行：视觉因素变化

环境来源于 DeepMind Control Suite、Gymnasium 经典控制、OGBench、Craftax、Arcade Learning Environment（100+ Atari 游戏）以及经典世界模型基准（Two-Room、PushT）。大多数环境都配备了一套因素变化——独立可控的视觉和物理参数（光照、纹理、动力学、形态），使得无需额外设置即可轻松评估对分布偏移的零样本泛化能力。添加新环境只需遵循 Gymnasium 接口标准。

完整环境列表

环境 ID	FoV 数量
swm/PushT-v1	16
swm/TwoRoom-v1	17
swm/OGBCube-v0	11
swm/OGBScene-v0	12
swm/HumanoidDMControl-v0	7
swm/CheetahDMControl-v0	7
swm/HopperDMControl-v0	7
swm/ReacherDMControl-v0	8
swm/WalkerDMControl-v0	8
swm/AcrobotDMControl-v0	8
swm/PendulumDMControl-v0	6
swm/CartpoleDMControl-v0	6
swm/BallInCupDMControl-v0	9
swm/FingerDMControl-v0	10
swm/ManipulatorDMControl-v0	8
swm/QuadrupedDMControl-v0	7
swm/CartPoleControl-v1	10
swm/MountainCarControl-v0	5
swm/MountainCarContinuousControl-v0	4
swm/AcrobotControl-v1	11
swm/PendulumControl-v1	9
swm/FetchReach-v3	8
swm/FetchPush-v3	11
swm/FetchSlide-v3	11
swm/FetchPickAndPlace-v3	11
swm/CraftaxClassicPixels-v1	—
swm/CraftaxClassicSymbolic-v1	—
swm/CraftaxPixels-v1	—
swm/CraftaxSymbolic-v1	—
ALE/* (100+ Atari 游戏)	—

求解器与基线

求解器	类型
交叉熵方法（CEM）	采样
改进 CEM（iCEM）	采样
模型预测路径积分（MPPI）	采样
预测采样（Predictive Sampling）	采样
梯度下降（SGD, Adam）	梯度
投影梯度下降（PGD）	梯度
增广拉格朗日（Augmented Lagrangian）	约束优化

基线	类型
DINO-WM	JEPA
PLDM	JEPA
LeWM	JEPA
GCBC	行为克隆
GCIVL	强化学习
GCIQL	强化学习

命令行界面

安装后，可以使用 swm 命令来检查/转换数据集、环境和检查点，无需编写代码：

swm datasets                                        # 列出缓存的数据集
swm inspect pusht_expert_train                      # 检查特定数据集
swm envs                                            # 列出所有已注册的环境
swm fovs PushT-v1                                   # 显示某个环境的因素变化
swm checkpoints                                     # 列出可用的模型检查点
swm convert pusht_expert_train --dest-format video  # 将数据集转换为其他格式

文档

完整文档位于 galilai-group.github.io/stable-worldmodel，包含 API 参考、教程和指南。

基于 `stable-worldmodel` 构建的项目

C-JEPA
LeWM

引用

@misc{maes_lld2026swm,
  title  = {stable-worldmodel: A Platform for Reproducible World Modeling Research and Evaluation},
  author = {Lucas Maes and Quentin Le Lidec and Luiz Facury and Nassim Massaudi and
            Ayush Chaurasia and Francesco Capuano and Richard Gao and Taj Gillin and
            Dan Haramati and Damien Scieur and Yann LeCun and Randall Balestriero},
  year   = {2026},
  eprint = {2605.21800},
  archivePrefix = {arXiv},
  primaryClass = {cs.LG},
  url    = {https://arxiv.org/abs/2605.21800},
}

问题

欢迎提出 issue —— 我们很乐意提供帮助。

开源项目galilai-group2026-05-29原文

打开互动版

stable-worldmodel

README