TD-GFN

Official implementation of "Beyond the Proxy: Trajectory-Distilled Guidance for Offline GFlowNet Training" (ICML 2026), by Ruishuo Chen, Xun Wang, Rui Hu, Zhuoran Li, and Longbo Huang (IIIS, Tsinghua University).

TL;DR

TD-GFN is a proxy-free framework for training GFlowNets from an offline trajectory dataset. We use inverse RL to distill dense, transition-level edge rewards from the data, then apply them indirectly — through DAG pruning and prioritized backward sampling — while keeping gradient updates anchored to ground-truth terminal rewards. The result: faster convergence and higher sample quality than both proxy-based and prior proxy-free baselines, with no risk of reward-model error propagation.

Method

1. Edge rewards via IRL. Using the GFlowNet ↔ entropy-regularized RL equivalence (Tiapkin et al., 2024), run max-causal-entropy IRL on a rebalanced offline dataset to recover an edge-level reward R_E : E → ℝ.
2. DAG pruning. Drop low-utility transitions based on R_E, yielding a compact subgraph.
3. Prioritized backward sampling + clean gradients. Sample trajectories backward with probabilities shaped by both terminal and edge rewards; the policy loss (FM / TB / SubTB / DB) uses only ground-truth terminal rewards — R_E never enters the gradient, which keeps the policy insulated from IRL estimation error.

Results

Molecule design (sEH binding)

On the 1,500-trajectory offline dataset, TD-GFN matches or exceeds proxy-based baselines' top-10 reward while needing substantially fewer training trajectories.

Hypergrid

Across dataset qualities, TD-GFN drives the empirical L1 error down faster and discovers more modes than COFlowNet, Dataset-GFN, CQL, IQL, BC, and rebalanced GAIL.

Biosequence (AMP)

On AMP sequence design, TD-GFN attains the highest top-100 reward and matches or exceeds the best top-100 diversity across all baselines (CQL, IQL, BC, Dataset-GFN, Proxy-GFN, COFlowNet, rebalanced GAIL).

See the paper for full ablations and additional metrics.

Repository

TD-GFN/
├── mols/    # Molecule design (sEH binding) — primary task
├── grid/    # Hypergrid (8², 20⁴, 256²)
├── bioseq/  # AMP sequence design
├── setup/   # One-shot conda installers
└── assets/  # Paper figures used here

The TD-GFN algorithm lives in mols/offline_mols.py (--obj=dac-fm / --obj=dac-qm) and mols/DAC.py. The dac-* naming is a development-period artifact — it's the TD-GFN method from the paper.

Installation

Each task gets its own conda env.

bash setup/install_mols.sh    # env 'mols'         (~5 min)
bash setup/install_grid.sh    # env 'gflownet-rl'  (~3 min)
bash setup/install_bioseq.sh  # env 'bioseq'      (~15 min; pulls LFS files)

Pass an alternate env name as the first arg if you want to override the default.

Prerequisites: conda, a CUDA-capable GPU (CPU is fine for grid), and ~5 GB free for the bioseq oracle (ProtT5 / AlBert; auto-downloaded from Hugging Face on first AMP run).

Data: Small offline datasets and the pretrained mols reward proxy ship with the repo. The empty bioseq/{clamp-gen-data,design-bench,Rostlab}/ directories are placeholders — install_bioseq.sh installs everything as proper Python packages, and Rostlab/ is just the Hugging Face runtime cache.

Reproducing experiments

Molecule design (primary)

conda activate mols && cd mols
WANDB_MODE=disabled bash run_main_td_gfn_fm.sh

Experiment	Script
TD-GFN (FM objective)	mols/run_main_td_gfn_fm.sh
TD-GFN with QM objective	mols/run_td_gfn_qm.sh
Ablation: FM ± DAG mask	mols/run_ablation_fm_mask.sh
Ablation: online-origin	mols/run_ablation_online_origin.sh
CQL + IQL (bundled w/ main)	mols/run_main_td_gfn_fm.sh
GAIL + BC imitation	mols/run_imitation_gail_bc.sh
BraVE	mols/run_brave.sh, mols/run_brave_full.sh
Online proxy-guided	mols/run_online.sh
QM-objective GFN	mols/run_qm.sh

Hypergrid

conda activate gflownet-rl && cd grid

# TD-GFN variants
python dac_gfn.py             --dataType=expert.pt
python dac_backward.py        --dataType=expert.pt
python dac_backward_target.py --dataType=expert_backward_target.pt
python dac_backward_uniform.py --dataType=expert_backward_uniform.pt

# Baselines: conservative.py (COFlowNet), cql.py, BC.py, online.py

Datasets: {expert,median,mixed,bad,random}[_8^2|_20^4][_backward_{target,uniform}].pt.

Biosequence (AMP)

conda activate bioseq && cd bioseq

bash dac.sh           # TD-GFN
bash repro.sh         # FM (cfn) + TB (gfn) baselines
bash online.sh        # online proxy-guided
bash bc_imitation.sh  # BC + DAC-imitation
bash offline_rl.sh    # CQL + IQL

Notes

• Wandb logs by default. Pass WANDB_MODE=disabled or run wandb disabled once to skip.
• The mols run scripts end with source activate mols; edit if your env has a different name.
• For smaller / debug runs, use --data_size=1000 (or 10000) together with the matching mols/mask_data/forward_mask-1000.pkl / forward_mask-10000.pkl shipped in the repo.

Citation

@inproceedings{chen2026tdgfn,
  title     = {Beyond the Proxy: Trajectory-Distilled Guidance for Offline GFlowNet Training},
  author    = {Chen, Ruishuo and Wang, Xun and Hu, Rui and Li, Zhuoran and Huang, Longbo},
  booktitle = {Proceedings of the 43rd International Conference on Machine Learning (ICML)},
  year      = {2026}
}

Acknowledgements

Builds on COFlowNet (repo scaffolding) and Jain et al. 2022 (bioseq library, under bioseq/LICENSE.md); AMP oracle uses the MJ10 forks of design-bench and clamp-gen-data.

License

MIT. Upstream files retain their original licenses (bioseq/LICENSE.md).