Diffusion Gym

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:

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, which currently needs to be installed manually:

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:

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, including tutorials and API references.

Citation

If this library is useful to you, consider citing the following work:

@inproceedings{jensen2026value,
  title={Value Matching: Scalable and Gradient-Free Reward-Guided Flow Adaptation},
  author={Cristian Perez Jensen and Luca Schaufelberger and Riccardo De Santi and Kjell Jorner and Andreas Krause},
  booktitle={The Fourteenth International Conference on Learning Representations},
  year={2026},
}