Modular PyTorch library for world models — many algorithms, one consistent API.

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.

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.

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.

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, 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

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

Loading

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.

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.