Documentation ~ Code ~ Paper

torchgfn

A Python Package for GFlowNets. Please cite this paper if you are using the library for your research!

Installing the package

The codebase requires python >= 3.10. To install the latest stable version with the core dependencies:

pip install torchgfn

torchgfn supports installation with multiple sets of dependencies, under the following tags:

• dev: dependencies required for development of the core library.
• scripts: dependencies needed to run examples in tutorials/examples/.
• all: everything.

and can be called by running

pip install torchgfn[scripts]

Or to install the latest release (from the master branch) with all dependencies in a Conda environment:

git clone https://github.com/GFNOrg/torchgfn.git
conda create -n gfn python=3.10
conda activate gfn
cd torchgfn
pip install -e ".[all]"

About this repo

This repo serves the purpose of fast prototyping GFlowNet (GFN) related algorithms. It decouples the environment definition, the sampling process, and the parametrization of the function approximators used to calculate the GFN loss. It aims to accompany researchers and engineers in learning about GFlowNets, and in developing new algorithms.

The library is shipped with many environments under the gym, including discrete environments (e.g., Discrete Energy Based Model, Hyper Grid, Graph Generation), and continuous environments (e.g., Box). The library is designed to allow users to define their own environments relatively easily. See here for more details.

Getting Started with Example Scripts & Notebooks

• Simple example: a concise description of how to train a gflownet using the library.
• Tutorials README for an overview of the included example scripts and notebooks drawn from the included gym environments. Where indicated, these scripts are intended to reproduce published results.

Contributing

Please see the Contributing Guidelines.

Components of the Library

• States, Actions, & Containers: The two core elements of torchgfn are the concepts of states, which are emitted by stateless environments, and actions, which transform states through the environment logic. These are encapsulated with metadata in containers.
• Modules, Estimators, & Samplers: The components of a GFlowNet policy (those which select actions given a state), are encapsulated in Estimators. If the policy is a neural network, that logic is captured in a Module. This estimator can then be used with a Sampler to produce states (in the form of trajectories or transitions).
• Losses: Each type of GFlowNet has a specific parameterization related to a specific loss. Covers TB, DB, SubTB, FM, ZVar, and RTB with guidance on when to use each.
• Off-Policy Training & Replay Buffers: How to train with replay buffers, expert data warm-starting, and mixed on/off-policy batches.
• Exploration Strategies: Epsilon-greedy, temperature scaling, noisy layers, local search, and variance annealing for mode discovery.
• Conditional GFlowNets: Learning families of distributions parameterized by external conditions.
• Recurrent & Non-Autoregressive Policies: LSTM/GRU policies, tree DAG simplifications, and order-invariant generation.
• Diffusion GFlowNets: Continuous diffusion sampling with Brownian motion estimators and RTB fine-tuning.
• Defining Environments: For most applications of torchgfn, the main challenge will be to define a stateless environment which will produce a valid sampler.
• Advanced Usage: Extending torchgfn with Custom GFlowNets: While torchgfn aims to support major usages of GFlowNets, we hope this will also serve as a platform for the community to extend the possible use cases of the technology. This guide details how one can extend the library for new research and applications.
• Full API Reference: A breakdown of all major components of both the core GFN library, the associated gym, and the included tutorials.