Metadata-Version: 2.4
Name: torchwm
Version: 0.3.1
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.13
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: ale-py>=0.11.2
Requires-Dist: einops>=0.8.2
Requires-Dist: fastapi[standard]>=0.116.0
Requires-Dist: gym>=0.26.2
Requires-Dist: gymnasium>=1.2.2
Requires-Dist: mlagents-envs>=0.28.0
Requires-Dist: moviepy>=2.2.1
Requires-Dist: myst-parser>=5.0.0
Requires-Dist: nbsphinx>=0.9.8
Requires-Dist: opencv-python>=4.12.0.88
Requires-Dist: plotly>=6.5.0
Requires-Dist: pre-commit>=4.5.0
Requires-Dist: pygame>=2.6.1
Requires-Dist: pytest>=9.0.2
Requires-Dist: pyyaml>=6.0.3
Requires-Dist: selenium>=4.41.0
Requires-Dist: sphinx>=9.1.0
Requires-Dist: sphinx-autodoc-typehints>=3.6.2
Requires-Dist: sphinx-copybutton>=0.5.2
Requires-Dist: sphinx-rtd-theme>=3.1.0
Requires-Dist: sphinxcontrib-bibtex>=2.6.5
Requires-Dist: sphinxext-opengraph>=0.13.0
Requires-Dist: tensorboard>=2.20.0
Requires-Dist: tensorboardx>=2.6.4
Requires-Dist: tqdm>=4.67.1
Requires-Dist: uvicorn>=0.35.0
Requires-Dist: wheel>=0.46.3
Requires-Dist: wandb>=0.16.0
Requires-Dist: scikit-learn>=1.8.0
Requires-Dist: umap-learn>=0.5.11
Requires-Dist: mypy>=1.19.1
Requires-Dist: pybullet>=3.2.7
Dynamic: license-file
Dynamic: requires-python

# TorchWM

TorchWM is a modular PyTorch library for world models and latent dynamics learning.
It includes practical implementations of Dreamer-style agents, PlaNet/RSSM utilities,
JEPA-style representation learning, and diffusion/transformer building blocks.

## Highlights

- Modular components for encoders, decoders, RSSMs, reward/value heads, and policies
- **Modular RSSM** with swappable encoder/decoder/backbone for research experiments
- Multiple environment backends: DMC, Gym/Gymnasium, Atari, MuJoCo, Unity ML-Agents
- Replay/memory utilities for both Dreamer and PlaNet-style training loops
- ViT + masking utilities for JEPA workflows
- Diffusion utilities (DDPM schedule + DiT model)

## Installation

Install from PyPI:

```bash
pip install torchwm
```

Install from source:

```bash
git clone https://github.com/ParamThakkar123/torchwm.git
cd torchwm
pip install -e .
```

Development extras:

```bash
pip install -e ".[dev]"
```

## Quick Start

### Train Dreamer on Gym

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

cfg = DreamerConfig()
cfg.env_backend = "gym"  # dmc | gym | unity_mlagents
cfg.env = "Pendulum-v1"
cfg.total_steps = 10_000

agent = DreamerAgent(cfg)
agent.train()
```

### Train Dreamer on Unity ML-Agents

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

cfg = DreamerConfig()
cfg.env_backend = "unity_mlagents"
cfg.unity_file_name = r"E:\UnityBuilds\MyEnv.exe"
cfg.unity_behavior_name = "MyBehavior"
cfg.unity_no_graphics = True
cfg.unity_time_scale = 20.0

agent = DreamerAgent(cfg)
agent.train()
```

### Train JEPA

```python
from world_models.models import JEPAAgent
from world_models.configs import JEPAConfig

cfg = JEPAConfig()
cfg.dataset = "imagefolder"
cfg.root_path = "./data"
cfg.image_folder = "train"
cfg.epochs = 10

agent = JEPAAgent(cfg)
agent.train()
```

## Documentation

- Sphinx source: `docs/source`
- Getting started guide: `docs/source/getting_started.md`
- Package overview: `docs/source/package_overview.md`
- API reference (autodoc): `docs/source/api_reference.rst`

Build HTML docs locally:

```bash
sphinx-build -b html docs/source docs/build/html
```

## Package Layout

- `world_models/models`: Agents and model architectures (`Dreamer`, `JEPAAgent`, `Planet`, ViT, diffusion)
- `world_models/models/modular_rssm`: Modular RSSM with swappable encoder/decoder/backbone
- `world_models/configs`: Config classes (`DreamerConfig`, `JEPAConfig`, `DiTConfig`)
- `world_models/envs`: Environment adapters and wrappers
- `world_models/training`: Script-style training entrypoints
- `world_models/datasets`: CIFAR10/ImageNet/ImageFolder data loaders
- `world_models/memory`: Replay and episodic memory implementations
- `world_models/utils`: Training/logging/distributed helper utilities

## Contributing

Contributions are welcome. See [CONTRIBUTING.md](CONTRIBUTING.md).

## License

MIT. See [LICENSE](LICENSE).

## Citation

```bibtex
@misc{Thakkar_GitHub_-_ParamThakkar123_torchwm,
author = {Thakkar, Param},
title = {{GitHub - ParamThakkar123/torchwm: A modular PyTorch library designed for learning, training, and deploying world models across various environments.}},
year = {2025},
url = {https://github.com/ParamThakkar123/torchwm}
}
```

Package: [torchwm on PyPI](https://pypi.org/project/torchwm/)
