torchgfn: A PyTorch GFlowNet library

TL;DR

torchgfn is a modular PyTorch library designed to facilitate rapid prototyping and testing of generative flow networks across diverse environments and training methods, supporting research and benchmarking.

Contribution

It introduces a flexible, decoupled architecture for GFlowNets, enabling easy experimentation with different components and quick replication of published results.

Findings

Provides a unified API for GFlowNet components

Supports rapid prototyping and testing of new features

Includes examples replicating published results

Abstract

The growing popularity of generative flow networks (GFlowNets or GFNs) from a range of researchers with diverse backgrounds and areas of expertise necessitates a library that facilitates the testing of new features (e.g., training losses and training policies) against standard benchmark implementations, or on a set of common environments. We present torchgfn, a PyTorch library that aims to address this need. Its core contribution is a modular and decoupled architecture which treats environments, neural network modules, and training objectives as interchangeable components. This provides users with a simple yet powerful API to facilitate rapid prototyping and novel research. Multiple examples are provided, replicating and unifying published results. The library is available on GitHub (https://github.com/GFNOrg/torchgfn) and on pypi (https://pypi.org/project/torchgfn/).