🤖 machine learning

Beyond the Proxy: Trajectory-Distilled Guidance for Offline GFlowNet Training

This paper introduces Trajectory-Distilled GFlowNet (TD-GFN), a proxy-free framework that leverages inverse reinforcement learning to extract dense edge rewards from offline data for guided exploration while relying exclusively on ground-truth terminal rewards to ensure robust training and outperform existing baselines in convergence speed and sample quality.

Original authors: Ruishuo Chen, Xun Wang, Rui Hu, Zhuoran Li, Longbo Huang

Published 2026-05-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ruishuo Chen, Xun Wang, Rui Hu, Zhuoran Li, Longbo Huang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to bake the perfect cake. In an ideal world, the robot would bake a cake, you would taste it, give it a score (a "reward"), and it would try again, learning from every mistake. This is how most AI learns today.

But in many real-world situations—like designing new medicines or creating biological sequences—you can't just "taste" every possibility. The experiments are too expensive, take too long, or require human experts who are in short supply. So, you are stuck with a "static cookbook": a pile of old recipes (data) that someone else wrote down, along with the final scores for those specific cakes. You can't ask for new scores; you can only look at what's already there.

This is the problem TD-GFN (Trajectory-Distilled GFlowNet) solves.

The Problem with the Old Way: The "Fake Judge"

Previously, when scientists tried to train AI using only these old cookbooks, they had to build a "proxy" or a "fake judge." This fake judge would look at a new, untried recipe and guess what the score would be.

The Flaw: If the fake judge is wrong (which happens often because it hasn't seen enough data), it gives bad advice. The robot follows this bad advice, makes a mistake, and the error spreads, making the whole system worse. It's like hiring a food critic who has never actually tasted the food to tell you how to cook.

The TD-GFN Solution: The "Map of the Kitchen"

Instead of building a fake judge to guess scores, TD-GFN looks at the paths the robot took to get to the final cakes in the old cookbook. It asks: "Which steps in these recipes were actually helpful, and which were dead ends?"

Here is how it works, step-by-step, using a simple analogy:

1. The "Reverse Engineering" (IRL)

Imagine you have a map of a city (the "DAG" or Directed Acyclic Graph) where every street leads to a destination. Some destinations are gold mines (high reward), and some are swamps (low reward).
TD-GFN uses a technique called Inverse Reinforcement Learning. Instead of asking "What is the score of this street?", it looks at the traffic patterns in the old data and asks, "If I were an expert trying to get to the gold mines, which streets would I have taken?"
It creates a heat map for every single street (edge) in the city. Some streets get a "high heat" score because they are crucial for reaching the gold; others get a "cold" score because they lead nowhere.

2. The "Road Closure" (Pruning)

Now, imagine you are the robot. You look at the heat map.

The Old Way: You try every street, hoping the fake judge tells you which ones are good.
The TD-GFN Way: You see that the "cold" streets are likely dead ends. So, you close them off (prune the graph). You don't even waste time thinking about them. You only keep the "hot" streets that lead toward the gold mines.
This makes your job much easier. You aren't guessing; you are navigating a streamlined map that only shows the promising paths.

3. The "Smart Backtracking" (Prioritized Sampling)

Finally, when the robot needs to learn, it doesn't just wander randomly. It uses a special trick: Backward Sampling.
Imagine you want to learn how to get to the gold mine. Instead of starting at the front door and guessing your way forward, you start at the gold mine and walk backward to the front door, but you do it smartly. You are more likely to take the "hot" streets (the ones the heat map said were important) and less likely to take the cold ones.
This ensures the robot spends its time studying the most valuable paths, learning much faster.

Why This is a Big Deal

The paper claims that by using this "Map of the Kitchen" approach, TD-GFN is:

Faster: It finds the best solutions (high-reward molecules or sequences) much quicker than other methods.
Smarter: It doesn't just copy the old recipes; it figures out the structure of what makes a good recipe and can invent new, even better ones that weren't in the original cookbook.
Safer: Because it doesn't rely on a "fake judge" to guess scores for new ideas, it avoids the trap of following bad advice. It only trusts the actual, known scores of the final results.

The Bottom Line

Think of TD-GFN as a master chef who doesn't need to taste every dish to know how to cook. Instead, they look at the history of successful dishes, figure out which specific ingredients and steps were the "secret sauce" (the edge rewards), remove the useless steps, and then teach the apprentice to focus only on the steps that matter. The result is a chef who learns faster, makes fewer mistakes, and creates better dishes than anyone else using just the old recipe book.

Technical Summary: Beyond the Proxy: Trajectory-Distilled Guidance for Offline GFlowNet Training

Problem Statement

Generative Flow Networks (GFlowNets) are powerful models for sampling compositional objects proportional to their rewards. However, in many practical applications—such as molecule discovery, biological sequence design, and recommendation systems—active interaction with the environment to query rewards is infeasible due to high experimental costs or the need for human evaluation. Consequently, GFlowNets must be trained using static offline datasets.

Current offline training paradigms typically rely on training a proxy reward model to approximate reward signals from the dataset, which the GFlowNet then queries for policy optimization. This approach faces significant challenges:

Data Scarcity and Cost: Constructing a reliable proxy model often requires large, diverse datasets and significant domain expertise.
Error Propagation: Reward queries for out-of-distribution (OOD) samples or states inaccurately estimated by the proxy can propagate errors through gradients, compromising the learned policy.
Limitations of Proxy-Free Methods: Existing proxy-free approaches (e.g., RO-GFlowNets, COFlowNet) often impose coarse constraints to align the policy with the dataset. These constraints can restrict generalization, inhibit effective exploration, and heighten sensitivity to data quality.

The central question addressed is: How can we effectively leverage offline trajectories to train GFlowNets when terminal rewards are insufficient to construct a reliable proxy model?

Methodology: Trajectory-Distilled GFlowNet (TD-GFN)

The authors propose TD-GFN, a novel proxy-free framework that extracts rich structural information from the environment's Directed Acyclic Graph (DAG) to guide policy training without relying on a proxy reward model. The methodology consists of three main phases:

1. Edge Reward Extraction via Inverse Reinforcement Learning (IRL)

Instead of predicting terminal rewards, TD-GFN utilizes IRL to distill transition-level edge rewards ( $R_E$ ) from offline trajectories.

Rebalancing: To address the issue that offline data may not represent expert behavior, the authors introduce a rebalancing strategy. Trajectories are sampled with probability proportional to their terminal rewards ( $P(\tau) \propto R(s_T)$ ), creating a pseudo-expert dataset.
Adversarial Training: Using a modified GAIL (Generative Adversarial Imitation Learning) framework, a discriminator $D_\phi$ distinguishes between transitions in the rebalanced dataset and those generated by a policy $\pi_\psi$ .
Reward Derivation: An unbiased edge-level reward function is extracted from the discriminator:
$R_E(s, s') = \log D_\phi(s, s') - \log(1 - D_\phi(s, s'))$
Theoretically, this reward approximates the log conditional incoming-flow score ( $\log P_B(s|s')$ ), quantifying the importance of specific transitions for learning an effective policy. Crucially, these rewards provide dense, intermediate guidance rather than terminal predictions.

2. Reward-Guided DAG Pruning

To ensure robustness and prevent error propagation from potentially inaccurate IRL estimates, TD-GFN does not use edge rewards to directly shape policy gradients. Instead, it uses them to prune the environment DAG.

Thresholding: Edges with rewards falling below a statistical threshold (mean minus $K$ times the standard deviation of the reward distribution) are removed.
Theoretical Guarantee: The authors prove that this pruning strategy preserves all "high-bottleneck" modes (terminal states reachable via paths with high edge scores) while removing low-utility transitions. This creates a compact, expressive action space that prioritizes high-value trajectories while insulating the policy from IRL approximation errors.

3. Prioritized Backward Sampling

Training is conducted on the pruned subgraph using a prioritized backward sampling mechanism.

Sampling Strategy: Terminal states are sampled from the intersection of the dataset and the pruned graph with probability proportional to their ground-truth rewards.
Backward Trajectory Construction: From each sampled terminal, trajectories are constructed in reverse toward the root using a learned backward policy $P_B$ , defined by the learned edge rewards:
$P_B(s_t | s_{t+1}) = \frac{\exp\{R_E(s_t, s_{t+1})\}}{\sum \exp\{R_E(s, s_{t+1})\}}$
Gradient Updates: The GFlowNet policy is optimized using these sampled trajectories and their ground-truth terminal rewards. By relying exclusively on recorded terminal rewards for gradient updates, the method avoids the error propagation issues inherent in proxy-based methods.

Key Contributions

A Novel Proxy-Free Paradigm: TD-GFN eliminates dependence on proxy models and OOD reward queries. It employs IRL to distill transition-level edge rewards from rebalanced datasets, providing structured guidance for policy training and generalization.
Robust Guidance Mechanism: The framework introduces two indirect guidance mechanisms: DAG pruning and prioritized backward sampling. This design ensures that gradient updates depend exclusively on ground-truth rewards, preventing error propagation while still leveraging dense structural guidance for efficient exploration.
State-of-the-Art Performance: Extensive experiments demonstrate that TD-GFN achieves superior convergence speed and sample quality compared to a broad range of baselines, including proxy-based methods (Proxy-GFN), offline RL algorithms (CQL, IQL), and other proxy-free GFlowNet variants (COFlowNet, Dataset-GFN).

Experimental Results

The authors evaluated TD-GFN on three benchmarks:

Hypergrid (Synthetic): On an 84-dimensional grid with sparse, multi-modal rewards, TD-GFN identified all 16 reward modes with fewer than 5,000 state visits (a six-fold improvement over baselines) and converged to the lowest Empirical L1 Error twice as fast as the best competitor (COFlowNet). It demonstrated strong robustness under noisy data (Mixed dataset) and data scarcity (10% of data).
Biosequence Design (Anti-Microbial Peptides): TD-GFN generated sequences with the highest rewards and greatest diversity among the top 100 sequences, outperforming both proxy-based and other offline methods.
Molecule Design: TD-GFN discovered more high-reward molecules and diverse modes than baselines while requiring significantly fewer trajectories to converge (2.75 × 10^4 vs. 43.7 × 10^4 for Proxy-GFN). It matched the performance of an Oracle-GFN (which uses ground-truth rewards directly) in identifying high-reward molecules.

Significance and Claims

The paper claims that TD-GFN establishes a more robust, stable, and efficient paradigm for offline GFlowNet training. By leveraging the structural information inherent in offline trajectories through IRL, the method bridges the gap between limited offline datasets and the complex exploration requirements of generative modeling.

The authors emphasize that TD-GFN:

Unlocks Potential: It enables the deployment of GFlowNets in complex domains where active reward queries are prohibitive.
Prevents Error Propagation: By decoupling the guidance mechanism (edge rewards) from the gradient update mechanism (terminal rewards), it avoids the pitfalls of proxy models.
Enhances Exploration: The dense structural guidance allows the model to explore high-reward regions effectively, even when they are sparsely represented or unobserved in the training data.

The work suggests that trajectory-level information, often overlooked in proxy-based methods, is a cornerstone for effective offline GFlowNet training, offering a scalable solution for real-world applications in drug discovery, biological design, and combinatorial optimization.