Metadata-Version: 2.4
Name: simulusion
Version: 1.0.0
Summary: A Modular Pytorch Based library for training world models
Author: Param Thakkar
License-Expression: MIT
Project-URL: Homepage, https://github.com/paramthakkar123/simulusion
Project-URL: Documentation, https://paramthakkar123.github.io/simulusion/
Project-URL: Issues, https://github.com/paramthakkar123/simulusion/issues
Project-URL: Changelog, https://github.com/paramthakkar123/simulusion/releases
Keywords: world-models,pytorch,mbrl
Classifier: Development Status :: 5 - Production/Stable
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Operating System :: OS Independent
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.1.0
Requires-Dist: torchvision>=0.16.0
Requires-Dist: click>=8.0.0
Requires-Dist: einops>=0.8.2
Requires-Dist: tqdm>=4.67.1
Requires-Dist: pyyaml>=6.0.3
Provides-Extra: gym
Requires-Dist: gymnasium[box2d]>=1.2.2; extra == "gym"
Requires-Dist: huggingface-hub>=0.23.0; extra == "gym"
Requires-Dist: ale-py>=0.10.1; extra == "gym"
Requires-Dist: autorom>=0.6.1; extra == "gym"
Provides-Extra: ml-agents
Requires-Dist: mlagents-envs<0.29.0,>=0.28.0; extra == "ml-agents"
Requires-Dist: protobuf>=5.29.6; extra == "ml-agents"
Provides-Extra: dmc
Requires-Dist: dm-control>=1.0.0; python_version < "3.13" and extra == "dmc"
Requires-Dist: labmaze-new>=0.0.2; python_version >= "3.13" and extra == "dmc"
Requires-Dist: gymnasium>=1.2.2; extra == "dmc"
Provides-Extra: dmc-uv
Requires-Dist: dm-control>=1.0.0; extra == "dmc-uv"
Requires-Dist: labmaze-new>=0.0.2; python_version >= "3.13" and extra == "dmc-uv"
Requires-Dist: gymnasium>=1.2.2; extra == "dmc-uv"
Provides-Extra: genie
Requires-Dist: h5py>=3.11.0; extra == "genie"
Requires-Dist: huggingface-hub>=0.23.0; extra == "genie"
Requires-Dist: omegaconf>=2.3.0; extra == "genie"
Provides-Extra: mujoco
Requires-Dist: mujoco>=3.2.0; extra == "mujoco"
Provides-Extra: robotics
Requires-Dist: gymnasium-robotics>=1.3.0; extra == "robotics"
Requires-Dist: mujoco<3.12,>=3.2.0; extra == "robotics"
Provides-Extra: brax
Requires-Dist: brax>=0.13.0; extra == "brax"
Requires-Dist: jax>=0.10.1; extra == "brax"
Provides-Extra: bsuite
Requires-Dist: bsuite>=0.3.5; extra == "bsuite"
Provides-Extra: viz
Requires-Dist: opencv-python>=4.12.0.88; extra == "viz"
Requires-Dist: imageio-ffmpeg>=0.5.1; extra == "viz"
Requires-Dist: umap-learn>=0.5.11; extra == "viz"
Requires-Dist: scikit-learn>=1.8.0; extra == "viz"
Requires-Dist: plotly>=6.5.0; extra == "viz"
Requires-Dist: h5py>=3.11.0; extra == "viz"
Provides-Extra: worldmodels
Requires-Dist: albumentations>=2.0.0; extra == "worldmodels"
Requires-Dist: cma>=4.0.0; extra == "worldmodels"
Requires-Dist: opencv-python>=4.12.0.88; extra == "worldmodels"
Requires-Dist: imageio-ffmpeg>=0.5.1; extra == "worldmodels"
Provides-Extra: ml
Requires-Dist: torchaudio>=0.13.0; extra == "ml"
Requires-Dist: hydra-core>=1.3.2; extra == "ml"
Requires-Dist: omegaconf>=2.3.0; extra == "ml"
Requires-Dist: wandb>=0.16.0; extra == "ml"
Requires-Dist: tensorboard>=2.20.0; extra == "ml"
Requires-Dist: nvidia-ml-py>=12.535.77; extra == "ml"
Requires-Dist: tensorboardx>=2.6.4; extra == "ml"
Provides-Extra: dev
Requires-Dist: pytest>=9.0.3; extra == "dev"
Requires-Dist: pytest-cov>=7.1.0; extra == "dev"
Requires-Dist: pytest-timeout>=2.3.1; extra == "dev"
Requires-Dist: pytest-xdist>=3.6.0; extra == "dev"
Requires-Dist: pre-commit>=4.5.0; extra == "dev"
Requires-Dist: mypy>=1.19.1; extra == "dev"
Requires-Dist: gitpython>=3.1.58; extra == "dev"
Requires-Dist: ruff==0.16.8; extra == "dev"
Provides-Extra: all
Requires-Dist: simulusion[brax,bsuite,dev,dmc,docs,gym,ml,ml-agents,mujoco,robotics,viz,worldmodels]; extra == "all"
Provides-Extra: docs
Requires-Dist: sphinx<9.0.0,>=8.0.0; extra == "docs"
Requires-Dist: sphinx-autodoc-typehints<1.26.0; extra == "docs"
Requires-Dist: sphinx-copybutton>=0.5.2; extra == "docs"
Requires-Dist: pydata-sphinx-theme>=0.15.0; extra == "docs"
Requires-Dist: sphinx-rtd-theme>=3.1.0; extra == "docs"
Requires-Dist: sphinxcontrib-bibtex>=2.6.5; extra == "docs"
Requires-Dist: sphinxcontrib-mermaid>=0.9.0; extra == "docs"
Requires-Dist: sphinx-thebe>=0.3.0; extra == "docs"
Requires-Dist: sphinxext-opengraph>=0.13.0; extra == "docs"
Requires-Dist: myst-parser<5.0.0,>=2.0.0; extra == "docs"
Requires-Dist: nbsphinx>=0.9.8; extra == "docs"
Requires-Dist: mistune>=3.3.0; extra == "docs"
Requires-Dist: nbconvert>=7.17.1; extra == "docs"
Dynamic: license-file

<h1 align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/ParamThakkar123/simulusion/main/docs/source/_static/simulusion-logo-dark.svg">
    <img alt="Simulusion" src="https://raw.githubusercontent.com/ParamThakkar123/simulusion/main/docs/source/_static/simulusion-logo-light.svg" height="72">
  </picture>
</h1>

<div align="center">
  <p>
    <a href="https://pypi.org/project/simulusion/"><img alt="PyPI version" src="https://badge.fury.io/py/simulusion.svg"></a>
    <a href="https://pypi.org/project/simulusion/"><img alt="PyPI downloads" src="https://img.shields.io/pypi/dm/simulusion.svg"></a>
    <a href="https://opensource.org/licenses/MIT"><img alt="License: MIT" src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>
    <a href="https://paramthakkar123.github.io/simulusion/"><img alt="Documentation" src="https://img.shields.io/badge/docs-link-blue.svg"></a>
    <a href="https://github.com/paramthakkar123/simulusion/actions/workflows/test.yml"><img alt="CI" src="https://github.com/paramthakkar123/simulusion/actions/workflows/test.yml/badge.svg"></a>
  </p>
  <p><strong>Modular PyTorch library for world models — many algorithms, one consistent API.</strong></p>
</div>

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

<table>
<tr>
<td width="50%" valign="top"><a href="https://paramthakkar123.github.io/simulusion/gallery.html#gallery-dreamer"><img src="https://raw.githubusercontent.com/ParamThakkar123/simulusion/main/docs/source/_static/gallery/dreamer_dream.gif" alt="Dreamer" width="100%"></a><br><b>Dreamer</b>: the real environment (left) beside the world model imagining it open loop (right).</td>
<td width="50%" valign="top"><a href="https://paramthakkar123.github.io/simulusion/gallery.html#gallery-diamond"><img src="https://raw.githubusercontent.com/ParamThakkar123/simulusion/main/docs/source/_static/gallery/diamond_dream.gif" alt="DIAMOND" width="100%"></a><br><b>DIAMOND</b>: Breakout in the emulator beside a diffusion model generating every frame itself.</td>
</tr>
<tr>
<td width="50%" valign="top"><a href="https://paramthakkar123.github.io/simulusion/gallery.html#gallery-genie"><img src="https://raw.githubusercontent.com/ParamThakkar123/simulusion/main/docs/source/_static/gallery/genie_replay.gif" alt="Genie" width="100%"></a><br><b>Genie</b>: a real Sonic clip beside Genie regenerating it from the first frame, with actions it learned from unlabelled video.</td>
<td width="50%" valign="top"><a href="https://paramthakkar123.github.io/simulusion/gallery.html#gallery-iris"><img src="https://raw.githubusercontent.com/ParamThakkar123/simulusion/main/docs/source/_static/gallery/iris_dream.gif" alt="IRIS" width="100%"></a><br><b>IRIS</b>: a policy boxing inside a Transformer world model made of discrete tokens.</td>
</tr>
<tr>
<td width="50%" valign="top"><a href="https://paramthakkar123.github.io/simulusion/gallery.html#gallery-planet"><img src="https://raw.githubusercontent.com/ParamThakkar123/simulusion/main/docs/source/_static/gallery/planet_dream.gif" alt="PlaNet" width="100%"></a><br><b>PlaNet</b>: a latent model predicting the future while a CEM planner searches inside it.</td>
<td width="50%" valign="top"><a href="https://paramthakkar123.github.io/simulusion/gallery.html#gallery-dit"><img src="https://raw.githubusercontent.com/ParamThakkar123/simulusion/main/docs/source/_static/gallery/dit_samples.png" alt="DiT" width="100%"></a><br><b>DiT</b>: CIFAR-10 samples from a class-conditional Diffusion Transformer, one row per class.</td>
</tr>
</table>

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/`](demos/).

## Quick Start

```bash
# 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/):

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

```python
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:

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

```mermaid
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.

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

## Documentation

- [Full Documentation](https://paramthakkar123.github.io/simulusion/)
- [Installation Guide](https://paramthakkar123.github.io/simulusion/installation.html)
- [Training Guide](https://paramthakkar123.github.io/simulusion/training_guide.html)
- [API Reference](https://paramthakkar123.github.io/simulusion/api_reference.html)

## Community

- [Issue Tracker](https://github.com/paramthakkar123/simulusion/issues)
- [Discussions](https://github.com/paramthakkar123/simulusion/discussions)
- [PyPI](https://pypi.org/project/simulusion/)
- [Contributing Guide](CONTRIBUTING.md)
- [Code of Conduct](CODE_OF_CONDUCT.md)

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