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.

See it in action

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/.

Quick Start

# 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:

# 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:

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()

Swap the algorithm, keep the code

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. Construction is unified across all registered models. A shared step-budget train() currently covers the Dreamer family — other agents use their own trainers (simulusion train … / JEPAAgent.train() / DiamondAgent.train()). The example reports which is which rather than assuming.

Features

• Unified interfaces across world-model algorithms
• 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

Architecture

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

Supported Algorithms

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.

Documentation

• Full Documentation
• Installation Guide
• Training Guide
• API Reference

Community

• Issue Tracker
• Discussions
• PyPI
• Contributing Guide
• Code of Conduct

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.