Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Latent Spherical Flow Policy for Reinforcement Learning with Combinatorial Actions

Lingkai Kong, Anagha Satish, Hezi Jiang, Akseli Kangaslahti, Andrew Ma, Wenbo Chen, Mingxiao Song, Lily Xu, Milind Tambe

TL;DR

This work tackles reinforcement learning with combinatorial action spaces by introducing LSFlow, a solver-augmented latent spherical flow policy. The method learns an expressive stochastic policy over a latent cost space on the unit sphere and delegates feasibility to a combinatorial optimizer, ensuring feasible actions by design. It combines flow matching in latent space with a von Mises–Fisher smoothing-based Bellman operator and a latent-space critic to stabilize learning and reduce solver calls, achieving strong improvements on public benchmarks and STI testing tasks. The approach promises scalable, expressive decision-making in constraint-heavy domains where combinatorial feasibility is strict and critical.

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.

Latent Spherical Flow Policy for Reinforcement Learning with Combinatorial Actions

TL;DR

This work tackles reinforcement learning with combinatorial action spaces by introducing LSFlow, a solver-augmented latent spherical flow policy. The method learns an expressive stochastic policy over a latent cost space on the unit sphere and delegates feasibility to a combinatorial optimizer, ensuring feasible actions by design. It combines flow matching in latent space with a von Mises–Fisher smoothing-based Bellman operator and a latent-space critic to stabilize learning and reduce solver calls, achieving strong improvements on public benchmarks and STI testing tasks. The approach promises scalable, expressive decision-making in constraint-heavy domains where combinatorial feasibility is strict and critical.

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.
Paper Structure (55 sections, 7 theorems, 59 equations, 6 figures, 4 tables, 1 algorithm)

This paper contains 55 sections, 7 theorems, 59 equations, 6 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Problem Definition
  4. Flow Matching
  5. Flow Policies
  6. Proposed Method
  7. Combinatorial Flow Policy Representation
  8. Efficient Policy Learning in Latent Cost Space
  9. On-Sphere Smoothing
  10. Practical Implementation
  11. Related Works
  12. Experiments
  13. Public Benchmarks
  14. Sexually Transmitted Infection Testing
  15. Ablation Study
  16. ...and 40 more sections

Key Result

Lemma 3.0

Fix any state $\bm{s}$ with $\mathcal{A}(\bm{s})\neq\emptyset$ and define $\bm{a}^{\star}(\bm{s},\bm{c})\in\mathop{\mathrm{arg\,min}}\limits_{\bm{a}\in\mathcal{A}(\bm{s})} \bm{c}^\top \bm{a}$. Then for any $\alpha>0$, the set of minimizers is unchanged: In particular, the solver mapping is positively scale invariant: any selection $\bm{a}^{\star}(\bm{s},\cdot)$ satisfies $\bm{a}^{\star}(\bm{s},\a

Figures (6)

  • Figure 1: Overall framework of LSFlow. To enable expressive stochastic policies under hard combinatorial constraints, we shift policy learning to a continuous spherical cost-direction space and let a combinatorial optimization (CO) solver enforce feasibility (\ref{['sec:policy-rep']}). Crucially, both policy and critic are learned directly in this latent cost space. We train the flow policy $\pi_\theta$ via weighted flow matching, using weights $w\propto \tilde{Q}_\phi(\bm{s},\bm{c})$ predicted by a latent-space critic (\ref{['sec:spherical-flow']}). To stabilize cost-space critic learning under solver-induced discontinuities, we apply vMF smoothing by sampling $\tilde{\bm{c}}\sim K_\kappa(\cdot\mid \bm{c})$ when constructing Bellman targets (\ref{['sec:smoothing-fixedpoint']}).
  • Figure 2: (a) The solver partitions the latent space on $\mathbb{S}^1$ into regions that map to different actions . We smooth this partition by averaging with a vMF kernel $K_{\kappa}(\cdot|\bm{c})$. (b) The corresponding value $\tilde{Q}(\bm{s},\bm{c})$ is originally piecewise constant across regions, but becomes a smooth function $\widetilde{Q}(\bm{s},\bm{c})$ after vMF smoothing.
  • Figure 3: Performance on sexually transmitted infection testing
  • Figure 4: Ablation study. Left: effect of training the flow directly on the sphere. Right: effect of vMF smoothing. Here $\kappa$ is the concentration parameter: larger $\kappa$ yields weaker smoothing, while smaller $\kappa$ yields stronger smoothing. None denotes no smoothing.
  • Figure 5: Graph used for the dynamic routing problem, based on the London tube network.
  • ...and 1 more figures

Theorems & Definitions (13)

  • Lemma 3.0: Positive scale invariance of the solver mapping
  • Proposition 3.0: Exact expressivity of solver-induced policies
  • Remark 3.1
  • Theorem 3.3: Contraction and smoothness of the fixed point
  • Lemma 5.0: Positive scale invariance of the solver mapping
  • proof
  • Proposition 6.0: Exact expressivity of solver-induced policies
  • proof
  • Theorem 7.1: Contraction and smoothness of the fixed point
  • proof : Proof of Theorem \ref{['thm:smooth-fixed-point']}
  • ...and 3 more