Show HN: Simulusion – A PyTorch library for training and running world models

Developer Param Thakkar released Simulusion 1.0, a PyTorch library formerly named TorchWM, on PyPI under the package name `simulusion`, unifying Dreamer, PlaNet, JEPA, IRIS, DIAMOND, DiT and Genie world models behind a single `create_config` / `create_model` / `make_env` API. The library ships built-in adapters for Gym/Gymnasium, DeepMind Control, MuJoCo, Brax, Atari and Unity ML-Agents, plus ONNX, TorchScript and TensorRT export, and every model includes a demo that trains from scratch on a 4 GB RTX 3050 laptop GPU in one to two hours. Trained Dreamer, DreamerV2, PlaNet, IRIS, DiT and I-JEPA checkpoints are published in the Simulusion demo models collection on the Hugging Face Hub.

Formerly TorchWM. The project was renamed to Simulusion for 1.0 and is published on PyPI as simulusion : pip install simulusion , then import simulusion . Simulusion brings the major world-model families together under a single PyTorch API. Train Dreamer, PlaNet, JEPA, IRIS, DIAMOND, DiT, and Genie agents through create config / create model / make env , or drop down to their encoders, decoders, and latent-dynamics backbones to compose your own architecture. Environment adapters Gym/Gymnasium, DeepMind Control, MuJoCo, Brax, Atari, Unity ML-Agents and ONNX / TorchScript / TensorRT export come built in. Every model ships with a demo that trains from scratch on a laptop GPU RTX 3050, 4 GB in one to two hours and records what it learned. These are short runs, not paper-scale results: the point is to show what each model does. | Dreamer : the real environment left beside the world model imagining it open loop right . | DIAMOND : Breakout in the emulator beside a diffusion model generating every frame itself. | | Genie : a real Sonic clip beside Genie regenerating it from the first frame, with actions it learned from unlabelled video. | IRIS : a policy boxing inside a Transformer world model made of discrete tokens. | | PlaNet : a latent model predicting the future while a CEM planner searches inside it. | DiT : CIFAR-10 samples from a class-conditional Diffusion Transformer, one row per class. | More clips, learning curves and the commands to reproduce each one are in the gallery https://paramthakkar123.github.io/simulusion/gallery.html ; the scripts live in demos/ https://github.com/ParamThakkar123/simulusion/blob/main/demos . The trained Dreamer, DreamerV2, PlaNet, IRIS, DiT and I-JEPA checkpoints are on the Hugging Face Hub in the Simulusion demo models https://huggingface.co/collections/simillusion/simulusion-demo-models-6ae0419992e0177bc1695362 collection; each model card shows how to load it. Install the core package from PyPI. This keeps environment integrations and experiment logging optional. pip install simulusion With extras pip install simulusion gym Gym/Gymnasium environments runnable quick start pip install simulusion dmc DeepMind Control Suite walker-walk, cheetah-run, ... On CPython 3.13 also run: python -m simulusion.install dmc pip install simulusion worldmodels Classic World Models ConvVAE + CMA-ES controller pip install simulusion ml-agents Unity ML-Agents pip install simulusion ml TensorBoard, W&B logging pip install simulusion viz Latent-space visualization plotly, UMAP pip install simulusion dev Testing and linting Or add it to a uv-managed project. uv add simulusion Simulusion depends on PyTorch but does not force a single PyTorch wheel index. If you need a specific PyTorch build, install or add the PyTorch packages with the index recommended for your platform by the PyTorch installation selector https://pytorch.org/get-started/locally/ : Example: CUDA 12.1 wheels. Choose a different index for CPU, ROCm, CUDA 11.x, CUDA 12.4+, or macOS. uv add torch torchvision torchaudio --index https://download.pytorch.org/whl/cu121 Use the friendly top-level API for the common path. The example below runs on a base pip install simulusion gym — no simulator downloads required: python import simulusion Trains a Dreamer agent on a Gymnasium task. Bump total steps for real runs. seed steps of random play come first and count towards total steps ; the final checkpoint is written to <logdir /ckpts/ when training finishes. agent = simulusion.create model "dreamer", env="Pendulum-v1", env backend="gym", seed steps=1 000, total steps=10 000, agent.train To train on DeepMind Control tasks such as walker-walk , install the DMC extra pip install simulusion dmc and use the default backend: agent = simulusion.create model "dreamer", env="walker-walk", total steps=1 000 000 agent.train Every algorithm in the table below is reachable through the same factory, so comparing them is a loop rather than a rewrite: python import simulusion for algo in "dreamer-v1", "dreamer-v2" : agent = simulusion.create model algo, env="Pendulum-v1", env backend="gym", total steps=20 000 agent.train examples/algorithm comparison.py runs exactly this and writes a comparison plot. Every family also trains on one shared Trainer from a YAML run config. Device, precision, resuming, evaluation and checkpoints work the same for all of them, and every run writes its resolved config.yaml and run.json git commit, package versions so it can be repeated: simulusion train --config dreamer model.env=cheetah-run trainer.precision=bf16 simulusion train --config iris game=ALE/Breakout-v5 simulusion train genie data.path=data/clips trainer.max steps=20000 simulusion train --config runs/<a previous run /config.yaml repeat it World models also train without their agent , from recorded episodes, including robot data converted from LeRobot, and any agent that speaks the WorldModel protocol can then plan with them: python simulusion datasets import lerobot datasets/pusht data/pusht simulusion train --config dreamer-wm data.path=data/pusht See the trainer guide https://paramthakkar123.github.io/simulusion/trainer guide.html and the data guide https://paramthakkar123.github.io/simulusion/episodes guide.html . - Unified interfaces across world-model algorithms: one WorldModel protocol, one Trainer , one YAML run config per experiment - One episode format for video, RL and robot data, with a LeRobot reader - Modular encoders, decoders, dynamics models, and backbones - Training and inference utilities for model-based reinforcement learning - Environment integrations for Gym/Gymnasium, Unity ML-Agents, MuJoCo, Brax, and robotics extras - Optional logging, visualization, development, and documentation extras flowchart LR subgraph API "simulusion API" CFG "create config " MDL "create model " ENV "make env " end subgraph CONFIGS "Configs" DC "DreamerConfig" JC "JEPAConfig" IC "IRISConfig" GC "GenieConfig" DIC "DiTConfig / DiamondConfig" end subgraph AGENTS "Agents / Models" DR "Dreamer / DreamerV1 / DreamerV2" JP "JEPAAgent" IR "IRISAgent" GN "Genie" DT "DiT / DIAMOND" end subgraph BACKBONES "Backbones" RSSM "RSSM / ModularRSSM" VIT "VisionTransformer" VQ "VQ-VAE / VideoTokenizer" ST "STTransformer" DIF "DDPM / DiT diffusion" end subgraph ENVS "Environments" GYM "Gym / Atari" DMC "DeepMind Control" MJ "MuJoCo" BR "Brax" UN "Unity ML-Agents" ROB "Robotics" more "..." end subgraph EXPORT "Export" ONNX "ONNX" TS "TorchScript" TRT "TensorRT" end CFG -- CONFIGS MDL -- AGENTS ENV -- ENVS AGENTS -- BACKBONES AGENTS -.- ENVS AGENTS -- EXPORT Every row is a registry entry — pass the name straight to simulusion.create model ... or simulusion.create config ... . Run simulusion.list models for the live list. | Name | Algorithm | Description | Key Features | |---|---|---|---| | dreamer | Dreamer | Model-based RL with latent dynamics alias for dreamer-v1 | Imagination, actor-critic | | dreamer-v1 | DreamerV1 | Latent imagination with Gaussian heads | Normal heads, standard KL | | dreamer-v2 | DreamerV2 | Discrete latents for pixel control | Symlog two-hot heads, balanced KL | | planet | PlaNet | Latent planning from pixels, no explicit policy | RSSM, CEM planner | | modular-rssm | ModularRSSM | Composable recurrent state-space model | Swappable priors/posteriors, custom heads | | iris | IRIS | Sample-efficient RL with Transformers | Discrete VAEs, world models | | jepa | JEPA | Self-supervised visual representations | Masked prediction, ViT | | dit | DiT | Diffusion Transformer workflows | Patch embeddings, diffusion backbones | | diamond | DIAMOND | Diffusion world model for pixel-control RL | EDM sampling, Atari imagination rollouts | | genie | Genie | Generative interactive environments from video | Latent actions, spatiotemporal transformer | | genie-small | Genie small | Development- and test-sized Genie | Same architecture, reduced width/depth | | genie-large | Genie large | Scaled-up Genie variant | Higher capacity dynamics + tokenizer | Simulusion follows semantic versioning https://semver.org/ as of 1.0.0. The public API — everything listed in the Public API reference https://paramthakkar123.github.io/simulusion/public api.html and re-exported from the top-level simulusion namespace — will not break within the 1.x line; anything removed gets a deprecation warning for at least one minor release first. Submodule internals not listed there may still change.