Jasmine: A Simple, Performant and Scalable JAX-based World Modeling Codebase

TL;DR

Jasmine is a high-performance, scalable JAX-based codebase for world modeling that enables efficient training, reproducibility, and benchmarking across diverse configurations and large datasets.

Contribution

It introduces Jasmine, a scalable, optimized, and reproducible world modeling framework built in JAX, supporting extensive benchmarking and large-scale training.

Findings

Achieves an order-of-magnitude faster reproduction of CoinRun.

Supports scalable training from single hosts to hundreds of accelerators.

Provides infrastructure for rigorous benchmarking and ablation studies.

Abstract

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.