Generative Augmented Flow Networks
Authors: Ling Pan
Dinghuai Zhang
Aaron Courville
Longbo Huang
Yoshua Bengio
Venue: ICLR 2023 notable top 25% (2023)
Keywords: Generative Flow Networks (GFlowNets)
Exploration

Abstract: The Generative Flow Network is a probabilistic framework where an agent learns a stochastic policy for object generation, such that the probability of generating an object is proportional to a given reward function. Its effectiveness has been shown in discovering high-quality and diverse solutions, compared to reward-maximizing reinforcement learning-based methods. Nonetheless, GFlowNets only learn from rewards of the terminal states, which can limit its applicability. Indeed, intermediate rewards play a critical role in learning, for example from intrinsic motivation to provide intermediate feedback even in particularly challenging sparse reward tasks. Inspired by this, we propose Generative Augmented Flow Networks (GAFlowNets), a novel learning framework to incorporate intermediate rewards into GFlowNets. We specify intermediate rewards by intrinsic motivation to tackle the exploration problem in sparse reward environments. GAFlowNets can leverage edge-based and state-based intrinsic rewards in a joint way to improve exploration. Based on extensive experiments on the GridWorld task, we demonstrate the effectiveness and efficiency of GAFlowNet in terms of convergence, performance, and diversity of solutions. We further show that GAFlowNet is scalable to a more complex and large-scale molecule generation domain, where it achieves consistent and significant performance improvement.

Published as a conference paper at ICLR 2023
GENERATIVE AUGMENTED FLOW NETWORKS
Ling Pant.2, Dinghuai Zhang Aaron Courville 1.2,4 Longbo Huang', Yoshua Bengio 1,2,4 Mila- Québec AI Institute UUniversité de Montréal TTinnhua University *CIFAR AI Chair
{penny. .iing..i.................. paangppannaae gmail com
ABSTRACT
The Generative Flow Network (Bengio et al. 2021b] GFlowNet) a probabilis- tic framework where agent learns a stochastic policy for object generation, such that the probability of generating an object is proportional to a given re- ward function. Its effectiveness has been shown in discovering high-quality and diverse solutions, compared to -awaadd-aximizzing reinforcement learning-baxe methods. Nonetheless, GFlow Nets only learn from rewards of the terminal states, which can limit its applicability. Indeed, intermediate rewards play critical role in learning, for example from intrinsic motivation to provide intermediate feed- back even in particularly challenging sparse reward tasks. Inspired by this, we propose Generative Augmented Flow Networks (GAFlowNets), novel learning framework to incorporate intermediate rewards into GFlowNets. We specify inter- mediate rewards by intrinsic motivation to tackle the exploration problem in sparse reward environments. GAFlowNets can leverage edge-based and state-based in- trinsic rewards in a joint way to improve exploration. Based on extensive exper- iments the GridWorld task, demonstrate the effectiveness and efficiency GAFlowNet in terms of convergence, performance, and diversity of solutions. We further show that GAFlowNet is scalable to a more complex and large-scale molecule generation domain, where achieves consistent and significant perfor- mance improvement.
INTRODUCTION
Deep reinforcement learning (RI.) has achieved significant progress in recent years with particular success in games (Mnih et al. 2013 Silver et al., (2016] Vinyals et al.] 2019). RL methods applied to the setting where a reward is only given at the end (i.e., terminal states) typically aim at maximizing that reward function for learning the optimal policy. However, diversity of the generated states is
desirable in a wide range of practical scenarios including molecule generation (Bengio et al. 2021a),
biological sequence design (Jain et al., (022b), recommender systems (Kunaver & Požr 2017))
dialogue systems (Zhang et al.] 2020). etc. For example, in molecule generation, the reward function used in in-silico simulations can be uncertain and imperfect itself (compared the more expensive in-vivo experiments). Therefore, it is not sufficient to only search the solution that maximizes the return. Instead, it is desired that we sample many high-reward candidates, which can be achieved by sampling them proportionally to the reward of each terminal state.
Interestingly, GetswNets (Bengio et 2021a learn a stochastic policy to sample composite ob- jects x E X' with probability proportional to the return R(x). The learning paradigm of GelowNets different from other RL methods, as it is explicitly aiming at modeling the diversity in the target distribution, i.e., all the modes of the reward function. This makes it natural for practical applica- tions where the model should discover objects that are both interesting and diverse, which a focus of previous GFlowNet works (Bengio et al. 202 a Malkin et al. 2022 Jain et al.) 2022b).
Yet, GFlowNets only learn from the reward of the terminal state, and do not consider intermediate rewards, which can limit its applicability, especially in more general RL settings. Rewards play critical role in learning (Silver et al. 2021). The tremendous success of RL largely depends on the reward signals that provide intermediate feedback. Even in environments with sparse rewards. RL agents can motivate themselves for efficient exploration by intrinsic motivation, which augments the sparse extrinsic learning signal with dense intrinsic reward at each step. Our focus in this paper is

---

Published as a conference paper at ICLR 2023
precisely introducing such intermediate intrinsic rewards in GFlowNets, since they can be applied even in settings where the extrinsic reward is sparse (say con-zero only on a few terminal states).
Inspired by this missing element of GFlowNets, we propose new GFlowNet learning framework that takes intermediate feedback signals into account to provide an exploration incentive during training. The notion of flow in GFlowNets (Bengio et al. 2021a b) refers to marginalized quantity that sums rewards over all downstream terminal states following a given state, while sharing that reward with other states leading to the same terminal states. Apart from the existing flows in the net- work, we introduce augmented flows as intermediate rewards. Our new framework is well-suited for sparse reward tasks by considering intrinsic motivation as intermediate rewards, where the training of GFlowNet can get trapped in a few modes, since it may be difficult for it to discover new modes based on those it visited (Bengio et al. 2021b).
We first propose an edge-based augmented flow, based on the incorporation of an intrinsic reward at each transition. However, we find that although it improves learning efficiency, it only performs local exploration and still lacks sufficient exploration ability to drive the agent to visit solutions with zero rewards. On the other hand, we find that incorporating intermediate rewards in a state-based man- ner (Bengio et al.) (02Tb)) can result in slower convergence and large bias empirically, although it can explore more broadly. Therefore, we propose a joint way to take both edge- and state- state-based augmented flows into account. Our method can improve the diversity of solutions and learning effi- ciency by reaping the best from both worlds. Extensive experiments on the GridWorld and molecule domains that are already used to benchmark GFlowNets corroborate the effectiveness of our pro- posed framework. The code is publicly available atht /ttpt:::://ghthapt http/aa..oooocccoooommmcaco www.liinnnggniiieeeeeee pan/GAFN
The main contributions of this paper are summarized as follows:
We propose a novel GFlowNet learning framework, dubbed Generative Augmented Flow Networks (GAFlowNet), to incorporate intermediate rewards, which are represented by augmented flows in the flow network.
We specify intermediate rewards by intrinsic motivation to deal with the exploration of state space for GFlowNets in sparse reward tasks. We theoretically prove that our augmented objective asymptotically yields an unbiased solution to the original formulation.
We conduct extensive experiments on the GridWorld domain, demonstrating the effective- ness of our method in terms of convergence, diversity, and performance. Our method is also general, being applicable to different types of GFlowNets. We further extend our method to the larger- scgle and more challenging molecule generation task, where our method achieves consistent and substantial improvements over strong baselines.
2 BACKGROUND
Consider a directed acyclic graph (DAG) G = (S, A), A), where S denotes the state space, and A represents the action space, which is a subset of x S. We denote the vertex So E S to be the initial state with no incoming edges, while the vertex sy without outgoing edges is called the sink state, and state-action pairs correspond to edges. The goal for GFlowNets is to learn a stochastic policy that can construct discrete objects E X' with probability proportional to the reward function R : X' + Rzo, i.e., π (x) & R(x). GFlowNets construct objects sequentially, where each step adds from the initial state to a terminal state 7 = (Sg ... + sn) a trajectory, where T E 'T with T denoting the set of trajectories. Bengio al. (2021a) detine a trajectory flow F T + Rzo. Let F(s) Erse F(( define a state flow for any state S, and F(s + 5') E F(T) defines F(r), the edge flow for any edge + s'. The trajectory flow induces a probability measure P(T) = Z- where Z = ETeT F(T) denotes the total flow. We then define the corresponding forward policy Ps(s'(s) = F(s-xx) and the backward policy Ps(sss') = r(x->x), The flows can be considered as the amount of water flowing through edges (like pipes) states (like tees connecting pipes) (Malkin ett al., 2022). with R(x) the amount of water through terminal state x, and (ss(s s)) the relative amount of water flowing in edges outgoing from s.
element to the construction. We call the resulting sequence of state transitions
F())
2

---

Published as a conference paper at ICLR 2023
2.1 GELOWNETS TRAINING CRITERION
We call flow consistent if it satisfies the flow matching constraint for all internal states s, i.e,, -s F(s" s) = F(s) = Ee Lي F(s + s'), which means that the incoming flows equal the outgoing flows. Bengio al. (02 prove that for a consistent flow F where the terminal flow is set to be the reward, the forward policy can sample objects with probability proportional to R(x).
Flow matching (FM), Bengio et al. (021a) propose approximate the edge flow by a model Fo(s,s') parameterized by following the FM objective, i.e., LFM(S) (log ((s- Fo(s", .) - log ((s-wa) s')s for non-terminal states. At terminal states, a similar objective encourages the incoming flow to match the corresponding reward. The objective is optimized using trajectories sampled from training policy with full support such as a tempered version of Pr, or a mixture of Pr, with a uniform policy U, i.e. To = (1 Pr, + c. U, This is similar e-greedy and entropy-regularize strategies in RL improve exploration. |Bengio et al. (2021)) prove that if we reach a global minimum of the expected loss function and the training policy Ty has full support, then Gel samples from the target distribution.
Detailed balance (DB). Bengio et al. (2021b) propose the DB objective to avoid the computa- tionally expensive summing operation over the parents children of states. For learning based on DB, we train a neural network with a state flow model Fg. a forward policy model PF, (-Is), and a backward policy model PB, (-Is) parameterized by 0. The optimization objective is to mini- mize (s distribution if global minimum of the expected loss is reached and Th has full support.
(log(Fe Ps) (s' s)) log( (s') PB, (sls')))² It also samples from the target
Trajectory balance (TB). Malkin et al. (2022) propose the TB objective for faster credit assignment and learning over longer trajectonies. The loss function for TB is (TB(T)
(log(Z) Zo IT== PF. (S)+1(s)) log( ([(x) [T= PB, (s)ss where Zg learnable scalar.
RELATED WORK
202 there has been an increasing interest in improving (Bengio et al., 20216] Malkin al. 2022), understanding, and applying this
GFlowNets. Since the proposal of GFlowNets (Bengio et al.
framework to practical scenarios. It is general-purpose high-level probabilistic inference frame-
work, and induces fruitful applications (Zhang et al. 2022a c Deleu et al. 2022 Jain et al. 2022a). Pan et al. (2023)) propose Forward-Looking GFlowNets by incorporating local credits, and makes it possible to be trained with incomplete trajectories (which is a requirement for previous GelowNet methods). Pan et al. (2023b) and Zhang et al. (2023) recently generalize GFlowNets to the more general stochastic environments with stochastic transition dynamics or rewards. However, previous works only consider learning based the terminal reward, which make it difficult to provide good training signal for intermediate states, especially when the reward is sparse (i.e., significantly non-zero in only tiny fraction of the terminal states).
Reinforcement learning (RL). Different from GFlowNets that aim to sample proportionally to the reward function, RL learns seward-maximizatin policy. Entropy-regulatiized RL (Attias 2003 Ziebart 2010 Haarnoja et al.) 2017 2018; Zhang et al.) (022b) introduces entropy regularization the learning objectives, which learns with the Boltzmann softmax policy and the log-sum-exp
operator in value function updates (Shhulmnn et al. 2017a, Asadi & Littman 2017 Pan et al., 2020). Although it can improve diversity, this is limited tree-structured DAGs. This because could only sample terminal state x in proportion to the sum of rewards over all trajectories leading to x. can fail on general (non-tree) DAGs (Bengio al. 2021a) for which the same terminal state can be obtained with potentially large number of trajectories (and very different number of trajectories for different x's).
Intrinsic motivation. There has been a line of research to incorporate intrinsic motivation (Pathak et al. 2017 7 Burda et al.) 2018 Zhang et al.) 2021) for improving exploration in RL. Yet, such ideas have not been explored with GFlowNer because the current mathematical framework of GelowNets only allows for terminal rewards, unlike the standard RL frameworks. This deficiency as well as the potential of introducing intrinsic intermediate rewards motivates this paper.
3

---

Published as a conference paper at ICLR 2023
4 GENERATIVE AUGMENTED FLOW NETWORKS
The potential difficulty in learning only from the terminal reward is related to the challenge of sparse rewards in RL, where most states do not provide an informative reward. We demonstrate the sparse reward problem for GFtowNets and reveal interesting findings based on the Grid task (as shown in Figure 4 with sparse rewards. Specifically, the agent only receives reward of +1 only when it reaches one of the 3 goals located around the corners of the world (except the starting state cornr) with size H H (with H E {64, 128)), and the reward is 0 otherwise. A more detailed description of the task can be found in Section 5.1 We evaluate the number of modes discovered by the GelowNet trained with TB, following Bengio al. (2021a). As summarized in Figure I GFlowNet training get trapped in subset of the modes. Therefore, it remains
a critical challenge for GFlowNets to efficiently learn when the reward signal is sparse and non-informative.
GFlewee
GAFlowNet
On the other hand, there has been recent progress with intrinsic moti- vation methods (Pathak et al., 2017 Burda et al.) 2018 to improve ploration of RL algorithms, where the agent learns from both a sparse extrinsic reward and dense intrinsic bonus at each step. Building this, we aim to address the exploration challenge of Gets by en- abling intermediate rewards in GFlowNets and thus intrinsic rewards.
oo found
Horizon
We now propose our learning framework, which is dubbed Generative Augmented Flow Network (GAFlowNet), take intermediate rewards into consideration.
Figure 1: Comparison of GFlowNet and our augmented (GAFlowNet) method Gridworld with sparse rewards.
4.1 EDGE- EDGE-BASED INTERMEDIATE REWARD AUGMENTATION
We start our derivation from the flow matching consistency constraint, to take advantage of insights brought by the water flow metaphor as discussed in Section By incorporating intermediate rewards r(st + $1+1) for transitions from states Sy to St-1 into the flow matching constraint, we obtain
Σ +s) F(x)) = Σ [F(s -> s+11 + r(( + St+1) 8
+
8++1
by considering extra flow r(s: S going of the transition St Stt1- Based on Eq. (1) we define the corresponding forward and backward policies
Pp(s{(s--1)= F(s.-1 + s; r(s, + s;)
PB(((-1(s() F(s, +))
F(s
F(s))
Combining these, we obtain the detailed balance objective with the incorporation of intermediate rewards as
F(s P(psss))p(((sss-)) (s(ss- r + + St).
(3)
Finally, we have our resulting edge-based reward augmented learning objective for trajectory balance
in Eq. (4) via telescoping calculation upon Eq. (3), where Sn, and Z Σ., 1-78; r(s;-1
8;) + Σx R(x) is the augmented total flow. Detailed derivation can be found in Appendix
n
P ( r(s: + S441) F(s++1)
Z
(S = R(x) t=0
II
II t=
(4)
We explain the semantics of Eq. 44 in Figure aa). For the transition from an internal state (yel- low circles) s, to the i-th next state s/+1. we associate s with a special state $ (red circle) with pseudo-exit. Specifically, from the state st, we choose associated next states $ with prob- ability (F(s; - st+1) + r(s: + Sx+1)) /F(s)) according to the forward policy in Eq. (2). At the associated next state 8141. we "virtually" choose the sink state (purple circles) s with probabil- ity r(s. - 1++1)/F(st).
we choose the next state S with probability F(st St+1)/F(ss Adding them together and multiplying these probabilities by the incoming flow F(s:), we have F(s, S(41) + r(s: + St+1). Therefore, considering all possible next states, we have the aug- mented flow consistency equation (incoming flow = outgoing flow) as in Eq. ( The intermediate rewards r(s: - 8t+1) similar to transitions into pseudo-exit that never taken but still attracts
4

---

Published as a conference paper at ICLR 2023
K(xy) R6X2)
Source state
1it FUx hes? F(4)
state
Interaa state
rating edge
Viriual state
Non Non-tomirating dig:
Sink
Terminating
Nan-cammnnainng slgs
Ring = Rog) R0n mo
(a)
(b)
Figure 2: (a) Edge-based reward augmentation can be seen as introducing an augmented flow of amount r(s + S(+1) towards pseudo-exit to the sink state (that we never actually take) at every transition step. (b) For tasks with sparse rewards, agents can easily get stuck few modes (e.g., Xg). Our proposed method motivates the agent to discover unexplored states and trajectories to find diverse sets modes (i.e., x12) by increasing the probability of visiting alternative transitions in proportion to all the transitions reachable from there. Note that we omit the sink state from terminating edges for simplicity.
larger probabilities into its ancestors the DAG. From the water analogy, it can be considered that the flow from to S (and thus the probability of choosing that transition) is augmented by all the intermediate rewards due to pseudo-exits in all the accessible downstream transitions.
In contrast to simply adding constant uniform probability every action (which is commonly used for exploration with GFlowNets). the pscudo-exit intermediate rewards have an effect that is not local. In addition, can specify non-uniform intermediate rewards intrinsic motivation r(s; + $++1) with novelty-based methods (Pathak et al. 2017: Burda et al. 2018) to better tackle exploration in sparse reward tasks. Although how accurately measure the novelty degree remains an open problem, it has been shown that random network distillation (RND) (Burda et al.) (2018) simple yet effective method for encouraging the agent visit states of interest. We use the scaled difference between the predicted features by a trainable state encoder and a random fixed state encoder as the intrinsic rewards based on RND, i.e., allø(s) Grandom (s) where is reward scaling parameter that controls the degree of exploration. The random distillation network trained by minimizing such differences. Therefore, the novelty measure generally smaller for more often seen states similar states. The overall training procedure is shown in Algorithm by substituting the augmented trajectory balance loss according to Eq. 44).
We now demonstrate the conceptual advantage of edge-based reward augmentation which specifies intermediate rewards by intrinsic motivation in sparse reward task for exploration in Figure depicts a flow network Markov decision process (MDP) with sparse rewards, where only R(xs) and R(x12) are and other terminal rewards are all 0. Consider the case where the agent had already discovered solution Xş with the red flow A. Since the rewards of most other solutions can easily get trapped in the mode of Xg. and thus fails to discover other solutions. Nonetheless, our edge-based reward augmentation could motivate the agent discover other paths (e.g., the blue flow B) to Xs, which can be beneficial for the agent to discover other solutions (e.g., x12) subsequently.
0, it
Following the evaluation scheme in (Bengio et al. 202 a0) we summarize the number of dis- covered modes and the empirical L1 error for GFlowNet and GAFlowNet with edge-based reward augmentation Figure 3 The figure also includes the base- and joint objec- tives introduced later sections (Sections 4.2 and 4.3) for completeness. The Li enror is defined as E [[p(x) π (x)]], where p(x) = R(x)/Z denotes the true reward distribution, and we estimate æ by repeated sampling and summarizing frequencies for visitation of each
C200
6.5 (noge)
CartonNi: olatel
CA owild bauld
cartznlfe Ban0
GM swild (61%)
Updits
(a)
(b)
Figure 3: Different reward augmentations pro-
posed Sections 4.1 4.3 (a) Diversity metric:
the number of modes found. (b) Distribution fit-
ness metric: empirical L2 emror.
5

---

Published as a conference paper at ICLR 2023
possible state As shown, GAFlowNet (edge) is able to discover more modes opposed to stan- dard GFlowNet and improves diversity. In addition, it learns more efficiently and leads to a smaller level of L1 emror
4.2 STATE- BASED INTERMEDIATE REWARD AUGMENTATION
Although the learning framework of edge- reward augmentation is able to improve the diversity of solutions found, it still fails discover all of the modes as shown in Figure 3 We hypothesize that this is due to its "local" exploration ability, where it is able to consider difterent paths to solution x; with non-zero rewards. However, it fails to sufficiently motivate the agent to globally explore solutions whose rewards may be zero. Therefore, it can still get trapped in a few modes, lacking sufficient exploration ability to discover other modes.
Different from the edge-based reward augmentation, Bengio al. (022 defines a trajectory return as the sum of intermediate rewards in state-based reward augmentation manner. Specifically, baste reward augmentation for trajectory balance yields the following criterion
Z
II t=
Pr (s:+1|s:)
R(x)=
+ Σ r(s: + s
II
PB (sr/s()))), 1=0
(5)
t=
and we also use RND for intrinsic rewards. Such an objective directly motivates the agent to explore different terminate states in more global way (e.g., X2, Xs. X7, in Figure 2(b), which are bene- ficial for discovering X12). As shown in Figure is able discover all the modes, exhibiting great diversity. However, it explicitly changes the underlying target probability distribution, and is directly and highly affected by the length of the trajectory. Therefore, this leads to much slower convergence demonstrated in Figure ()).
4.3 JOINT INTERMEDIATE REWARD AUGMENTATION
As discussed above, the state- reward augmentation is effective in improving diversity but fails to fit the target distribution efficiently. On the other hand, edge-based reward augmentation performs more efficiently, but lacks sufficient exploration ability which cannot discover all the modes.
Therefore, we propose joint method take both state and edge-based intermediate reward aug- mentation into account to reap the best from both worlds. Specifically, we redefine the trajectory return as the sum of the terminal reward and the intrinsic reward for the terminal state only. This can be considered as we augment the extrinsic terminal reward with its curiosity degree. On the other hand, we include intrinsic rewards for internal states according to the edge-based reward augmenta- tion. This integration inherits the merits of both state and edge-based reward augmentation, which makes possible to improve exploration in more global way while learning more efficiently.
n-1
n-
r(s: +141))
Z II (S(S [R(x) + r(s)) II t=
PB PB(S
F(((+1)
Our resulting flow consistency constraint is shown in Eq. 66 where Z is the augmented total flow Ex R(x) + Ex. r(s,-1 + s,). Our new optimization objective (GAPloNNet Eq. which is trained by Algorithm
log 2 III
Pr Pr(s441)ss) ==
llog [I((X) + rss)))
II 1=
Ps P + r(s: - Sht1) (
F(8
In Theorem [ we theoretically justify that the resulting joint augmentation method leads to an unbiased solution to the original formulation asymptotically. The proof can found in Appendix B Note that we employ RND (Burda et al. 2018) for the intrinsic rewards, which decrease as the agent has more knowledge about the state.
Theorem 1.
Suppose that VT, LGAFlOwNea (7) and Vx, R(x) r(x) > 0. When edge-based (2) If state-based intrinsic
intrinsic rewards converge to O, we have that (1) P(x) = +(())
R(x)+r(x)
rewards converge to 0, then P(x) is an unbiased sample distribution.

---

Published as a conference paper at ICLR 2023
Algorithm Generative Augmented Flow Networks.
1: Initialize the forward and backward policies Pp, PB. learnable parameter Z, and state flow F
2: Initialize the random fixed target network & and the predictor network Φ
3: for each training step ( to 7 do
4: Collect a batch of B trajectories T = {so - .... s based on the forward policy PF
5: Compute intrinsic rewards for each sample in the batch of trajectories based the random
target network and the predictor network ø
6: Update the GAFlowNet model according to the augmented trajectory balance loss in Eq.
7:
Update the predictor network & by minimizing [[x(s) x(s)]l2
As shown in Figure 3 the joint method is able to discover all of the modes. In addition, it converges to the smallest level of L1 eπror, and is more efficient than state-based and edge-based formulations, which validates its effectiveness in practice.
EXPERIMENTS
We conduct comprehensive experiments to understand the effectiveness of our method and inves- tigate the following key questions: i) How does GAFlowNet compare against previous baselines? ii) What are the effects of state and edge-bases based flow augmentation, the form of the intrinsic reward mechanism, and critical hyperparameters? iii) Can it scale larger-scale and more complex tasks?
5.1 GRIDWORLD
We first conduct a series of experiments based on GridWorld with sparse re- wards (Figure The task is the same introduced in (Bengio al. (2221)), except that the reward function is sparse as described in Section[4] which makes it much harder due the challenge of exploration. With larger value of the size H, it requires the agent to plan in longer horizon and learn from sparse re- ward signals. Actions include operations to increase one coordinate in (Ben- gio et al. 2021a), and stop operation indicating termination to guarantee that the underlying MDP is directed acyclic graph. We compare GAFlowNet against strong baselines including Metropolis-HasiingssMMMM 2020), PPO (Schulman et al., 201 7b), and GFlowNet (Malkin et al.] 2022). We also include variant of PPO with intrinsic rewards based on the same intrinsic motivation mechanism using RND (Burda et al. 2018). All baselines are implemented based on the open- source coddi Each algorithm is run for five random seeds, and report their mean and standard deviation. A detailed description of the hyperparameters and setup can be found Appendix C.1
Apirtt
(Dai et al.)
Figure 4: The GridWorld task.
5.1.1 PERFORMANCE COMPARISON
We conduct experiments small, medium, and large GridWorlds with increasing sizes H. Full results other values of H can be found in Appendix C.3, To investigate the effectiveness of GAFlowNet, we first compare it against basclines in terms of the empirical L1 error as computed in Section 4 As shown in the first row in Figure S MCMC and PPO fail to converge due to smaller L2 error, it still underperforms GFlowNets by a large margin. GAFlowNet converges fastest and to the smallest level of L1 error, which shows that our method is effective to both explore efficiently and converge to sampling goals with probability proportional to the extrinsic reward function even if the reward signals are sparse.
the particularly sparse rewards. Although RPO-RND
The number of modes that each method discovers during the course of training is shown in the second row in Figure 5 Although incorporating PPO with intrinsic rewards improves the number of discovered modes compared to that of PPO in larger-scale tasks, still plateaus quickly. On the other hand, GFlowNets can get trapped in a few modes, while GAFlowNet is able discover all of the modes efficiently. A detailed comparison in terms of performance can be found in Ap- pendix C.4 Visualization of the learned structures be found in Appendix C22 Besides the
https://github... httpss//gittobb.c////////////r//

---

Published as a conference paper at ICLR 2023
Small
Midium
2414
OFliwher
GRiceelet
GPive GAFiceeee
dillinaNet
Gallnahw
shot
1200
196
166 nhe
1990
205
500
mie mie
medee fourd
MCMC
MCMC
NCM
Pro
PeQ
PHD
FPO-RND Gillwket DATonell
PRO RMO
PRD RND
CFlcenter GAllosheee
Giiinaket
Updste
Update
Updara
Figure 5: Comparison of GeFsowNets and baselines in Grid with increasing sizes correspond- ing to each column (left: small, middle: medium, right: large). The first and second rows correspond empirical L1 error and the number discovered modes, respectively.
particularly challenging sparse reward tasks, we show that GAFlowNets also provide consistent per- formance improvement in tasks with dense reward functions as in (Bengio et al. 202 Malkin
al. 2022). Results and discussions can be found in Appendix C.5 due to space limitations.
5.1.2 ABLATION STUDY
We now provide an in-depth ablation study on the important components and hyperparameters of GAFlowNet in the large Gril task. We also study the effect of different mechanisms of intrin- sic rewards besides RND, where results can be found in Appendix C.6
The effect of state-based and edge-based flow augmentation. In Figure a). we investigate the ef- feet edge-based, state-based, and joint intrinsic rewards. As discussed in Section incorporating intrinsic rewards for the trajectory in state-based manner can result in slower convergence, which has large L error. On the other hand, augmenting the TB objective with intrinsic rewards in an way still fails to motivate the agent to visit states with zero rewards. In contrast, the joint augmentation mechanism is effective in both diversity and performance, achieving the smallest level of I. error in experiments. It is also worth noting that only incorporating the intrinsic reward for the terminal state using state- based augmentation less efficient, which implies the importance of both edge-based and terminal state-based intrinsic rewards.
The effect of the coefficient of intrinsic rewards. In practice, we scale intrinsic rewards by coefficient. Figure b]b) illustrates the effect of the coefficient of the intrinsic rewards. A small coefficient does not improve the performance, while a too large coefficient converges slower. There exists an intermediate value that provides the best trade off.
58 Sruttu 204] Co/
Seoeter (990
GATIoWNetPMM
(a)
(b)
(a)
(b)
Figure 6: Ablation study. (a) The effect of state- Figure 7: Empirical L1 error of GFlowNet and and bdge-based intrinsic rewards. (b) The effect of the coefficient of intrinsic rewards.
GAFlowNet based on (a) DB and (b) FM.
5.1.3 VERSATILITY
We demonstrate that proposed framework is versatile by building it upon the other two GFlowNet objectives based on the detailed balance (DB) (Bengio et al.] (0221)) and flow matching (FM) (Bengio et al. 2021a) criteria. Comparison of empirical L, error averaged over increasing
8

---

Published as a conference paper at ICLR 2023
As demonstrated. GAFlowNet also significantly improves training convergence DB and FM, which provides consistent improvement gains.
sizes H are summarized in Figure 7
5.2
MOLECULE GENERATION
5.2.1
EXPERIMENTAL SETUP
We now investigate the effectiveness of our method in larger-scale tasks, by evaluating it on the more challenging molecule generation task (Bengio et al. (021a) as depicted in Figure B(a). A molecule represented by graph, which consists of vocabulary of building blocks. The agent sequentially generates the molecule by choosing where to attach a block and also which block to attach at each step considering chemical validity constraints. There is also an exit action indicating whether the agent decides to stop the generation process. This problem is challenging with large state (about 10 and action (around 100 to 2000) spaces. The agent aims to discover diverse molecules with high rewards, i.e. low binding energy to the soluble epoxide hydrolase (sEH) protein. We use a pretrained proxy model to compute this binding energy. We consider a sparse reward function here, where the agent only obtains non-zero reward if the corresponding molecule succeeds to meet a target score, and the reward is O otherwise. A detailed description of the environment is in Appendix C.I.I We compare our method with previous GelowNet results (Bengio al. 2021a), PPO (Schulman et al. 20176), PPO with intrinsic rewards based on RND, and MARS (Xie et al.) 2020 All baselines are run with three random seeds in (Bengio al. 2021a). More details for the setup can be found in Appendix(C.I) C.1.2
5.2.2 PERFORMANCE COMPARISON
We follow the evaluation metric in (Bengio et al. 2021a) and investigate our method in both perfor- mance and diversity. Figure 8 b) demonstrates the average reward of the top-10 unique molecules generated by each method. The number of modes discovered by each method with rewards above 7.5 is summarized in Figure 8 c).. We compute the average pairwise Tanimoto similarities for the
top-10 10 samples in Figure 8 d). Additional comparison results can be found in Appendix C.7
MoRS
MAR5
MARS
CAP lowleweet
CAP lowleet
CAF
Sietevi bx
Stmee vidied octo")
Statee diie (410")
(b)
(c)
(d)
Figure Molecule generation task. (a) The environment. (b) Average reward of the top-10 molecules. (c) The number of modes with R > 7.5. (d) Tanimoto similarity (lower is better).
As shown, MARS fails to perform well given sparse rewards since most of the reward signals non-informative. On the other hand, PPO and its variant with intrinsic rewards are better at finding higher-quality solutions than MARS, but suffer both from high similarities of the samples. The unaugmented GelowNet is better at discovering more diverse molecules, but does not perform well terms of solution quality. GAFlowNet significantly outperforms baseline methods in performance and diversity. We also visualize the top-10 molccules generated by GFlowNet and GAFlowNet in run in Appendix As shown, GAFlowNet able to generate diverse and nigh-quality molecules efficiently, which demonstrates consistent and significant performance improvement.
6 CONCLUSION
In this paper, we propose a new learning framework, GAFlowNet. for GFlowNet to incorporate intermediate rewards. We specify intermediate rewards by intrinsic motivation to tackle the explo- ration problem of GFlowNets in sparse reward tasks, where it can get trapped in a few modes. We conduct extensive experiments to evaluate the effectiveness of GAFlowNets, which significantly outperforms strong baselines in terms of diversity, convergence, and performance when the rewards very sparse. GAFlowNet is also scalable complex tasks like molecular graph generation.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENT
The authors would like to thank Emmanuel Bengio, Qingpeng Cai, and anonymous reviewers for their comments and insightful discussions. Aaron Courville thanks the support of CIFAR. Longho Huang is supported in part by the Technology and Innovation Major Project the Ministry of Science and Technology of China under Grant 2020AAA0108400 and 2020AAA0108400 the Ts- inghua University Initiative Scientific Research Program, and Tsinghua Precision Medicine Foun- dation 10001020109. Yoshua Bengio acknowledges the funding from CIFAR, Samsung, IBM and Microsoft.
REPRODUCIBILITY STATEMENT
All details for our experiments are in Appendix with detailed description the task, hyper- parameters, network architectures for baselines, and setup. Our implementation for all baselines and environments based on open-source repositories. The proof of Theorem can be found in Appendix B The code is publicly available at wtttts:://gttuub..
ng Paag--ann/AANN
REFERENCES
Kavosh Asadi and Michael L Littman. An alternative softmax operator for reinforcement learning. In International Conference on Machine Learning, pp. 243-252. PMLR, 2017.
Hagai Attias. Planning by probabilistic inference. In AISTATS, 2003.
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. Advances in Neural Information Processing Systems, $4:27381-27394 2021a.
Yoshua Bengio, Tristan Deleu, Edward Hu, Salem Lahlou, Mo Tiwari, and Emmanuel Bengio. Gflownet foundations. arXiv preprint (ttiv:2/11.092200 2021b.
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. In International Conference on Learning Representations, 2018.
Hanjun Dai, Rishabh Singh, Bo Dai, Charles Sutton, and Dale Schuurmans. Learning discrete energy-based models via auxiliary-variabl local exploration. Advances in Neural Information Processing Systems, 13:10443-10455 2020.
Tristan Deleu, Ant' onio G' ois, Chris C. Emezue, Mansi Rankawat, Simon Lacoste-Julien, Stefan Bauer, and Yoshua Bengio. Bayesian structure learning with generative flow networks. ArXiv, abs/2202. 2022.
Justin Gilmer, Samuel Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning. pp- 1263-1272. PMLR, 2017.
Tuomas Haaroja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International conference on machine learning, pp. 1352-1361. PMLR, 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learning, pp. 1861-1870. PMLR, 2018.
Moksh Jain, Emmanuel Bengio, Alex Gareía, Jarrid Rector- Rector-Brooks, Brooks, Bonaventure F. Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang. Micheal Kilgour, Dinghuai Zhang, Lena Simine, Payel Das, and Yoshua Bengio. Biological sequence design with gflownets. ICMI, 2022a.
Moksh Jain, Emmanuel Bengio, Alex Hemnandez-Gariia. Jarrid Rector- Rector-Brooks, Bonaventure FP Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, et al. Biological sequence design with gflownets. In International Conference on Machine Learning, pp. 9786-9801. PMLR, 2022h.
10

---

Published as a conference paper at ICLR 2023
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6990, 2014.
Matevž Kunaver and Tomaž Požrl. Diversity in recommender systems-a survey. Knowledge-based systems, 1262154-16,, 2017.
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment gflownets. arXiv preprint arXiv:2201.13255 2022.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 5 518(7540):529-53,, 2015.
Ling Pan, Qingpeng Cai, and Longbo Huang. Softmax deep double deterministic policy gradients. Advances in Neural Information Processing Systems, 33: 1767-11777, 2020.
Ling Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Better training of gflownets with local credit and incomplete trajectories. arXiv preprint arXiv:2302.01667, 2023a.
Ling Pan, Dinghuai Zhang, Moksh Jain, Longbo Huang, and Yoshua Bengio. Stochastic generative flow networks. arXiv preprint arXiv:2302.09463 2023b.
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International conference on machine learning, pp. 2778-2787. PMLR, 2017.
John Schulman, Xi Chen, and Pieter Abbeel. Equivalence between policy gradients and soft q- learning. arXiv preprint arXiv:1704.06444, 2017a.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017b.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. Nature, 529(7587):484 2016.
David Silver, Satinder Singh, Doina Precup, and Richard Sutton. Reward enough. Artificial Intelligence, 299:103533 103535, 2021.
Oriol Vinyals, Igor Babuschkin, Wojciech M Czameeki,, Michaël Mathieu, Andrew Dudzik, Juny- oung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev. et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782):550-354 2019.
Yutong Xie, Chence Shi, Hao Zhou, Yuwei Yang, Weinan Zhang, Yong Yu, and Lei Li. Mars: Markov molecular sampling for multi-objective drug discovery. In International Conference on Learning Representations, 2020.
Dinghuai Zhang, Ricky T. Q. Chen, Nikolay Malkin, and Yoshua Bengio. Unifying generative models with gflownets and beyond. 2022a.
Dinghuai Zhang, Aaron C. Courville, Yoshua Bengio, Qinqing Zheng, Amy Zhang, and Ricky T. Q- Chen. Latent state marginalization as a low-cost approach for improving exploration. ArXiv, 0bs/2210.00999. 2022b.
Dinghuai Zhang, Nikolay Malkin, Z. Liu, Alexandra Volokhova, Aaron C. Courville, and Yoshua Bengio. Generative flow networks for discrete probabilistic modeling. In ICMI. 2022c.
Dinghuai Zhang, Ling Pan, Ricky TQ Chen, Aaron Courville, and Yoshua Bengio. Distributional gflownets with quantile flows. arXiv preprint arXiv:2302.05599 2023.
Tianjun Zhang, Huazhe Xu, Xiaolong Wang, Yi Wu, Kurt Keutzer, Joseph E Gonzalez, and Yuan- dong Tian. Noveld: A simple yet effective exploration criterion. Advances in Neural Information Processing Systems, 34:25217-25223 2021.

---

Published as a conference paper at ICLR 2023
Yichi Zhang, Zhijian Ou, and Zhou Yu. Trsk-tri dialog systems that consider multiple appro- priate responses under the same context. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 9604-9611, 2020.
Brian D. Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. 2010.
12

---

Published as a conference paper at ICLR 2023
A DERIVATION OF EDGE-BASED, STATE-BASED, AND JOINT INTERMEDIATE REWARD AUGMENTATION
For the bagee-baeed intermediate reward augmentation, with the incorporation of intermediate rewards as in Eq. ( following Section 4.1
obtain the detailed balance objective
F(s = ((s + r(s + +s))
(8)
Therefore, we have that
P(50)P((S((s)) = F(s:)) r(so -+ -1)
(9)
t= n ((s..)))P((slssa-)))) = F(ss,)PP( sm- + rsn + s.)
By accumulative multiplication on both sides, we get that
121
II (=(S4+1(ss) F(s1) ... F(S,) PB(S((S(+1) r(s: + St+1)
F(so) ... F(Sn-1)
F(s++1)
(10)
Therefore, we obtain the corresponding edge-based reward augmented formulation for trajectory balance as
Pn Pn(Se + r(s: - s(+1) F(s(+1)
F(Sn)
Pr(s (8.) = (snn)
II x=
(11)
1= 0
where F(su) = Z = R(x) Eo. r(s-1 + s;), and F(s.) R(x).
The state-based reward augmented formulation be obtained similarly by following Appendix
D in (Bengio et al.) (2221)) The joint reward augmented formulation is obtained by combining edge-based and state-based reward augmentation.
B PROOF OF THEOREM
Theorem 1. Suppose that VT, LGAFFooNN (T) = O, and Vx, R(x) + +r(x) > 0. When edge-based
intrinsic rewards converge
have that (1) P(x) [R( R(x)+((x)
(2) If state-based intrinsic
rewards converge to 0, then P(x) is an unbiased sample distribution.
Proof. By definition, we have that
Fo(T) Z
Ps(s t=0
II
(12)
Since VT, LGAFlowNod (T) we have that
n-1 Z
r(s: + s2+1)
II Pr (s(s = (R(x) + rx()) t=
II Pn(selse+1)
(13)
F(s
Therefore, obtain that
Po(T) Fo(T)
R(x) + r(x)
n-1
(=(s(() + r(s: - 8141)
II
(14)
Z
F(s.+1)
t=0
When edge-based intrinsic rewards converge to 0 and does not vanish, we have that
R(x) + +r(x)
Po(x)
Σ
Po(T)
Σ II Pa
(15)
Z
=X ==
13

---

Published as a conference paper at ICLR 2023
Due to the law of total probability, have that
n-1
Σ
II
=1
(16)
r= (350
Therefore, Po(x)- ((x)) As Ex Po(x)
L, we also get that == Ex (R(x)) + r(x)).
(H(x))
Therefore, we have Part (1) that
R(x) - r(x) Px(x) = Σx R(x) + rx)]
(17)
Based on the above analysis, P(x) is an unbiased estimation when state-based intrinsic rewards converge to 0, and we have Part (2).
C EXPERIMENTAL DETAILS
C.I EXPERIMENTAL SETUP
C.I.I TASK
The molecule generation task We adopt a pretrained proxy model for the reward, which is trained on dataset of 300, 000 molecules that randomly generated provided in (Bengio et al. (021a). For the original dense reward function, the agent receives a reward based on the normalized score. Here, use sparse reward function, where the agent only obtains the original non-zero reward the normalized score succeeds to meet target score (7.0), and the reward is O otherwise. As described in Section [.2.1 the agent can choose one of the blocks attach from the basic building blocks vocabulary (with a size of 105).
C.1.2 BASELINE
All baseline methods are implemented based the open-source implementation as described in the main text, where we follow the default hyperparameters and setup as in (Bengio et al. 2021a The code will be released upon publication of the paper.
World Specifically, in GridWorld, the GFlowNet model is feedforward network consisting of two hidden layers with 256 hidden units per layer using LeakyRel.U activation. We train all models based samples from parallel of 16 rollouts in the environment. We leverage random network distillation (RND) (Burda et al.) 2018)) as the intrinsic reward mechanism, where the random target network and the predictor network are both feedforward networks consisting of two hidden layers with 256 hidden units per layer using LeakyReLU activation. We train the GFlowNet model and RND jointly based on the Adam (Kingma Ba 2014) optimizer with a learning rate of 0.001 for the policy models (Pr and Pn)) and 0.1 for Z.
Molecule generation For the molecule generation task, we use a reward proxy provided in (Ben- gio al., 2021a). As the molecule represented an atom graph, we use Message Passing Neural Networks (MPNN) (Gilmer et al., 2017) as the network architecture for all models. Note that we build our method upon GFlowNet based on the flow matching criterion in the molecule generation task, since is the most competitive version in this task in terms of finding high-quality and diverse candidates.
For GAFlowNet, the only hyperparameter that requires tuning is the coefficient of intrinsic wards, where we use same value for state-based and edge-basee augmentation. We tune a in (0.001, 0.005, 0.01, 0.05, 0.1, 0.53 with grid search. Specifically, values of horizon except for e (16, 64). where we set a to be 0.005. For the molecule generation task, is set to be 0.1.
0.001 for Grid with all
14

---

Published as a conference paper at ICLR 2023
C.2 VISUALIZATION OF THE LEARNED STRUCTURES
In this section, we visualize the learned forward policies of GelowNets and GAFlowNets. The results are demonstrated in Figure 9 in the small GridWorld, and the borders (the last column and the last row) of the grid are omitted in the figure for better readability. The length of the arrows is proportional the likelihood of the corresponding actions under the forward policy PF. while the size of the gray circle is proportional to the probability of the termination action (stop). As shown, GFlowNet can get trapped few modes, while GAFlowNets is able discover all the modes.
(a)
(b)
Figure 9: Visualization of the learned forward policies of (a) GelowNets and (b) GAFlowNets in GridWorld (horizon H 8).
The corresponding distribution the samples collected by GFlowNets and GAFlowNets during the course of training in the small GridWorld task is shown in Figure [O] where GAFlowNets leads to more diverse solutions. Visualization of the top molecules discovered by GFlowNets and GA FlosowNets be found in Appendix C.8
Figure 10: Distribution of the samples during the course of training from (a) GFlowNets and (b)
GAFlowNets GridWorld (horizon = 8).
C.3 FULL RESULTS IN GRIDWORLD
We show in Figure 11 the full comparison results in GridWorld with increasing sizes H{8, 16, 32, 64, 128). As shown, GAFlowNet significantly outperforms baselines in empirical L error and the number of modes found.
C.4 PERFORMANCE COMPARISON
Apart from evaluating our method based on the metrics (the number of modes discovered by each method and empirical L1 error) as (Bengio et al.) 2021a), we are also interested in its perfor- mance after each update. Here, evaluate the performance of baselines after each update (instead of throughout the training process) as the evaluation scheme of RL algorithms. Figure 12 demon- strates the performance for the top-5 solutions among batch of 16 parallel rollouts of each method after each update for Worldorld with sizes H e (8, 16, 32, 64, 128]. As shown, although PPO is more efficient than PPO, both of them underperform GelowNet by large margin (especially with
15

---

Published as a conference paper at ICLR 2023
IP RND
PPO R&D
IPo RMD
PTO MD
1400
CAt swhst
Car swhlt
Upars
(a)
(b)
(c)
(d)
H-61
Ind RND
(Po RMD
GAn swhst
CAr onfic
Gar swhlll
CAP onficl
Updkte
(f)
(h)
4000
2P0 RMD
RMG
CHonnN
cot swhol
Got onfill
Lodada
Updara
ii)
Figure 11: Comparison of GAFlowNets and baselines in GridWorld with increasing sizes H {8, 16, 32, 64, 128}. (a), (c), (e). (g), (i) correspond to the empirical 1₁ error. (b), (d), (f), (h), (i)
correspond to the number of modes discovered by each method.
a larger value of H). We find that GAFlowNet significantly outperforms baseline methods, and also performs more efficiently than GFlowNet.
Hw 14
52
HCMC
HCNC P90
PRo POO-MD
Fnc MVD
TTC AND oifewter
(a)
(c)
Ho 64
H- 126
NCMC IPO IPD
HCMC ITO ITO-RAD Gfiielie CAr onNil.
Typ-4
Typ
GAR onNill
(d)
(e)
Figure 12: Top- K performance of baselines after each update for horizon H E {8, 16, 32, 64, 128).
16

---

Published as a conference paper at ICLR 2023
C.5 ADDITIONAL RESULTS ON GRIDWORLDS WITH DENSE REWARDS
In this section, we evaluate the performance of GAFlowNets in variant of GridWorld with dense reward functions. We employ the same reward function (Eq. (88) as in (Bengio et al. 2 Malkin et al. 2022) with Rg = -e - 3 (the hardest variant of the grid considered in (Bengio et al,
2021a Malkin et al., 2022)) which is shown Figure a]). Note that when Rg closer 0,
the problem becomes more difficult as it results in region of the state space that is undesirable to explore (Bengio et al. 2021a).
R(x) Ro + II T[ (0.25 < 12//-- 0.5l) + II [0.3 < | 0.5| < 0.4)
(18)
Comparison results of GAFlowNets and GFlowNets in empirical L1 error in GridWorlds with differ- ent sizes H E {8, 16, 32, 64, ]28} are shown in Figures (3)b)-(f). As demonstrated, GAFlowNets also greatly improve the learning efficiency of GFlowNets. This is because GAFlowNets exhibit greater exploration ability, and are more efficient in choosing training trajectories that the agent is not very familiar with yet and may have a high reward (Bengio et al. (021b). Therefore, it achieves consistent performance improvement.
Gffowiler
Gftawiec
(d)
(e)
Figure 13: (a) The dense reward function. (b)-(f) Empirical 7., error of GAFlowNets and GFlowNets in GridWorld with dense reward functions for horizon H E {8, 16, 32, 64, 128).
C.6 ADDITIONAL ABLATION STUDY OF GAFLOWNET
We investigate the effect of different types of intrinsic rewards in- cluding Intrinsic Curiosity Module (ICM) (Pathak al. 2017). Novelty Difference (NovelD) (Zhang et al. 2021). and Random Network Distillation (RND) (Burda al. 2018) in Figure 14 with fine-tuned coefficients for intrinsic rewards. We also include base- line with constant intrinsic rewards in GAFlowNet, which mimics the behavior of greedy exploration typically used in reinforcement learning algorithms. As demonstrated, GAFlowNet not sensi- tive to the forms of intrinsic rewards, but RND enables the fastest convergence in our simulations. It also validates the effectiveness of the novelty-based methods from the comparison of GAFlowNet and GAFlowNet with constant intrinsic rewards. This because novelty-based methods are more efficient than blindly wandering in the maze.
20301
SANEE E2 "HEWWO" Types of Intrinsic rewards
Figure 14: Ablation study different types of intrinsic re- wards.
17

---

Published as a conference paper at ICLR 2023
C.7 ADDITIONAL PERFORMANCE COMPARISON ON THE MOLECULE GENERATION TASK
Following the evaluation metrics in (Bengio et al. a)2 besides the results in Figure g in the main text, also evaluate the average reward of the 1op-100 molecules and the number of modes with R > 8.0 discovered by each method. As demonstrated in Figure [5| GAFlowNet achieves
consistent and significant performance improvement over previous baselines.
2000 1o
HARS
PPO
PPO
Averaage
FPO-RND
PPO RMD
Gilowftet
GPlowhie
States visited (x100)
States visited (x20*)
(a)
(b)
Figure 15: Molecule generation task. (a) The average reward of the top-100 molecules. (b) The number modes with > 8.0.
C.8 FULL VISUALIZATION OF TOP-10 MOLECULES
Figure 16 demonstrates the top-10 molecules generated by GFlowNet and GAFlowNet, where GAFlowNet discovers more diverse and higher-quality solutions.
(a) GFlowNet. Tanimoto similarity is 0.410 and the average reward is 7.747.
(b) GAFlowNet. Tanimoto similarity is 0.356 and the average reward is 8.120.
Figure 16: Full visualization top- molecules generated by GFlowNet and GAFlowNet.
18


Reviews:
[{"review_id": "oTuANQUjtFf", "reviewer": "Reviewer_WVnH", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper introduces  Generative Augmented Flow Networks (GAFlowNets) that incorporate intermediate rewards into GFlowNets. Original GFlowNets  learn from the reward of the terminal state, and do not consider intermediate rewards. That is a problem in environments with sparse rewards. For this problem taking intermediate feedback signals into account helps to provide an exploration. (Bengio et al., 2021b) already proposed to incorporate intermediate rewards in a state- based manner. This approach  can result in slower convergence and large bias empirically, although it can explore more broadly. Incorporating the edge-based augmentation for efficiency, this paper proposes a joint way to take both edge-based and state-based augmented flows into account. Experiments on the GridWorld and molecule domains show the effectiveness of the proposed framework. \n", "summary_of_the_review": "This paper formulate the problem by clarifying existing unsolved  problems of GFlowNets and proposes a novel approach. This paper proposes a joint way to take both edge-based and state-based augmented flows into account.The paper is enough described and organized to support claims with good experimental results.\n", "strength_and_weaknesses": "The strength of this paper is a novel learning framework to incorporate intermediate rewards into GFlowNets. Incorporating intermediate rewards into GFlowNets has been already proposed. This paper formulate the problem by clarifying existing unsolved  problems and proposes a novel approach.\nThe algorithm is described in details. Enough experiments with further ablation study.\nThis is not weakness, but more evaluations of differences of learned structures between GFlowNets and GAFlowNets would support claims more clearly", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "dThQP39bYZ", "reviewer": "Reviewer_eLGw", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes to incorporate intrinsic reward augmentation in GFlowNets to improve exploration in sparse reward settings. It extends the intrinsic rewards to both state and edge intrinsic rewards improving on convergence and mode discovery in GflowNets in Gridworld experiments and more complex molecule generation tasks.\n", "summary_of_the_review": "The paper tackles a problem of lack of intrinsic rewards in GFlowNets and provides a joint edge and state based flow intrinsic reward that captures well different modes and converges faster. \n", "strength_and_weaknesses": "Strengths: this paper is clear and well written, it provides a self contained and comprehensive literature review, it motivates the problem well and provides compelling results.\nWeaknesses:\nCan the authors provide citations/derivation for eq 4,5 and 6 in particular the telescoping derivation?\nThe experiments mention a coefficient parameter alpha that is never mentioned in the theoretical presentation of the model? Where is this used?\nHow important is the choice of intrinsic motivation function, is it enough to recover more structured forms of exploration? How far ahead can the flow information capture relevant future rewards? \nWhat are the main limitations of the GAFlowNets in exploration, can the authors frame the proposed approach in terms of where would this kind of exploration shine vs where would it fail?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "N9suuE0Klr", "reviewer": "Reviewer_ka6N", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes an extension of Generative Flow Networks (GEFlowNets). The main idea is to incorporate intermediate rewards as well as terminal rewards into the objective, which was missing in GEFlowNets. The paper proposes to combine two variants: edge-based augmentation and state-based augmentation and shows that adding intrinsic rewards leads to the same solution when intrinsic rewards converge to 0 over time. The empirical results on a grid world and a molecule generation task show that the proposed method significantly outperforms GEFlowNets and other RL baselines including PPO when combined with RND-like intrinsic rewards. Besides, it discovers more diverse set of solutions compared to the previous approaches thanks to the additional intrinsic rewards. ", "summary_of_the_review": "GEFlowNets are an interesting alternative to existing RL algorithms, which has been receiving attention. This paper extends GEFlowNets nicely to support intermediate rewards, which sounds like an exciting direction towards solving RL problems with GEFlowNets. Although the results are limited to the intrinsic rewards on specific types of problems, the comparison to the prior work GEFlowNets looks quite promising and convincing. I am willing to increase the score if the authors reflect (or provide a reasonable response to) some of my comments above.", "strength_and_weaknesses": "[Strength]\n\n* The proposed idea of incorporating intermediate rewards to GEFlowNets is novel.\n* The result convincingly demonstrates that the proposed method significantly outperforms the previous approach (GEFlowNets) both in terms of performance and the diversity of solutions.\n* The paper is well-written.\n\n[Weaknesses]\n\n* **Lack of result on standard RL benchmark:** To my understanding, the proposed method is general enough to be applied to any RL problems as it can incorporate intermediate rewards, and the introduction was also written in such a way. However, all of the experiments were conducted on the \"reward-at-termination\" types of environments, while intermediate rewards were given by intrinsic rewards (not from the environment). It would be much more comprehensive to evaluate the proposed method on a standard RL benchmark (e.g., Atari) that contain \"intermediate extrinsic\" rewards. \n\n* **Lack of clarity about intermediate rewards / intrinsic rewards:** Related to the above point, I felt that the paper often assumes that intermediate rewards = intrinsic rewards without a clear distinction, which I found confusing. In page 5, for example, it says \"our edge-based reward augmentation could motivate the agent to discover other paths ...\". Does the diversity of the solution come from the incorporation of intermediate rewards (i.e., the edge-based augmentation) or the combination of the edge-based augmentation and a particular intrinsic rewards designed for exploration? I believe that the paper meant the latter. I would be good to make this distinction clear throughout the paper.\n\n* **Lack of clarity about edge-based augmentation vs state-based augmentation:** I found it difficult to understand the different behaviors between the edge-based augmentation and the state-based augmentation. Why does the edge-based augmentation tend to explore \"locally\" while the state-based augmentation explores more \"globally\"? I could not figure it out from the equations immediately. Further clarifying this point would be helpful for justifying the joint approach.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "D4_2ujf6Dh1", "reviewer": "Reviewer_kXmH", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper argues that intermediate rewards are also crucial in the learning of Generative Flow Network, and proposes a novel GFlowNet learning framework, dubbed GAFlowNet, to incorporate intermediate rewards. The paper specifies intermediate rewards by intrinsic motivation to deal with the exploration of state space for GFlowNets in sparse reward tasks. The experiment results on the GridWorld domain and the larger-scale and more challenging molecule generation task show the effectiveness of the method in terms of convergence, diversity, and performance.", "summary_of_the_review": "Overall, this paper proposes an interesting method. But, why intermediate rewards play a critical role in RL  is not clear.  It is unclear that why the improved Generative Augmented Flow Networks (GAFlowNets) can work better than the previous Generative Flow Network. \n\n\nThe rebuttal is good, and I thus increase the score.\n ", "strength_and_weaknesses": "Strength:\n\n(1)The authors focus on the intermediate rewards in the learning of Generative Flow Network, and propose a novel GFlowNet learning framework, dubbed GAFlowNet, to incorporate intermediate rewards.\n\n (2) The paper proposes three different kinds of augmentation methods, namely, edge-based intermediate reward augmentation, state-based intermediate reward augmentation, and joint intermediate reward augmentation. And, the advantages and disadvantages of different augmentation methods are analyzed intuitively and clearly.\n\n (3) In general, the experiments are comprehensive and reasonable.\n\n (4) In general, the paper is well-organized and easy to follow.\n\nWeakness:\n\n(1)  The author just considers combining intermediate rewards to further improve the performance of GFlowNet in particularly challenging sparse reward tasks, which is not novel enough.\n\n(2) More novelty degree measure methods can be considered to further refine the experiments.\n\nMinor errors:\n(1) In Figure.5, the first and second rows correspond to the number of discovered modes and empirical L1 error respectively, which doesn’t correspond to the descriptions in the caption of Figure.5 and the analysis of the experiment in 5.1.1.\n\n(2) The position of the figures can be adjusted to make the paper read more fluently, for example, Figure.4 has been mentioned earlier in the paper.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%