To add more details:

• The policy classes are currently under ./gflownet/policy/.
• Currently, the base class in [base.py] implements an MLP.
• There are two more classes implementing tree architectures that I am not sure if they are well maintained (Michał might know)
• There is also a class for the State Flow model in state_flow.py used by the DB loss.
• An inconvenient aspect of the current implementation is how the combination of loss and architecture is handled in the configuration files. Currently, there needs to be a configuration file for every architecture/loss combination.

The new implementation should satisfy, in my opinion, the following desiderata:

• It should be easy to create new architectures for the policy models, say a CNN, GNN, Transformer, etc.
• The use of policy should not be entangled with the loss.
• There could be generic architectures (MLP, CNN, etc.) whose specific main properties can be specified via the config. For example, the number of layers and units of an MLP.
• Additionally, it should be easy to add a new class with a specific architecture that would be available to use with any loss.