🤖 AI Summary
Existing GFlowNet research lacks a unified, scalable PyTorch implementation framework, hindering the development of novel training objectives, integration with diverse environments, and reproducible benchmarking. To address this, we introduce the first modular, object-oriented open-source GFlowNet library built on PyTorch. Our method standardizes environment interfaces and sampler abstractions, enables plug-and-play loss functions—including trajectory balance (TB), detailed balance (DB), and unnormalized balance (UB)—and decouples state-space representation, action policies, and flow parameterizations to facilitate customization and composability. The framework successfully reproduces multiple state-of-the-art results across canonical benchmarks, substantially lowering the barrier for algorithm validation and extension. The codebase is publicly released and has been widely adopted by the research community.
📝 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 which facilitates the testing of new features such as training losses that can be easily compared to standard benchmark implementations, or on a set of common environments. torchgfn is a PyTorch library that aims to address this need. It provides users with a simple API for environments and useful abstractions for samplers and losses. Multiple examples are provided, replicating and unifying published results. The code is available in https://github.com/saleml/torchgfn.
