torchwm 0.3.2

A Modular Pytorch Based library for training world models

TorchWM

Modular PyTorch Library for World Models

---

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})
# Returns standardized tensors for inference

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.