torchgfn: A PyTorch GFlowNet library

TL;DR

Generative Flow Networks (GFNs) define probabilistic samplers over structured spaces using an unnormalized reward $R$ to generate objects along a DAG-shaped MDP. The paper presents torchgfn, a PyTorch library that enforces a strict separation of concerns among Environments, States, Actions, Containers, Estimators, Samplers, and GFlowNets, enabling modular, rapid prototyping across discrete, continuous, and graph GFNs. The $v2$ release adds graph-generation support, conditional GFNs, custom samplers, recurrent policies via PolicyMixin, and diffusion sampling, accompanied by extensive tutorials and documentation. The work emphasizes reproducibility and community-driven development, aiming to establish torchgfn as a standard platform for foundational GFN research.

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/).

Figures (2)

• Figure 1: A schematic of the modular torchgfn repository structure. States and Actions are top-level abstractions used to interface between the stateless Envs. Containers are generic objects used by the remainder of the library and contain States-Actions sequences drawn from Samplers, which interface (via a PolicyMixin) directly with the parametrization of an instantiated GFlowNet subclass (e.g., TrajectoryBalanceGFlowNet) which implements a specific training objective. Parametrization is handled by one or more Estimators (wrappers for pytorch's nn.Module) which represent the function approximators (e.g., policies $\pi$, state flows), and are trained using the GFlowNet's loss. A set of Envs is made available through the Gym module.
• Figure 2: torchgfn repository dependency graph as of v2.3.0. States and Actions are top-level abstractions used to interface between the stateless Env and a GFlowNet. Containers are generic objects used by the remainder of the library and contain States-Actions sequences drawn from Samplers via an Estimator (wrappers for pytorch's nn.Module) which represent the function approximators (e.g., policies, state flows), which are contained with a GFlowNet subclass (e.g., TrajectoryBalanceGFlowNet) that implements a specific training objective. Containers are utilized by both Samplers and GFlowNets to carry the relevant Transition or Trajectory-level information. A set of Environments is made available through the Gym module.