Jasmine：一个简单、高效且可扩展的 JAX 世界建模代码库
机器学习
2025-11-03 v1 人工智能
摘要
尽管世界模型日益被视为克服数据稀缺问题的关键路径（例如在机器人领域），但开源世界模型训练基础设施仍处于初级阶段。我们介绍了 Jasmine，这是一个基于 JAX 的高性能世界模型代码库，能够从单机扩展到数百个加速器，且几乎无需代码修改。Jasmine 通过在数据加载、训练和检查点方面进行性能优化，使其对 CoinRun 案例研究的复现速度提升了一个数量级，相较于先前的开源实现。该代码库保证完全可复现的训练过程，并支持多样化的分片配置。通过将 Jasmine 与精选的大规模数据集组合，我们构建了跨模型家族和架构消解的严谨基准测试管道。
引用
@article{arxiv.2510.27002,
title = {Jasmine: A Simple, Performant and Scalable JAX-based World Modeling Codebase},
author = {Mihir Mahajan and Alfred Nguyen and Franz Srambical and Stefan Bauer},
journal= {arXiv preprint arXiv:2510.27002},
year = {2025}
}
备注
Blog post: https://pdoom.org/jasmine.html