Latent Spherical Flow Policy for Reinforcement Learning with Combinatorial Actions

AI summary

Overview

Research area: Reinforcement learning (RL) with combinatorial action spaces, combining generative modeling (spherical flow matching) with combinatorial optimization solvers.

Technical level: Advanced. The paper includes formal definitions, two stated theorems with proofs referenced to appendices, and assumes familiarity with Markov decision processes, flow matching, mixed-integer programming, and von Mises–Fisher distributions.

One-sentence scope: The paper introduces LSFlow, a method that learns a stochastic policy on a sphere of cost directions in a continuous latent space and delegates feasibility to a combinatorial optimization solver, with a smoothed Bellman operator for stable value learning.

What This Paper Is About

In many sequential decision problems the set of valid actions is not a simple list of numbers but a structured, constrained object such as a subset, a matching, or a route, and the number of feasible options grows exponentially with problem size. This makes it impractical to write down a policy over all valid actions directly. The paper's goal is to give combinatorial RL the expressive, multimodal stochastic policies that diffusion and flow models have brought to continuous control, while still guaranteeing that every action the policy emits is feasible.

Key Contributions

1. A solver-augmented spherical flow policy. The authors propose LSFlow, described as the first flow- or diffusion-based policy framework for RL with combinatorial action spaces. It learns a stochastic distribution over unit-norm cost directions on a sphere and passes each sample through a combinatorial optimization (CO) solver that returns a guaranteed-feasible structured action.
2. Efficient training in latent cost space with a smoothed Bellman operator. The critic is trained directly on (s, c) pairs so the solver is invoked once per environment step rather than repeatedly inside the policy update loop. To handle the piecewise-constant, discontinuous value landscape induced by the solver, the authors introduce a smoothed Bellman operator using a von Mises–Fisher (vMF) kernel and prove contraction and consistency properties.
3. Theory on expressivity, contraction, and consistency. Lemma 3.1 shows the solver mapping is positively scale invariant (only the direction of c matters). Proposition 3.2 shows that with a suitable distribution over the sphere, the solver-induced construction can represent any stochastic policy over the feasible set. Theorem 3.5 shows the smoothed operator is a γ-contraction with a unique fixed point that is infinitely differentiable in c. Theorem 3.6 shows the smoothed fixed point converges to the unsmoothed value function as the kernel concentration κ → ∞.
4. Empirical validation on public benchmarks and a real-world application. LSFlow is evaluated on four combinatorially constrained tasks from a public benchmark suite plus a real-world sexually transmitted infection (STI) testing application, with an accompanying code release.

Main Findings

• Overall reward improvement. LSFlow achieves the highest reward on all four public benchmark tasks, improving over the strongest baseline (SRL) by an average of 20.6%. The authors attribute this to the expressiveness of a stochastic policy class, in contrast to SRL's deterministic structured policy, which they say can limit representational capacity and exploration.
• Per-task rewards (Table 1, higher is better). Dynamic Scheduling: LSFlow 28.85 ± 1.48 vs. SRL 24.50 ± 1.09, SEQUOIA 24.00 ± 1.44, DQN-Sampling 16.89 ± 0.95, Greedy 15.12 ± 1.56, Random 10.12 ± 0.84. Dynamic Routing: LSFlow 28.51 ± 2.60 vs. SRL 25.32 ± 1.49, SEQUOIA 20.99 ± 1.92, DQN-Sampling 18.29 ± 1.49, Greedy 11.58 ± 0.46, Random 11.11 ± 0.36. Dynamic Assignment: LSFlow 35.93 ± 2.71 vs. SEQUOIA 32.99 ± 3.45, SRL 28.46 ± 0.82, DQN-Sampling 22.25 ± 2.27, Greedy 20.15 ± 1.00, Random 13.28 ± 0.80. Dynamic Intervention: LSFlow 17.21 ± 0.39 vs. DQN-Sampling 15.16 ± 0.25, SRL 13.36 ± 1.93, SEQUOIA 10.84 ± 0.89, Greedy 10.55 ± 0.67, Random 8.91 ± 0.08.
• Average reward and training time. Average reward: LSFlow 27.62, SRL 22.91, SEQUOIA 22.21, DQN-Sampling 18.15, Greedy 14.35, Random 10.85. Average training time in hours: DQN-Sampling 0.52, LSFlow 1.29, SRL 3.89, SEQUOIA 9.11 (Random and Greedy are not reported with a time).
• Computational efficiency. LSFlow is approximately 3.0× faster than SRL in average training time. The authors state this arises because SRL differentiates through the combinatorial solver.
• Fixed-point guarantees. Theorem 3.5 establishes that the vMF-smoothed Bellman operator is a γ-contraction on bounded functions with the sup norm, admits a unique fixed point, and produces a value function that is infinitely differentiable in the cost direction for every state.
• Consistency of smoothing. Theorem 3.6 shows that as κ → ∞, the smoothed fixed point converges to the unsmoothed value function for cost directions outside the solver-switching boundary B_s, and almost surely when the direction is drawn from the policy. Remark 3.7 notes that a global sup-norm convergence guarantee is not achievable in general, because the smoothed fixed point is C^∞ while the unsmoothed one is typically piecewise constant with jumps.
• Deliberate bias–variance trade-off. Because the smoothed operator does not share the fixed point of the standard Bellman operator, finite κ introduces bias. The authors frame this as a deliberate trade-off against variance reduction and say they study it empirically (Section 5.3).
• Flow matching preferred over diffusion. Remark 3.3 gives two reasons: spherical flow matching achieves substantially better performance than spherical diffusion models (as shown in Chen and Lipman, 2024), and spherical diffusion models are considerably more complex to implement in practice.
• Real-world STI testing. LSFlow is evaluated on a real-world STI testing application, shown in Figure 3. Specific numerical results for this application are not reported in the provided content.

Methodology in Plain English

The central move is to stop asking the policy network to output actions directly. Instead, the network outputs a direction in a continuous space, and a combinatorial optimization solver turns that direction into a feasible action.

Concretely, the method works in two stages. First, a flow model learns a state-conditional distribution over cost vectors c in R^m. Each cost vector defines a linear objective for a solver, which returns the feasible action minimizing c^T a over the constraint set A(s). Because scaling a cost vector by any positive number leaves the solver's answer unchanged, only the direction of c matters, so the policy is defined on the unit sphere S^(m-1). The flow is integrated with a projected ODE that keeps samples on the sphere, and the base distribution is fixed on the sphere (for example, uniform).

Training the policy uses a proposal-and-weight scheme. At each iteration, cost directions are sampled from the current policy, each proposal is weighted by an exponential of the value it induces, w(s,c) ∝ exp((1/λ) Q(s, a*(s,c))), and the flow model is refitted to these weighted samples. The authors connect this to KL-regularized policy improvement with a trust region, similar in spirit to PPO and TRPO.

The solver is expensive, so rather than calling it inside the policy update, the authors learn a critic directly on the latent cost space, Q̃_φ(s,c) := Q(s, a*(s,c)). This makes the policy update depend only on (s, c) pairs.

The remaining problem is that the map from cost vector to action is piecewise constant: the solver returns the same action throughout a region and then jumps at a boundary. This makes the value landscape jagged and destabilizes Bellman backups. The fix is to smooth over directions with a von Mises–Fisher kernel, K_κ(c̃|c) ∝ exp(κ c^T c̃), both when choosing the action for the current step and when forming the bootstrap target. In practice, the target averages the critic over J perturbed next-step directions. During data collection, the agent samples a center direction, perturbs it, executes the resulting solver action, and stores the center direction with the transition.

Why This Matters

Impact on research. The paper connects two lines of work that have largely developed separately: expressive generative policies (diffusion and flow matching, mostly used in continuous control) and solver-based RL for combinatorial actions. It provides a theoretical bridge by proving that a spherical distribution over cost directions is expressive enough to represent any stochastic policy over a finite feasible set, and that a smoothed Bellman operator restores smoothness while remaining consistent in the limit. This gives researchers a template for combining generative policy classes with hard constraints rather than relaxing them.

Real-world applications (drawn from the paper's task suite and application):

• Dynamic scheduling: repeatedly selecting a limited set of service requests and assigning them to feasible time windows, where completing or delaying service changes which requests remain.
• Dynamic routing: planning a bounded-length route on a graph (for example, delivery), where visited nodes determine both immediate reward and future states.
• Dynamic resource assignment: repeatedly matching limited-capacity resources to incoming demands, where today's allocation affects future availability.
• Dynamic intervention and public health: selecting a constrained subset of interventions to apply to a population, including the paper's real-world STI testing application.

Industry relevance. Any operation that must repeatedly pick a feasible structured decision under uncertainty is a candidate: logistics and delivery routing, workforce and service scheduling, matching marketplaces, and public health resource allocation. The reported computational profile is relevant here, since LSFlow's average training time (1.29 hours) is lower than SRL (3.89 hours) and SEQUOIA (9.11 hours) on the evaluated tasks. The solver bottleneck is addressed by amortizing solver calls through a latent-space critic, which matters for deployments where each solver call is costly.

Future Directions

• Quantifying the convergence rate. Theorem 3.6 establishes consistency as κ → ∞, but the authors explicitly state that making the convergence quantitative would require additional regularity of the policy near the solver-switching boundary, and they leave a precise rate to future work.
• Choosing κ in practice. Because the smoothed operator does not share the fixed point of the standard Bellman operator, finite κ introduces bias. The paper says this bias–variance trade-off is studied empirically in Section 5.3, leaving open how best to select or schedule κ in new domains.
• Reducing solver dependence further. The critic is learned in latent cost space specifically to avoid repeated solver calls during policy optimization, but the solver is still invoked once per environment step to execute an action. Further amortization, especially in domains where solving is the dominant cost, is a natural extension.
• Extending and comparing across combinatorial and discrete policy classes. The paper notes concurrent work (Ma et al., 2025b) on discrete diffusion policies for standard discrete action spaces, which does not directly enforce hard feasibility constraints. Comparing or unifying these approaches, and testing spherical flow policies on a wider range of combinatorial structures, remains open.

Target Audience

This paper is most useful to reinforcement learning researchers working on large or structured action spaces, and to researchers at the intersection of RL and combinatorial optimization who are interested in feasibility-by-design methods. It will also interest practitioners in public health, logistics, and operations who need policies that never emit infeasible decisions, and graduate students who want a worked example of pairing a generative policy with a solver plus the accompanying smoothness theory. Readers should be comfortable with MDP notation, ordinary differential equations, and constrained optimization; the truncated content does not include the appendices, so the full proofs and the complete ablation results are not visible here.

Note: this summary is based on the provided (truncated) paper content. Details such as the appendix proofs, the full hyperparameter settings, the Section 5.3 ablation numbers, and the numerical results of the STI testing application are referenced in the text but not reported in the content supplied.

Authors’ abstract

Reinforcement learning (RL) with combinatorial action spaces remains challenging because feasible action sets are exponentially large and governed by complex feasibility constraints, making direct policy parameterization impractical. Existing approaches embed task-specific value functions into constrained optimization programs or learn deterministic structured policies, sacrificing generality and policy expressiveness. We propose a solver-induced \emph{latent spherical flow policy} that brings the expressiveness of modern generative policies to combinatorial RL while guaranteeing feasibility by design. Our method, LSFlow, learns a \emph{stochastic} policy in a compact continuous latent space via spherical flow matching, and delegates feasibility to a combinatorial optimization solver that maps each latent sample to a valid structured action. To improve efficiency, we train the value network directly in the latent space, avoiding repeated solver calls during policy optimization. To address the piecewise-constant and discontinuous value landscape induced by solver-based action selection, we introduce a smoothed Bellman operator that yields stable, well-defined learning targets. Empirically, our approach outperforms state-of-the-art baselines by an average of 20.6\% across a range of challenging combinatorial RL tasks.