Galilai Group发布Stable World Model

这是什么

stable-worldmodel 是由 galilai-group 维护的一个开源 Python 平台，旨在为“世界模型”（World Model）的研究与评估提供可复现的基础设施。该项目提供了一个统一的接口，覆盖了世界模型研究的三个核心阶段：数据收集、模型训练以及基于模型预测控制（Model-Predictive Control, MPC）的评估。

该项目不仅标准化了环境交互协议，还内置了常见的基线（Baselines）和规划求解器（Planning Solvers）的参考实现，使研究人员能够将代码重心集中在模型架构和目标函数本身的创新上，而非重复造轮子。目前该项目在 GitHub 上已获得约 1060 个 Star，主语言为 Python。

解决的问题

世界模型研究长期面临以下痛点，stable-worldmodel 试图通过工程化手段予以解决：

1. 研究碎片化与不可复现性：不同论文通常使用各自定制的数据格式、环境封装和评估脚本，导致结果难以横向对比或复现。
2. 基础设施负担过重：研究人员往往需要花费大量时间处理数据加载、格式转换、环境初始化以及 MPC 求解器的集成，而非专注于核心算法。
3. 泛化能力评估困难：在标准环境中评估模型时，缺乏系统化的机制来测试模型在分布外（Out-of-Distribution）场景下的零样本泛化能力。

核心功能

1. 统一的数据与模型接口

项目通过 swm 库提供了一套标准化的 API，支持从数据收集到评估的全流程：

• 数据收集：支持通过 World 对象收集演示数据，自动处理环境交互。
• 数据加载与转换：内置格式注册表（Format Registry），支持多种后端格式（如 LanceDB、Video/MP4、HDF5等）。用户可以使用 swm.data.load_dataset 自动检测格式，或使用 swm.data.convert 进行一次性格式迁移。
• 模型训练：提供训练脚本参考，如 scripts/train/lewm.py 实现了 LeWM，scripts/train/prejepa.py 复现了 DINO-WM。

2. 标准化的环境套件

项目集成了广泛的标准环境，包括：

• DeepMind Control Suite
• Gymnasium Classic Control
• OGBench
• Craftax
• Arcade Learning Environment（100+ Atari 游戏）
• 经典世界模型基准（Two-Room, PushT）

每个环境都预置了一组“变化因子”（Factors of Variation, FoV），即独立可控的视觉和物理参数（如光照、纹理、动力学、形态等）。这使得评估模型在分布偏移下的零样本泛化能力变得简单直接，无需额外配置。

3. 模型预测控制（MPC）评估

项目内置了 CEMSolver 等规划求解器，支持基于世界模型的 MPC 策略。研究人员可以快速构建 WorldModelPolicy，设置规划 horizon 和采样数量，直接对模型进行性能评估（如成功率 Success Rate）。

4. CLI 工具支持

安装后提供 swm 命令行工具，无需编写代码即可执行常见任务：

• swm datasets：列出缓存的数据集。
• swm inspect <name>：检查特定数据集。
• swm envs：列出所有注册的环境。
• swm fovs <env_name>：查看环境的变体因子。
• swm checkpoints：列出可用的模型检查点。
• swm convert：转换数据集格式。

亮点 / 与同类相比

• 格式灵活性与向后兼容： 不同于许多项目硬编码单一数据格式，stable-worldmodel 采用插件式的格式注册机制。用户可以选择 Lance（默认，高性能）、Video 或其他格式。写入模式支持 append（追加，默认）、overwrite（覆盖）和 error，重新运行收集脚本会自动扩展数据集而非报错，极大提升了实验迭代的便利性。

• 性能基准透明化： 项目提供了详细的吞吐量与存储基准测试（通过 scripts/benchmark/compare_h5_lance.py 复现），展示了不同后端在 PushT 数据集上的表现，帮助研究者根据硬件资源（如 H200 GPU）选择最优配置。

• 聚焦核心贡献： 通过封装数据 I/O、环境交互和 MPC 求解细节，项目强制规范了研究流程。参考实现（如 LeWM 和 DINO-WM）确保了基线的一致性，使得不同工作之间的对比更加公平和严谨。

• 活跃的 API 演进： 项目处于积极开发中，API 可能在次要版本间发生变化，但这也意味着它能快速适配最新的深度学习框架和环境标准。

适合谁用 / 上手

适合人群

• 世界模型研究人员：需要快速搭建实验环境，验证新的世界模型架构或训练目标。
• 强化学习工程师：希望利用世界模型进行模型预测控制（MPC）或离线强化学习的研究者。
• 复现研究者：希望复现 LeWM、DINO-WM 等前沿工作，但缺乏统一数据管道和评估标准的团队。

上手指南

1. 安装： 推荐使用 pip 安装。基础版本仅包含核心库，若需训练、环境和数据格式支持，建议安装全量版本：

   pip install 'stable-worldmodel[all]'
   

   若需支持 LeRobot 数据集，需使用 Python 3.12+ 并安装额外依赖：

   pip install 'stable-worldmodel[lerobot]'
   

   开发模式可使用 uv 管理虚拟环境：

   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
   

2. 配置存储路径： 数据集和检查点默认存储在 ~/.stable_worldmodel/。可通过设置环境变量 $STABLEWM_HOME 更改存储位置。

3. 快速开始示例：

   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. 使用 MPC 进行评估
   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}%")
   

4. 文档与引用： 完整文档位于 galilai-group.github.io/stable-worldmodel。 引用该项目时请使用以下 BibTeX：

   @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 Bale