Source-linked AI summary
stable-worldmodel: A Platform for Reproducible World Modeling Research and Evaluation
Lucas Maes, Quentin Le Lidec, Luiz Facury, Nassim Massaudi, Ayush Chaurasia, Francesco Capuano, Richard Gao, Taj Gillin, Dan Haramati, Damien Scieur, Yann LeCun, Randall Balestriero
TL;DR
World-model research lacks reliable shared baselines, reproducible evaluation, and efficient data handling. stable-worldmodel unifies data, model, solver, and benchmark tools, and its experiments show that current models remain brittle under even mild distribution shifts.
Problem
World-model research lacks reliable baselines and reproducible evaluation, while fragmented implementations and data-loading bottlenecks impede shared research workflows.
Source: §1 · p. 1, §3.1 · pp. 3–4Method
stable-worldmodel combines a Lance-based data layer, reproducible world-model baselines and planning solvers, and a benchmark suite with controllable factors of variation.
Source: §1 · p. 2Results
Current world models remain far from robust zero-shot generalization, with even mild distribution shifts causing substantial degradation in planning performance.
Source: §4.2 · p. 8, §6 · p. 9Takeaways & Limitations
Abstract
from arXiv · showhide
World models are central to building agents that can reason, plan, and generalize beyond their training data. However, research on world models is currently fragmented, with disparate codebases, data pipelines, and evaluation protocols hindering reproducibility and fair comparison. Current practice is further limited by three key bottlenecks: fragile one-off codebases, slow video data loading, and the lack of standardized generalization benchmarks. We present stable-worldmodel (swm), an open-source platform for standardized and reproducible world modeling research and evaluation. It delivers (1) a high-performance Lance-based data layer with native support and conversion tools for MP4, HDF5, and LeRobot datasets, (2) clean, well-tested implementations of modern world model baselines and planning solvers, and (3) a broad suite of environments and tasks extended with controllable visual, geometric, and physical factors of variation for systematic in-silico evaluation of dynamics understanding, control performance, representation quality, and out-of-distribution generalization. By unifying the full pipeline under a single, scalable framework, \texttt{swm} dramatically reduces research overhead and accelerates trustworthy progress toward reliable world models.
1 Introduction
World models learn predictive dynamics from data, but the field lacks reliable baselines, standardized benchmarks, and reproducible evaluation. stable-worldmodel (swm) addresses these gaps with shared data, model, solver, and evaluation tools.
1 Introduction
- Predictive decision-making has roots in 1960s–1970s control theory; recent neural world models instead learn predictive models from raw data without explicitly deriving dynamics equations.Source: §1 · p. 1
- The field lacks reliable baselines, standardized benchmarks, and reproducible evaluation protocols, while fragile, repeated reimplementation can introduce hidden bugs and inconsistent results.Source: §1 · pp. 1–2
- swm is an open-source, modular test-bed spanning data collection, training, and evaluation, including control, representation probing, out-of-distribution, and zero-shot generalization.Source: §1 · p. 2
- swm provides a high-performance Lance-based data layer, reproducible implementations of modern world-model baselines and planning solvers, and benchmarks with controllable visual, geometric, and physical variation.Source: §1 · p. 2
2 Background: the World Model Pipeline
World models encode observations into latent states and predict how those states evolve under actions. Planning then uses the learned dynamics to select actions, with reinforcement learning and model predictive control offering different optimization approaches.
Training World Models
- An encoder maps an observation o_t to a latent state s_t, and a predictor learns to estimate the next state from the current state and action.Source: §2 · p. 2
- Latent representations should retain task-relevant features and capture underlying dynamics to support accurate long-horizon prediction, planning, temporal reasoning, and out-of-distribution generalization.Source: §2 · pp. 2–3
Planning with World Models
- Planning seeks a policy that minimizes expected cumulative task cost under learned dynamics, with the cost measuring progress toward the goal state.Source: §2 · p. 3
- Reinforcement learning optimizes a policy over trajectories, while model predictive control optimizes finite-horizon predicted trajectories at test time and uses the world model directly.Source: §2 · p. 3
3 stable-worldmodel: A Unified Framework for World Model Research
stable-worldmodel (swm) is an open-source platform for reproducible world-model research, standardizing data collection, evaluation, and control while leaving model architectures and training code flexible. It addresses implementation fragmentation, data-loading bottlenecks, and evaluation gaps through shared interfaces, data and planning tools, and a broad benchmark suite.
3.1 Challenges in World Modeling Research
- Repeatedly reimplementing baselines across siloed codebases can introduce inconsistencies, undermining fair comparison and reproducibility and making it harder to isolate methodological gains.The paper notes that CEM has been independently reimplemented in at least five recent papers.Source: §3.1 · p. 3
- Frame-based storage supports fast random sampling but incurs I/O and storage overhead, while compressed video reduces disk footprint but makes later-frame access slow.Both formats can limit scaling and leave GPUs underused when data loading cannot keep the accelerator utilized.Source: §3.1 · pp. 3–4
- Standard benchmarks often stay close to training conditions, so in-distribution success may not reveal whether learned dynamics generalize to controlled changes or support counterfactual reasoning and long-horizon planning.The paper motivates zero-shot and out-of-distribution testing to probe changes in appearance, geometry, object properties, memory requirements, and physical parameters.Source: §3.1 · p. 4
3.2 Design Principles and Interface
- swm standardizes data collection, evaluation, and control while allowing researchers to retain their preferred model architectures and training code.Its World, Policy, and Solver abstractions connect environment interaction, action selection, and finite-horizon planning; new environments can be added if they follow the Gymnasium API.Source: §3.2 · pp. 4–5
3.3 System components
- Lance achieves the highest loading throughput in the reported benchmarks, while offering a good trade-off in compression; swm also supports MP4, HDF5, and LeRobot data with conversion tools.The comparisons cover Push-T and Two-Room, and include local and remote storage conditions.Source: Figure 3 · p. 5, §3.3 · p. 6
- The data pipeline supports offline and online collection with random, expert, checkpointed, or mixed policies, and can apply visual, geometric, or physical factors of variation.Online training can alternate data collection with model training or evaluation by repeating the platform loop.Source: §3.3 · pp. 5–6
- swm implements goal-conditioned reinforcement-learning baselines and latent world models, alongside tested sampling-based and gradient-based planning solvers.When paired with their original world models, the implementations reproduce the planning success rates reported in DINO-WM and PLDM.Source: §3.3 · p. 6
3.4 Benchmarking World Models
- The benchmark suite spans low-dimensional and 2D control, 3D visual environments, Atari, MuJoCo, and partially observable tasks, with controlled visual and physical variations for robustness and zero-shot evaluation.The platform reports success rate as its primary control and planning metric and supports online evaluation from randomly sampled or trajectory-selected boundary conditions.Source: §3.4 · p. 7
4 Case Study: Planning Performance and Zero-Shot Evaluation of WMs
The Push-T case study uses swm’s unified pipeline to assess planning and zero-shot robustness, finding that prediction error aligns poorly with success and that visual shifts substantially degrade planning.
4 Case Study: Planning Performance and Zero-Shot Evaluation of WMs
- The unified evaluation pipeline assesses planning and zero-shot generalization across diverse settings with minimal additional code, supporting standardized, reproducible diagnosis of model limitations.The case studies are presented as a way to turn ad-hoc analyses into reproducible experiments.Source: §4 · p. 7, §4.2 · p. 8
4.1 In-distribution Planning Performance
- On Push-T, implementations recover baseline performance consistent with prior reports under identical expert-data and planning configurations, while TD-MPC2 performs poorly offline.The authors conjecture that TD-MPC2’s offline result stems from out-of-distribution actions fooling its predictor.Source: §4.1 · p. 7
4.2 World Models Are Still Brittle Under Distribution Shift
- Across four Push-T shift regimes, prediction-error and planning-success distributions overlap for LeWM and PLDM, although stronger OOD shifts raise prediction error.The regimes progress from expert training and validation data to random-policy trajectories, then random trajectories with visual and other variations.Source: §4.2 · p. 7, Figure 4 · p. 8
- Mild visual perturbations sharply reduce most models’ success, while distractor robustness decays quadratically beyond small counts across baselines.The analysis varies visual distractor count and factors such as color, shape, size, and background.Source: §4.2 · pp. 7–8, Figure 5 · p. 8
5 Related Work
Related work includes world-model benchmarks, reinforcement-learning benchmarks, and complementary development frameworks, but the authors identify gaps in training infrastructure, breadth, and world-model support.
Benchmarking World Models
- World-model research spans generative, reward-prediction, and self-supervised approaches, while standalone codebases contribute to duplicated implementations and inconsistent evaluation.WorldTest and WorldMark standardize aspects of evaluation, but their training infrastructure remains fragmented; EB-JEPA is architecture-restricted and lacks scalable research components.Source: §5 · pp. 8–9
Benchmarking RL
- RL datasets often target model-free methods in fixed environments, while existing visual-robustness benchmarks are tied to particular environment families and visual perturbations.swm reuses benchmark environments and adds systematic variation for robustness, zero-shot adaptation, and generalization evaluation.Source: §5 · p. 9
Complementary Frameworks
- General-purpose RL and robotics frameworks cover adjacent needs, but the authors position swm as adding world-model components such as planning and evaluation rather than replacing them.The paper describes integrations with stable-baselines3 and LeRobot and aims for a unified reference library.Source: §5 · p. 9
6 Conclusion and Future Work
The conclusion presents swm as a modular, reproducible world-modeling platform and argues that current models still lack robust zero-shot generalization; the appendix documents its setup and core abstractions.
6 Conclusion and Future Work
- swm combines high-performance data management, planning solvers, and customizable benchmarks to reduce friction across the world-modeling pipeline.The authors identify visual and physical robustness and long-horizon consistency as areas needing improvement, and suggest architectural advances and systematic scaling may help generalization.Source: §6 · p. 9
Appendix
- The appendix walks through installation, core abstractions, environments, interventions, baselines, evaluation, planning solvers, command-line tools, and open research problems.It also gives contribution guidance and additional experimental results.Source: Appendix · p. 14
A Installation & Setup
- The setup guidance recommends uv and Linux, while noting that macOS, Windows, and a Conda setup are also possible.The library is designed for both local development and large-scale distributed training.Source: A Installation & Setup · p. 14
B Core Abstractions
- World wraps vectorized environments, while Policy and Solver abstractions connect action selection with planning and model-predictive control.A model implementing get_cost can be used as a controller through the Solver and MPCPolicy abstractions.Source: B Core Abstractions · pp. 14–15
- Factors of Variation expose controllable environment properties, and the Dataset API supports multiple formats and conversion within a modular pipeline.Wrappers can modify observations or default behavior, including by applying visual transformations.Source: B Core Abstractions · p. 15
C Environments and Factors of Variation (FoV)
swm integrates a broad set of environments with its World interface and distinguishes native environments with simulator-level variation controls from external environments that rely on wrappers or adapters.
C Environments and Factors of Variation (FoV)
- The environment suite spans control, manipulation, navigation, open-world, and Atari tasks, with Gymnasium integration supporting high-throughput collection and evaluation.Listed families include DeepMind Control Suite, OGBench, Classic Control, Fetch-Suite, Craftax, ALE, and swm’s PushT and TwoRoom environments.Source: C Environments and Factors of… · p. 15
C.1 Environment Registry
- Native swm environments expose structured state and typed variation spaces for reset-time interventions on scene, geometry, and physical properties.These controls are accessed through the Gymnasium options argument.Source: §C.1 · pp. 15–16
- External environments run without modification and lack native variation spaces, so interventions use visual wrappers or adapters; unsupported Gymnasium environments have no FoV unless added.The platform can load registered Gymnasium environments by default, but the intervention machinery is not automatically available for them.Source: C Environments and Factors of… · p. 15, §C.1 · p. 16
C.2 Factors of Variation (FoV) & Wrappers
swm supports interventions through editable simulator factors and observation-boundary visual wrappers, then provides shared implementations of goal-conditioned and latent world-model baselines. The mechanisms and baselines are designed to support controlled evaluation and comparison.
C.2 Factors of Variation (FoV) & Wrappers
- Editable environments use native factors of variation, while black-box environments use observation-boundary visual wrappers; both mechanisms feed the same evaluation pipeline and can compose.Factors can be sampled randomly or fixed, and their values remain constant through an episode; resampling wrappers can redraw perturbations at each reset.Source: §C.2 · pp. 16–17
D Available Baselines
- swm implements six offline baselines across goal-conditioned reinforcement learning and latent world models, with shared training structure and datasets for fair comparisons.The baselines use a common Hydra-configured training entry point and can share trajectories collected before training.Source: D Available Baselines · pp. 17–18
D.1 Goal-Conditioned Reinforcement Learning
- Goal-conditioned baselines include GCBC imitation and IQL/IVL value-learning variants; their policies use DINOv2 embeddings and act without a planning solver.GCBC reproduces expert actions from a state and goal; IQL learns value and critic functions before extracting a policy through advantage-weighted regression.Source: §D.1 · pp. 19–20
D.2 Latent World Models
- Latent-model baselines share an encoder-predictor interface and MPC rollout planning, spanning frozen-encoder DINO-WM, joint JEPA variants, and reward-driven TD-MPC2.PLDM uses a multi-term anti-collapse objective, whereas LeWorldModel uses SIGReg as its single regularizer.Source: §D.2 · pp. 21–22
E Evaluation Protocols
swm evaluates goal-conditioned agents using target observations and provides both episodic and dataset-driven protocols. Its planning solvers optimize action sequences using model-predicted costs across sampling-based and gradient-based approaches.
E Evaluation Protocols
- Goal-observation tasks give agents a start observation and target frame, then assess whether they reach the goal within a fixed interaction budget.The same evaluation question applies to latent-model planners and goal-conditioned policies, despite their different interfaces.Source: E Evaluation Protocols · p. 23
- Episodic evaluation samples tasks through the environment’s reset randomizer, while dataset-driven evaluation selects starts and later target observations from recorded trajectories.Dataset targets are reachable under the recorded behavior by construction, and the protocol can vary goal horizon and dataset distribution independently.Source: E Evaluation Protocols · pp. 23–24
- Both protocols cap episodes by environment steps; reported metrics include success rate, time-to-goal among successes, and wall-clock latency per planning step.The step cap is intended to keep wall-clock evaluation cost comparable across methods with different per-step computation.Source: E Evaluation Protocols · p. 24
F Planning Solvers
- Planning solvers minimize model-evaluated action-sequence costs, covering sampling-based and gradient-based methods for continuous and discrete action spaces.The objective uses the current latent state and a planning horizon of H steps.Source: F Planning Solvers · p. 24
F.1 Sampling-Based Solvers
The platform provides sampling- and gradient-based planning solvers, alongside scripts and interfaces for extending its research workflow.
F.1 Sampling-Based Solvers
- Sampling-based solvers require sufficiently large candidate populations to explore high-dimensional action spaces, though swm parallelizes rollouts for GPU batching.They do not use gradient information.Source: §F.1 · p. 24
- CEM refits a diagonal Gaussian over action sequences using the lowest-cost candidates, while Predictive Sampling returns the lowest-cost randomly perturbed nominal plan without iterative refinement.Predictive Sampling instead relies on random search and repeated MPC replanning over time.Source: §F.1 · p. 25
- MPPI updates a nominal sequence through soft cost-weighted samples, while iCEM adds colored noise, elite retention, and momentum to improve sample efficiency under small planning budgets.MPPI can restrict weighting to the top-k candidates; iCEM's colored noise correlates perturbations across time.Source: §F.1 · pp. 26–27
- Categorical CEM maintains per-step categorical action distributions and updates them from elite frequencies, avoiding a continuous relaxation for discrete-action tasks.Smoothing and momentum are optional.Source: §F.1 · p. 27
F.2 Gradient-Based Solvers
- Gradient-based solvers optimize action sequences using first-order information, but require differentiable costs and can face poorly conditioned gradients over long horizons.The Lagrangian solver adds differentiable inequality constraints with inner action updates and outer multiplier updates; PGD relaxes discrete actions to probability vectors, while GRASP optimizes intermediate virtual states.Source: §F.2 · pp. 28–30
G Scripts & Command Line Interface
- Hydra-configured scripts cover data collection, expert-policy training, and training and evaluation of imitation-learning, offline-RL, and world-model baselines.The CLI also reports dataset metadata and lists environments, factors of variation, and cached checkpoints.Source: G Scripts & Command Line Interface · p. 31
H Extending stable-worldmodel
- swm invites community additions of environments, models, and solvers, and plans further baseline benchmarking and a public leaderboard.The leaderboard is intended to track the state of the art and guide baseline choice.Source: H Extending stable-worldmodel · p. 31
I Experimental setup
The experiments use dataset-driven Push-T evaluations, with controlled start-goal offsets, specified trajectory budgets, and CEM planning unless otherwise noted.
Tasks and dataset
- All experiments use the DINO-WM expert dataset on Push-T, drawing each start and goal from one trajectory with a fixed 25-step offset; shift sweeps also vary the source dataset.The shift regimes are expert train, expert validation, random policy, and random policy with factors of variation.Source: I Experimental setup · p. 32
Evaluation budget and metrics
- Episodes are capped at 50 environment steps, with 50 trajectories for in-distribution comparisons and 256 per regime for prediction-versus-success analyses.Success rate follows the DINO-WM definition.Source: I Experimental setup · p. 32
Planning
- Unless noted otherwise, CEM plans with 30 iterations, 300 candidates, 30 elites, and a five-model-step horizon, executing the optimized sequence open-loop.Models use frameskip 5, making the horizon 25 environment steps; runs take a few minutes on one L40S GPU.Source: I Experimental setup · p. 32
Baselines
- Robustness analyses primarily compare PLDM and LeWM, with DINO-WM additionally reported in Table 5b; baseline hyperparameters follow the original publications.PLDM and LeWM were selected for strong in-distribution planning and low per-step overhead.Source: I Experimental setup · p. 32
J Additional Results
The supplementary results expose factor-specific generalization failures, data-loading trade-offs, weak prediction-error discrimination of planning success, and differences between online and offline control.
Prediction error and planning success
- Prediction MSE rises as evaluation shifts from expert training to random-policy data with all factors perturbed, yet successful and failed PLDM plans have heavily overlapping error distributions.The overlap means low prediction error does not guarantee planning success.Source: J Additional Results · p. 32
Background-color robustness
- LeWM planning success stays high near white and green backgrounds but collapses at higher-intensity red, blue, or purple shifts.The green background matches the target anchor's color in the original environment.Source: J Additional Results · p. 32
Factor-wise generalization
- 52.0% success rate: PLDM is notably robust to agent-shape variation, while changing target-anchor position or agent size reduces planning success across all baselines.Table 4 breaks down zero-shot generalization by individual factors on random-policy OOD trajectories.Source: J Additional Results · p. 32, Table 4 · p. 33
Data loading
- Over 5.0k samples per second locally and 3.4k when streaming from S3: Lance delivers high throughput, while MP4 uses the least storage at 220 MB but impairs random access.Lance offers a trade-off between compression and high-speed random access.Source: J Additional Results · p. 32, Figure 8 · p. 33
Online and offline control
- TD-MPC2 reaches competitive asymptotic rewards and matches SAC sample efficiency online, while its actor drifts outside the training-data manifold and performs poorly offline on Push-T.Figure 13 contrasts expert-policy rollouts that stay within the data distribution with TD-MPC2 actor trajectories that drift outside it.Source: J Additional Results · p. 33
K Research Opportunities
The paper identifies visual and physical robustness, longer-horizon planning, and world-model scaling as open research directions supported by swm's evaluation and data infrastructure.
Achieving performant zero-shot world models
- Current approaches remain brittle under simple visual perturbations, making visual and physical robustness a central research direction for zero-shot world models.The paper states that closing this robustness gap is required for safe real-world deployment and will demand substantial further progress.Source: K Research Opportunities · p. 34, Figure 11 · p. 35
Unlocking long-horizon planning
- Prediction errors compound over rollouts, while most published evaluations use short in-distribution horizons; controllable goal offsets and multiple solvers support testing longer-horizon planning.The paper points to horizon scaling, hierarchical planning, and multi-step planning as research directions.Source: K Research Opportunities · p. 34
Scaling World Models
- Whether scaling laws and emergent behaviors apply to world models remains open; Lance data loading and standardized pipelines support comparisons across data, parameter, and compute scales.The data layer is designed to alleviate I/O bottlenecks during multimodal training.Source: K Research Opportunities · pp. 34–35