Skip to content

Baselines

The baselines/ directory collects reproducible training scripts for canonical GFlowNet objectives and a few closely related algorithms. Each script follows the CleanRL-style single-file pattern, relies on Hydra for configuration management, and logs to the shared tmp/ directory unless overridden.

Implementation matrix

The table below summarizes which method-environment combinations currently ship with the repository. Cells marked with ✅ point to a ready-to-run script in baselines/; 🚧 means experimental and unverified script that may not work; ❌ indicates that the pairing has not been implemented yet.

Method / Environment Hypergrid BitSeq TFBind-8 QM9 Small AMP GFP
Detailed Balance (DB) [1] ✅ ✅ ✅ ✅ ✅ 🚧
Trajectory Balance (TB) [2] ✅ ✅ ✅ ✅ ✅ ❌
Sub-Trajectory Balance (SubTB) [3] ✅ ✅ ✅ ❌ ❌ ❌

Special-purpose scripts

Some environments require bespoke objectives or training tweaks. These live alongside the standard baselines:

  • baselines/soft_dqn_hypergrid.py – Online SoftDQN baseline for Hypergrid [4].
  • baselines/mdb_dag.py – Modified Detailed Balance for Bayesian network structure learning [5].
  • baselines/fldb_phylo.py – Forward-Looking Detailed Balance for phylogenetic tree generation [6].
  • baselines/tb_ising.py – Energy-based modeling of the Ising system using a TB objective [7].

How to run a baseline

  1. Install the package with baseline extras: pip install -e '.[baselines]'.
  2. Launch the desired script, optionally overriding Hydra configs. For example:
python baselines/db_hypergrid.py num_train_steps=1_000 logging.tqdm_print_rate=100

Each script documents the relevant configuration group in its module docstring. Point logging.log_dir or logging.checkpoint_dir to persistent storage when running long jobs.

References

  1. Bengio, Y. et al. (2023). GFlowNet Foundations.
    Journal of Machine Learning Research (JMLR).
  2. Malkin, N. et al. (2022). Trajectory Balance: Improved Credit Assignment in GFlowNets.
    Advances in Neural Information Processing Systems (NeurIPS).
  3. Madan, K. et al. (2023). Learning GFlowNets from Partial Episodes for Improved Convergence and Stability.
    International Conference on Machine Learning (ICML).
  4. Tiapkin, D. et al. (2024). Generative Flow Networks as Entropy-Regularized RL.
    International Conference on Artificial Intelligence and Statistics (AISTATS).
  5. Deleu, T. et al. (2022). Bayesian Structure Learning with Generative Flow Networks.
    Conference on Uncertainty in Artificial Intelligence (UAI).
  6. Pan, L. et al. (2023). Better Training of GFlowNets with Local Credit and Incomplete Trajectories.
    International Conference on Machine Learning (ICML).
  7. Zhang, D. et al. (2022). Generative Flow Networks for Discrete Probabilistic Modeling.
    International Conference on Machine Learning (ICML).