## Code for Hypergrid environment in Our Paper

## Requirements
This code requires PyTorch `torch` (to compile the cuda kernels):

- The experiments were conducted with Python 3.8.13 and PyTorch 1.9.0 (CUDA 11.1) 

- Python libraries:
```
pip install -r requirements.txt
```

## Running Instructions


- To train GFlowNets baseline with trajectory balance loss, run:
```
python trajectory_balance.py --save_path "tb.pkl.gz" --ndim 4 --horizon 8 --r 1e-3 --n_train_steps 625001
```

- To train GFlowNets while minimizing the path regularization via OT (reg_coef > 0), run:
```
python trajectory_balance_ot.py --save_path "tb_min_ot.pkl.gz" --ndim 4 --horizon 8 --r 1e-3 --n_train_steps 625001 --reg_coef 0.02
```

- To train GFlowNets while minimizing the upper bound of the path regularization via OT (reg_coef > 0), run:
```
python trajectory_balance_ot_upperbound.py --save_path "tb_ot_upperbound.pkl.gz" --ndim 4 --horizon 8 --r 1e-3 --n_train_steps 625001 --reg_coef 0.02
```

- To train GFlowNets while maximizing the path regularization via OT (reg_coef < 0), run:
```
python trajectory_balance_ot.py --save_path "tb_max_ot.pkl.gz" --ndim 4 --horizon 8 --r 1e-3 --n_train_steps 625001 --reg_coef -0.02
```

- The commands are similar when our hypergrid environment has ndim = 7 
