GFlowNet Trajectory Saver

This project provides an implementation of GFlowNet for saving trajectories generated by the model during training. The primary focus is on extracting and saving trajectories using the save_gflownet_trajectories function. This allows for detailed analysis and further use of the trajectories for research or model improvement.

---

Publication

This repository accompanies the paper:

"Solution Augmentation for ARC Problems Using GFlowNet: A Probabilistic Exploration Approach"
Accepted at Transactions on Machine Learning Research (TMLR), 2025.

OpenReview | Project Page (Constructing)

Authors: Sanha Hwang, Seungpil Lee, Sejin Kim, Sundong Kim

Citation

If you find this work useful, please cite our paper:

@article{hwang2025solution,
  title={Solution Augmentation for ARC Problems Using GFlowNet: A Probabilistic Exploration Approach},
  author={Hwang, Sanha and Lee, Seungpil and Kim, Sejin and Kim, Sundong},
  journal={Transactions on Machine Learning Research},
  year={2025}
}

---

Branches

• main: Stable branch containing the core GFlowNet implementation for ARC-style environments and trajectory saving.
• LLM-Experiment: Dedicated branch for downstream LLM fine-tuning experiments described in the paper (e.g., Section 4.3/4.4). This branch contains code for preparing trajectory datasets and integrating them with LLaMA-based models.

---

Acknowledgements

This work was supported by NRF (Reinforcement Learning-Based Program Synthesis Techniques for Solving Abstract Visual Reasoning Problems by Identifying Patterns and Combining Concepts; RS-2024-00451162, Developing Abstraction and Reasoning Capability for AI Model; RS-2024-00454000), IITP (Enhancing AI Model Reliability Through Domain-Specific Automated Value Alignment Assessment; RS-2024-00445087, Development of Artificial Complex Intelligence for Conceptually Understanding and Inferring like Human; RS-2023-00216011), and GIST (KH0330, Future-leading Specialized Research Project) grants funded by the Ministry of Science and ICT, Korea.

Furthermore, This implementation was initially inspired by augustwester/gflownet, which served as a reference during the early stages of development. Over time, the codebase has diverged significantly to accommodate ARC-specific environments, trajectory saving utilities, and LLM downstream integration.

---

Table of Contents

• Overview
• Installation
• Usage
  • Command-line Arguments
  • Saving Trajectories
• Project Structure
• Configuration
• Examples
• Contact

---

Overview

The save_gflownet_trajectories functionality allows users to:

• Generate trajectories using a GFlowNet model.
• Save these trajectories in a JSON format for post-processing.
• Customize the number of trajectories, environment settings, and file paths via command-line arguments.

---

Installation

Prerequisites

• Python 3.8+
• PyTorch
• numpy
• gymnasium
• wandb (optional for logging)

Setup

Clone the repository and install the required dependencies:

$ git clone https://github.com/GIST-DSLab/GFN_to_ARC.git
$ cd gfn
$ pip install -r requirements.txt

---

Usage

Command-line Arguments

Argument	Description	Default Value
--save_trajectories	File path to save the trajectories (JSON format)	None
--num_trajectories	Number of trajectories to generate	100
--env_mode	Environment mode (entire or partial)	entire
--prob_index	Problem index for ARC tasks	178
--num_actions	Number of available actions in the environment	5

---

Saving Trajectories

To save trajectories using the save_gflownet_trajectories function, execute the following command:

python main.py --save_trajectories "trajectories.json" --num_trajectories 100 \
               --env_mode "entire" --prob_index 178 --num_actions 5

Example Explanation

• --save_trajectories: Specifies the output JSON file for the trajectories.
• --num_trajectories: Sets the number of trajectories to generate.
• --env_mode: Configures the environment mode (entire includes all tasks).
• --prob_index: Selects the specific ARC problem index.
• --num_actions: Defines the number of actions available in the environment.

The command generates 100 trajectories for the specified environment and saves them to trajectories.json.

---

Project Structure

project/
├── main.py               # Main execution script
├── train.py              # Training and trajectory saving logic
├── replay_buffer.py      # Replay Buffer for off-policy learning
├── config.py             # Configuration file
├── utils.py              # Utility functions
├── gflow/
│   ├── gflownet_target.py
│   ├── utils.py
├── ARCenv/
│   ├── wrapper.py
├── arcle/
│   ├── loaders.py
└── policy_target.py      # Policy model definitions

---

Configuration

The config.py file provides a central location for modifying default settings, such as:

Key	Description	Default Value
CUDANUM	CUDA device index	0
DEVICE	Device configuration (CPU/GPU)	cuda:0
TASKNUM	ARC Task problem index	178
ACTIONNUM	Number of actions in the environment	5
WANDB_USE	Enable WandB logging	True
REPLAY_BUFFER_CAPACITY	Maximum size of the replay buffer	10000

---

Examples

Save 50 Trajectories

python main.py --save_trajectories "output_trajectories.json" --num_trajectories 50

Change Environment Mode

python main.py --save_trajectories "trajectories_partial.json" --env_mode "partial"

---

Contact

For further information or questions, please contact:

• Name: Sanha Hwang
• Email: hsh6449@gmail.com
• GitHub: GitHub Profile

---