stable-worldmodel：可复现世界模型研究与评估平台
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"成功率: {results['success_rate']:.1f}%")
在 `scripts/train/` 目录下提供了参考实现：`lewm.py` 实现了 LeWM，`prejepa.py` 复现了 DINO - WM。同时，还展示了在 H200 GPU 上使用 Push - T LanceDB 数据集训练 LeWM 时的 GPU 利用率。
数据格式
数据的记录、加载和转换都通过一个小型的格式注册表进行。可以根据需求选择合适的后端，也可以注册自己的格式。数据格式如下：
| 格式 | 磁盘布局 | 适用场景 |
|---|---|---|
| lance | LanceDB 表（按情节连续的扁平行） | 默认格式，便于追加数据，支持快速索引读取 |
| hdf5 | 单个 `.h5` 文件（每列一个数据集） | 便于携带的单文件制品 |
| folder | `.npz` 列 + 每步一张 JPEG 图片 | 便于检查和部分键流式传输 |
| video | `.npz` 列 + 每集一个 MP4 文件（使用 decord） | 长情节、紧凑的图像存储 |
| lerobot | `lerobot://`（只读适配器） | 直接在 LeRobot Hub 数据集上进行训练/评估 |
代码示例如下：
world.collect("data/pusht.lance", episodes=100) # 默认使用 lance 格式
world.collect("data/pusht_video", episodes=100, format="video") # 使用视频格式
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 数据集。测试结果如下：
| 格式 | 数据源 | 缓存情况 | 样本/秒 | 每步耗时（ms） |
|---|---|---|---|---|
| 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 | 视觉因素数量 |
|---|---|
| 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 游戏） | - |
求解器与基线模型
求解器类型
- 交叉熵方法（Cross - Entropy Method, CEM）：采样方法
- 改进的 CEM（Improved CEM, iCEM）：采样方法
- 模型预测路径积分（Model Predictive Path Integral, MPPI）：采样方法
- 预测采样（Predictive Sampling）：采样方法
- 梯度下降（Gradient Descent, SGD, Adam）：梯度方法
- 投影梯度下降（Projected Gradient Descent, PGD）：梯度方法
- 增广拉格朗日约束优化（Augmented Lagrangian Constrained Opt）
基线模型类型
- 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](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，将竭诚提供帮助。那么，在使用过程中会遇到哪些具体问题呢？
