TorchWM

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

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

Quick Start

# Install the core package from PyPI.
# This keeps environment integrations and experiment logging optional.
pip install torchwm

# With extras
pip install torchwm[gym]       # Gym/Gymnasium environments (runnable quick start)
pip install torchwm[dmc]       # DeepMind Control Suite (walker-walk, cheetah-run, ...)
pip install torchwm[procgen]   # Procgen benchmark environments
pip install torchwm[ml-agents] # Unity ML-Agents
pip install torchwm[ml]        # TensorBoard, W&B logging
pip install torchwm[viz]       # FastAPI visualization
pip install torchwm[dev]       # Testing and linting

# Or add it to a uv-managed project.
uv add torchwm

TorchWM 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 torchwm[gym] — no simulator downloads required:

import torchwm

# Trains a Dreamer agent on a Gymnasium task. Bump `total_steps` for real runs.
agent = torchwm.create_model(
    "dreamer",
    env="Pendulum-v1",
    env_backend="gym",
    total_steps=5_000,
)
agent.train()

To train on DeepMind Control tasks such as walker-walk, install the DMC extra (pip install torchwm[dmc]) and use the default backend:

agent = torchwm.create_model("dreamer", env="walker-walk", total_steps=1_000_000)
agent.train()

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["torchwm 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

Algorithm	Description	Key Features
Dreamer	Model-based RL with latent dynamics	Imagination, actor-critic
JEPA	Self-supervised visual representations	Masked prediction, ViT
IRIS	Sample-efficient RL with Transformers	Discrete VAEs, world models
DiT	Diffusion Transformer workflows	Patch embeddings, diffusion backbones
DIAMOND	Diffusion world model for pixel-control RL	EDM sampling, Atari imagination rollouts

Documentation

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

Community

• Issue Tracker
• Discussions
• PyPI

TorchWM is under active development. APIs may change between versions.