
### Overview

Below we provide the general overview of our code base and list relevant aspects of the code.

[src/gflownet/trainer.py](src/gflownet/trainer.py) script that contains general classes for task, algorithm, and training pipeline.

[src/gflownet/tasks](src/gflownet/tasks) is a directory that contains scripts for the runnable trainer files for the graph, sequence, and hypergrid environments.

[src/gflownet/tasks/basic_graph_task.py](src/gflownet/tasks/basic_graph_task.py) contains the main elements of the graph generation task code, including state space generation, reward computations, dataset generation, task and trainer instantiations, p(x) and p(x; \theta) computation, dsitirbutional eveluations, and more.

[src/gflownet/envs/graph_building_env.py](src/gflownet/envs/graph_building_env.py) handle graph building environment which induces a DAG state space, compatible with GFlowNet. Supports forward and backward actions, with a parents function that list parents of forward actions.

[src/gflownet/data/sampling_iterator.py](src/gflownet/data/sampling_iterator.py) handles data sampling pipeline for training and evaluation, for both online and offline training. 

[src/gflownet/algo/trajectory_balance.py](src/gflownet/algo/trajectory_balance.py) contains trajectory balance class, the primary algorithm class used for training GFlowNets in this work. This class also contains the implementation for SubTB and DB objectives.  

[src/gflownet/models](src/gflownet/models) directory containing scripts that define models for the graph [src/gflownet/models/graph_transformer.py](src/gflownet/models/graph_transformer.py), sequence [src/gflownet/models/seq_transformer.py](src/gflownet/models/seq_transformer.py) tasks.

Further details are indluded in [implementation_notes.md](implementation_notes.md)

### Running Code

To run the code for the graph generation tasks, the first step is to generate and cach the state space of all possible graphs. To do this, run:

`python basic_graph_task.py --recompute-all ./data/basic_graphs 7`

## Single run

To train a single model, you can use [expts/task_single_run_gfn.py](expts/task_single_run_gfn.py) with various parameter settings:

`python expts/task_single_run_gfn.py`

## Example experiment 

To reproduce a set of experiments, for instance, distilling to P_F and F_sa for the graph generation task, use [expts/task_distilled_flows.py](expts/task_distilled_flows.py).


