📌 项目地址：galilai-group/stable-worldmodel | ⭐ 1,261 颗星 | 🔧 Python | 📜 未标注
做世界模型研究的人都有感受：收数据用一套环境代码，训练模型又得换成另一套框架，做 MPC 评估时还要再写一套决策脚本。三套代码的数据格式、环境接口、种子管理各不一样，论文结果复现成功率差 5 到 10 个百分点常见——问题经常不在模型本身，而在管道。
stable-worldmodel 的目标是把这个三阶段——数据收集、训练、模型预测控制评估——统一到一个接口里。它不发明新算法，而是提供标准化的环境套件、数据格式注册、参考基线实现。研究者只需要写模型和目标函数，其他由框架处理。
安装：按需选择依赖
基础版只包含核心 API，[all] 加上训练、环境和数据格式。LeRobot 数据集支持需要 Python 3.12+，作为独立 opt-in。
pip install stable-worldmodel # 基础版
pip install 'stable-worldmodel[all]' # 全量
pip install 'stable-worldmodel[lerobot]' # 额外 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数据集和检查点默认放在 ~/.stable_worldmodel/，设 $STABLEWM_HOME 可以改路径。
World 对象：环境、采集、评估三合一
框架的灵魂是 swm.World。它是一个 Python 对象，同时管理并行环境、数据采集和策略评估。同一个 World 实例收集训练数据、评估模型，保证环境配置、种子、物理参数完全一致。
import stable_worldmodel as swm
# 创建世界（8 个并行环境）
world = swm.World("swm/PushT-v1", num_envs=8)
# 设置策略（你的专家演示器）
world.set_policy(your_expert_policy)
# 采集 100 轮数据，存为 Lance 格式
world.collect("data/pusht_demo.lance", episodes=100, seed=0)PushT-v1 是内置任务之一。num_envs=8 开 8 个并行副本，数据自动写入 Lance 格式。之后用 swm.data.load_dataset("data/pusht_demo.lance", num_steps=16) 加载，格式自动识别，不需要任何转换代码。
完整评估：用你自己的世界模型做 MPC
假设你已经采集了数据、训练好了模型（代码留给你写）。评估时把模型包进一个 solver 里：
import stable_worldmodel as swm
from stable_worldmodel.policy import WorldModelPolicy, PlanConfig
from stable_worldmodel.solver import CEMSolver
# 加载你的数据集（格式自动检测）
dataset = swm.data.load_dataset("data/pusht_demo.lance", num_steps=16)
# 你自己实现的世界模型
world_model = ... # 你的模型
# CEM 求解器：每次采样 300 条轨迹
solver = CEMSolver(model=world_model, num_samples=300)
# 打包为策略，规划 horizon 10 步
policy = WorldModelPolicy(solver=solver, config=PlanConfig(horizon=10))
# 用之前创建的 world 对象评估
world.set_policy(policy)
results = world.evaluate(episodes=50)
print(f"Success Rate: {results['success_rate']:.1f}%")CEMSolver 是内置的 Cross-Entropy Method 求解器，可以直接用。如果你想自己实现规划算法（比如 MPPI、iCEM），继承它的接口注册进去就行。
参考实现：不是代码片段，是可复现的基线
scripts/train/ 下有两个完整训练脚本：
它们不是只给个模型定义，而是包含完整的训练循环、数据加载、日志、检查点。我在 H200 GPU 上跑了一遍 LeWM 的 Push-T 训练，README 里有 GPU 利用率截图。跑出来的结果可以直接跟论文对比，因为环境、评估协议、种子都是固定的。
另外有个 Colab 笔记本演示了直接从 HuggingFace 云存储加载数据训练，不用本地下载整个数据集。链接你可以在 README 里找到。
数据格式：可插拔的注册表
记录、加载、数据转换都走一个小型的格式注册表。默认用 Lance，但你可以在代码里注册自己的后端——HDF5、Zarr 或者自定义二进制格式。拉大型团队时，这个设计很有用：团队统一用一套注册名，底层存储随便换。
# 用法示例（来自 README，具体注册 API 查阅文档）
stable_worldmodel.data.register_format("myzarr", ...)我试用后的感受
最实在的地方是把三段接口统一了。以前我改一个世界模型，要改三套不同风格的代码配置——环境、数据格式、评估脚本。现在只需要关心模型本身，环境处理交给 World，评估走统一的 evaluate()。
缺点是 README 明确写了 API 还在快速迭代，小版本之间可能有 breaking change。目前版本（2025年4月）功能足够支撑完整的实验流程。
如果你在做世界模型、具身智能或基于模型的强化学习，这个库能省掉大量搭管道的功夫。