Metadata-Version: 2.4
Name: torchwm
Version: 0.4.2
Summary: A Modular Pytorch Based library for training world models
Author: 
License-Expression: MIT
Keywords: world-models,pytorch,mbrl
Classifier: Programming Language :: Python :: 3
Classifier: Operating System :: OS Independent
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=1.13.0
Requires-Dist: torchvision>=0.14.0
Requires-Dist: torchaudio>=2.10.0
Requires-Dist: einops>=0.8.2
Requires-Dist: pyyaml>=6.0.3
Requires-Dist: tqdm>=4.67.1
Requires-Dist: opencv-python>=4.12.0.88
Requires-Dist: requests>=2.32.0
Requires-Dist: click>=8.0.0
Requires-Dist: gym>=0.26.2
Requires-Dist: gymnasium>=1.2.2
Requires-Dist: moviepy>=2.2.1
Requires-Dist: h5py>=3.11.0
Requires-Dist: plotly>=6.5.0
Requires-Dist: ale-py>=0.11.2
Requires-Dist: selenium>=4.41.0
Requires-Dist: scikit-learn>=1.8.0
Requires-Dist: umap-learn>=0.5.11
Requires-Dist: huggingface-hub>=1.14.0
Requires-Dist: playwright>=1.60.0
Requires-Dist: jax>=0.10.1
Provides-Extra: gym
Requires-Dist: huggingface-hub>=0.23.0; extra == "gym"
Requires-Dist: pygame>=2.6.1; extra == "gym"
Requires-Dist: autorom>=0.6.1; extra == "gym"
Provides-Extra: ml-agents
Requires-Dist: mlagents-envs>=0.28.0; extra == "ml-agents"
Provides-Extra: mujoco
Requires-Dist: mujoco>=3.2.0; extra == "mujoco"
Provides-Extra: robotics
Requires-Dist: gymnasium-robotics>=1.3.0; extra == "robotics"
Provides-Extra: brax
Requires-Dist: brax>=0.13.0; extra == "brax"
Provides-Extra: viz
Requires-Dist: fastapi[standard]>=0.116.0; extra == "viz"
Requires-Dist: uvicorn>=0.35.0; extra == "viz"
Requires-Dist: myst-parser<5.0.0,>=2.0.0; extra == "viz"
Requires-Dist: nbsphinx>=0.9.8; extra == "viz"
Provides-Extra: ml
Requires-Dist: tensorboard>=2.20.0; extra == "ml"
Requires-Dist: tensorboardx>=2.6.4; extra == "ml"
Requires-Dist: wandb>=0.16.0; extra == "ml"
Provides-Extra: dev
Requires-Dist: pytest>=9.0.2; extra == "dev"
Requires-Dist: pytest-cov>=7.1.0; extra == "dev"
Requires-Dist: pre-commit>=4.5.0; extra == "dev"
Requires-Dist: mypy>=1.19.1; extra == "dev"
Provides-Extra: all
Requires-Dist: torchwm[brax,dev,docs,gym,ml,ml-agents,mujoco,robotics,viz]; extra == "all"
Provides-Extra: docs
Requires-Dist: sphinx<9.0.0,>=8.0.0; extra == "docs"
Requires-Dist: sphinx-autodoc-typehints<1.26.0; extra == "docs"
Requires-Dist: sphinx-copybutton>=0.5.2; extra == "docs"
Requires-Dist: sphinx-rtd-theme>=3.1.0; extra == "docs"
Requires-Dist: sphinxcontrib-bibtex>=2.6.5; extra == "docs"
Requires-Dist: sphinxcontrib-mermaid>=0.9.0; extra == "docs"
Requires-Dist: sphinx-thebe>=0.3.0; extra == "docs"
Requires-Dist: sphinxext-opengraph>=0.13.0; extra == "docs"
Requires-Dist: myst-parser<5.0.0,>=2.0.0; extra == "docs"
Requires-Dist: nbsphinx>=0.9.8; extra == "docs"
Dynamic: license-file
Dynamic: requires-python

# TorchWM

<div align="center">
  <img src="docs/images/torchwm-logo.svg" alt="TorchWM Logo" height="80">
  <p>
    <a href="https://pypi.org/project/torchwm/"><img alt="PyPI version" src="https://badge.fury.io/py/torchwm.svg"></a>
    <a href="https://pypi.org/project/torchwm/"><img alt="PyPI downloads" src="https://img.shields.io/pypi/dm/torchwm.svg"></a>
    <a href="https://opensource.org/licenses/MIT"><img alt="License: MIT" src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>
    <a href="https://paramthakkar123.github.io/torchwm/"><img alt="Documentation" src="https://img.shields.io/badge/docs-link-blue.svg"></a>
  </p>
  <p><strong>Modular PyTorch Library for World Models</strong></p>
</div>

---

## ⚡ Quick Start

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

### Installation

```bash
# 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:

```python
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:

```python
from torchwm import DreamerAgent, DreamerConfig

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

### Creating Environments and Operators

```python
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](https://paramthakkar123.github.io/torchwm/)

### Get Started
- [Installation](https://paramthakkar123.github.io/torchwm/getting_started.html)
- [Training Guide](https://paramthakkar123.github.io/torchwm/training_guide.html)
- [Inference Guide](https://paramthakkar123.github.io/torchwm/inference_guide.html)

### User Guides
- [Operators Guide](https://paramthakkar123.github.io/torchwm/operators_guide.html)
- [Environments Guide](https://paramthakkar123.github.io/torchwm/environments_guide.html)
- [Package Overview](https://paramthakkar123.github.io/torchwm/package_overview.html)

### Algorithms
- [Dreamer](https://paramthakkar123.github.io/torchwm/dreamer.html)
- [JEPA](https://paramthakkar123.github.io/torchwm/jepa.html)
- [IRIS](https://paramthakkar123.github.io/torchwm/iris.html)
- [DiT](https://paramthakkar123.github.io/torchwm/dit.html)

## 🤝 Community

- 🐛 [Issue Tracker](https://github.com/paramthakkar123/torchwm/issues)
- 💬 [Discussions](https://github.com/paramthakkar123/torchwm/discussions)
- 📧 [PyPI](https://pypi.org/project/torchwm/)

---

```{important}
TorchWM is under active development. APIs may change between versions.
```

```{toctree}
:hidden:
:maxdepth: 1
:caption: Get Started

getting_started
installation
```

```{toctree}
:hidden:
:maxdepth: 1
:caption: User Guides

operators_guide
training_guide
inference_guide
environments_guide
package_overview
```

```{toctree}
:hidden:
:maxdepth: 1
:caption: Algorithms

dreamer
jepa
iris
dit
```

```{toctree}
:hidden:
:maxdepth: 1
:caption: Reference

api_reference
configs_reference
```

```{toctree}
:hidden:
:maxdepth: 1
:caption: Development

contributing
benchmarks
```

```{toctree}
:hidden:
:maxdepth: 1
:caption: User Guides

operators_guide
training_guide
inference_guide
environments_guide
package_overview
```

```{toctree}
:hidden:
:maxdepth: 1
:caption: Algorithms

dreamer
jepa
iris
dit
```

```{toctree}
:hidden:
:maxdepth: 1
:caption: Reference

api_reference
configs_reference
```

```{toctree}
:hidden:
:maxdepth: 1
:caption: Development

contributing
benchmarks
```

---

## ⚡ Quick Start

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

### Installation

```bash
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:

```python
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:

```python
from torchwm import DreamerAgent, DreamerConfig

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

### Creating Environments and Operators

```python
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](https://paramthakkar123.github.io/torchwm/)
- 🐛 [Issue Tracker](https://github.com/paramthakkar123/torchwm/issues)
- 💬 [Discussions](https://github.com/paramthakkar123/torchwm/discussions)
- 📧 [PyPI](https://pypi.org/project/torchwm/)

---

```{important}
TorchWM is under active development. APIs may change between versions.
```
