Jasmine: A Simple, Performant and Scalable JAX-based World Modeling Codebase
Fuente:
arXiv
Guardado en:
| Autores principales: | , , , |
|---|---|
| Formato: | Preprint |
| Publicado: |
2025
|
| Materias: | |
| Acceso en línea: | |
| Etiquetas: |
Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
|
| _version_ | 1866912679105396736 |
|---|---|
| author | Mahajan, Mihir Nguyen, Alfred Srambical, Franz Bauer, Stefan |
| author_facet | Mahajan, Mihir Nguyen, Alfred Srambical, Franz Bauer, Stefan |
| contents | While world models are increasingly positioned as a pathway to overcoming data scarcity in domains such as robotics, open training infrastructure for world modeling remains nascent. We introduce Jasmine, a performant JAX-based world modeling codebase that scales from single hosts to hundreds of accelerators with minimal code changes. Jasmine achieves an order-of-magnitude faster reproduction of the CoinRun case study compared to prior open implementations, enabled by performance optimizations across data loading, training and checkpointing. The codebase guarantees fully reproducible training and supports diverse sharding configurations. By pairing Jasmine with curated large-scale datasets, we establish infrastructure for rigorous benchmarking pipelines across model families and architectural ablations. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2510_27002 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Jasmine: A Simple, Performant and Scalable JAX-based World Modeling Codebase Mahajan, Mihir Nguyen, Alfred Srambical, Franz Bauer, Stefan Machine Learning Artificial Intelligence While world models are increasingly positioned as a pathway to overcoming data scarcity in domains such as robotics, open training infrastructure for world modeling remains nascent. We introduce Jasmine, a performant JAX-based world modeling codebase that scales from single hosts to hundreds of accelerators with minimal code changes. Jasmine achieves an order-of-magnitude faster reproduction of the CoinRun case study compared to prior open implementations, enabled by performance optimizations across data loading, training and checkpointing. The codebase guarantees fully reproducible training and supports diverse sharding configurations. By pairing Jasmine with curated large-scale datasets, we establish infrastructure for rigorous benchmarking pipelines across model families and architectural ablations. |
| title | Jasmine: A Simple, Performant and Scalable JAX-based World Modeling Codebase |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2510.27002 |