Learning Shortest Paths with Generative Flow Networks

In plain words

Finding the shortest route from one place to another is easy when you can see the whole map. It gets hard when the map is too big to store, as with a Rubik's cube, where the positions cannot all be kept in memory. The usual answer is to train a model that estimates how far each position is from the goal and to search with those estimates.

This paper takes a different route. It trains a model that does not estimate distances at all. Instead, it learns to walk backwards from any position toward the goal, choosing moves so that the walk is always as short as possible. The authors first prove that a model which keeps its expected walk length as small as possible can only put probability on shortest paths. They then turn the search problem into a training problem for a generative model that builds random walks, and train it with a loss that also pushes the total "flow" through the graph down.

At test time the learned walker is combined with beam search, which keeps a small set of the most promising partial paths at each step. On a synthetic permutation puzzle and on Rubik's cubes, the authors report solution lengths that are close to those of a strong specialised method, often with a smaller search budget.

TL;DR

• Learns a stochastic policy whose optimum is exactly the set of shortest paths to a goal in an unweighted directed graph §3.1 §3.2.
• Training uses trajectory balance on every prefix of short sampled trajectories, plus a flow regulariser weighted by lambda §3.3.1.
• Inference is beam search over the learned backward policy; width 1 is greedy decoding §3.4.
• Rubik's cube results are competitive with CayleyPy Cube and are reported with smaller beam widths for 2x2x2 and 4x4x4. All numbers are author-reported tab. 1.
• The public code (MIT) covers the 2x2x2, 3x3x3 and Swap tasks only. It does not contain the 4x4x4 setup or the baseline.

Problem

Classical search (Dijkstra's algorithm, A*) gives exact or near-exact paths when the graph can be explored and a good heuristic exists. Designing such heuristics is hard in high-dimensional spaces, and some state spaces are too large to store §1. Learning-based methods estimate how far each state is from the goal and use those estimates to steer a search at test time §1 §2.2. The paper keeps the aim modest: it does not claim to replace exact solvers where good heuristics exist, but to offer a route to problems where heuristics are unavailable or costly to design §2.2.

The method assumes a finite directed graph with unit edge costs and a path from every vertex to the goal vertex §3.2. Expected trajectory length of a non-acyclic GFlowNet controls sampling efficiency, and prior work did not analyse what minimising it implies §1 §2.1.

Method

The method has four parts: a reduction of shortest paths to a GFlowNet, a theorem about which policies minimise trajectory length, a training objective, and beam search at test time.

1. Reduction to a GFlowNet §3.2. Remove all edges leaving the goal vertex. States are the graph vertices plus a sink state s_f. The initial state s_0 is the goal. Transitions are the reversed graph edges, plus a transition from every state to s_f. The reward is uniform, R(s) = 1 on all non-sink states §3.2.

2. Trajectory length and flow §2.1 eq. 3. The expected trajectory length equals the normalised total state flow, so minimising it is the same as minimising total flow §2.1. Lemma 3.2 gives a lower bound on expected length for any backward policy with finite expectation eq. 4 §3.1. Lemma 3.3 shows a backward policy reaching that bound always exists eq. 5 §3.1.

3. Main theorem §3.1 (Theorem 3.4). A backward policy that meets reward matching minimises expected length if and only if every trajectory that ends in a terminal state x with length different from the shortest distance to x has probability zero. In words: the minimiser puts all its mass on shortest paths. The backward policy is not unique: at each state it may split probability arbitrarily among parents that lie on shortest paths §3.1.

4. Training §3.3.1 alg. 1. The forward policy starts at s_0 and samples fixed-length prefixes of N_max steps. The stop action is masked during these rollouts. Every prefix, together with a stop transition, is treated as a complete trajectory. The trajectory-balance residual is computed on every prefix and summed, and the flow regulariser is added for each prefix. Because each state has a transition to s_f, the flow is F(s) = R(s) / P_F(s_f | s), so no separate flow network is needed. The regulariser is not applied to the source state s_0 §3.3.1. The forward and backward policies share one network with separate linear heads app. D.

5. Beam search at test time §3.4. Start from the state to be solved and extend each beam entry with every backward action. Score entries by the sum of log backward probabilities along the path, keep the top W entries, and stop when s_0 is reached. Width 1 with the most probable action is greedy decoding §3.4. Duplicate states are removed at each step app. D.

The training objective, with the paper's notation, is:

\mathcal{L}_{\mathrm{regTB}}(\theta,\tau)=\sum_{i=0}^{N_{\max}}\Big\{\mathcal{L}_{\mathrm{TB}}(\theta,\tau_{0:i})+\frac{\lambda}{P_F(s_f\mid s_i,\theta)}\Big\} eq. 6

\mathcal{L}_{\mathrm{TB}}(\theta,\tau_{0:i})=\Big(\log\frac{P_F(s_f\mid s_i,\theta)\prod_{t=1}^{i}P_F(s_t\mid s_{t-1},\theta)}{\tfrac{1}{|V|}\prod_{t=1}^{i}P_B(s_{t-1}\mid s_t,\theta)}\Big)^2 §3.3.1

The normalising constant is Z = |V|, the number of states, because the reward is uniform. It is therefore not learned §3.3.1. The paper's eq. 6 sums from i = 0, but the text excludes the source flow; see Paper vs code.

The training step, in the paper's form:

# Flow-regularised trajectory balance step  [alg. 1] [§3.3.1]
# Input:  policy theta (shared PF and PB heads), lam, N_max
# Output: loss; gradients are averaged over a batch of rollouts  [alg. 1]

def regtb_loss(theta, start_state, lam, n_max):
    # Roll out N_max forward steps from s0 with the stop action masked  [§3.3.1]
    states, actions = sample_forward(theta, start_state, n_max)
    # states has n_max + 1 entries and states[0] is s0  [§3.3.1]
    log_pf_prefix = 0.0
    log_pb_prefix = 0.0
    loss = 0.0
    for i in range(n_max + 1):
        # Stop probability at the prefix end closes the trajectory  [§3.3.1]
        pf_stop = P_F_stop(states[i], theta)
        # Trajectory balance residual for the prefix; log Z = log |V|  [§3.3.1]
        residual = (log_Z + log_pf_prefix + log(pf_stop) - log_pb_prefix) ** 2
        loss += residual
        # Flow regulariser on every state except the source  [§3.3.1]
        if i > 0:
            loss += lam / pf_stop
        if i < n_max:
            # Extend prefix by one forward and one backward step  [§3.3.1]
            log_pf_prefix += log(P_F(actions[i], states[i], theta))
            log_pb_prefix += log(P_B(states[i], states[i + 1], theta))
    return loss

Beam search over the backward policy, with the paper's stopping rule:

# Beam search on the backward policy  [§3.4] [app. D]
# Input:  start state, beam width W, step budget T
# Output: length of the first completed path, or None

def beam_search(start_state, width, budget):
    # Each beam entry holds a state and its accumulated log-probability  [§3.4]
    beam = [(start_state, 0.0)]
    for step in range(budget):
        # Extend every entry with every backward action  [§3.4]
        candidates = [
            (apply(action, state), score + log_PB(action, state))
            for state, score in beam
            for action in backward_actions
        ]
        # Drop duplicate states before pruning  [app. D]
        candidates = remove_duplicate_states(candidates)
        # Keep the W highest-scoring entries  [§3.4]
        beam = top_k(candidates, width)
        # Stop when the goal s0 is reached  [§3.4]
        if any(state == goal for state, _ in beam):
            return step + 1
    return None

Hyperparameters

Values are as stated in the paper. Code values are in Paper vs code.

Results

Author-reported. Setup: Swap puzzle with n = 15 (about 1.3e12 states) and n = 20 (about 2.4e18 states), evaluated on 500 random permutations §4.1 fig. 2; 2x2x2, 3x3x3 and 4x4x4 Rubik's cubes against CayleyPy Cube, using its official implementation with its published settings and a model of matched size, keeping the better of the two §4.2 app. D.2. Table 1 reports average solution length and solve rate across beam widths. The paper's "x" marks are not defined in the paper; where the solve rate is non-zero the average-length cell is still "x", so its meaning there is unclear.

Source: tab. 1. The 2x2x2 optimum of 10.64 was checked by BFS, per the caption tab. 1.

Other reported results, all author-reported:

• Swap: after enough training, both greedy and width-4 beam search produce exact shortest paths for every test permutation, and the sampled policy is close to optimal on average §4.1 fig. 2. The expected backward walk length from a uniformly random state is, after training, within 0.6% (n = 15) and 2.4% (n = 20) of the optimum §3.3.2 fig. 2.
• Regulariser ablation: larger lambda gave better models until they failed to find valid paths. The paper's rule of thumb is to pick the largest lambda that still finds valid paths after a short run §4.3 fig. 3.
• 2x2x2: greedy decoding (W = 1) gives valid paths for the whole test set. The method needs a beam 16 times smaller than CayleyPy for optimal solutions §4.2.
• 4x4x4: valid paths for the whole test set with beam 2^12, where CayleyPy needs at least 2^18. The method is ahead at every width from 2^3 to 2^18 and slightly behind at 2^21 §4.2 tab. 1.
• Runtime on one H200: 3x3x3 at W = 2^18 takes 1.74 s with a 25M-parameter network, versus 6.19 s for CayleyPy with 4M parameters. 4x4x4 at W = 2^21 takes 50.7 s versus 330.6 s. The gap is attributed to one forward pass giving logits for all neighbours, where the baseline needs one pass per neighbour §4.2.
• Weighted graphs (appendix only): on a weighted Swap variant with n = 10, lambda = 1e-3, N_max = 20 and 50k steps, the model found all shortest paths at width 128 in a test set checked with Dijkstra's algorithm. The authors note it needed more search than the unweighted case app. C fig. 4.

Paper vs code

Code: repository at commit 4c2dc375f894b1259164a17ae56fec85524d2377, MIT licence, copyright held by the first author. The paper links it directly.

• Loss weighting. Trajectory balance is averaged over the N_max + 1 prefixes, while the regulariser is summed over prefixes [code: train.py#L239] train.py. Paper eq. 6 sums both. The code objective therefore equals, up to a constant factor 1/(N_max + 1), eq. 6 with lambda replaced by (N_max + 1) times lambda (and without the i = 0 term; see below). The appendix lambda values are only meaningful with the code's weighting. This is a derivation from the code, not a statement by the authors.
• Source flow. Eq. 6 sums from i = 0 eq. 6, but the text says the source flow is not regularised §3.3.1. The code excludes index 0 [code: train.py#L239].
• Swap N_max. Paper: 50 for both n = 15 and n = 20 app. D.1. README commands: 35 for n = 15 and 60 for n = 20 README.md [code: README.md#L58].
• 3x3x3 lambda. Paper: 1e-7 app. D.2. README: 5e-7 README.md.
• Iterations. Paper: 500,000 for 2x2x2 app. D.2; README uses 510,000 README.md. Paper: 1,000,000 for 3x3x3 app. D.2; README uses 2,000,000 README.md.
• Evaluation budget. Table 1 counts solves within 100 steps tab. 1. README evaluation commands use a 20-step budget for 2x2x2 and 30 for 3x3x3 README.md README.md. Solutions longer than the budget count as unsolved, so these commands do not reproduce the 100-step solve rate.
• Beam stopping rule. The paper stops when the goal is reached §3.4. The code runs for a fixed budget of eval_max_length steps and reports the shortest length among beams that reached s_0 train.py train.py.
• Forward action mask. The code masks some forward actions on top of the stop mask, which the paper does not mention. A forward action a is disabled in a state equal to the backward permutation for a [code: train.py#L163-L169]. For Swap, that state is a neighbour of the goal and applying a returns the goal. For Rubik's cubes, the comparison uses colours, dividing the permutation by 4 (2x2x2) or 9 (3x3x3) [code: train.py#L166]. The mask is applied when sampling train.py and in the loss train.py. Because PF is renormalised over the unmasked actions, the mask also changes the flow values 1 / P_F(s_f | s) that the regulariser uses.
• Architecture. The code adds an input linear layer before the six blocks and feeds x6 + x0 to the output heads train.py train.py. The paper describes only the six blocks app. D.
• Rubik states. The code represents a Rubik's cube by a colour per sticker, with the goal at index // 4 (2x2x2) or // 9 (3x3x3) train.py. A solved state is a colour match, not a specific permutation. The paper says states are permutations §4.
• Partition function. Z is hard-coded: 3,674,160 x 24 for 2x2x2, 43252003274489856000 for 3x3x3 and log(n!) for Swap train.py train.py train.py. Changing the state encoding requires changing these.
• 4x4x4. The code has tasks for rubik2, rubik3 and swap only [code: train.py#L527] train.py train.py. The 4x4x4 results in Table 1 cannot be reproduced from this repository.
• Not in the repository. The CayleyPy baseline, the weighted-graph experiment and the 4x4x4 test set are not included.

Pitfalls

• Lambda transfers badly. The paper's lambda spans 1e-2 to 1e-13 across tasks, because flows near the goal scale with |V| app. D.2. The code uses float32 by default and nothing in the repository enables 64-bit mode. Tune lambda per task with the paper's rule of thumb §4.3 fig. 3.
• Trajectory length is far shorter than solutions. N_max is 50 for Swap with n = 20, while random permutations there have about 95 inversions on average §4.1. The 4x4x4 solutions average more than 50 moves with N_max = 30 §4.1. Balance is only enforced within N_max steps of s_0, so the generalisation to longer solutions is empirical §3.3.2.
• The objective is an approximation. The regularised objective replaces the constraint with a penalty. For fixed lambda the minimiser need not satisfy balance exactly, and the bias is not quantified §3.3.2.
• Beam width and budget interact. The reported lengths depend on W and on the step budget, not only on the policy. The paper's 2x2x2 and 3x3x3 results use widths from 1 to 2^18 tab. 1. The code's min-over-beams reporting train.py can report a shorter length than the first arrival.
• The table's "x" is undefined. See Results. Compare cells with the solve rate, not just the average length.
• Training-time evaluation samples the test set. When the test set is larger than the evaluation batch, a random subset is drawn [code: train.py#L366-L368], so reported numbers can vary with the seed.
• Unpinned dependencies. requirements.txt lists jax, equinox and optax without versions. tb_loss uses jnp.cumulative_sum with include_initial train.py, so check the JAX version before reproducing.
• Checkpoints. The repository includes weights for rubik2 and rubik3 only README.md README.md. No Swap checkpoint is included. These were not evaluated here.
• Test sets were not inspected. The shapes of the .npy test sets were not checked against the sizes stated in the paper, because running Python was not approved in this session.
• Version. This card follows the metadata version (v2, updated 2026-10-13). The code is a single shallow clone at the commit above.

Lineage

• Builds on: Brunswic et al. 2024 (AAAI), "A Theory of Non-acyclic Generative Flow Networks", which introduced flow regularisation for non-acyclic GFlowNets §2.1. Morozov et al. 2025 (ICML), "Revisiting Non-Acyclic GFlowNets in Discrete Environments", which extended the theory §2.1. The paper changes two things: trajectories are bounded in length, and trajectory balance replaces detailed balance §3.3.
• Compared against: CayleyPy Cube (Chervov et al. 2025, NeurIPS, "A machine learning approach that beats Rubik's cubes") tab. 1. The text also names DeepCubeA (Agostinelli et al. 2019) and EfficientCube (Takano et al. 2023) as earlier baselines §4.2.
• Superseded by: not known from the content on this wiki.