torchwm 0.4.0

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

from world_models.models import DreamerAgent
from world_models.configs import DreamerConfig

cfg = DreamerConfig()
cfg.env = "walker-walk"
cfg.total_steps = 1_000_000

agent = DreamerAgent(cfg)
agent.train()

Using Inference Operators

from world_models.inference.operators import DreamerOperator

op = DreamerOperator(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

from world_models.models import DreamerAgent
from world_models.configs import DreamerConfig

cfg = DreamerConfig()
cfg.env = "walker-walk"
cfg.total_steps = 1_000_000

agent = DreamerAgent(cfg)
agent.train()

Using Inference Operators

from world_models.inference.operators import DreamerOperator

op = DreamerOperator(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.