TorchWM Documentation#
TorchWM is a modular PyTorch library for world models, latent-dynamics planning, and representation learning. Train Dreamer, JEPA, IRIS, DiT, Genie, and DIAMOND agents with a unified API.
import torchwm
agent = torchwm.create_model("dreamer", action_size=6)
agent.train(env_name="walker-walk", total_steps=100_000)
Get Started
User Guides
- Public API Quick Reference
- Using Operators for Inference
- Training Guide
- Inference Guide
- Evaluation Guide
- Memory & Replay Buffers
- Environments Guide
- Environment Backends
- NuPlan Dataset
- TorchWM CLI
- Package Overview
- Controllers and Policies
- World Models Study Guide
- Tutorial: Plug TorchWM world models into RL libraries
- Modular RSSM
- Vision Components
- Datasets
- Loss Functions
- Plugin Registry
- World Models Deep Dive (Ha & Schmidhuber, 2018)
Algorithms
Development