Flow Matching for Offline Reinforcement Learning with Discrete Actions

Fairoz Nower Khan; Nabuat Zaman Nahim; Ruiquan Huang; Haibo Yang; Peizhong Ju

cs.LGFeb 5, 2026

Flow Matching for Offline Reinforcement Learning with Discrete Actions

Authors: Fairoz Nower Khan, Nabuat Zaman Nahim, Ruiquan Huang, Haibo Yang, Peizhong Ju

Organizations: Department of Computer Science, University of Kentucky · Department of Computing and Information Sciences, Rochester Institute of Technology

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 across diverse settings and benchmarks, including high-dimensional control, multi-agent games, and dynamically changing preferences over multiple objectives, while outperforming traditional offline RL methods in practical multi-modal decision-making scenarios. 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 & tables

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: Multimodal decision-making on a 4-goal gridworld. All baselines converge to one goal; QDFM recovers all four from the same offline data.

Dataset	Env	AWAC	AWBC	BCQ	CQL	GreedyQ	BRAC	BoltzQ	QDFM (Ours)
Expert	HalfCheetah	2.49 $\pm$ 2.47	-0.71 $\pm$ 0.45	10.33 $\pm$ 5.49	-1.19 $\pm$ 0.01	-0.31 $\pm$ 0.00	-0.03 $\pm$ 0.01	4.84 $\pm$ 2.15	1.32 $\pm$ 1.09
Medium	HalfCheetah	19.69 $\pm$ 16.28	2.22 $\pm$ 1.23	20.23 $\pm$ 8.21	-0.58 $\pm$ 0.81	-0.23 $\pm$ 0.06	0.42 $\pm$ 1.65	4.62 $\pm$ 0.58	6.02 $\pm$ 2.00
Expert	Hopper	0.95 $\pm$ 0.00	1.03 $\pm$ 0.03	4.95 $\pm$ 4.75	0.90 $\pm$ 0.00	0.90 $\pm$ 0.00	0.81 $\pm$ 0.01	0.69 $\pm$ 0.01	17.74 $\pm$ 7.61
Medium	Hopper	1.54 $\pm$ 0.03	7.12 $\pm$ 0.16	36.18 $\pm$ 26.22	1.13 $\pm$ 0.02	1.11 $\pm$ 0.02	2.82 $\pm$ 0.04	1.84 $\pm$ 0.03	37.79 $\pm$ 16.11
Expert	Walker2d	1.53 $\pm$ 0.70	-0.20 $\pm$ 0.02	3.26 $\pm$ 2.92	-0.31 $\pm$ 0.02	-0.30 $\pm$ 0.03	1.97 $\pm$ 0.37	0.06 $\pm$ 0.20	4.85 $\pm$ 2.34
Medium	Walker2d	-0.18 $\pm$ 0.02	2.48 $\pm$ 2.42	9.30 $\pm$ 5.93	6.65 $\pm$ 0.25	6.59 $\pm$ 0.25	3.15 $\pm$ 1.98	-0.09 $\pm$ 0.01	14.01 $\pm$ 9.88

Table 1: Performance on discretized MuJoCo benchmark. We report mean

\pm

standard deviation over 5 seeds. Evaluation numbers are normalized as suggested by Fu et al. (2020) .

Table 4

Map	Dataset	BC	MACQL	CFCQL	OMAR	OMIGA	QDFM
2s3z	medium-replay	16.71 $\pm$ 0.09	15.30 $\pm$ 1.27	16.79 $\pm$ 0.27	15.10 $\pm$ 1.31	16.27 $\pm$ 0.73	17.06 $\pm$ 0.39
5m_vs_6m	expert	12.13 $\pm$ 0.23	12.17 $\pm$ 1.47	15.05 $\pm$ 0.67	11.32 $\pm$ 0.43	14.95 $\pm$ 1.52	14.99 $\pm$ 0.19
	medium	11.21 $\pm$ 0.08	11.26 $\pm$ 0.20	11.59 $\pm$ 1.23	11.01 $\pm$ 1.84	11.48 $\pm$ 0.73	11.58 $\pm$ 1.55
	medium-replay	6.89 $\pm$ 0.11	6.66 $\pm$ 0.13	6.76 $\pm$ 0.11	6.83 $\pm$ 0.42	6.74 $\pm$ 0.13	6.94 $\pm$ 0.25
	poor	1.30 $\pm$ 0.25	1.54 $\pm$ 0.14	1.41 $\pm$ 0.05	3.43 $\pm$ 0.15	1.24 $\pm$ 0.02	1.67 $\pm$ 0.29

Table 4: Offline MARL performance on SMAC benchmarks (Return

\uparrow

). Results are averaged over 5 seeds.

Appendix figures & tables27 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 3 — Figure 3: Intuition for the CTMC formulation. (a) The rate field specifies how probability mass flows between actions. (b) The jump-to-endpoint target field moves mass directly toward the endpoint. (c) Inference simulates the learned CTMC in discrete time via stay-or-jump updates.

Method	Train Time (s)	Act Time (ms)
QDFM	52.41	7.27
Energy-based	52.93	14.87
Diffuse-based	6913.83	7.50

Appendix

Table 5: Runtime comparison between the jump-to-endpoint rate field used by QDFM and alternative valid CTMC rate-field constructions.

Dataset	QDFM	single-a	uniform
Hopper-Medium	$\mathbf{35.81\pm 18.80}$	$0.75\pm 0.01$	$1.82\pm 0.03$
Walker2d-Medium	$\mathbf{7.90\pm 5.65}$	$0.94\pm 0.81$	$3.90\pm 2.98$

Appendix

Table 6:

\hat{\mu}

ablation on MuJoCo datasets.

Table 9

Dataset	Environment	AWAC	AWBC	BCQ	CQL	GreedyQ	Ours
Expert	HalfCheetah	29.14 $\pm$ 306.94	-367.81 $\pm$ 55.60	1002.09 $\pm$ 682.03	-428.26 $\pm$ 1.14	-319.23 $\pm$ 0.40	-123.18 $\pm$ 93.52
Medium	HalfCheetah	2164.79 $\pm$ 2021.35	-5.00 $\pm$ 152.87	2231.27 $\pm$ 1019.49	-351.66 $\pm$ 100.36	-308.52 $\pm$ 6.93	330.95 $\pm$ 319.08
Expert	Hopper	10.63 $\pm$ 0.13	13.30 $\pm$ 0.87	140.93 $\pm$ 154.72	8.88 $\pm$ 0.08	8.88 $\pm$ 0.10	467.86 $\pm$ 277.28
Medium	Hopper	29.92 $\pm$ 0.97	211.61 $\pm$ 5.18	1157.34 $\pm$ 853.44	16.48 $\pm$ 0.68	15.96 $\pm$ 0.57	1193.29 $\pm$ 711.09
Expert	Walker2d	71.89 $\pm$ 32.31	-7.50 $\pm$ 0.93	151.19 $\pm$ 133.83	-12.77 $\pm$ 0.89	-12.20 $\pm$ 1.21	211.31 $\pm$ 101.13
Medium	Walker2d	-6.85 $\pm$ 0.72	115.42 $\pm$ 111.01	428.60 $\pm$ 272.39	307.01 $\pm$ 11.26	303.94 $\pm$ 11.29	574.62 $\pm$ 586.73

Appendix

Table 9: Performance comparison at

K_{\text{act}}=256

. Mean

\pm

standard deviation across seeds. Best results per row (highest return) are in bold .

Dataset	Environment	AWAC	AWBC	BCQ	CQL	GreedyQ	Ours
Expert	HalfCheetah	-232.53 $\pm$ 74.42	-337.04 $\pm$ 37.28	-384.38 $\pm$ 21.48	-263.23 $\pm$ 95.31	-415.24 $\pm$ 44.37	-191.31 $\pm$ 97.18
Medium	HalfCheetah	318.29 $\pm$ 458.39	-5.80 $\pm$ 388.09	-198.59 $\pm$ 82.80	-442.33 $\pm$ 69.31	-143.53 $\pm$ 148.09	309.99 $\pm$ 495.70
Expert	Hopper	7.06 $\pm$ 0.41	8.42 $\pm$ 0.19	-0.77 $\pm$ 0.34	105.71 $\pm$ 31.11	1.04 $\pm$ 0.19	173.45 $\pm$ 91.35
Medium	Hopper	37.42 $\pm$ 1.32	198.77 $\pm$ 76.68	43.95 $\pm$ 4.19	29.17 $\pm$ 0.80	28.52 $\pm$ 0.58	286.67 $\pm$ 101.14
Expert	Walker2d	79.51 $\pm$ 78.24	12.27 $\pm$ 10.29	444.78 $\pm$ 221.17	68.31 $\pm$ 81.92	5.67 $\pm$ 20.73	254.65 $\pm$ 191.24
Medium	Walker2d	319.66 $\pm$ 218.78	46.59 $\pm$ 102.14	293.99 $\pm$ 72.79	31.02 $\pm$ 107.75	90.36 $\pm$ 2.39	124.95 $\pm$ 202.83

Appendix

Table 10: Performance comparison at

K_{\text{act}}=16

. Mean

\pm

standard deviation across seeds. Best results per row (highest return) are in bold .

Dataset	$M=16$	$M=32$	$M=64$	$M=128$
Hopper-Medium	3158.7 $\pm$ 637.2	3583.4 $\pm$ 359.6	4719.3 $\pm$ 104.4	7247.1 $\pm$ 118.8
Walker2d-Medium	3152.7 $\pm$ 607.6	3588.1 $\pm$ 361.6	4664.4 $\pm$ 9.5	7346.0 $\pm$ 12.4

Appendix

Table 11: Ablation on

M

: training time (seconds).

Dataset	$M=16$	$M=32$	$M=64$	$M=128$
Hopper-Medium	8.86 $\pm$ 2.06	8.86 $\pm$ 1.96	8.93 $\pm$ 2.12	8.86 $\pm$ 2.01
Walker2d-Medium	8.91 $\pm$ 2.22	8.91 $\pm$ 2.18	10.02 $\pm$ 0.09	10.08 $\pm$ 0.10

Appendix

Table 12: Ablation on

M

: action generation time (ms).

Dataset	$N=5$	$N=10$	$N=20$	$N=40$
Hopper-Medium	2286.7 $\pm$ 13.3	2381.3 $\pm$ 21.6	2429.6 $\pm$ 43.9	2637.3 $\pm$ 7.7
Walker2d-Medium	2314.8 $\pm$ 28.9	2382.0 $\pm$ 77.5	2480.2 $\pm$ 62.5	2683.2 $\pm$ 84.8

Appendix

Table 13: Ablation on

N

: training time (seconds).

Dataset	$N=5$	$N=10$	$N=20$	$N=40$
Hopper-Medium	1.89 $\pm$ 0.02	3.60 $\pm$ 0.23	6.29 $\pm$ 0.20	12.69 $\pm$ 0.84
Walker2d-Medium	1.92 $\pm$ 0.03	3.48 $\pm$ 0.18	6.57 $\pm$ 0.48	12.59 $\pm$ 0.75

Appendix

Table 14: Ablation on

N

: action generation time (ms).

$\|\mathcal{A}\|$	$M=32$	$M=64$	$M=128$
64	1527	1469	1481
256	1411	1448	1578
512	1420	1510	1869

Appendix

Table 15: Training time (seconds) while varying

|\mathcal{A}|

and

M

on Hopper-Medium.

Dataset	GreedyQ	CQL	BCQ	AWAC	IQL	QDFM
HalfCheetah-M	1380.1 $\pm$ 13.2	1729.0 $\pm$ 24.8	2364.8 $\pm$ 85.3	2606.1 $\pm$ 16.6	3620.3 $\pm$ 11.2	2209.4 $\pm$ 19.7
HalfCheetah-E	1351.0 $\pm$ 10.5	1751.5 $\pm$ 42.2	2324.4 $\pm$ 9.8	2606.7 $\pm$ 32.6	3611.9 $\pm$ 38.1	2232.0 $\pm$ 27.9
Hopper-M	1403.7 $\pm$ 38.0	1761.5 $\pm$ 26.5	2364.6 $\pm$ 69.1	2602.9 $\pm$ 27.8	3633.2 $\pm$ 24.5	2208.3 $\pm$ 16.6
Hopper-E	1378.8 $\pm$ 30.3	1722.0 $\pm$ 26.1	2416.9 $\pm$ 68.9	2607.3 $\pm$ 61.2	3674.4 $\pm$ 134.1	2214.6 $\pm$ 31.4
Walker2d-M	1374.8 $\pm$ 10.2	1742.9 $\pm$ 36.0	2389.5 $\pm$ 74.3	2598.9 $\pm$ 13.9	3683.5 $\pm$ 102.1	2233.5 $\pm$ 17.3
Walker2d-E	1354.9 $\pm$ 8.1	1720.9 $\pm$ 14.2	2401.2 $\pm$ 61.1	2619.7 $\pm$ 5.2	3617.7 $\pm$ 33.7	2249.1 $\pm$ 26.9

Appendix

Table 16: Same-hardware training time (seconds).

Dataset	GreedyQ	CQL	BCQ	AWAC	IQL	QDFM
HalfCheetah-M	0.22 $\pm$ 0.00	0.21 $\pm$ 0.01	0.40 $\pm$ 0.03	0.22 $\pm$ 0.00	0.21 $\pm$ 0.00	6.28 $\pm$ 0.04
HalfCheetah-E	0.21 $\pm$ 0.00	0.23 $\pm$ 0.03	0.39 $\pm$ 0.00	0.22 $\pm$ 0.02	0.22 $\pm$ 0.00	6.38 $\pm$ 0.02
Hopper-M	0.29 $\pm$ 0.00	0.30 $\pm$ 0.02	0.45 $\pm$ 0.05	0.28 $\pm$ 0.01	0.23 $\pm$ 0.00	6.26 $\pm$ 0.02
Hopper-E	0.33 $\pm$ 0.01	0.35 $\pm$ 0.04	0.48 $\pm$ 0.09	0.34 $\pm$ 0.03	0.24 $\pm$ 0.03	6.28 $\pm$ 0.09
Walker2d-M	0.26 $\pm$ 0.02	0.24 $\pm$ 0.01	0.44 $\pm$ 0.01	0.24 $\pm$ 0.01	0.23 $\pm$ 0.01	6.39 $\pm$ 0.08
Walker2d-E	0.23 $\pm$ 0.00	0.23 $\pm$ 0.00	0.46 $\pm$ 0.05	0.24 $\pm$ 0.00	0.24 $\pm$ 0.00	6.35 $\pm$ 0.12

Appendix

Table 17: Same-hardware action generation time (milliseconds).

Dataset	QDFM ( $\hat{\mu}$ )	single-a	uniform
Hopper-Medium	35.81 $\pm$ 18.80	0.75 $\pm$ 0.01	1.82 $\pm$ 0.03
Walker2d-Medium	7.90 $\pm$ 5.65	0.94 $\pm$ 0.81	3.90 $\pm$ 2.98

Appendix

Table 18: Ablation on the initialization distribution

\hat{\mu}

. Using a data-informed initialization significantly improves performance compared to single-action or uniform initialization.

Table 21

Table 22

Figure 8: Multi-Objective Fork: state visitation heatmaps under preference conditioning. Each panel shows state visitation counts aggregated over rollouts from the learned CTMC policy under a different preference vector $\omega$ : left objective ( $\omega=(1,0)$ ), balanced ( $\omega=(0.5,0.5)$ ), and right objective ( $\omega=(0,1)$ ). The key qualitative signature is trap avoidance : the trap state remains unvisited (dark) even under balanced preferences, while trajectories concentrate along the left or right corridor depending on $\omega$ . This indicates that the learned discrete flow maintains multiple modes instead of collapsing to an “average” (trap-entering) behavior, while still allowing controllable traversal between modes via $\omega$ .

Figure 11: Two-agent matrix game (offline). Mean Pareto points (point in objective space is Pareto if no other achievable point improves one objective without degrading another) obtained by sweeping the preference vector $\omega$ , averaged over three seeds, curve traces the trade-off between objective 1 (task reward) and objective 2 (safety/cost). Our factorized CTMC policy with centralized multi-objective guidance (QDFM) produces a broad, preference-controllable Pareto front, whereas independent and centralized behavior cloning yield collapsed fronts with weak dependence on $\omega$ with the endpoints for $\omega=(1,0)$ and $\omega=(0,1)$ coinciding, reflecting the fact that these methods are insensitive to preference conditioning.

Method	HV $\uparrow$	Spread $\uparrow$	Coord@ $\omega{=}(0,1)$
Independent BC	30.145 $\pm$ 0.704	0.043 $\pm$ 0.041	0.748 $\pm$ 0.025
Centralized BC	32.786 $\pm$ 0.284	0.039 $\pm$ 0.015	0.824 $\pm$ 0.012
QDFM (ours)	47.667 $\pm$ 0.529	0.616 $\pm$ 0.072	0.283 $\pm$ 0.028

Appendix

Table 23: In a two-agent offline matrix game we compare 2D hypervolume (higher is better), spread (higher indicates broader coverage here), and Coord which is the coordination rate at

\omega=(0,1)

(lower indicates the policy deliberately trades coordination for safety).

Map	Dataset	BC	MACQL	CFCQL	OMAR	OMIGA	QDFM
2s3z	medium-replay	0.57 $\pm$ 0.03	0.42 $\pm$ 0.10	0.52 $\pm$ 0.05	0.45 $\pm$ 0.10	0.48 $\pm$ 0.12	0.59 $\pm$ 0.04
3m	medium	0.18 $\pm$ 0.02	0.31 $\pm$ 0.09	0.31 $\pm$ 0.03	0.13 $\pm$ 0.00	0.33 $\pm$ 0.03	0.29 $\pm$ 0.08
5m_vs_6m	expert	0.43 $\pm$ 0.03	0.41 $\pm$ 0.17	0.51 $\pm$ 0.07	0.29 $\pm$ 0.06	0.49 $\pm$ 0.15	0.48 $\pm$ 0.09
5m_vs_6m	medium	0.17 $\pm$ 0.01	0.16 $\pm$ 0.04	0.20 $\pm$ 0.12	0.17 $\pm$ 0.15	0.19 $\pm$ 0.08	0.19 $\pm$ 0.06

Appendix

Table 24: Win rate comparison on SMAC benchmarks.

Figure 12: Multi-goal gridworld analysis. (a) State visitation heatmaps at $K{=}4$ . (b) Mode coverage vs. number of goals. (c) Trap entry rate vs. number of goals. (d) Mean return vs. number of goals. All values are mean $\pm$ std over 5 seeds.

	$K=2$		$K=3$		$K=4$		$K=5$
Method	Return	Mode Cov	Return	Mode Cov	Return	Mode Cov	Return	Mode Cov
CQL	10.00	0.50	10.00	0.33	10.00	0.25	10.00	0.20
IQL	10.00	0.50	10.00	0.33	10.00	0.25	10.00	0.20
AWAC	10.00	0.50	10.00	0.33	10.00	0.25	10.00	0.20
BCQ	$-$ 4.00 $\pm$ 5.5	0.00	$-$ 4.00 $\pm$ 5.5	0.00	$-$ 6.00 $\pm$ 5.5	0.00	2.00 $\pm$ 4.5	0.04
BRAC	6.00 $\pm$ 8.9	0.40	10.00	0.33	6.00 $\pm$ 8.9	0.20	10.00	0.20

Appendix

Table 25: Multi-goal gridworld results. Mean

\pm

std over 5 seeds, 200 eval episodes each. Mode coverage is the fraction of distinct goals reached by a single trained policy.

Experiment group	GPU type	Memory	Typical run time per seed
Discretized MuJoCo	A100 80GB class	32 to 64GB	0.6 to 2.1 hours
Multi-objective gridworlds	CPU or one GPU	16 to 32GB	Under one hour
Resource Gathering	CPU or one GPU	16 to 32GB	One to a few hours
Multi-goal gridworld	One GPU	8 to 16GB	Under 10 minutes
SMAC	One GPU with CPU workers	32 to 64GB	Map dependent

Appendix

Table 26: Compute resources for the reported experiments.

\mathcal{L}_{\mathrm{DFM}}(\theta)=\frac{1}{B}\sum_{i=1}^{B}\big\|\!\Big(u_{t_{i}}(\cdot,A_{i,t}\mid Z_{i})-u_{\theta,t_{i}}(\cdot,A_{i,t}\mid s_{i},\omega_{i})\Big)\big\|_{2}^{2}.

Appendix

Algorithm 3 Q-weighted discrete flow matching for offline RL

Explore similar work

Discrete Flow Matching for Offline-to-Online Reinforcement Learning
May 12, 2026Fairoz Nower Khan, Nabuat Zaman Nahim, Peizhong JuFlow Matching Reinforcement Learning
Department of Computer Science, University of Kentucky
Entropy-Regularized Adjoint Matching for Offline Reinforcement Learning
May 7, 2026Abdelghani Ghanem, Mounir GhoghoFlow Matching Offline RL
College of Computing Mohammed VI Polytechnic University Rabat 11103, Morocco
From Suboptimal to Optimal: One-Step Latent Flow Policies for Offline Reinforcement Learning
Oct 12, 2026Shenghua Lu, Zixuan Yang, Rongxu Zhao +3Offline RL Latent Action Learning
Jilin University

May 12, 2026J/K move · Enter open · S save

Discrete Flow Matching for Offline-to-Online Reinforcement Learning

Fairoz Nower Khan, Nabuat Zaman Nahim, Peizhong Ju

Department of Computer Science, University of Kentucky

Flow Matching Reinforcement Learning RL Fine-Tuning Offline-To-Online RL Discrete Flow Matching

Open paper AskPDF arXiv

Many reinforcement learning (RL) tasks have discrete action spaces, but most generative policy methods based on diffusion and flow matching are designed for continuous control. Meanwhile, generative policies usually rely heavily on offline datasets and offline-to online RL is itself challenging, as the policy must improve from new interaction without losing useful behavior learned from static data. To address those challenges, we introduce DRIFT, an online fine-tuning method that updates an offline pretrained continuous time Markov chain (CTMC) policy with an advantage-weighted discrete flow matching loss. To preserve useful pretrained knowledge, we add a path-space penalty that regularizes the full CTMC trajectory distribution, rather than only the final action distribution. For large discrete action spaces, we introduce a candidate-set approximation that updates the actor over a small subset of actions sampled from reference-policy rollouts and uniform exploration. Our theoretical analysis shows that focusing on a small set of high probability actions can reduce computation while keeping the approximation error small, making our method practical for large action spaces. We evaluate on macro-action MinAtar games where offline-pretrained DRIFT outperforms its online-only counterpart across all five games and achieves the highest mean return on three and on oracle free Jericho, DRIFT achieves a mean normalized score of 7.3% with a lightweight GRU encoder, outperforming the evaluated pretrained language-model baselines. Ablations demonstrate the importance of appropriately weighted path-space regularization and reference-guided candidate construction.

Abstract

Explore similar work

Discrete Flow Matching for Offline-to-Online Reinforcement Learning

Entropy-Regularized Adjoint Matching for Offline Reinforcement Learning

From Suboptimal to Optimal: One-Step Latent Flow Policies for Offline Reinforcement Learning