Paper: Expected flow networks in stochastic environments and two-player zero-sum games (ICLR-2024)

Reviewer comment (Novelty):
It is mentioned that "We propose expected flow networks, ... generalizing GFlowNets on tree-structured state spaces." In the simpler case of trees instead of non-tree DAGs, why don't you just apply maximum entropy reinforcement learning (MaxEnt RL) approaches? It has been extensively demonstrated in previous GFlowNets papers that when the DAG degenerates to a tree, it is indeed equivalent to MaxEnt RL approaches (e.g., Soft Q-learning), which can natually handle stochastic worlds. Therefore, it makes little sense for only EFlowNets to be only evaluated in tree-structured state spaces.

Author response:
You are right that GFlowNets = MaxEnt RL in **deterministic** environments, but this is not the case in stochastic environments. To demonstrate this point, we train a soft version of DQN using the exact same architecture as AFlowNets on tic-tac-toe with self-play. After a non-negligible amount of hyperparameter tuning, we obtain an agent that learns a meaningful policy (winning >99% of the time against a uniform agent) but is still noticeably worse than AFlowNets and AlphaZero both in terms of training time and the policy it seems to converge to (worse % of optimal moves and still loses occasionally to a uniform agent). See the new results in Appendix E.2.

Reviewer comment (Evaluation):
The performance and explanation of EFlowNets does not make sense. As described in the paper, this paper aims to extend (Pan et al., 2023; Bengio et al., 2023) to handle more general cases considering stochasticity in the transition dynamics, i.e., extend GFlowNets -- which aims to sample proportionally to a reward function \( R(x) \), i.e., \( \pi(x) \propto R(x) \), to the settings with stochastic transition dynamics. However, the solution learned by EFlowNets is incorrect as demonstrated in Figure 1(a). Although it is able to solve \( p \) (with a solution of around \( 0.66 \)) in this case while \( p \) cannot find satisfiable solutions in Stoch-GFlowNets, it leads to sampling \( x_1 \) with probability \( p(x_1) \approx 0.33 \neq 2/10 \), \( p(x_2) \approx 0.33 \neq 3/10 \), \( p(x_3) \approx 0.306 \neq 1/10 \), and \( p(x_2) \approx 0.034 \neq 4/10 \). Although \( p \) is solvable in EFlowNets, the learned sampling policy fails to realize the goal of GFlowNets (sampling proportionally from the reward function), and learns a meaningless sampling policy. Consequently, EFlowNets do not address the limitations of Stoch-GFlowNets in some cases where satisfiable solutions are unavailable. However, this claim is made throughout the paper, which is an overstatement in this context.

Author response:
This is precisely the problem in stochastic environments: **learning a policy that samples proportionally to reward is in general impossible**. As illustrated in Fig. 1(a), even in an incredibly small tree, it is impossible to sample proportionally to reward (as you cannot change the environment dynamics). StochGFNs aim to sample from the reward distribution, but this is not always possible, as we show in that example. (This issue with StochGFNs has as a consequence that the global optimum may depend on the training policy or not be unique – it may be impossible to achieve zero loss.) 

**EFlowNets solve a fundamentally different problem** – sampling from the expectation over the environment’s transitions. This is one reason we invented a new name (**expected** flow networks), rather than calling them variants of GFlowNets, as generative flow networks should solve the generative modeling problem of sampling from the reward distribution.

The advantage of the EFlowNet formulation is *satisfiability*: in the overwhelming majority of cases where sampling proportional to reward is not possible, the EFlowNet constraints are still satisfiable and the set of global optima of the loss does not depend on the choice of full-support training policy (as well as other benefits mentioned in the paper).

In addition, we would like to politely push back against the assertion that the learned policy is meaningless. We believe considering your expected reward from a state is an intuitive way to deal with random transitions and show connections between EFlowNets and marginalized quantal response policies (see Appendix A).

Reviewer comment (Theory):
The statement seems incorrect -- "Stochastic GFlowNets directly apply the training algorithms applicable to deterministic GFlowNets (e.g., DB) to the augmented DAG \( G \), with the only modification being that the forward policy \( P_F \) is free to be learned only on agent edges, while on environment edges it is fixed to the transition function." In stochastic environments, the transitions dynamic is unknown, and after I checked the Stochastic GFlowNets paper, I found that it indeed learns the forward policy and the underlying transition dynamics (which is not fixed, unknown and is learned by experiences).

Author response:
We agree this sentence is misleading and it will be amended. The point being made is that they use the deterministic GFlowNet algorithms with a learned policy for agent edges and a transition function for environment edges. This transition function can either be taken as given or learned through some maximum likelihood objective. Similarly, EFlowNets can either take the transition function as given or learn it (in the TFBind task, it is learned, just as in StochGFNs).

Reviewer comment (Theory):
In the "Violated desiderata in stochastic GFlowNets" section, a very important property is missed (which is satisfied by Stochastic GFlowNets with solvable cases but EFlowNets unfortunately violated) -- the correct sampling behavior, which means that GFlowNets function well in stochastic environments which sample proportionally to the reward function. Failures to satisfy this basic and fundamental requirement will lead to meaningless solutions.

Author response:
This is precisely the problem in stochastic environments: **learning a policy that samples proportionally to reward is in general impossible**. As illustrated in Fig. 1(a), even in an incredibly small tree, it is impossible to sample proportionally to reward (as you cannot change the environment dynamics). StochGFNs aim to sample from the reward distribution, but this is not always possible, as we show in that example. (This issue with StochGFNs has as a consequence that the global optimum may depend on the training policy or not be unique – it may be impossible to achieve zero loss.) 

**EFlowNets solve a fundamentally different problem** – sampling from the expectation over the environment’s transitions. This is one reason we invented a new name (**expected** flow networks), rather than calling them variants of GFlowNets, as generative flow networks should solve the generative modeling problem of sampling from the reward distribution.

The advantage of the EFlowNet formulation is *satisfiability*: in the overwhelming majority of cases where sampling proportional to reward is not possible, the EFlowNet constraints are still satisfiable and the set of global optima of the loss does not depend on the choice of full-support training policy (as well as other benefits mentioned in the paper).

In addition, we would like to politely push back against the assertion that the learned policy is meaningless. We believe considering your expected reward from a state is an intuitive way to deal with random transitions and show connections between EFlowNets and marginalized quantal response policies (see Appendix A).

Reviewer comment (Novelty):
EFlowNets bear great similarity to (Yang et al., 2023), however, the latter aims to extend Decision Transformer to stochastic environments in the realm of RL. Taking the expectation is correct in RL, since the objective of RL is to maximize the expected discounted future rewards. In addition, EFlowNets also bear similarity to Section 3.3.2 about stochastic rewards in (Bengio et al., 2023), which is mentioned about stochasticity in the rewards (not the transition dynamics), which is a case studied in (Zhang et al., 2023). Therefore, it does not make sense to apply the methodology about learning an expectation as in RL (Yang et al., 2023) to GFlowNets with the hope to tackle stochasticity in transition dynamics as discussed in (Bengio et al., 2023).

Author response:
We are not sure about why the referenced works imply it is inappropriate to model uncertainty in transition dynamics in the way we do. Actually, [Bengio et al.] already mention that the modeling in [Pan et al., 2023] that aims to sample proportionally to reward is not in general satisfiable, which is one consideration that motivates us to consider a different sampling problem in stochastic environments – the one solved by EFlowNets.

Reviewer comment (Experiments):
It is mentioned that trajectory balance is used for branch-adjusted AFlowNets. However, most practical games in multi-agent systems are actually partially observable (e.g., StarCraft II), which renders stochasticity in the transition dynamics. TB leads to large variance as mentioned in the text and in experiments in (Madan et al., 2023). Therefore, it seems to be more appropriate to apply this with DB or SubTB (since one of the claims for SubTB is to improve credit assignment).

Author response:
We thank the reviewer for pointing out the potential flaws of TB, but we found that the improved credit assignment of the TB objective already yielded strong results (in spite of the potential variance) when compared to the EDB constraints. (We also note that our TB for AFlowNets is distinct from the TB in GFlowNets.) 

This parallels the development of GFlowNets, where TB was first proposed as an improvement over DB for improved credit assignment [Malkin et al., 2022] and then SubTB was proposed for problems where TB has excessively high variance [Madan et al., 2023]. We agree that SubTB could potentially help further, especially in problems with longer trajectories, and believe this is an interesting future area of research!

Reviewer comment (Theory):
It is mentioned that "This data defines a fully observed sequential game", which can limit its practical applicability in a wider range of applications with partial observation.

Author response:
We agree this is a limitation of the work, but the fully observable setting is widely studied in the RL / adversarial games literature. We already mention the extension to partially observable states as possible future work in the conclusion and continue to think this is an interesting future direction.

Reviewer comment (Novelty):
A very relevant paper extending GFlowNets to the multi-agent setting (Li et al., 2023) is not cited and discussed. Since both papers study a very similar setting, it is worth comparing the approach with FCN and CFN in (Li et al., 2023), or at least discussing thoroughly. "Generative Multi-Flow Networks: Centralized, Independent and Conservation. Yinchuan Li, Haozhi Wang, Shuang Luo, Yunfeng Shao, Jianye Hao. 2023."

Author response:
We thank the reviewer for pointing this paper out. The paper in question ([a rejected ICLR’23 submission](https://openreview.net/forum?id=OTIhUlChVaT)) discusses a setting that is not applicable to turn-taking zero-sum games, as it considers multiple agents taking actions simultaneously. We can mention this paper in the related work, but its relevance to our problems is quite limited and experimental comparison is not appropriate.

Reviewer comment (Experiments):
It is unclear why experiments in Section 4.1 employ very large \( \beta = 10, 30 \) which is different from (Jain et al., 2022), which leads to the case of a tree with very peaky rewards. In addition, it is worth investigating how the approach behaves under the traditional \( L_1 \) error metric, which measures how well EFlowNets learn compared with other approaches.

Author response:
In the StochGFN paper, the beta parameter was artificially set to induce a less flat sampling distribution. We experiment with larger betas **in addition to** to the baseline value as this makes the problem of sampling from the reward distribution unsatisfiable (since addition of randomness induces smoothness in the sampling distribution, while the reward is peaky).

Reviewer comment (Theory):
Why the two-play games correspond to a tree instead of a graph? There can be many parent states for a state actually.

Author response:
Indeed most two-player games are graphs instead of trees; however, any DAG can be converted to a tree-structured environment by augmenting states with history. See the end of Section 2.1 for further details.

Reviewer comment (Writing):
The paper still requires solid background knowledge on generative flow networks and it is not very friendly to non-expert.

Author response:
> Paper is not friendly to non-experts.

We concede that the paper is fairly dense, but believe we have given the minimal background on GFlowNets that is required for understanding our generalizations. We would be glad to consider any recommendations which would make the paper more beginner-friendly for those less familiar with GFlowNets!

Reviewer comment (Evaluation):
From Figure 3, it seems that the proposed method is much more computation heavy, can you show the computation time required for each step of each method? Also how does the computation time scale with the size of the problem for each method?

Author response:
> From Figure 3, it seems that the proposed method is much more computation heavy, can you show the computation time requires for each step of each method? Also how does the computation time scale with the size of the problem for each method?

It is important to compare the adversarial algorithms under a common framework. When developing the training scripts, we took inspiration from AlphaZero-like methods [1]. In this training framework, an iteration is defined by the following steps:
1. Collect training data by interacting with the environment. AlphaZero employs MCTS to construct training examples; we can instead just sample whole trajectories so long as the sampling policy has full support.
2. Train the neural networks on the training data for a number of epochs. (Comparing epoch counts is somewhat misleading as AlphaZero trains policy and value networks on single transitions, while we train a policy network and scalar $\log Z$ on whole trajectories.)
3. Evaluate performance. Pit agents against baselines or old agents [1].

AlphaZero and AFlowNets are trained for 100 iterations each. The data points in Figure 3 represent the performance measured every five or so iterations. An important note on the seeming performance differential: AlphaZero collects a small number of examples on each iteration and adds them to a buffer [1, 2]. Our implementation did not do this, instead we sample about 10 times more examples on every iteration and fully replace our old examples. Also note that our implementation collects whole trajectories, not individual transitions. As such, the data collection process is quite computationally expensive, but can be made more efficient in future implementations by implementing a replay buffer similar to AlphaZero.

The presence of a buffer alone, however, does not fully capture the performance differential. An important factor to consider is the use of MCTS in AlphaZero, a recursive algorithm which runs 25 or more game simulations for each step in a trajectory [2]. The more complex the game, the more simulations need to be run to maintain good performance. 

Assuming that the number of trajectories collected for our method and the number of transitions collected for AlphaZero (once the buffer is full) are equal, the training process takes roughly the same amount of time to complete.

We can roughly state the worst-case computational complexity of AlphaZero data collection as $O(NMT)$. We can roughly state the computational complexity of AFlowNet data collection as $O(NT)$. The variables are:
- $N$: number of examples to collect,
- $M$: number of MCTS simulations,
- $T$: maximum length of a trajectory.

The computational complexity of collecting data for AlphaZero is expected to grow substantially with more complex games as $M$ must be increased. We believe that the computational complexity of training given a similar number of training examples is similar for both methods. The computational complexity of evaluation should also be relatively similar. As such, we expect the computational complexity of data collection to dominate. We list data collection times below for various settings of $M$ and $N$ in tic-tac-toe. $T = 9$ in the worst case.
- AlphaZero: 17 seconds to collect examples, 4 seconds to train for 10 epochs. $M=25$, $N = 5224$.
- AlphaZero: 24 seconds to collect examples, 4 seconds to train for 10 epochs. $M=50$, $N = 5736$.
- AlphaZero: 42 seconds to collect examples, 4 seconds to train for 10 epochs. $M=100$, $N = 6216$.
- AlphaZero: 72 seconds to collect examples, 4 seconds to train for 10 epochs. $M=200$, $N = 5864$.
- AFlowNet: 8 seconds to collect examples, 4 seconds to train for 10 epochs. $N = 6144$.

It is clear from the above that, for the same number of examples, the AFlowNet data generation process is quicker than AlphaZero’s data generation process. Furthermore, the training time is roughly similar between the two methods. We did not directly compare evaluation time because the evaluation process is quite different between the two methods, and serves more as a benchmark of progress rather than a necessary part of training. Thus, as the game gets more complex, we expect AFlowNets to have a faster data generating process than AlphaZero, allowing our method to train on more data in the same amount of time.

Reviewer comment (Writing):
The differences between stochastic environments and common environments for GFlowNet may not be stated clearly.

Author response:
We would like to thank the reviewer for their insightful comments on the paper. We would like to quickly articulate the novelty and power of our constructions for applications in stochastic and adversarial settings. Existing GFlowNet literature is relatively sparse when it comes to environment stochasticity. Those methods that do exist emphatically stick to the deterministic GFlowNet construction despite this construction providing undesirable sampling behavior or unsatisfiable constraints, which are not appropriate to our settings (see Figure 1(a)). 

In this paper, we depart from the deterministic constructions of GFlowNets in favor of a new paradigm of flow networks: Expected Flow Networks (EFlowNets) and Adversarial Flow Networks (AFlowNets). Our new constraints, expected detailed balance (EDB) and especially the branch-adjusted trajectory balance (TB), are an algorithmic novelty. They have nontrivial proofs of correctness, are grounded in the notions of quantal response policies (Appendix A), and allow for the first time the inference of stochastic equilibria in zero-sum games without self-play.

We respond to your questions and concerns in the following paragraphs.

> Can you provide more clear description of the differences between stochastic environments and common environments?

In common (deterministic) environments, a transition from one state to another is wholly determined by the action taken at the starting state. For example, if we have a state $s_0$ with two corresponding actions: $a_1$ and $a_2$; which transition to states $s_1$ and $s_2$ respectively, by committing action $a_2$, the agent is certain to transition to $s_2$. Under our notation, this corresponds to the environmental transition $P_{\text{env}}(s_2|(s_0,a_2)) = 1$. In other words: when we take an action, we know where we will end up. 

In contrast, stochastic environments feature stochasticity in their transitions. If we extend the above example so that $P_{\text{env}}(s_2|s_0) \leq 1$, the action $a_2$ may not lead to $s_2$ every time, depending on the environment’s transition function. In AFlowNets, the environment’s transition function is seen as an adversary’s policy, $P_{\text{adv}}(-)$, **allowing us to model adversarial games as a special type of stochastic environment**.

The distinction between deterministic and stochastic environments is important for GFlowNets in particular since the classical GFlowNet theory assumes that transitions are deterministic (constructive actions applied to an object give known results).

Reviewer comment (Writing):
Can you provide more clear description of the differences between stochastic environments and common environments?

Author response:
We would like to thank the reviewer for their insightful comments on the paper. We would like to quickly articulate the novelty and power of our constructions for applications in stochastic and adversarial settings. Existing GFlowNet literature is relatively sparse when it comes to environment stochasticity. Those methods that do exist emphatically stick to the deterministic GFlowNet construction despite this construction providing undesirable sampling behavior or unsatisfiable constraints, which are not appropriate to our settings (see Figure 1(a)). 

In this paper, we depart from the deterministic constructions of GFlowNets in favor of a new paradigm of flow networks: Expected Flow Networks (EFlowNets) and Adversarial Flow Networks (AFlowNets). Our new constraints, expected detailed balance (EDB) and especially the branch-adjusted trajectory balance (TB), are an algorithmic novelty. They have nontrivial proofs of correctness, are grounded in the notions of quantal response policies (Appendix A), and allow for the first time the inference of stochastic equilibria in zero-sum games without self-play.

We respond to your questions and concerns in the following paragraphs.

> Can you provide more clear description of the differences between stochastic environments and common environments?

In common (deterministic) environments, a transition from one state to another is wholly determined by the action taken at the starting state. For example, if we have a state $s_0$ with two corresponding actions: $a_1$ and $a_2$; which transition to states $s_1$ and $s_2$ respectively, by committing action $a_2$, the agent is certain to transition to $s_2$. Under our notation, this corresponds to the environmental transition $P_{\text{env}}(s_2|(s_0,a_2)) = 1$. In other words: when we take an action, we know where we will end up. 

In contrast, stochastic environments feature stochasticity in their transitions. If we extend the above example so that $P_{\text{env}}(s_2|s_0) \leq 1$, the action $a_2$ may not lead to $s_2$ every time, depending on the environment’s transition function. In AFlowNets, the environment’s transition function is seen as an adversary’s policy, $P_{\text{adv}}(-)$, **allowing us to model adversarial games as a special type of stochastic environment**.

The distinction between deterministic and stochastic environments is important for GFlowNets in particular since the classical GFlowNet theory assumes that transitions are deterministic (constructive actions applied to an object give known results).

Reviewer comment (Novelty):
What is the motivation of extending GFlowNet to stochastic environments and adversarial environments? And what is the general advantage of GFlowNet-based methods compared to RL methods in both environments?

Author response:
> What is the motivation of extending GFlowNet to stochastic environments and adversarial environments? And what is the general advantage of GFlowNet-based methods compared to RL methods in both environments?

Since GFlowNets are theorized as operating in deterministic environments, it is natural to extend them to stochastic environments. Some interesting examples of real-world stochasticity are 1) demand-response for electricity grids where the variation of the energy supply comes from the weather [1], 2) market making for securities [2], and 3) supply chain management [3].

The main motivation for extending GFlowNets specifically to stochastic and adversarial environments is their ability to sample a diverse set of high-reward objects. In stochastic environments, our goal is to maintain a diverse sampling of high-reward outcomes despite the inherent unpredictability. In the case of adversarial environments, we aim to produce agents which can play with a variety of winning strategies, rather than overfitting to a single mode that produces high rewards but is also exploitable. In both cases, it is natural to sample states according to their expected reward rather than the raw reward as there is no longer any guarantee that this reward will be reached given a set of actions taken in an environment.

Both of these aims are in stark contrast to the usual paradigm of reinforcement learning, which is to maximize rewards. Existing literature has demonstrated that GFlowNets are capable of sampling a more diverse set of solutions than comparable RL methods [4, 5].

Reviewer comment (Novelty):
How do you think about the potential of AFlowNet compared to RL methods considering the limitation in the proposed algorithms currently?

Author response:
> How do you think about the potential of AFlowNet compared to RL methods considering the limitation in the proposed algorithms currently?

We believe that our method has considerable advantages in zero-sum games such as tic-tac-toe and Connect-4, especially compared to standard RL methods (e.g., PPO, Deep Q-Learning). The latter struggle to learn strong policies through self-play in these multi-agent settings. To demonstrate this, after significant tuning, we manage to train a DQN and SoftDQN agent (through self-play) on tic-tac-toe that can fairly consistently beat the uniform agent. Even so, they take considerably longer to train and converge to a suboptimal solution. **We have added results with DQN and SoftDQN agents to Appendix E.2.**