Flow Matching for Offline Reinforcement Learning with Discrete Actions

TL;DR

This paper extends flow matching to offline RL with discrete actions by introducing Q-weighted discrete flow matching (QDFM) built on Continuous-Time Markov Chains (CTMCs). It replaces continuous ODE-based flows with CTMC dynamics, enabling principled policy improvement via a Boltzmann-style, Q-guided objective and allowing multi-objective and multi-agent extensions through preference conditioning and factorized transitions. The authors prove gradient equivalence between the weighted conditional loss and a guided marginal objective, ensuring recovery of the KL-regularized optimal policy under suitable assumptions. Empirically, QDFM demonstrates strong performance on discretized MuJoCo tasks, multi-objective discrete benchmarks, and a two-agent coordination game, while offering flexible inference through action quantization and robust preference-driven behavior. The framework provides a principled, scalable path to discrete, multi-objective offline RL with potential extensions to online and decentralized multi-agent settings.

Abstract

Generative policies based on diffusion models and flow matching have shown strong promise for offline reinforcement learning (RL), but their applicability remains largely confined to continuous action spaces. To address a broader range of offline RL settings, we extend flow matching to a general framework that supports discrete action spaces with multiple objectives. Specifically, we replace continuous flows with continuous-time Markov chains, trained using a Q-weighted flow matching objective. We then extend our design to multi-agent settings, mitigating the exponential growth of joint action spaces via a factorized conditional path. We theoretically show that, under idealized conditions, optimizing this objective recovers the optimal policy. Extensive experiments further demonstrate that our method performs robustly in practical scenarios, including high-dimensional control, multi-modal decision-making, and dynamically changing preferences over multiple objectives. Our discrete framework can also be applied to continuous-control problems through action quantization, providing a flexible trade-off between representational complexity and performance.

Figures (9)

• Figure 1: Motivating example a) Offline dataset exhibits two disjoint modes corresponding to different objectives (Goal A and Goal B); the central region is a trap state never observed in the data. b) A continuous treatment of the decision space can induce interpolation between the modes and potential collapse toward the mean, causing trajectories to enter the trap.
• Figure 2: Walker2d-v5: sensitivity to CTMC rate scaling $\alpha$ for discretizations $K_{\mathrm{act}}\in\{16,32\}$. We report mean episodic return across $5$ seeds (100 episodes/seed); error bars indicate standard deviation. All runs use the same Euler simulation horizon with $K_{\mathrm{steps}}=15$.
• Figure 3: Effect of the number of Euler steps used to simulate the CTMC at inference time on Walker2d-v5. Performance remains stable across a wide range of discretization steps.
• Figure 4: Toy two-objective CartPole. Sweeping the preference parameter $\omega$ induces a smooth traversal of the Pareto front without retraining. Points correspond to different $\omega$ values and are connected for visualization. This curve demonstrates that our method learns not only two extreme behaviors but also intermediate trade-offs.
• Figure 5: Toy preference sweep on CartPole-v1 using the preference-conditioned CTMC policy. We report task return $R_1$, smoothness objective $R_2$ (negative number of action switches), scalarized return $R_\omega$, and the fraction of episodes achieving $R_1\ge 475$. All results use $\alpha=5$, $K_{\mathrm{steps}}=20$, and $k_{\pi}=20$ with 200 evaluation episodes.

Theorems & Definitions (9)

• Theorem 4.1: Gradient Equivalence
• Corollary 4.2: Recovery of the KL-Regularized Optimal Policy
• Proposition A.1: Affine invariance of the Bregman gradient
• proof
• Proposition A.2: Chain rule
• proof
• proof
• Theorem A.5: Recovery of the Boltzmann Policy
• proof