diffusiongym 2.0__tar.gz

diffusiongym-2.0/.editorconfig

@@ -0,0 +1,16 @@
+# http://editorconfig.org/#file-format-details
+root = true
+
+[*]
+charset = utf-8
+end_of_line = lf
+indent_size = 4
+indent_style = space
+insert_final_newline = true
+trim_trailing_whitespace = true
+
+[*.md]
+trim_trailing_whitespace = false
+
+[Makefile]
+indent_style = tab

diffusiongym-2.0/.envrc

@@ -0,0 +1,2 @@
+watch_file pixi.lock
+eval "$(pixi shell-hook -e dev)"

diffusiongym-2.0/.github/workflows/docs.yml

@@ -0,0 +1,22 @@
+name: Deploy Docs
+
+on:
+  push:
+    branches: [main]
+
+permissions:
+  contents: read
+  pages: write
+  id-token: write
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: prefix-dev/setup-pixi@v0.8.1
+      - run: pixi run -e dev sphinx-build -b html docs docs/_build/html
+      - uses: actions/upload-pages-artifact@v3
+        with:
+          path: docs/_build/html
+      - uses: actions/deploy-pages@v4

diffusiongym-2.0/.github/workflows/test.yml

@@ -0,0 +1,32 @@
+name: Test
+
+on:
+  pull_request: {}
+  push:
+    branches: master
+
+jobs:
+  test:
+    strategy:
+      matrix:
+        python-version: ['3.13']
+        os: [ubuntu-latest]
+
+    name: Python ${{ matrix.os }} ${{ matrix.python-version }}
+    runs-on: ${{ matrix.os }}
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - uses: prefix-dev/setup-pixi@v0.8.8
+        with:
+          pixi-version: v0.48.0
+          cache: true
+          environments: dev
+
+      - run: pixi run fmt
+      - run: pixi run lint

diffusiongym-2.0/.gitignore

@@ -0,0 +1,188 @@
+# Created by https://www.toptal.com/developers/gitignore/api/python
+# Edit at https://www.toptal.com/developers/gitignore?templates=python
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+### Python Patch ###
+# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
+poetry.toml
+
+# ruff
+.ruff_cache/
+
+# LSP config files
+pyrightconfig.json
+
+# End of https://www.toptal.com/developers/gitignore/api/python
+
+# Pixi
+.pixi
+*.ipynb
+
+docs/api/generated/
+.DS_Store
+*.png
+jobs/
+slurm-*
+*.pt
+scripts/

diffusiongym-2.0/.pre-commit-config.yaml

@@ -0,0 +1,27 @@
+exclude: '.pixi/'
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v5.0.0  # this is optional, use `pre-commit autoupdate` to get the latest rev!
+  hooks:
+    - id: check-yaml
+    - id: check-toml
+    - id: end-of-file-fixer
+    - id: trailing-whitespace
+
+- repo: local
+  hooks:
+  - id: ruff
+    name: ruff-format
+    stages: [pre-commit, pre-push]
+    language: system
+    entry: pixi run fmt
+    types: [python]
+    pass_filenames: false
+
+  - id: ruff
+    name: ruff-check
+    stages: [pre-commit, pre-push]
+    language: system
+    entry: pixi run lint
+    types: [python]
+    pass_filenames: false

diffusiongym-2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Cristian Perez Jensen
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

diffusiongym-2.0/PKG-INFO

@@ -0,0 +1,98 @@
+Metadata-Version: 2.4
+Name: diffusiongym
+Version: 2.0
+Summary: Diffusion Gym
+Project-URL: Homepage, https://github.com/cristianpjensen/diffusiongym
+Project-URL: Issues, https://github.com/cristianpjensen/diffusiongym/issues
+Author: Cristian Perez Jensen
+License: MIT License
+        
+        Copyright (c) 2025 Cristian Perez Jensen
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Requires-Python: <3.12,>=3.8
+Requires-Dist: datasets==2.11.0
+Requires-Dist: dgl==2.4
+Requires-Dist: diffusers==0.33.1
+Requires-Dist: meeko==0.6.1
+Requires-Dist: numpy<2,>=1.24
+Requires-Dist: open-clip-torch<3,>=2.26
+Requires-Dist: pandas<3,>=2
+Requires-Dist: polars<2,>=1.34.0
+Requires-Dist: prody<3,>=2.4
+Requires-Dist: pyarrow==11.0.0
+Requires-Dist: pydantic<3,>=1.10
+Requires-Dist: pyyaml<7,>=6
+Requires-Dist: rdkit-stubs<0.9,>=0.6
+Requires-Dist: rdkit<2025,>=2023.9.4
+Requires-Dist: torch==2.3.1
+Requires-Dist: torchdata<0.9,>=0.7
+Requires-Dist: torchvision==0.18.1
+Requires-Dist: tqdm<5,>=4.66
+Requires-Dist: transformers>=4.36
+Requires-Dist: typing-extensions<5,>=4.9
+Description-Content-Type: text/markdown
+
+# Diffusion Gym
+
+<div align="center">
+  <img src="docs/_static/teaser.gif" width="100%" />
+</div>
+
+<p align="center">
+<a href="https://github.com/cristianpjensen/diffusiongym/blob/master/LICENSE"><img alt="License" src="https://img.shields.io/github/license/cristianpjensen/diffusiongym"></a>
+<a href="https://github.com/astral-sh/ruff"><img alt="Code style: ruff" src="https://img.shields.io/badge/code%20style-ruff-000000.svg"></a>
+</p>
+
+`diffusiongym` is a library for reward adaptation of any pre-trained flow model on any data modality.
+
+## Installation
+
+In order to install *diffusiongym*, execute the following command:
+```bash
+pip install diffusiongym
+```
+
+*diffusiongym* requires PyTorch 2.3.1, and there may be other hard dependencies. Please open an issue if
+installation fails through the above command.
+
+Molecule environments depend on [FlowMol](https://github.com/cristianpjensen/FlowMol),
+which currently needs to be installed manually:
+```bash
+pip install git+https://github.com/cristianpjensen/FlowMol.git@8f4c98cbe68111e4e63480b250d925b6d960d3bc
+```
+
+Some image rewards depend on the clip package, which needs to be installed manually as well:
+```bash
+pip install git+https://github.com/openai/CLIP.git
+```
+
+## High-level overview
+
+Diffusion and flow models are largely agnostic to their data modality. They only require that the underlying data type supports a small set of operations. Building on this idea, *diffusiongym* is designed to be fully modular. You only need to provide the following:
+ * Data type `YourDataType` that implements `DDProtocol`, which defines some functions necessary for interacting with it as a flow model.
+ * Base model `BaseModel[YourDataType]`, which defines the scheduler, how to sample $p_0$, how to compute the forward pass, and how to preprocess and postprocess data.
+ * Reward function `Reward[YourDataType]`.
+
+Once these are defined, you can sample from the flow model and apply reward adaptation methods, such as Value Matching.
+
+## Documentation
+
+Much more information can be found in [the documentation](https://cristianpjensen.github.io/diffusiongym/), including tutorials and API references.

diffusiongym-2.0/README.md

@@ -0,0 +1,46 @@
+# Diffusion Gym
+
+<div align="center">
+  <img src="docs/_static/teaser.gif" width="100%" />
+</div>
+
+<p align="center">
+<a href="https://github.com/cristianpjensen/diffusiongym/blob/master/LICENSE"><img alt="License" src="https://img.shields.io/github/license/cristianpjensen/diffusiongym"></a>
+<a href="https://github.com/astral-sh/ruff"><img alt="Code style: ruff" src="https://img.shields.io/badge/code%20style-ruff-000000.svg"></a>
+</p>
+
+`diffusiongym` is a library for reward adaptation of any pre-trained flow model on any data modality.
+
+## Installation
+
+In order to install *diffusiongym*, execute the following command:
+```bash
+pip install diffusiongym
+```
+
+*diffusiongym* requires PyTorch 2.3.1, and there may be other hard dependencies. Please open an issue if
+installation fails through the above command.
+
+Molecule environments depend on [FlowMol](https://github.com/cristianpjensen/FlowMol),
+which currently needs to be installed manually:
+```bash
+pip install git+https://github.com/cristianpjensen/FlowMol.git@8f4c98cbe68111e4e63480b250d925b6d960d3bc
+```
+
+Some image rewards depend on the clip package, which needs to be installed manually as well:
+```bash
+pip install git+https://github.com/openai/CLIP.git
+```
+
+## High-level overview
+
+Diffusion and flow models are largely agnostic to their data modality. They only require that the underlying data type supports a small set of operations. Building on this idea, *diffusiongym* is designed to be fully modular. You only need to provide the following:
+ * Data type `YourDataType` that implements `DDProtocol`, which defines some functions necessary for interacting with it as a flow model.
+ * Base model `BaseModel[YourDataType]`, which defines the scheduler, how to sample $p_0$, how to compute the forward pass, and how to preprocess and postprocess data.
+ * Reward function `Reward[YourDataType]`.
+
+Once these are defined, you can sample from the flow model and apply reward adaptation methods, such as Value Matching.
+
+## Documentation
+
+Much more information can be found in [the documentation](https://cristianpjensen.github.io/diffusiongym/), including tutorials and API references.

diffusiongym-2.0/diffusiongym/__init__.py

@@ -0,0 +1,59 @@
+"""Diffusion Gym package."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("diffusiongym")
+except PackageNotFoundError:
+    __version__ = "0.0.0"
+
+from diffusiongym.base_models import BaseModel
+from diffusiongym.environments import (
+    EndpointEnvironment,
+    Environment,
+    EpsilonEnvironment,
+    Sample,
+    ScoreEnvironment,
+    VelocityEnvironment,
+)
+from diffusiongym.make import construct_env, make
+from diffusiongym.registry import base_model_registry, reward_registry
+from diffusiongym.rewards import DummyReward, Reward
+from diffusiongym.schedulers import (
+    ConstantNoiseSchedule,
+    CosineScheduler,
+    DiffusionScheduler,
+    MemorylessNoiseSchedule,
+    NoiseSchedule,
+    OptimalTransportScheduler,
+    Scheduler,
+)
+from diffusiongym.types import D, DDMixin, DDTensor
+from diffusiongym.utils import train_base_model
+
+__all__ = [
+    "BaseModel",
+    "ConstantNoiseSchedule",
+    "CosineScheduler",
+    "D",
+    "DDMixin",
+    "DDTensor",
+    "DiffusionScheduler",
+    "DummyReward",
+    "EndpointEnvironment",
+    "Environment",
+    "EpsilonEnvironment",
+    "MemorylessNoiseSchedule",
+    "NoiseSchedule",
+    "OptimalTransportScheduler",
+    "Reward",
+    "Sample",
+    "Scheduler",
+    "ScoreEnvironment",
+    "VelocityEnvironment",
+    "base_model_registry",
+    "construct_env",
+    "make",
+    "reward_registry",
+    "train_base_model",
+]

diffusiongym-2.0/diffusiongym/base_models/__init__.py

@@ -0,0 +1,6 @@
+"""Base models for flow matching and diffusion."""
+
+from .base import BaseModel
+from .one_dim_gmm import OneDimensionalBaseModel
+
+__all__ = ["BaseModel", "OneDimensionalBaseModel"]

diffusiongym-2.0/diffusiongym/base_models/base.py

@@ -0,0 +1,163 @@
+"""Abstract base class for base models used in flow matching and diffusion models."""
+
+from abc import ABC, abstractmethod
+from typing import Any, Generic, Literal, Optional
+
+import torch
+from torch import nn
+
+from diffusiongym.schedulers import Scheduler
+from diffusiongym.types import D
+
+OutputType = Literal["epsilon", "endpoint", "velocity", "score"]
+
+
+class BaseModel(ABC, nn.Module, Generic[D]):
+    """Abstract base class for base models used in flow matching and diffusion."""
+
+    output_type: OutputType
+
+    def __init__(self, device: Optional[torch.device]):
+        super().__init__()
+
+        if device is None:
+            device = torch.device("cpu")
+
+        self.device = device
+
+    @property
+    @abstractmethod
+    def scheduler(self) -> Scheduler[D]:
+        """Base model-dependent scheduler used for sampling."""
+
+    @abstractmethod
+    def sample_p0(self, n: int, **kwargs: Any) -> tuple[D, dict[str, Any]]:
+        """Sample n data points from the base distribution p0.
+
+        Parameters
+        ----------
+        n : int
+            Number of samples to draw.
+        **kwargs : dict
+            Additional keyword arguments.
+
+        Returns
+        -------
+        samples : D
+            Samples from the base distribution p0.
+        kwargs : dict
+            Additional keyword arguments.
+        """
+
+    @abstractmethod
+    def forward(self, x: D, t: torch.Tensor, **kwargs: Any) -> D:
+        """Forward pass of the base model.
+
+        Parameters
+        ----------
+        x : D
+            Input data.
+        t : torch.Tensor, shape (n,)
+            Time steps, values in [0, 1].
+
+        Returns
+        -------
+        output : D
+            Output of the model.
+        """
+
+    def preprocess(self, x: D, **kwargs: Any) -> tuple[D, dict[str, Any]]:
+        """Preprocess data and keyword arguments for the base model.
+
+        Parameters
+        ----------
+        x : D
+            Input data to preprocess.
+        **kwargs : dict
+            Additional keyword arguments to preprocess.
+
+        Returns
+        -------
+        output : D
+            Preprocessed data.
+        kwargs : dict
+            Preprocessed keyword arguments.
+        """
+        return x, kwargs
+
+    def postprocess(self, x: D) -> D:
+        """Postprocess samples x_1 (e.g., decode with VAE).
+
+        Parameters
+        ----------
+        x : D
+            Input data to postprocess.
+
+        Returns
+        -------
+        output : D
+            Postprocessed output.
+        """
+        return x
+
+    def train_loss(
+        self,
+        x1: D,
+        xt: Optional[D] = None,
+        t: Optional[torch.Tensor] = None,
+        pred: Optional[D] = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        """Compute loss for a single batch training step.
+
+        Parameters
+        ----------
+        x1 : D
+            Target data points.
+        xt : Optional[D], default=None
+            Noisy data points at time t. If None, will be sampled.
+        t : Optional[torch.Tensor], shape (len(x1),), default=None
+            Time steps. If None, will be sampled.
+        pred : Optional[D], default=None
+            Model predictions. If None, will be computed by the model.
+        **kwargs : dict
+            Keyword arguments
+
+        Returns
+        -------
+        loss : torch.Tensor, shape (len(x1),)
+            Computed loss for the training step.
+        """
+        if t is None:
+            t = torch.rand(len(x1), device=x1.device)
+
+        assert t.shape == (len(x1),)
+
+        alpha = self.scheduler.alpha(x1, t)
+        beta = self.scheduler.beta(x1, t)
+
+        if xt is None:
+            x0 = x1.randn_like()
+            xt = alpha * x1 + beta * x0
+        else:
+            assert len(xt) == len(x1)
+            x0 = (xt - alpha * x1) / beta
+
+        if pred is None:
+            pred = self.forward(xt, t, **kwargs)
+
+        target = None
+        if self.output_type == "velocity":
+            alpha_dot = self.scheduler.alpha_dot(x1, t)
+            beta_dot = self.scheduler.beta_dot(x1, t)
+            target = alpha_dot * x1 + beta_dot * x0
+        elif self.output_type == "score":
+            target = -x0 / beta
+        elif self.output_type == "endpoint":
+            target = x1
+        elif self.output_type == "epsilon":
+            target = x0
+        else:
+            raise ValueError(f"Unknown output type: {self.output_type}")
+
+        return ((pred - target) ** 2).aggregate("mean")