A test bench for world models that actually stays still

What it does

stable-worldmodel wraps the three stages of world-model research—data collection, training, and model-predictive control evaluation—into one Python interface. It bundles a couple dozen environments (DM Control, OGBench, Atari, Craftax, plus classics like Push-T and Two-Room) and ships reference implementations for baselines including LeWM and DINO-WM. The goal is letting researchers swap in a new model or loss without rewriting the surrounding plumbing.

The interesting bit

Most environments come with tunable “factors of variation”—lighting, textures, dynamics, morphology—that let you stress-test zero-shot generalization without any extra setup. The data layer is also unusually flexible: a small format registry supports LanceDB, HDF5, folder-per-episode, video, and direct LeRobot Hub streaming, with one-shot conversion between them. Benchmarks in the repo show LanceDB running ~3.4× faster than local HDF5 and ~350× faster than HDF5-over-S3, while taking a third the disk space.

Key highlights

• Unified World object for collecting, training, and evaluating with MPC
• Reference training scripts for LeWM and PreJEPA (DINO-WM reproduction)
• Pluggable data backends with autodetect and append-by-default semantics
• 100+ Atari games plus robotics, control, and navigation tasks via Gymnasium interface
• Direct training from HuggingFace object storage without local download (Colab notebook provided)

Caveats

• APIs may change between minor versions; library is in active development
• LeRobot support requires Python 3.12+ and is a separate opt-in extra

Verdict Worth a look if you’re actually training world models and tired of maintaining your own benchmark harness. Probably overkill if you just need a single environment and a custom dataloader.