galilai-group/stable-worldmodel

解决的问题

提供一个统一且可复现的世界模型研究平台，简化工作流的三个主要阶段：数据收集、模型训练和使用模型预测控制（MPC）进行评估。它消除了研究人员编写数据处理和规划求解器样板代码的需要，使他们能够专注于模型架构和目标函数。

如何工作

该平台将多个关键组件整合到一个界面中：

• 数据管理：格式注册表支持多种后端（LanceDB、HDF5、文件夹、视频和LeRobot），用于记录和加载数据集，并提供内置工具实现格式间的一次性迁移。
• 环境套件：提供来自DeepMind Control Suite、Gymnasium、OGBench和Atari等来源的大量标准化环境，其中许多环境包含“变化因素”（factors of variation），用于测试模型在不同视觉或物理设置下的泛化能力。
• 规划求解器：包含采样型求解器（如CEM、MPPI）和基于梯度的求解器（如SGD、Adam）的参考实现，用于将世界模型预测转化为动作。
• 统一API：提供简单的Python API，允许用户使用专家策略收集数据、训练模型，并使用WorldModelPolicy与求解器组合进行评估。

适用人群

需要在多样化环境中对结果进行基准测试和复现的世界模型、具身智能和模型预测控制领域的AI研究人员和工程师。

主要亮点

• 丰富的环境库：涵盖从简单控制（CartPole）到复杂机器人（Fetch）和Atari游戏的广泛任务。
• 高性能数据I/O：利用LanceDB实现快速索引读取和可追加存储，优化性能。
• 内置基线：包含基于JEPA的模型（LeWM、DINO-WM）以及其他行为克隆或强化学习基线的参考实现。
• 泛化能力测试：集成的“变化因素”使评估模型在分布偏移下的零样本泛化能力变得简单。