torchwm 0.4.1

A Modular Pytorch Based library for training world models

TorchWM

Modular PyTorch Library for World Models

---

⚡ Quick Start

Train Dreamer agents in just 3 lines of code. TorchWM provides a unified interface for training and deploying world models.

Installation

# Core dependencies
pip install torchwm

# With extras
pip install torchwm[gym]       # Additional gym 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

Training a Dreamer Agent

Use the friendly top-level API for the common path:

import torchwm

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

The lower-level research modules are still available when you need direct control:

from torchwm import DreamerAgent, DreamerConfig

cfg = DreamerConfig()
cfg.env = "walker-walk"
agent = DreamerAgent(cfg)

Creating Environments and Operators

import torchwm

env = torchwm.make_env("CartPole-v1", backend="gym")
op = torchwm.get_operator("dreamer", image_size=64, action_dim=6)
processed = op.process({"image": image, "action": action})

🚀 Features

• 🎯 Unified Interface: Consistent API across all world model algorithms
• 🔧 Modular Components: Swappable encoders, decoders, and backbones
• 🚀 High Performance: Optimized for both training and inference
• 🌍 Multi-Environment: Support for DMC, Gym, Unity, and custom environments
• 📊 Rich Monitoring: Integrated logging with Weights & Biases and TensorBoard
• 🧠 Research Ready: Easy experimentation with different architectures

🧠 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
Diamond	Diffusion + RL for continuous control	EDM sampling, value functions

📖 Documentation

📖 Full Documentation

Get Started

• Installation
• Training Guide
• Inference Guide

User Guides

• Operators Guide
• Environments Guide
• Package Overview

Algorithms

• Dreamer
• JEPA
• IRIS
• DiT

🤝 Community

• 🐛 Issue Tracker
• 💬 Discussions
• 📧 PyPI

---

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

:hidden:
:maxdepth: 1
:caption: Get Started

getting_started
installation

:hidden:
:maxdepth: 1
:caption: User Guides

operators_guide
training_guide
inference_guide
environments_guide
package_overview

:hidden:
:maxdepth: 1
:caption: Algorithms

dreamer
jepa
iris
dit

:hidden:
:maxdepth: 1
:caption: Reference

api_reference
configs_reference

:hidden:
:maxdepth: 1
:caption: Development

contributing
benchmarks

:hidden:
:maxdepth: 1
:caption: User Guides

operators_guide
training_guide
inference_guide
environments_guide
package_overview

:hidden:
:maxdepth: 1
:caption: Algorithms

dreamer
jepa
iris
dit

:hidden:
:maxdepth: 1
:caption: Reference

api_reference
configs_reference

:hidden:
:maxdepth: 1
:caption: Development

contributing
benchmarks

---

⚡ Quick Start

TorchWM provides a unified interface for training and deploying world models.

Installation

pip install torchwm
# or with uv
uv add torch torchvision torchaudio

Training a Dreamer Agent

Use the friendly top-level API for the common path:

import torchwm

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

The lower-level research modules are still available when you need direct control:

from torchwm import DreamerAgent, DreamerConfig

cfg = DreamerConfig()
cfg.env = "walker-walk"
agent = DreamerAgent(cfg)

Creating Environments and Operators

import torchwm

env = torchwm.make_env("CartPole-v1", backend="gym")
op = torchwm.get_operator("dreamer", image_size=64, action_dim=6)
processed = op.process({"image": image, "action": action})

Features

• 🎯 Unified Interface: Consistent API across all world model algorithms
• 🔧 Modular Components: Swappable encoders, decoders, and backbones
• 🚀 High Performance: Optimized for both training and inference
• 🌍 Multi-Environment: Support for DMC, Gym, Unity, and custom environments
• 📊 Rich Monitoring: Integrated logging with Weights & Biases and TensorBoard
• 🧠 Research Ready: Easy experimentation with different architectures

🧠 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
Diamond	Diffusion + RL for continuous control	EDM sampling, value functions

🤝 Community

• 📖 Documentation
• 🐛 Issue Tracker
• 💬 Discussions
• 📧 PyPI

---

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