DynaGFN: Model-Based Planning for Generative Flow Networks

This repository contains the implementation of "DynaGFN: Model-Based Planning for Generative Flow Networks".

Installation

This project uses uv for dependency management. To install:

1. Install uv (if not already installed):

curl -LsSf https://astral.sh/uv/install.sh | sh

2. Install dependencies:

uv sync

This will install all dependencies specified in pyproject.toml and uv.lock.

Usage

1. Experiments with the Deceptive Hypergrid

Navigate to the adaptive-teacher/grid/ directory:

cd adaptive-teacher/grid/

Run experiments with different world models (A, B, C) using the following commands:

python trainer_dyna_better.py --seed 42 --agent tb_dyna --world_model A --ndim 2 --horizon 256 --run_name "dyna_wmA_tb" --logger wandb --plot &
python trainer_dyna_better.py --seed 42 --agent tb_dyna --world_model B --ndim 2 --horizon 256 --run_name "dyna_wmB_tb" --logger wandb --plot &
python trainer_dyna_better.py --seed 42 --agent tb_dyna --world_model C --ndim 2 --horizon 256 --run_name "dyna_wmC_tb" --logger wandb --plot &

Additional examples with advanced features:

python trainer_dyna_better.py --seed 42 --agent tb_dyna --world_model A --use_gafn --ri_scale 0.3 --run_name "dyna_wmA_gafn_riscale0.3" --ndim 2 --horizon 256 --logger wandb --plot &
python trainer_dyna_better.py --seed 42 --agent tb_dyna --world_model A --use_teacher --run_name "dyna_wmA_teacher" --ndim 2 --horizon 256 --logger wandb --plot &
python trainer_dyna_better.py --seed 42 --agent tb_dyna --world_model A --use_buffer --buffer_pri reward --run_name "dyna_wmA_PRT" --ndim 2 --horizon 256 --logger wandb --plot &
python trainer_dyna_better.py --seed 42 --agent tb_dyna --world_model A --use_teacher --use_buffer --buffer_pri reward --run_name "dyna_wmA_teacher_PRT" --ndim 2 --horizon 256 --logger wandb --plot &

For more details on running experiments, see adaptive-teacher/grid/README.md.

2. Combinatorial Optimization Problems

CVRP (Capacitated Vehicle Routing Problem)

Navigate to the cvrp_nls/ directory:

cd cvrp_nls/

Train DynaGFN for CVRP:

python train_dyna_modular.py $N

For more details on running and other options, see cvrp_nls/README.md.

TSP (Traveling Salesman Problem)

Navigate to the tsp_nls/ directory:

cd tsp_nls/

Train DynaGFN for TSP:

python train_dyna.py $N

For more details on running and other options, see tsp_nls/README.md.

3. Discovery Tasks

Navigate to the adaptive-teacher/discovery/ directory:

cd adaptive-teacher/discovery/

Run DynaGFN experiments:

# DynaGFN with Model A
python rundynawb.py --setting qm9str --world_model_type A

For additional setup instructions and more details, see adaptive-teacher/discovery/README.md.

Acknowledgements

We thank the authors of GFACS (Ant Colony Sampling with GFlowNets for Combinatorial Optimization) and Adaptive Teacher for their valuable codebases, which this work builds upon. This repository contains code adapted from both projects.