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

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; the scripts live in demos/. The trained Dreamer, DreamerV2, PlaNet, IRIS, DiT and I-JEPA checkpoints are on the Hugging Face Hub in the Simulusion demo models collection; each model card shows how to load it.

pip install simulusion

pip install simulusion[gym]       # Gym/Gymnasium environments (runnable quick start)
pip install simulusion[dmc]       # DeepMind Control Suite (walker-walk, cheetah-run, ...)
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

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:

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:

import simulusion

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:

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:

simulusion datasets import lerobot datasets/pusht data/pusht
simulusion train --config dreamer-wm data.path=data/pusht

See the trainer guide and the data guide.

• Unified interfaces across world-model algorithms: one WorldModel protocol, oneTrainer , 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.

Simulusion follows semantic versioning as of 1.0.0. The public API — everything listed in the Public API reference 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.