A Unified Toolbox for JEPA and World Model Research: stable-worldmodel

World‑model researchers often encounter fragmented data formats and heterogeneous environment interfaces, requiring them to rebuild baseline pipelines for each experiment. The open‑source platform stable-worldmodel integrates these components.

Four Core Components

Datasets : Supports Lance, HDF5, MP4, LeRobot and automatically detects formats.

Baselines : Bundles state‑of‑the‑art implementations such as DINO‑WM, LeWM, PLDM, GCBC, GCIVL, and GCIQL.

Environments : Provides roughly 150 environments covering DeepMind Control Suite, Gymnasium, OGBench, Craftax, over 100 Atari games, and classic benchmarks like PushT and Two‑Room.

Solvers : Includes CEM, iCEM, MPPI, Predictive Sampling, as well as SGD/Adam, PGD, and Augmented Lagrangian optimizers.

Data‑storage and read‑performance comparison on the PushT dataset shows: HDF5 occupies 43 GB, LanceDB 13 GB, and a compressed video format 496 MB. Locally, LanceDB reads three times faster than HDF5, achieving over 4,800 samples / s. When stored on S3, HDF5 without caching drops to 9.1 samples / s (≈7 s per step), while LanceDB maintains above 3,000 samples / s.

Three‑Step Workflow

import stable_worldmodel as swm
from stable_worldmodel.policy import WorldModelPolicy, PlanConfig
from stable_worldmodel.solver import CEMSolver

# 1. Collect data
world = swm.World("swm/PushT-v1", num_envs=8)
world.set_policy(your_expert_policy)
world.collect("data/pusht_demo.lance", episodes=100, seed=0)

# 2. Load and train
dataset = swm.data.load_dataset("data/pusht_demo.lance", num_steps=16)
world_model = ...

# 3. Evaluate with MPC
solver = CEMSolver(model=world_model, num_samples=300)
policy = WorldModelPolicy(solver=solver, config=PlanConfig(horizon=10))
world.set_policy(policy)
results = world.evaluate(episodes=50)

Most environments also provide a set of “factors of variation” (independent parameters for lighting, texture, dynamics, and morphology) that enable zero‑shot generalisation tests without additional experimental setups.

Installation requires a single command: pip install stable-worldmodel.

Project repository: https://github.com/galilai-group/stable-worldmodel

Paper: https://arxiv.org/abs/2605.21800

The library remains under active development, so APIs may change between minor releases.