Latent Spherical Flow Policy for Reinforcement Learning with Combinatorial Actions

Technical Summary: Latent Spherical Flow Policy for Reinforcement Learning with Combinatorial Actions

1. Problem Definition

The paper addresses the challenge of applying Reinforcement Learning (RL) to environments with combinatorial action spaces. In these settings, actions are structured decisions (e.g., subset selection, routing, scheduling) governed by hard feasibility constraints. The feasible set of actions, $\mathcal{A}(s)$, typically grows exponentially with problem size, making direct policy parameterization over all valid actions intractable.

Existing approaches face a trade-off:

• Solver-based value optimization: Embeds value networks into Mixed-Integer Programs (MIPs) to extract greedy actions. This often requires specific value architectures and scales poorly with network complexity.
• Optimization-induced decision layers: Learns structured policies end-to-end but typically yields deterministic policies, limiting expressiveness and exploration in complex landscapes.

While modern generative models (diffusion and flow-based) have shown success in continuous control by representing rich, multimodal stochastic distributions, they cannot be directly applied to combinatorial RL because they do not inherently enforce discrete feasibility constraints. Naive adaptation leads to infeasible actions.

2. Methodology: LSFLOW

The authors propose LSFLOW (Latent Spherical Flow Policy), a framework that decouples policy expressiveness from feasibility enforcement. The core idea is to learn a stochastic policy in a compact continuous latent space and delegate feasibility to a combinatorial optimization (CO) solver.

A. Solver-Induced Spherical Flow Policy

• Latent Representation: Instead of outputting actions directly, the policy learns a distribution over cost vectors $\mathbf{c} \in \mathbb{R}^m$.
• Solver Mapping: A CO solver maps a sampled cost vector to a feasible action via a linear objective: $\mathbf{a}^*(s, \mathbf{c}) = \arg\min_{\mathbf{a} \in \mathcal{A}(s)} \mathbf{c}^\top \mathbf{a}$.
• Spherical Constraint: Due to the positive scale invariance of the solver (the optimal action depends only on the direction of $\mathbf{c}$, not its magnitude), the latent space is restricted to the unit sphere $\mathbb{S}^{m-1}$.
• Generative Model: The policy $\pi_\theta(\mathbf{c} | s)$ is modeled using spherical flow matching. It learns a time-dependent vector field on the sphere to transport a base distribution to a target distribution of cost directions. This allows the policy to represent any stochastic distribution over feasible actions without explicitly enumerating them.

B. Efficient Latent-Space Training
To avoid the computational bottleneck of calling the CO solver repeatedly during policy updates, LSFLOW trains both the actor and the critic directly in the latent cost space.

• Critic: A value network $\tilde{Q}_\phi(s, \mathbf{c})$ is trained to approximate $Q(s, \mathbf{a}^*(s, \mathbf{c}))$. This amortizes solver costs, requiring only one solver call per environment step rather than per policy proposal.
• Policy Update: The policy is updated using a weighted spherical flow-matching objective. Proposals are reweighted based on the critic's estimated value, shifting the distribution toward cost directions that yield higher-value feasible actions.

C. Smoothed Bellman Operator
A critical challenge is that the mapping from cost direction to action is piecewise constant and discontinuous (due to the solver), leading to unstable value learning.

• vMF Smoothing: The authors introduce a von Mises–Fisher (vMF) kernel to smooth the Bellman operator. Instead of evaluating the value at a single direction $\mathbf{c}$, the target is computed by averaging over a neighborhood of directions sampled from the vMF kernel.
• Theoretical Guarantees: The authors prove that this smoothed operator is a contraction mapping with a unique fixed point that is infinitely differentiable ($C^\infty$) with respect to the cost direction. This ensures stable, low-variance learning targets. They further show that as the kernel concentration parameter $\kappa \to \infty$, the smoothed value converges to the true value function almost everywhere.

3. Key Contributions

1. First Flow/Diffusion Framework for Combinatorial RL: LSFLOW is the first policy framework to transfer the expressiveness of flow-based generative models to combinatorial action spaces while guaranteeing feasibility via a solver.
2. Efficient Latent-Space Optimization: The method introduces a training paradigm where both the actor and critic operate in the latent cost space, avoiding repeated solver calls during optimization.
3. Stabilized Learning via Smoothing: The paper proposes a smoothed Bellman operator using vMF kernels to mitigate the instability caused by solver-induced discontinuities, providing theoretical proofs of contraction and smoothness.

4. Experimental Results

The authors evaluate LSFLOW on public benchmarks and a real-world application:

• Public Benchmarks: Tested on four dynamic tasks (Scheduling, Routing, Assignment, Intervention) from Xu et al. (2025). LSFLOW outperformed state-of-the-art baselines (including SEQUOIA and SRL) by an average of 20.6% in reward.
• Real-World Application: Applied to a Sexually Transmitted Infection (STI) testing task involving network-based screening for Gonorrhea, HIV, Syphilis, and Chlamydia. LSFLOW demonstrated higher detection efficiency (fraction of positive cases found) compared to baselines, particularly in low-to-mid budget regimes.
• Efficiency: LSFLOW was approximately 3.0× faster in training time than the deterministic Structured RL (SRL) baseline, as it avoids differentiating through the solver or solving embedded MIPs during training.

5. Significance and Claims

The paper claims that LSFLOW offers a principled and scalable route to RL with combinatorial actions. By shifting the burden of feasibility to a solver and learning an expressive stochastic policy in a latent spherical space, the method overcomes the limitations of previous approaches that sacrifice either generality (solver-based methods) or expressiveness (deterministic methods).

The authors emphasize that their approach successfully transfers the benefits of modern generative policies (multimodality, rich exploration) to constrained domains. They acknowledge limitations, noting that smoothing introduces a bias-variance trade-off and that inference still requires solver calls, suggesting future work on adaptive smoothing and learned guidance to further reduce runtime.