$ show galilai-group/stable-worldmodel
Stable-worldmodel is a research framework for goal-conditioned world models in reinforcement learning, providing a full pipeline from data collection and storage to model training and planning-based evaluation. The codebase is primarily a research platform with strong software engineering practices—pluggable data formats (Lance, HDF5, video, folder), a command-line interface (swm), and a Hydra-based configuration system for reproducible experiments. It includes implementations of multiple world model architectures (LeWM, PLDM, PreJEPA/DINO-WM, TD-MPC2, GCRL), planning solvers (CEM, MPPI, iCEM, gradient-based, Lagrangian), and environment wrappers with domain randomization across 30+ tasks (DeepMind Control, Gymnasium, PushT, OGBench, Craftax, robotics). The repository demonstrates high craftsmanship in its data engineering: the data format registry extends cleanly, the Lance format supports cloud-native streaming from Hugging Face Buckets, and utilities like GoalDataset, HistoryBuffer, and ReplayBuffer solve recurring RL data management problems. The project is actively maintained with extensive test coverage, comprehensive documentation, and a scalable training pipeline using stable-pretraining and Lightning. While the primary goal is research (testing hypotheses about world model generalization under variation), the codebase includes application-level production quality—CI/CD, CLI, checkpoint management, and robust error handling—making it a rare hybrid that serves both researchers and engineers.
The LLM identified 5 files worth highlighting for reuse:
Implements the Lance dataset reader and writer, an append-friendly columnar format optimized for RL data. Supports episode metadata side tables, JPEG-compressed images, dot-rename column mapping, and HTTP ranged reads from Hugging Face Buckets (no local download). The pattern (detect/open_reader/open_writer) follows the format registry and can be extracted for any episodic data storage system.
Core dataset abstraction providing window-based sampling with frameskip, goal augmentation (current, random, geometric future, uniform future), and episode-scoped metadata. Integrates with format readers and handles heterogeneous columns (images, actions, proprio). The GoalDataset pattern is reusable for any offline RL or imitation learning project needing structured goal-conditioned data loading.
Cross-Entropy Method (CEM) solver with configurable horizon, action repeats, elites, and warm-start from a repeated previous plan. Supports batched planning across multiple environments, callback hooks for logging, and optional actor-model initialization. The solver implements the Costable protocol and can be reused in any MPC pipeline with action spaces.
Wrapper for the Half-Cheetah environment that adds a programmatic variation_space with factors like body color, joint friction, and mass. Demonstrates the variation pattern used across all DMControl environments: resetting space parameters to modify MuJoCo XML on-the-fly. This pattern is reusable for any MuJoCo-based environment requiring systematic domain randomization.
Extends Gymnasium spaces (Box, Discrete, Dict) with init_value, sampling_order, and constrain_fn for rejection sampling. Used by all environments to ensure valid variation combinations (e.g., door fits in maze, agent not inside wall). The constraint-aware sampling pattern is reusable for any RL or simulation library needing structured, valid state generation.