Skip to content

GFNOrg/torchgfn

Repository files navigation

Python Documentation Status

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 main 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

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 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 captures 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.
  • 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