Package: torchgfn
Documentation~Code~Papertorchgfn: a Python package for GFlowNetsPlease citethis paperif you are using the library for your researchInstalling the packageThe codebase requires python >= 3.10To install the latest stable version:pipinstalltorchgfnOptionally, to run scripts:pipinstalltorchgfn[scripts]To install the cutting edge version (from themainbranch):gitclonehttps://github.com/saleml/torchgfn.git
condacreate-ngfnpython=3.10
condaactivategfncdtorchgfn
pipinstall.About this repoThis repo serves the purpose of fast prototypingGFlowNet(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.Currently, the library is shipped with three environments: two discrete environments (Discrete Energy Based Model and Hyper Grid) and a continuous box environment. The library is designed to allow users to define their own environments. Seeherefor more details.Scripts and notebooksExample scripts and notebooks for the three environments are providedhere. For the hyper grid and the box environments, the provided scripts are supposed to reproduce published results.Standalone exampleThis example, which shows how to use the library for a simple discrete environment, requirestqdmpackage to run. Usepip install tqdmor install all extra requirements withpip install .[scripts]orpip install torchgfn[scripts].importtorchfromtqdmimporttqdmfromgfn.gflownetimportTBGFlowNet# We use a GFlowNet with the Trajectory Balance (TB) lossfromgfn.gymimportHyperGrid# We use the hyper grid environmentfromgfn.modulesimportDiscretePolicyEstimatorfromgfn.samplersimportSamplerfromgfn.utilsimportNeuralNet# NeuralNet is a simple multi-layer perceptron (MLP)if__name__=="__main__":# 1 - We define the environmentenv=HyperGrid(ndim=4,height=8,R0=0.01)# Grid of size 8x8x8x8# 2 - We define the needed modules (neural networks)# The environment has a preprocessor attribute, which is used to preprocess the state before feeding it to the policy estimatormodule_PF=NeuralNet(input_dim=env.preprocessor.output_dim,output_dim=env.n_actions)# Neural network for the forward policy, with as many outputs as there are actionsmodule_PB=NeuralNet(input_dim=env.preprocessor.output_dim,output_dim=env.n_actions-1,torso=module_PF.torso# We share all the parameters of P_F and P_B, except for the last layer)# 3 - We define the estimatorspf_estimator=DiscretePolicyEstimator(module_PF,env.n_actions,is_backward=False,preprocessor=env.preprocessor)pb_estimator=DiscretePolicyEstimator(module_PB,env.n_actions,is_backward=True,preprocessor=env.preprocessor)# 4 - We define the GFlowNetgfn=TBGFlowNet(init_logZ=0.,pf=pf_estimator,pb=pb_estimator)# We initialize logZ to 0# 5 - We define the sampler and the optimizersampler=Sampler(estimator=pf_estimator)# We use an on-policy sampler, based on the forward policy# Policy parameters have their own LR.non_logz_params=[vfork,vindict(gfn.named_parameters()).items()ifk!="logZ"]optimizer=torch.optim.Adam(non_logz_params,lr=1e-3)# Log Z gets dedicated learning rate (typically higher).logz_params=[dict(gfn.named_parameters())["logZ"]]optimizer.add_param_group({"params":logz_params,"lr":1e-1})# 6 - We train the GFlowNet for 1000 iterations, with 16 trajectories per iterationforiin(pbar:=tqdm(range(1000))):trajectories=sampler.sample_trajectories(env=env,n_trajectories=16)optimizer.zero_grad()loss=gfn.loss(env,trajectories)loss.backward()optimizer.step()ifi%25==0:pbar.set_postfix({"loss":loss.item()})ContributingBefore the first commit:pipinstall-e.[dev,scripts]pre-commitinstall
pre-commitrun--all-filesRunpre-commitafter staging, and before committing. Make sure all the tests pass (By runningpytest). Note that thepytesthook ofpre-commitonly runs the tests in thetesting/folder. To run all the tests, which take longer, runpytestmanually.
The codebase usesblackformatter.To make the docs locally:cddocs
makehtml
openbuild/html/index.htmlDetails about the codebaseDefining an environmentSeehereStatesStates are the primitive building blocks for GFlowNet objects such as transitions and trajectories, on which losses operate.An abstractStatesclass is provided. But for each environment, aStatessubclass is needed. AStatesobject
is a collection of multiple states (nodes of the DAG). A tensor representation of the states is required for batching. If a state is represented with a tensor of shape(*state_shape), a batch of states is represented with aStatesobject, with the attributetensorof shape(*batch_shape, *state_shape). Other
representations are possible (e.g. a state as a string, anumpyarray, a graph, etc...), but these representations cannot be batched, unless the user specifies a function that transforms these raw states to tensors.Thebatch_shapeattribute is required to keep track of the batch dimension. A trajectory can be represented by a States object withbatch_shape = (n_states,). Multiple trajectories can be represented by a States object withbatch_shape = (n_states, n_trajectories).Because multiple trajectories can have different lengths, batching requires appending a dummy tensor to trajectories that are shorter than the longest trajectory. The dummy state is the $s_f$ attribute of the environment (e.g.[-1, ..., -1], or[-inf, ..., -inf], etc...). Which is never processed, and is used to pad the batch of states only.For discrete environments, the action set is represented with the set ${0, \dots, n_{actions} - 1}$, where the $(n_{actions})$-th action always corresponds to the exit or terminate action, i.e. that results in a transition of the type $s \rightarrow s_f$, but not all actions are possible at all states. For discrete environments, eachStatesobject is endowed with two extra attributes:forward_masksandbackward_masks, representing which actions are allowed at each state and which actions could have led to each state, respectively. Such states are instances of theDiscreteStatesabstract subclass ofStates. Theforward_maskstensor is of shape(*batch_shape, n_{actions}), andbackward_masksis of shape(*batch_shape, n_{actions} - 1). Each subclass ofDiscreteStatesneeds to implement theupdate_masksfunction, that uses the environment's logic to define the two tensors.ActionsActions should be though of as internal actions of an agent building a compositional object. They correspond to transitions $s \rightarrow s'$. An abstractActionsclass is provided. It is automatically subclassed for discrete environments, but needs to be manually subclassed otherwise.Similar toStatesobjects, each action is a tensor of shape(*batch_shape, *action_shape). For discrete environments for instances,action_shape = (1,), representing an integer between $0$ and $n_{actions} - 1$.Additionally, each subclass needs to define two more class variable tensors:dummy_action: A tensor that is padded to sequences of actions in the shorter trajectories of a batch of trajectories. It is[-1]for discrete environments.exit_action: A tensor that corresponds to the termination action. It is[n_{actions} - 1]fo discrete environments.ContainersContainers are collections ofStates, along with other information, such as reward values, or densities $p(s' \mid s)$. Two containers are available:Transitions, representing a batch of transitions $s \rightarrow s'$.Trajectories, representing a batch of complete trajectories $\tau = s_0 \rightarrow s_1 \rightarrow \dots \rightarrow s_n \rightarrow s_f$.These containers can either be instantiated using aStatesobject, or can be initialized as empty containers that can be populated on the fly, allowing the usage of theReplayBufferclass.They inherit from the baseContainerclass, indicating some helpful methods.In most cases, one needs to sample complete trajectories. From a batch of trajectories, a batch of states and batch of transitions can be defined usingTrajectories.to_transitions()andTrajectories.to_states(), in order to train GFlowNets with losses that are edge-decomposable or state-decomposable.  These exclude meaningless transitions and dummy states that were added to the batch of trajectories to allow for efficient batching.ModulesTraining GFlowNets requires one or multiple estimators, calledGFNModules, which is an abstract subclass oftorch.nn.Module. In addition to the usualforwardfunction,GFNModules need to implement arequired_output_dimattribute, to ensure that the outputs have the required dimension for the task at hand; and some (but not all) need to implement ato_probability_distributionfunction.DiscretePolicyEstimatoris aGFNModulethat defines the policies $P_F(. \mid s)$ and $P_B(. \mid s)$ for discrete environments. Whenis_backward=False, the required output dimension isn = env.n_actions, and whenis_backward=True, it isn = env.n_actions - 1. Thesennumbers represent the logits of a Categorical distribution. The correspondingto_probability_distributionfunction transforms the logits by masking illegal actions (according to the forward or backward masks), then return a Categorical distribution. The masking is done by setting the corresponding logit to $-\infty$. The function also includes exploration parameters, in order to define a tempered version of $P_F$, or a mixture of $P_F$ with a uniform distribution.DiscretePolicyEstimator`` withis_backward=False`` can be used to represent log-edge-flow estimators $\log F(s \rightarrow s')$.ScalarModuleis a simple module with required output dimension 1. It is useful to define log-state flows $\log F(s)$.For non-discrete environments, the user needs to specify their own policies $P_F$ and $P_B$. The module, taking as input a batch of states (as aStates) object, should return the batched parameters of atorch.Distribution. The distribution depends on the environment. Theto_probability_distributionfunction handles the conversion of the parameter outputs to an actual batchedDistributionobject, that implements at least thesampleandlog_probfunctions. An example is providedhere, for a square environment in which the forward policy has support either on a quarter disk, or on an arc-circle, such that the angle, and the radius (for the quarter disk part) are scaled samples from a mixture of Beta distributions. The provided example shows an intricate scenario, and it is not expected that user defined environment need this much level of details.In allGFNModules, note that the input of theforwardfunction is aStatesobject. Meaning that they first need to be transformed to tensors. However,states.tensordoes not necessarily include the structure that a neural network can used to generalize. It is common in these scenarios to have a function that transforms these raw tensor states to ones where the structure is clearer, via aPreprocessorobject, that is part of the environment. More on thishere. The default preprocessor of an environment is the identity preprocessor. Theforwardpass thus first calls thepreprocessorattribute of the environment onStates, before performing any transformation. Thepreprocessoris thus an attribute of the module. If it is not explicitly defined, it is set to the identity preprocessor.For discrete environments, aTabularmodule is provided, where a lookup table is used instead of a neural network. Additionally, aUniformPBmodule is provided, implementing a uniform backward policy. These modules are providedhere.SamplersASamplerobject defines how actions are sampled (sample_actions()) at each state, and trajectories  (sample_trajectories()), which can sample a batch of trajectories starting from a given set of initial states or starting from $s_0$. It requires aGFNModulethat implements theto_probability_distributionfunction. For off-policy sampling, the parameters ofto_probability_distributioncan be directly passed when initializing theSampler.LossesGFlowNets can be trained with different losses, each of which requires a different parametrization, which we call in this library aGFlowNet. AGFlowNetis aGFNModulethat includes one or multipleGFNModules, at least one of which implements ato_probability_distributionfunction. They also need to implement alossfunction, that takes as input either states, transitions, or trajectories, depending on the loss.Currently, the implemented losses are:Flow MatchingDetailed Balance (and it's modified variant).Trajectory BalanceSub-Trajectory Balance. By default, each sub-trajectory is weighted geometrically (within the trajectory) depending on its length. This corresponds to the strategy definedhere. Other strategies exist and are implementedhere.Log Partition Variance loss. Introducedhere