flowgym 1.0

Flow Gym

Flow Gym

flowgym is a library for reward adaptation of any pre-trained flow model on any data modality.

Installation

In order to install flowgym, execute the following command:

pip install flowgym

If you want access to pre-trained image or molecular generation models, specify them as options:

pip install flowgym[images]
pip install flowgym[molecules]

If you want to run value matching on your environment, you will need to install it:

pip install value_matching

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, flowgym is designed to be fully modular. You only need to provide the following:

• Data type that implements DataProtocol, which defines basic arithmetic operations, factory methods, and gradient methods.
• 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.