Learning Shortest Paths with Generative Flow Networks

Nikita Morozov; Ian Maksimov; Daniil Tiapkin; Sergey Samsonov

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

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 $s_{0}$ are removed in our construction — it is an absorbing state with respect to $\mathcal{P}_{\mathrm{B}}$ , while there are forward and backward transitions in both directions between $s_{1}$ and $s_{3}$ . Each state has a forward transition to $s_{f}$ , corresponding to the stop action, while $\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=15

and

n=20

for three evaluation protocols: 1) the path is faithfully sampled from

\mathcal{P}_{\mathrm{B}}

, 2) the path is obtained by a greedy evaluation of

\mathcal{P}_{\mathrm{B}}

, 3) the path is generated with beam search of a small width

W=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=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=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

Ours

CayleyPy

Ours

CayleyPy

Ours

CayleyPy

Beam size

W

Avg. length

Solve rate

Avg. length

Solve rate

Beam size

W

Avg. length

Solve rate

Avg. length

Solve rate

Beam size

W

Avg. length

Solve rate

Avg. length

Solve rate

2^{0}

11.62

1.00

0.00

2^{0}

0.471

0.000

2^{3}

0.07

0.00

2^{2}

11.24

1.00

0.00

2^{3}

0.984

0.001

2^{6}

0.27

0.00

2^{4}

10.78

1.00

0.00

2^{6}

25.33

1.000

0.687

2^{9}

0.74

0.00

2^{6}

10.64

1.00

0.06

2^{9}

23.49

1.000

24.34

1.000

2^{12}

64.09

1.00

0.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=2^{18}

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Weighted Swap Puzzle results for

n=10

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

Your GFlowNet Secretly Learns an Optimal Transport Plan
Jun 4, 2026Ian Maksimov, Nikita Morozov, Denis Belomestny +1Generative Flow Networks Flow-Based Generative Modeling
1HSE University · 2Duisburg-Essen University
Expected flow networks in stochastic environments and two-player zero-sum games
Oct 4, 2023Marco Jiralerspong, Bilun Sun, Danilo Vucetic +4Game-Playing Agents Zero-Sum Games
Mila – Qu´ebec AI Institute, Universit´e de Montr´eal
Sample Complexity Bounds for Stochastic Shortest Path with a Generative Model
Apr 17, 2026Jean Tarbouriech, Matteo Pirotta, Michal Valko +1Markov Decision Processes Policy Learning
Facebook AI Research Paris & Inria Lille · Facebook AI Research Paris · DeepMind Paris

Jun 4, 2026J/K move · Enter open · S save

Your GFlowNet Secretly Learns an Optimal Transport Plan

Ian Maksimov, Nikita Morozov, Denis Belomestny, Sergey Samsonov

1HSE University · 2Duisburg-Essen University

Generative Flow Networks Flow-Based Generative Modeling Optimal Transport

Open paper AskPDF arXiv

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.

Abstract

Graph-Based Planning

Explore similar work

Your GFlowNet Secretly Learns an Optimal Transport Plan

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

Sample Complexity Bounds for Stochastic Shortest Path with a Generative Model