cs.LGMar 2, 2026

Learning Shortest Paths with Generative Flow Networks

Authors: Nikita Morozov, Ian Maksimov, Daniil Tiapkin, Sergey Samsonov

Organizations: HSE University · CMAP, CNRS, École polytechnique · LMO, Université Paris-Saclay

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

In this paper, we present a novel learning framework for finding shortest paths in graphs utilizing Generative Flow Networks (GFlowNets). First, we examine theoretical properties of GFlowNets in non-acyclic environments in relation to shortest paths. We prove that, if the total flow is minimized, forward and backward policies traverse the environment graph exclusively along shortest paths between the initial and terminal states. Building on this result, we reduce shortest-path finding in an arbitrary unweighted graph to finding a non-acyclic GFlowNet with minimal total flow, and approximate it in practice by training with flow regularization. We experimentally demonstrate the performance of our method in pathfinding in permutation environments and in solving Rubik's Cubes of various sizes. For the latter problems, our approach shows competitive results with state-of-the-art machine learning approaches designed specifically for this task in terms of the solution length, while requiring smaller search budget at test-time.

Figures & tables

Figure 1: Visualization of the GFlowNet environment in the proposed construction using Rubik’s Cube as an example. Forward transitions are denoted by blue arrows, backward ones are denoted by red arrows. Note that backward transitions from the goal state s0s_{0} are removed in our construction — it is an absorbing state with respect to PB\mathcal{P}_{\mathrm{B}} , while there are forward and backward transitions in both directions between s1s_{1} and s3s_{3} . Each state has a forward transition to sfs_{f} , corresponding to the stop action, while PB(s∣sf)\mathcal{P}_{\mathrm{B}}(s\mid s_{f}) corresponds to the reward distribution, which we set as uniform.
Figure 2: The results for Swap Puzzle n=15n=15 and n=20n=20 for three evaluation protocols: 1) the path is faithfully sampled from PB\mathcal{P}_{\mathrm{B}} , 2) the path is obtained by a greedy evaluation of PB\mathcal{P}_{\mathrm{B}} , 3) the path is generated with beam search of a small width W=4W=4 . The results for greedy evaluation are not presented for checkpoints where it failed to find a valid solution for every example in the test set. For n=20n=20 , 100k training iterations took 15 minutes to run on an NVIDIA H200 GPU.
Figure 3: Effect of the regularization coefficient λ\lambda on the obtained solution lengths on Swap Puzzle n=15n=15 and 2x2x2 Rubik’s cube. In both cases, choosing a larger λ\lambda led to models that fail to find any valid paths.
 2x2x2 Rubik’s Cube 3x3x3 Rubik’s Cube 4x4x4 Rubik’s Cube
 OursCayleyPy OursCayleyPy OursCayleyPy
Beam size WWAvg. lengthSolve rateAvg. lengthSolve rateBeam size WWAvg. lengthSolve rateAvg. lengthSolve rateBeam size WWAvg. lengthSolve rateAvg. lengthSolve rate
202^{0}11.621.00x0.00202^{0}x0.471x0.000232^{3}x0.07x0.00
222^{2}11.241.00x0.00232^{3}x0.984x0.001262^{6}x0.27x0.00
242^{4}10.781.00x0.00262^{6}25.331.000x0.687292^{9}x0.74x0.00
262^{6}10.641.00x0.06292^{9}23.491.00024.341.0002122^{12}64.091.00x0.16
Table 1: Comparison with CayleyPy Cube ( Chervov et al.,, 2026 ) . We present average lengths of solutions found by both methods for different beam search budgets. Solve rate corresponds to the fraction of examples in the test set for which the model has found a valid path to the goal within 100 steps. For the 2x2x2 Rubik’s Cube, the average length of 10.64 is optimal for the test set, validated by the BFS algorithm. We focus on comparison to CayleyPy Cube as it was shown to thoroughly outperform previous methods in both solution lengths and runtime efficiency. For reference, the primary established RL-based approach DeepCubeA obtains the avg. length of 21.50 in 3x3x3 cube with the search budget of W=218W=2^{18} .
Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Weighted Swap Puzzle results for n=10n=10 .
More in this topic

Graph-Based Planning

Graph-Based Planning represents states, tasks, or environments as nodes and edges to search for feasible, low-cost action sequences.

See all 99 papers

Explore similar work

CardsList
Jun 4, 2026

Your GFlowNet Secretly Learns an Optimal Transport Plan

Generative Flow Networks (GFlowNets) are a framework for sampling structured objects via stochastic trajectories in a directed graph. In this work, we establish a theoretical connection between non-acyclic GFlowNets and optimal transport (OT). We show that fixing the initial flow distribution in a minimum-flow GFlowNet reduces its objective to a Kantorovich OT problem with graph-induced shortest path costs. At the optimum, the learned GFlowNet policy therefore encodes an optimal transport plan from the source distribution to the target distribution: we show that sampling trajectories from the minimum-flow GFlowNet recovers the corresponding optimal coupling. Our formulation enables applying the GFlowNet learning framework to OT problems on large graphs via edge flows and neural parameterization. Experiments confirm agreement with exact OT solvers and demonstrate that GFlowNets can learn high-quality transport plans.

Ian Maksimov, Nikita Morozov, Denis Belomestny +1

1HSE University · 2Duisburg-Essen University

Generative Flow NetworksFlow-Based Generative ModelingOptimal Transport
AskPDF
Oct 4, 2023

Expected flow networks in stochastic environments and two-player zero-sum games

Generative flow networks (GFlowNets) are sequential sampling models trained to match a given distribution. GFlowNets have been successfully applied to various structured object generation tasks, sampling a diverse set of high-reward objects quickly. We propose expected flow networks (EFlowNets), which extend GFlowNets to stochastic environments. We show that EFlowNets outperform other GFlowNet formulations in stochastic tasks such as protein design. We then extend the concept of EFlowNets to adversarial environments, proposing adversarial flow networks (AFlowNets) for two-player zero-sum games. We show that AFlowNets learn to find above 80% of optimal moves in Connect-4 via self-play and outperform AlphaZero in tournaments.

Marco Jiralerspong, Bilun Sun, Danilo Vucetic +4

Mila – Qu´ebec AI Institute, Universit´e de Montr´eal

Game-Playing AgentsZero-Sum GamesGenerative Flow NetworksFlow-Based Generative ModelingProtein Design
AskPDF
Apr 17, 2026

Sample Complexity Bounds for Stochastic Shortest Path with a Generative Model

We study the sample complexity of learning an εε-optimal policy in the Stochastic Shortest Path (SSP) problem. We first derive sample complexity bounds when the learner has access to a generative model. We show that there exists a worst-case SSP instance with SS states, AA actions, minimum cost cmin⁡c_{\min}, and maximum expected cost of the optimal policy over all states B⋆B_{\star}, where any algorithm requires at least Ω(SAB⋆3/(cmin⁡ε2))Ω(SAB_{\star}^3/(c_{\min}ε^2)) samples to return an εε-optimal policy with high probability. Surprisingly, this implies that whenever cmin⁡=0c_{\min} = 0 an SSP problem may not be learnable, thus revealing that learning in SSPs is strictly harder than in the finite-horizon and discounted settings. We complement this lower bound with an algorithm that matches it, up to logarithmic factors, in the general case, and an algorithm that matches it up to logarithmic factors even when cmin⁡=0c_{\min} = 0, but only under the condition that the optimal policy has a bounded hitting time to the goal state.

Jean Tarbouriech, Matteo Pirotta, Michal Valko +1

Facebook AI Research Paris & Inria Lille · Facebook AI Research Paris · DeepMind Paris

Markov Decision ProcessesPolicy LearningSample-Efficient RL
AskPDF