arXiv CSopen access

Discrete Flow Matching for Offline-to-Online Reinforcement Learning

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

machine learning, deep learning, neural networks

Discrete Flow Matching for Oﬄine-to-Online Reinforcement Learning Fairoz Nower Khan, Nabuat Zaman Nahim, and Peizhong Ju

arXiv:2605.12379v1 [cs.LG] 12 May 2026

Department of Computer Science, University of Kentucky

Abstract Many reinforcement learning (RL) tasks have discrete action spaces, but most generative policy methods based on diﬀusion and ﬂow matching are designed for continuous control. Meanwhile, generative policies usually rely heavily on oﬄine datasets and oﬄine-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 ﬁne-tuning method that updates an oﬄine pretrained continuous-time Markov chain (CTMC) policy with an advantage-weighted discrete ﬂow matching loss. To preserve useful pretrained knowledge, we add a path-space penalty that regularizes the full CTMC trajectory distribution, rather than only the ﬁnal 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 the candidate-set error is controlled by missing target probability mass, and the induced CTMC generator error decreases as the candidate set covers more high-probability actions. Experiments on prevailing discrete action RL task show that our method provides stable oﬄine-to-online improvement across all tasks, achieving the highest average score on Jericho with a simple GRU encoder while outperforming methods that use pretrained language models. Controlled experiments further conﬁrm that the path-space penalty remains bounded during ﬁne-tuning and that the CTMC generator adapts to shifted rewards faster than deterministic baselines. The candidate-set mechanism is supported by a stability analysis showing that the generator error decreases exponentially with candidate coverage.

Introduction

Oﬄine-to-online RL asks a simple but important question: how can a policy learned from static data keep improving through real interaction? Oﬄine RL trains agents from logged data without costly or risky exploration [Levine et al., 2020, Kumar et al., 2020, Kostrikov et al., 2022], but the learned policy is limited by the coverage and quality of that dataset. If good actions or important states are missing from the data, oﬄine training alone cannot discover them. Oﬄine-to-online ﬁne-tuning addresses this limitation by using the oﬄine policy as a strong initialization and then improving it with online experience [Nair et al., 2020, Lee et al., 2022, Nakamoto et al., 2023, Zhang et al., 2023]. The challenge is to improve the policy without forgetting the useful behavior structure learned during oﬄine pre-training. Diﬀusion [Ho et al., 2020, Song et al., 2021] and ﬂow-matching [Lipman et al., 2023, Liu, 2022] policies capture multimodal action distributions in oﬄine data [Wang et al., 2023, Chen et al., 2023, Kang et al., 2023, Zhang et al., 2025a, Park et al., 2025], and can guide exploration during oﬄine-toonline ﬁne-tuning. However, these methods rely on continuous dynamics that do not naturally apply 1

G)al visits / 50 episodes

10.5 10.0

Reward

9.5 9.0 8.5 8.0 7.5 Sub7.0

optimal

Optimal

0.0

0.5

1.0

Coverage |cand|/||

(b) Multim)dal )verage

vs. DQN mode- ollapse DRIFT (o.rs) DQN

50 40 30 20 10 0

G1 G2 (+10) (+10)

G3 (+6)

Steps t) s0it hed g)al ↓

(a) Candidate-sephase -ransi-ion (|| = 128)

( ) P)st-swi-ch adap-a-ion speed 1.9 × faster 2.69

3 2

2.63

1.42

1 0

DRIFT ()urs)

DQN

PPO

Figure 1: Three empirical motivations for our method. (a) The candidate-set approximation shows sharp performance phase transition at 16% action coverage. (b) Our CTMC policy maintains multimodal coverage, while DQN mode-collapses to single goal. (c) In a goal-switch environment, our stochastic policy adapts to new reward signal 1.9× faster than DQN and PPO. to discrete spaces, leaving out domains such as recommendation, combinatorial control, scheduling, and token-level decision making. Discrete ﬂow matching [Gat et al., 2024, Campbell et al., 2024] replaces continuous ﬂows with CTMC generators making generative policies possible in discrete action spaces [Khan et al., 2026], but online ﬁne-tuning such policies remains open. During the online stage, the generator must move toward newly discovered high-reward actions, preserve useful pretrained behavior, and scale when the action space is large. Existing oﬄine discrete ﬂow methods do not address online adaptation, while oﬄine-to-online methods are mostly designed for continuous spaces. Our method We present an oﬄine-to-online ﬁne-tuning algorithm for discrete action spaces, improving online rewards while staying close to the pretrained generator. First, we design a discrete ﬂow matching actor update toward a reward-weighted target policy [Peters et al., 2010b, Ziegler et al., 2019]. Second, we use a path-space trust region that regularises the full CTMC trajectory rather than only the terminal action distribution [Schulman et al., 2015, Zhang et al., 2025b]. Third, for large action spaces, we update only over a candidate set built from reference-policy rollouts and uniform exploration, and show that the approximation error decreases as the candidate set covers more high-probability actions. We provide a stability analysis (Theorem 4.3) showing that the generator error under the candidate-set restriction decays exponentially with budget, and validate the resulting phase-transition empirically (Figure 1a). Our motivating experiments (Figure 1) show our generative policy maintains multimodal goal coverage where DQN [Mnih et al., 2015] mode-collapses, our stochastic generator adapts 1.9× faster than DQN and PPO [Schulman et al., 2017] after a reward shift, and the candidate-set approximation reaches optimal performance at just 16% action coverage in a 128-action gridworld, matching our coverage bound in Proposition 4.1. Together, these observations identify a common regime where standard discrete-action RL struggles: structured action spaces with multimodal actions, shifting rewards, or large action spaces where exhaustive evaluation is infeasible. This work is a ﬁrst step in that direction which we start by evaluating how our method performs on standard benchmarks, comparing with baselines on Jericho text games, MinAtar, discretised MuJoCo tasks,

and combinatorial gridworlds.

Related Work

Oﬄine-to-online RL and discrete-action RL. Oﬄine-to-online methods stabilize the transition from static data to live interaction through advantage weighting [Nair et al., 2020], balanced replay with pessimistic ensembles [Lee et al., 2022], calibrated value learning [Nakamoto et al., 2023], policy expansion [Zhang et al., 2023], and adaptive constraints [Li et al., 2026]. These approaches target continuous control or use standard policy classes. In discrete-action RL, DQN [Mnih et al., 2015], Double DQN [Van Hasselt et al., 2016], Rainbow [Hessel et al., 2018], A2C [Mnih et al., 2016], and PPO [Schulman et al., 2017] provide strong online baselines but do not model multimodal action distributions. Generative policies for RL. Diﬀusion and ﬂow-based policies capture multimodal behavior in oﬄine data. Diﬀusion-QL [Wang et al., 2023], IDQL [Hansen-Estruch et al., 2023], eﬃcient diffusion policies [Kang et al., 2023], energy-weighted ﬂow matching [Zhang et al., 2025a], and Flow Q-Learning [Park et al., 2025] all improve oﬄine RL through expressive action modeling with value guidance. Online ﬁne-tuning of such policies has been explored in continuous control [Zhang et al., 2025b, McAllister et al., 2025], but these methods require continuous action representations and do not extend to discrete decisions. Discrete generative models and trust-region ﬁne-tuning. Discrete diﬀusion and ﬂow models replace Euclidean dynamics with Markov processes over ﬁnite state spaces for sequences [Gat et al., 2024], molecular design [Campbell et al., 2024], and biological generation [Holderrieth et al., 2025, Wang et al., 2025], but none address oﬄine-to-online RL with value learning. Our path-space regularization builds on trust-region policy improvement [Peters et al., 2010b, Schulman et al., 2015, 2017] and KL-regularized ﬁne-tuning from human feedback [Ziegler et al., 2019, Ouyang et al., 2022]. We diﬀer in that we regularize the full CTMC path measure rather than only the terminal action distribution.

Preliminaries

3.1

Markov Decision Process and Oﬄine Data

We consider a Markov decision process (MDP) M = (S, A, r, P, γ), where S is the state space, A is a ﬁnite discrete action space with K = |A| actions, r : S × A → R is the reward function, P is the transition kernel, and γ ∈ (0, 1) is the P∞ discount factor. The state-action value function is π t Q (s, a) = Eπ t=0 γ r(st , at ) | s0 = s, a0 = a . An oﬄine dataset Doﬀ = {(si , ai , ri , s′i )}N i=1 is collected by an unknown behavior policy. The KLregularized optimal policy admits the well-known Boltzmann form [Todorov, 2006, Haarnoja et al., 2017, Lu et al., 2023]: π ⋆ (a | s) ∝ πref (a | s) exp A(s, a)/β , where A(s, a) = Q(s, a) − V (s) is the advantage, πref is a reference policy, and β > 0 is a temperature that trades oﬀ reward maximization against proximity to πref .

3.2

Discrete Flow Matching via CTMCs

Discrete Flow Matching (DFM) [Campbell et al., 2024, Gat et al., 2024] extends the ﬂow-matching paradigm from continuous Euclidean spaces to ﬁnite discrete spaces by replacing continuous vector 3

ﬁelds used in standard ﬂow matching with Continuous-Time Markov Chains (CTMCs). A CTMC over the action space A is speciﬁed by a time-dependent rate matrix (or generator) ut (i → j | s) satisfying the constraints for every source action i: X ut (i → j | s) ≥ 0 ∀ j 6= i, ut (i → j | s) = 0. (1) j∈A

The ﬁrst constraint ensures nonnegative oﬀ-diagonal transition rates and the second ensures probability conservation. The marginal Pprobability mass function pt evolves according to the Kolmogorov d pt (j | s) = i∈A pt (i | s) ut (i → j | s). forward equation: dt DFM learns a parametric rate model uθ by regressing it against a tractable conditional target generator derived from a prescribed probability path connecting a source distribution p0 to a target distribution p1 . At inference time, actions are sampled by simulating the learned CTMC from t = 0 to t = 1 using Euler discretization.

3.3

Oﬄine Pre-Training via Advantage-Weighted DFM

Our ﬁne-tuning method only requires an initial reference generator uref and does not depend on a speciﬁc pre-training procedure. In our experiments, we obtain uref using an advantage-weighted DFM pre-training stage that combines standard DFM with value-based guidance. Speciﬁcally, we construct a target policy π̃pre (a | s) ∝ exp(Āoﬀ (s, a)/β), where Āoﬀ is computed from oﬄine advantage estimates, and train a CTMC generator to ﬂow from a uniform prior to π̃pre using the DFM objective. This produces a reference generator that captures high-value behaviors from the ofﬂine data before online interaction begins. Importantly, this pre-training stage is not part of the deﬁnition of our ﬁne-tuning algorithm. Any pretrained discrete-action policy can be used as an initialization, as it can be represented as, or distilled into, a CTMC generator. For example, given a pretrained policy πpre (a | s) from behavior cloning, oﬄine RL, PPO, DQN, or another generative policy, we can ﬁt uθ by training a DFM bridge from a simple prior p0 to πpre (· | s). The resulting generator is then used as uref and ﬁne-tuned by the same online procedure in Section 4. Thus, our contribution is the CTMC-

Algorithm 1 DRIFT: Discrete ﬂow matching with Regularised Information-theoretic Fine-Tuning (sketched version, full algorithm in Appendix H) Require: Pre-trained generator uref , oﬄine data Doﬀ , temperature β, KL weight α, CTMC sub-steps M , candidate sizes Nroll , Nrand , refresh interval K, ﬂow-time truncation δ ∈ (0, 1) 1: Initialize uθ ← uref , critics Qφ1 , Qφ2 , value Vψ , targets Qφ− , Vψ− , buﬀer B k 2: for n = 1 to Nsteps do 3: Simulate CTMC under uθ from X0 ∼ Unif(A) for M substeps; execute a = XM ; store (s, a, r, s′ , d) in B 4: Sample mixed minibatch: (1−ρ)B from B and ρB from Doﬀ 5: Update critics: Qφk ← arg min LQk 6: Update value: Vψ ← arg min LV using lagged π̃ − (Eq. 4) 7: 8: 9: 10: 11: 12: 13: 14: 15:

for each state sb in actor batch do Build Acand (sb ) via Nroll rollouts of frozen uref + Nrand uniform actions Compute smoothed πref (Eq. 2) and clipped advantage Ā Set target policy π̃(a | sb ) ∝ πref (a | sb ) exp(Ā(sb , a)/β) (Eq. 3) Construct bridge pt (Eq. 5) and target generator u∗t by independent coupling (Eq. 6) Compute ℓDFM (sb ) by sampling t ∼ Unif(0, 1 − δ) and i ∼ pt (Eq. 7). c b ) against frozen uref (Eq. 8) Estimate path-space KL(s end for Update uθ by minimizing Lactor = LDFM + α LbKL (Eq. 9)

− − Soft-update targets: φ− k ← (1 − τ )φk + τ φk , ψ ← (1 − − τ )ψ +τ ψ 17: If n mod K = 0: set uref ← uθ 18: end for 19: return Fine-tuned generator uθ

16:

based oﬄine-to-online ﬁne-tuning mechanism, while the initialization can be obtained from any suitable pretrained policy. Full details of the pre-training procedure are given in Appendix A.

Our Method: Online ﬁne-tuning of Discrete Flow Matching Policy

Oﬄine-pretrained policies are limited by the coverage and quality of their data [Levine et al., 2020], while naive online ﬁne-tuning can forget the useful structure learned oﬄine. We introduce DRIFT (Discrete ﬂow matching with Regularised Information-theoretic Fine-Tuning), which improves a pretrained CTMC generator with online rewards while preserving its behavior through a path-space trust region. Overview. Starting from a reference generator uref obtained via advantage-weighted DFM pretraining (Section 3.3), DRIFT performs iterative online learning. Each iteration consists of four steps, summarized in Algorithm 1: (1) collect online transitions by rolling out the current CTMC generator in the environment (Line 3), (2) update critics and a value network from a mixed replay buﬀer (Lines 5–6), (3) construct a reward-weighted target policy and its corresponding DFM target generator over a candidate action set (Lines 8–11), (4) update the generator by minimizing a combined DFM loss plus path-space KL penalty (Lines 12–15). The reference generator is periodically refreshed to track improving performance (Line 17).

4.1

Environment Interaction and Replay

Actions are generated by simulating the current CTMC generator uθ forward in ﬂow time. Starting from X0 ∼ Uniform(A), we run M Euler sub-steps with step size ∆t = 1/M (Algorithm 2 in Appendix B). At each sub-step, the process either stays at the current action or jumps to a new one with probability proportional to the learned transition rates. The terminal action a = XM is executed in the environment, and the resulting transition (s, a, r, s′ , d) is stored in a replay buﬀer B. To prevent the ﬁne-tuning distribution from drifting too far from the pre-trained policy, each training minibatch mixes online transitions from B with a fraction ρ of oﬄine transitions from Doﬀ (Line 4), following oﬄine-to-online practice [Nakamoto et al., 2023, Lee et al., 2022].

4.2

Candidate Set Construction and Target Policy

Evaluating the target policy over the entire action space is intractable when K = |A| is large. We restrict computation to a candidate subset Acand (s) for each state (Line 8), by combining: 1. Reference rollouts. We run the frozen reference generator uref for Nroll independent CTMC simulations and collect the terminal actions, recording count(a) for each action. 2. Uniform exploration. We draw Nrand actions uniformly from A, ensuring coverage of potentially high-reward actions outside the pre-trained support. Smoothed reference policy. The reference distribution over the candidate set is estimated from rollout frequencies with additive smoothing (Line 9): πref (a | s) ∝

ǫ count(a) + , Nroll |Acand (s)| 5

a ∈ Acand (s).

(2)

Target policy. The target policy π̃ is the solution to the KL-regularized policy improvement problem [Ziegler et al., 2019, Peters et al., 2010a] (Line 10): πref (a | s) exp Ā(s, a)/β . (3) π̃(a | s) = P ′ ′ a′ ∈Acand (s) πref (a | s) exp Ā(s, a )/β

We extend this target to the full action space by setting π̃(a | s) = 0, a ∈ / Acand (s). This policy concentrates probability on high-advantage actions while remaining anchored to πref .

4.3

Critic and Value Network Updates

We maintain two independent critic networks Qφ1 , Qφ2 and a value network Vψ , together with frozen target copies Qφ− and Vψ− . k

Critic update. Each critic is trained by minimizing the squared Bellman error Pagainst a one-step TD target computed from the frozen value network (Line 5): LQk (φk ) = B1 B b=1 Qφk (sb , ab ) − 2 ′ sg[rb + γ Vψ− (sb )(1 − db )] , where sg[·] denotes stop-gradient. Using the frozen value network Vψ− in the TD target avoids the cost of a full CTMC rollout at every critic step and reduces variance [Haarnoja et al., 2018]. Value update. The value network is updated toward the expected critic value under a fully frozen lagged target policy π̃ − (Line 6). The lagged target policy is constructed from frozen networks only where the advantages are computed as A− (s, a) = mink Qφ− (s, a) − Vψ− (s), then normalized and k

clipped to Ā− (s, a) = clip(Â− (s, a), −c, c), giving π̃ − (a | s) ∝ πref (a | s) exp Ā− (s, a)/β ,

a ∈ Acand (s).

(4)

We compute both π̃ − and the value target on Acand (s) so that the value backup is consistent with the same restricted action distribution used by the actor update, while avoiding an expenP sive summation over the full action space. The value loss is then: LV (ψ) = B1 B b=1 Vψ (sb ) − hP i2 − sg a π̃ − (a | sb ) mink Qφ− (sb , a) . All target networks are soft-updated: φ− k ← (1 − τ )φk + τ φk , k

ψ − ← (1 − τ )ψ − + τ ψ. The advantage used in the actor update (Section 4.4) takes the pessimistic minimum of the two critics to suppress overestimation [Fujimoto et al., 2018]: A(s, a) = mink Qφk (s, a) − Vψ− (s). Advantages are normalized per-state togive β a consistent scale across states, then clipped to A(s,a)−µs [−c, c]: Ā(s, a) = clip σs +ε , −c, c , where µs , σs are the mean and standard deviation of A(s, ·) over Acand (s).

4.4

Actor Update: DFM Loss and Path-Space KL

The actor update has two components: a discrete ﬂow matching loss that steers the generator toward π̃, and a path-space KL penalty that prevents catastrophic forgetting. Linear bridge and target generator. We deﬁne a linear interpolation from the full-action uniform prior p0 = Uniform(A) to the target policy: pt (a | s) = (1 − t) p0 (a) + t π̃(a | s), 6

ṗt (a | s) = π̃(a | s) − p0 (a).

(5)

For the candidate-set update, we view π̃(· | s) as a distribution on the full action space by setting π̃(a | s) = 0 for a ∈ / Acand (s). Because π̃ spreads mass across multiple actions, the Kolmogorov forward equation does not uniquely determine a generator consistent with ṗt . We resolve this via the independent coupling transport [Campbell et al., 2024, Holderrieth et al., 2025], which constructs a valid target generator u∗t by routing outﬂow from each losing action to all gaining actions proportionally to their demand (Line 11). Concretely, for i 6= j: u∗t (i → j | s) =

(ṗt (i))− (ṗt (j))+ , pt (i | s) Zt

Zt =

(ṗt (a))+ .

(6)

Details of this construction and mass conservation proof is given in Appendix C.1. Source actions outside the candidate set. Because the source distribution is deﬁned on the full action space, p0 = Uniform(A), while the target π̃(· | s) is supported only on Acand (s), a sampled intermediate action i ∼ pt (· | s) may lie outside Acand (s), especially for small t. The independent-coupling construction handles this naturally. For i ∈ / Acand (s), π̃(i | s) = 0,

ṗt (i | s) = π̃(i | s) − p0 (i) = −

1 . |A|

Discrete ﬂow matching loss. For each state sb , we sample a ﬂow time t ∼ Uniform(0, 1 − δ) and an intermediate action i ∼ pt (· | sb ), then regress the learned rates against the target rates (Line 12): B 2 1 XX LDFM (θ) = (7) uθ (ib → j, tb | sb ) − u∗tb (ib → j | sb ) . B b=1 j6=ib

Path-space KL regularization. Unlike the policy KL commonly used in RL, which compares two action distributions at a state, our KL compares two CTMC path distributions, so it constrains how probability mass moves over ﬂow time. We penalize deviations across the entire trajectory of the CTMC, including both the jump destinations and the holding times (Line 13). For CTMCs, the KL divergence between path measures Puθ and Puref admits a tractable decomposition via the Radon–Nikodym derivative [Kipnis and Landim, 2013, Zhang et al., 2025b]:

KL Puθ k Puref = EPuθ

k:jumps

Z 1 0

log

uθ (Xt− → Xtk , tk | s) k

uref (Xt− → Xtk , tk | s) k

# λref (Xt , t | s) − λθ (Xt , t | s) dt ,

(8)

where λθ (i, t | s) = j6=i uθ (i → j, t | s) is the total exit rate. This is a Monte Carlo plug-in surrogate where gradients ﬂow only through the evaluated uθ and λθ along the simulated path, so 7

the term acts as a practical regularizer toward uref rather than as an unbiased gradient estimator of the path-space KL. For each sampled state sb , we estimate this quantity with a Monte Carlo plug-in estimator LbKL (sb , θ) along a simulated CTMC path.

Combined ﬁne-tuning objective. 1 PB

where LbKL (θ) = B

4.5

The generator is then updated by minimizing

Lactor (θ) = LDFM (θ) + αLbKL (θ),

(9)

b=1 LKL (sb , θ) and α > 0 controls the strength of the path-space trust region.

Rate Network Parameterization

The CTMC generator uθ constraints (Eq. 1) are enforced by construction. The network outputs unconstrained logits gθ (i → j, t | s) for j 6= i, which are passed through a softplus to obtain nonnegative oﬀ-diagonal rates: uθ (i → j, t | s) = softplus(gθ (i → j, t | s)). The diagonal is set to P uθ (i → i, t | s) = − j6=i uθ (i → j, t | s), guaranteeing a valid generator at every parameter setting. The rate model is a lightweight MLP with two hidden layers of size 256.

4.6

Theoretical Analysis

We analyze the error caused by replacing the full action space with a candidate set. For the analysis, let π̃(· | s) denote the ideal target policy on the full action space, and deﬁne its candidate-restricted renormalization as π̃(a | s) 1{a ∈ Acand (s)} π̃cand (a | s) = P . ′ a′ ∈Acand (s) π̃(a | s) P Proposition 4.1 (Coverage Error). Let ǫ(s) = a∈A / cand (s) π̃(a | s) be the target mass excluded by the candidate set. If ǫ(s) < 1, then kπ̃cand − π̃k1 = 2ǫ(s). If the candidate set is built from Nroll independent samples from πref and Nrand independent uniform samples, then X E[kπ̃cand − π̃k1 ] = 2 π̃(a | s) (1 − πref (a | s))Nroll (1 − 1/K)Nrand . a∈A

Thus, the expected excluded mass decreases exponentially in both Nroll and Nrand . Proposition 4.2 (Mass Conservation). Assume pt (i) > 0 for actions with outgoing mass and P Zt = a (ṗt (a))+ > 0. Then the independent-coupling generator u∗t (i → j) =

(ṗt (i))− (ṗt (j))+ , pt (i) Zt

i 6= j,

satisﬁes the Kolmogorov forward equation. be the target generators induced by π̃ and Theorem 4.3 (Generator Stability). Let u∗t and u∗,cand t π̃cand , respectively. Assume both bridges use the same source distribution p0 = Unif(A). Fix δ ∈ (i) ≥ p > 0, Zt , Ztcand ≥ Z > 0. Then, for (0, 1) and suppose that for all t ∈ [0, 1 − δ], pt (i), pcand t every source action i, X ∗,cand C ut (i → j) − u∗t (i → j) ≤ 2 2 kπ̃cand (· | s) − π̃(· | s)k1 , p Z j6=i

Table 1: Jericho text-game results (raw score, single seed, 300K steps). Best score per game in bold among our methods. † Published scores from Yao et al. [2020]. Prior work†

Our experiments Game (max)

DRIFT

DRRN

DQN

IQL

AWAC

CALM

KG-A2C

Deephome (300) Detective (360) Enchanter (400) Ludicorp (150) Omniquest (50) Pentari (70) Temple (35) Zork1 (350) Zork3 (7) Ztuu (100)

35.8 291.5 20.0 13.8 8.0 26.8 8.0 25.0 2.5 5.0

6.0 50.0 20.0 6.0 10.0 25.0 5.0 35.0 3.0 5.0

6.0 70.0 20.0 6.0 10.0 25.0 8.0 25.0 0.0 5.0

6.0 50.0 20.0 6.0 5.0 25.0 8.0 25.0 0.0 5.0

3.5 66.5 17.0 11.4 5.5 11.75 7.4 1.25 0.0 2.0

1.0 289.7 19.1 10.1 6.9 0.0 0.0 30.4 0.5 3.7

1.0 207.9 12.1 17.8 3.0 50.7 7.6 34.0 0.0 9.2

Avg. norm. (%)

23.2

15.3

12.1

10.6

8.3

12.6

19.2

for a universal constant C. Therefore, as the candidate set covers more target-policy mass, the induced CTMC generator approaches the full-action generator. All proofs are provided in Appendix C.

Experiments

We evaluate DRIFT on three benchmarks: Jericho text games [Hausknecht et al., 2020], MinAtar [Young and Tian, 2019], and discretised D4RL MuJoCo tasks [Fu et al., 2020]. All methods share the same oﬄine data, pre-training budget, and online ﬁne-tuning steps within each benchmark. We report mean scores over 5 seeds unless otherwise noted. Full experimental details are provided in Appendix F.

5.1

Baselines

We compare against nine oﬄine-to-online methods: CQL [Kumar et al., 2020], Cal-QL [Nakamoto et al., 2023], IQL [Kostrikov et al., 2022], AWAC [Nair et al., 2020], DQN [Mnih et al., 2015], PPO [Schulman et al., 2017], Rainbow [Hessel et al., 2018], PEX [Zhang et al., 2023] and SPA [Li et al., 2026]. For methods designed for continuous control, we implement faithful discrete-action adaptations (Appendix F.6). On Jericho, we compare against DRRN [He et al., 2016] and report published scores for CALM [Yao et al., 2020] and KG-A2C [Ammanabrolu and Hausknecht, 2020].

5.2

Jericho Text Games

We evaluate DRIFT on ten text-based games from the Jericho benchmark [Hausknecht et al., 2020] to test whether the oﬄine-to-online ﬁne-tuning framework transfers to a qualitatively diﬀerent domain. In these games the agent receives textual observations and must select a text command from a variable-size set of valid actions at each step. The policy architecture uses GRU text encoders shared across all methods (see Appendix F.3 for details). DRIFT achieves the highest average normalized score (23.2%) across all ten games, outperforming every baseline including the results for CALM (12.6%) and KG-A2C (19.2%). This is notable because CALM uses a ﬁne-tuned GPT-2 language model for action generation while DRIFT uses a 9

Table 2: MinAtar oﬄine→online results (mean over 5 seeds). Best online score per game in bold. † Methods that fail to improve over their oﬄine score. Game

DRIFT

Cal-QL

CQL

IQL

AWAC

DQN

Rainbow

PPO

PEX

SPA

Breakout Asterix Freeway Seaquest Space Inv.

0.43→17.10 0.50→1.17 1.26→25.00 0.15→1.84 2.23→42.64

1.19→14.61 0.67→1.04 19.89→25.71 0.46→0.86 0.27→54.09

0.62→14.69 0.58→1.15 22.97→26.50 0.42→0.95 0.00→52.56

0.01→17.01 0.66→1.13 2.54→23.56 0.12→0.90 0.00→41.60

0.50→0.51† 0.38→0.53† 15.85→15.69† 0.28→0.20† 4.52→4.52†

0.38→14.25 0.55→1.06 2.99→25.38 0.24→0.89 0.00→49.04

0.29→12.48 0.55→0.91 0.31→24.57 0.38→0.58 0.19→35.52

—→0.52 —→0.50 —→0.13 —→0.13 —→2.96

0.32→16.73 0.53→1.05 1.89→22.98 0.37→1.00 0.00→48.58

0.19→0.24† 0.44→0.53† 0.95→0.76† 0.33→0.35† 0.00→0.00†

Average

0.91→17.55

4.50→19.26

4.92→19.17

0.67→16.84

4.31→4.29

0.83→18.12

0.34→14.81

—→0.85

0.62→18.07

0.38→0.38

simple GRU encoder. DRIFT achieves its strongest gains on games with deeper puzzle chains. On Deephome it scores 35.8 compared to 6.0 for almost every other baseline, a 6× improvement. On Detective it solves the majority of the game (291.5 out of 360) and on Ludicorp it more than doubles DQN, DRRN, and IQL while exceeding the published CALM score. However, DRRN outperforms DRIFT on Zork1 (35.0 vs 25.0) and Omniquest (10.0 vs 8.0), and four games show ﬂoor eﬀects where most methods converge to the same score within 300K steps. Path-space regularization helps prevent forgetting in sparse-reward settings by anchoring the ﬁne-tuned generator to the pretrained reference, allowing DRIFT to preserve successful strategies while exploring and notably, it is the only method with a positive score on every game.

5.3

MinAtar

We evaluate DRIFT against nine oﬄine-to-online baselines on MinAtar [Young and Tian, 2019] using the same oﬄine dataset and 300K online steps. As shown in Table 2, DRIFT improves on all ﬁve games, achieves the best score on Breakout, Asterix, and Seaquest, and remains competitive on Freeway and Space Invaders. Unlike AWAC and SPA, which collapse to oﬄine scores on four games, DRIFT avoids catastrophic forgetting through path-space KL regularization to the pretrained reference. Although its generative overhead is less beneﬁcial in this small-action regime, DRIFT’s consistent gains show reliable ﬁne-tuning. We also report macro-action MinAtar results (|A|=216) in Appendix F.5.

5.4

D4RL Continuous Control (Discretised)

We evaluate on three MuJoCo locomotion tasks from D4RL (Hopper, Walker2d, and HalfCheetah) each with two dataset qualities (medium, expert). Continuous action spaces are discretised into k=22 clusters via k-means on the oﬄine dataset and scores follow the standard D4RL normalisation. Details are in Appendix F.4. DRIFT improves over its oﬄine initialisation on all six tasks, achieving the best online score on Hopper (both qualities), Walker-medium, and HalfCheetah-expert. The largest gains appear on Hopper, where DRIFT improves by +47.8 (medium) and +46.7 (expert) normalised points. DRIFT’s online improvement magnitude (+24.0 avg) is competitive with CQL (+10.8) and Cal-QL (+11.7), despite starting from a weaker oﬄine initialisation (0.4 vs. 8.0 average oﬄine score). Ablation Study We run single-factor ablations in a controlled tabular setting (10 states, 50 actions, 50K online steps) and provide full results and analysis in Appendix D. Path-space KL (α=0.01) outperforms terminal KL and no regularization, lower temperature (β=0.1) speeds convergence, frequent reference refreshes (K=50) avoid stale targets, and higher oﬄine mixing (ρ=0.75) stabilizes critic learning. These settings do not directly transfer to MinAtar, where the original hyperparameters (β=0.5, α=0.1, K=500, ρ=0.25) work better, showing the need for environment-speciﬁc 10

Table 3: D4RL discretised results: oﬄine→online normalised score (mean, 5 seeds). Best online score per task in bold. All methods use k-means discretisation (k=22). Task

DRIFT

CQL

IQL

AWAC

Cal-QL

DQN

PPO

PEX

SPA

Hopper-med Hopper-exp Walker-med Walker-exp Cheetah-med Cheetah-exp

0.1→47.9 0.4→47.1 1.0→15.9 0.1→14.8 0.2→11.1 0.5→9.7

28.8→35.8 8.8→29.3 2.5→12.6 4.0→10.7 3.3→15.8 0.7→8.6

0.4→27.9 0.1→17.2 0.3→13.0 −0.2→10.1 0.2→7.8 0.2→6.8

26.0→25.3 11.5→11.2 4.2→5.8 9.4→9.0 13.1→12.7 1.1→1.0

28.8→44.1 8.8→25.7 2.5→11.7 4.0→12.4 3.3→15.8 0.7→8.5

0.4→23.7 −0.3→21.7 1.8→12.3 0.2→6.5 2.1→16.9 0.1→4.7

—→3.1 —→3.7 —→7.2 —→7.0 —→4.9 —→5.9

0.4→43.1 0.1→39.5 −0.3→11.8 0.2→7.2 0.2→14.6 0.2→7.7

0.4→0.4 0.3→0.3 1.8→1.8 0.2→0.2 2.1→2.1 0.1→0.1

Average

0.4→24.4

8.0→18.8

0.2→13.8

10.9→10.8

8.0→19.7

0.7→14.3

—→5.3

0.1→20.6

0.8→0.8

tuning.

Conclusion, Limitations and Future Work

DRIFT extends generative policies to discrete action spaces and provides the ﬁrst oﬄine-to-online ﬁne-tuning method for CTMC policies, with the largest gains in domains where action spaces are large, variable, or multimodal. DRIFT ﬁne-tunes a pretrained CTMC generator using rewardweighted discrete ﬂow matching and a path-space penalty that regularizes full CTMC trajectories rather than only terminal actions. A candidate-set mechanism scales the method to large action spaces using reference rollouts and uniform samples, with provable coverage guarantees. Experiments on MinAtar, discretized D4RL, and Jericho show reliable improvement: DRIFT achieves the best score on three of ﬁve MinAtar games and the highest average normalized score on Jericho (23.2%), outperforming baselines and published CALM and KG-A2C results. Path-space regularization is the key stabilizer, making DRIFT the only method that improves on every game across all benchmarks. Several limitations remain. Path-space KL prevents forgetting but may slow exploration when the oﬄine reference is poor, suggesting adaptive schedules for α. The candidate-set mechanism is theoretically supported and validated in ablations, but still needs large-scale testing in genuinely large action spaces such as real-time strategy games or combinatorial optimization. Extending DRIFT to these domains, as well as to multi-objective and multi-agent settings with factorized CTMC generators, can be a prominent future direction.

References Prithviraj Ammanabrolu and Matthew Hausknecht. Graph constrained reinforcement learning for natural language action spaces. In International Conference on Learning Representations, 2020. Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative ﬂows on discrete state-spaces: Enabling multimodal ﬂows with applications to protein co-design. In International Conference on Machine Learning, pages 5453–5512. PMLR, 2024. Huayu Chen, Cheng Lu, Chengyang Ying, Hang Su, and Jun Zhu. Oﬄine reinforcement learning via high-ﬁdelity generative behavior modeling. In The Eleventh International Conference on Learning Representations, 2023. Justin Fu, Aviral Kumar, Oﬁr Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep datadriven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.

Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International conference on machine learning, pages 1587–1596. PMLR, 2018. Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky TQ Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete ﬂow matching. Advances in Neural Information Processing Systems, 37:133345–133385, 2024. Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International conference on machine learning, pages 1352–1361. PMLR, 2017. Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Oﬀ-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pages 1861–1870. Pmlr, 2018. Philippe Hansen-Estruch, Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idql: Implicit q-learning as an actor-critic method with diﬀusion policies. arXiv preprint arXiv:2304.10573, 2023. Matthew Hausknecht, Prithviraj Ammanabrolu, Marc-Alexandre Côté, and Xingdi Yuan. Interactive ﬁction games: A colossal adventure. In Proceedings of the AAAI Conference on Artiﬁcial Intelligence, volume 34, pages 7903–7910, 2020. Ji He, Jianshu Chen, Xiaodong He, Jianfeng Gao, Lihong Li, Li Deng, and Mari Ostendorf. Deep reinforcement learning with a natural language action space. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1621–1630, 2016. Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Proceedings of the AAAI conference on artiﬁcial intelligence, volume 32, 2018. Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diﬀusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840–6851, 2020. Peter Holderrieth, Marton Havasi, Jason Yim, Neta Shaul, Itai Gat, Tommi Jaakkola, Brian Karrer, Ricky TQ Chen, and Yaron Lipman. Generator matching: Generative modeling with arbitrary markov processes. In The Thirteenth International Conference on Learning Representations, 2025. Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Eﬃcient diﬀusion policies for oﬄine reinforcement learning. Advances in Neural Information Processing Systems, 36:67195–67212, 2023. Fairoz Nower Khan, Nabuat Zaman Nahim, Ruiquan Huang, Haibo Yang, and Peizhong Ju. Flow matching for oﬄine reinforcement learning with discrete actions. arXiv preprint arXiv:2602.06138, 2026. Claude Kipnis and Claudio Landim. Scaling limits of interacting particle systems. Springer Science & Business Media, 2013. Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Oﬄine reinforcement learning with implicit q-learning. In International Conference on Learning Representations, 2022. Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for oﬄine reinforcement learning. Advances in neural information processing systems, 33:1179–1191, 2020. Seunghyun Lee, Younggyo Seo, Kimin Lee, Pieter Abbeel, and Jinwoo Shin. Oﬄine-to-online reinforcement learning via balanced replay and pessimistic q-ensemble. In Conference on Robot Learning, pages 1702– 1712. PMLR, 2022. Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Oﬄine reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643, 2020.

Songlin Li, Wei Xiao, Hao Wu, Xiaodan Zhang, Daolong An, and Shuai Lü. State proﬁciency-based adaptive ﬁne-tuning for oﬄine-to-online reinforcement learning. In Proceedings of the AAAI Conference on Artiﬁcial Intelligence, volume 40, pages 23169–23176, 2026. Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2023. Yaron Lipman, Marton Havasi, Peter Holderrieth, Neta Shaul, Matt Le, Brian Karrer, Ricky TQ Chen, David Lopez-Paz, Heli Ben-Hamu, and Itai Gat. Flow matching guide and code. arXiv preprint arXiv:2412.06264, 2024. Qiang Liu. Rectiﬁed ﬂow: arXiv:2209.14577, 2022.

A marginal preserving approach to optimal transport.

arXiv preprint

Cheng Lu, Huayu Chen, Jianfei Chen, Hang Su, Chongxuan Li, and Jun Zhu. Contrastive energy prediction for exact energy-guided diﬀusion sampling in oﬄine reinforcement learning. In International Conference on Machine Learning, pages 22825–22855. PMLR, 2023. David McAllister, Songwei Ge, Brent Yi, Chung Min Kim, Ethan Weber, Hongsuk Choi, Haiwen Feng, and Angjoo Kanazawa. Flow matching policy gradients. arXiv preprint arXiv:2507.21053, 2025. Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529–533, 2015. Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pages 1928–1937. PmLR, 2016. Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with oﬄine datasets. arXiv preprint arXiv:2006.09359, 2020. Mitsuhiko Nakamoto, Simon Zhai, Anikait Singh, Max Sobol Mark, Yi Ma, Chelsea Finn, Aviral Kumar, and Sergey Levine. Cal-ql: Calibrated oﬄine rl pre-training for eﬃcient online ﬁne-tuning. Advances in Neural Information Processing Systems, 36:62244–62269, 2023. James R Norris. Markov chains. Number 2. Cambridge university press, 1998. Long Ouyang, Jeﬀrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35:27730–27744, 2022. Seohong Park, Qiyang Li, and Sergey Levine. Flow q-learning. In International Conference on Machine Learning. PMLR, 2025. Jan Peters, Katharina Mulling, and Yasemin Altun. Relative entropy policy search. In Proceedings of the AAAI Conference on Artiﬁcial Intelligence, volume 24, pages 1607–1612, 2010a. Jan Peters, Katharina Mülling, and Yasemin Altun. Relative entropy policy search. In AAAI, 2010b. Walter Rudin. Principles of Mathematical Analysis. McGraw-Hill, 3 edition, 1976. John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pages 1889–1897. PMLR, 2015. John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.

Zijing Shi, Yunqiu Xu, Meng Fang, and Ling Chen. Self-imitation learning for action generation in textbased games. In Proceedings of the 17th Conference of the European Chapter of the Association for Computational Linguistics, pages 703–726, 2023. Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic diﬀerential equations. In International Conference on Learning Representations, 2021. Emanuel Todorov. Linearly-solvable markov decision problems. Advances in neural information processing systems, 19, 2006. Jens Tuyls, Shunyu Yao, Sham M Kakade, and Karthik R Narasimhan. Multi-stage episodic control for strategic exploration in text games. In International Conference on Learning Representations, 2022. Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q-learning. In Proceedings of the AAAI conference on artiﬁcial intelligence, volume 30, 2016. Chenyu Wang, Masatoshi Uehara, Yichun He, Amy Wang, Avantika Lal, Tommi Jaakkola, Sergey Levine, Aviv Regev, Tommaso Biancalani, et al. Fine-tuning discrete diﬀusion models via reward optimization with applications to dna and protein design. In The Thirteenth International Conference on Learning Representations, 2025. Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diﬀusion policies as an expressive policy class for oﬄine reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023. Shunyu Yao, Rohan Rao, Matthew Hausknecht, and Karthik Narasimhan. Keep calm and explore: Language models for action generation in text-based games. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 8736–8754, 2020. Kenny Young and Tian Tian. Minatar: An atari-inspired testbed for thorough and reproducible reinforcement learning experiments. arXiv preprint arXiv:1903.03176, 2019. Haichao Zhang, Wei Xu, and Haonan Yu. Policy expansion for bridging oﬄine-to-online reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023. Shiyuan Zhang, Weitong Zhang, and Quanquan Gu. Energy-weighted ﬂow matching for oﬄine reinforcement learning. In The Thirteenth International Conference on Learning Representations, 2025a. Tonghe Zhang, Chao Yu, Sichang Su, and Yu Wang. Reinﬂow: Fine-tuning ﬂow matching policy with online reinforcement learning. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025b. Daniel M Ziegler, Nisan Stiennon, Jeﬀrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoﬀrey Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv:1909.08593, 2019.

Oﬄine Pre-Training Details

Our Stage 1 pre-training produces a reference generator uref from the oﬄine dataset Doﬀ . This procedure is an advantage-weighted DFM variant.

Critic pre-training. We pre-train two critic networks Qφ1 , Qφ2 and a value network Vψ using standard TD learning on Doﬀ : yb = rb + γ Vψ− (s′b ) (1 − db ),

LQk (φk ) =

2 1 X Qφk (sb , ab ) − sg[yb ] . B b

The value network is trained via an advantage-weighted target. This diﬀers from the Boltzmann soft-value backup with in-support sampling.

Generator pre-training. We compute a target policy π̃pre (a | s) ∝ exp(Āoﬀ (s, a)/β) over the full action space using the pre-trained critics, where Āoﬀ is the clipped, normalized pessimistic advantage. The generator is then trained via standard DFM to ﬂow from p0 = Uniform(A) to π̃pre using the independent coupling (Eq. 6). The pre-trained generator captures high-value behaviors from the oﬄine data and provides a strong initialization for online ﬁne-tuning.

CTMC Simulation Procedure

Algorithm 2 describes the Euler discretization used to sample actions from the learned CTMC generator. This procedure is used both during environment interaction (Line 3 of Algorithm 1) and when constructing candidate sets via reference rollouts (Line 8).

Algorithm 2 CTMC Euler Simulation Require: State s, generator u, sub-steps M , step size ∆t = 1/M 1: Sample X0 ∼ Uniform(A) 2: for m = 0 to M − 1 do 3: Set tm ← m ∆t and current action P i ← Xm 4: Compute total exit rate λ ← j6=i u(i → j, tm | s) 5: if λ = 0 then 6: Xm+1 ← i {no transition possible} 7: else 8: With probability 1 − λ ∆t: Xm+1 ← i {stay} 9: With probability λ ∆t: sample Xm+1 ∼ u(·, tm | s)/λ {jump} 10: end if 11: end for 12: return Terminal action XM The scheme corresponds to the standard naive Euler discretization of a CTMC [Lipman et al., 2024, Eq. (6.6)] and incurs O(∆t) local error in transition probabilities. Validity requires λ ∆t ≤ 1; in practice this is ensured by choosing M suﬃciently large.

Proofs and Derivations

C.1

Independent Coupling: Construction and Mass Conservation

Construction. Given the linear bridge pt (a | s) = (1 − t) p0 (a) + t π̃(a | s) with p0 = Uniform(A), its time derivative is ṗt (a) = π̃(a | s) − p0 (a). We decompose ṗt into its positive and negative parts: X (ṗt (a))+ = max(ṗt (a), 0), (ṗt (a))− = max(−ṗt (a), 0), Zt = (ṗt (a))+ . a

The independent coupling assigns rates proportional to the product of outﬂow from i and inﬂow to j: u∗t (i → j | s) =

(ṗt (i))− (ṗt (j))+ , pt (i | s) Zt

i 6= j,

with u∗t (i → j | s) = 0 whenever pt (i | s) = 0 or Zt = 0. This construction is valid because it produces nonnegative oﬀ-diagonal rates and the diagonal is set by conservation.

Interpretation in the ﬁne-tuning context. The independent coupling spreads each source’s outﬂow across all sinks proportionally to demand. In the ﬁne-tuning setting, this means that actions gaining probability under π̃ receive mass from all actions losing probability, rather than from a single greedy source. This produces a dense, stable supervision signal that reduces susceptibility to mode collapse during ﬁne-tuning.

Proof of Proposition 4.2 (Mass Conservation). Substituting the generator into the Kolmogorov forward equation, the net ﬂow into state j is: X X Net(j) = pt (i) u∗t (i → j) − pt (j) u∗t (j → k) i6=j

k6=j

(ṗt (j))− X (ṗt (k))+ . Zt Zt k6=j i6=j P P P negative parts are equal: a (ṗt (a))+ = a (ṗt (a))− = Zt . Therefore Since a ṗt (a) = 0, total positive and P P + + − − k6=j (ṗt (k)) = Zt − (ṗt (j)) . Substituting and simplifying: i6=j (ṗt (i)) = Zt − (ṗt (j)) and (ṗt (j))+ Zt − (ṗt (j))− − (ṗt (j))− Zt − (ṗt (j))+ Net(j) = = (ṗt (j))+ − (ṗt (j))− = ṗt (j). Zt =

C.2

X (ṗt (i))− (ṗt (j))+

−

Path-Space KL Derivation

The KL divergence between CTMC path measures is a standard result from the theory of Markov jump processes [Norris, 1998, Kipnis and Landim, 2013]. We provide a self-contained derivation for completeness. Let Puθ and Puref denote the path measures of CTMCs with generators uθ and uref , respectively, both starting from the same initial distribution. A CTMC path over [0, 1] consists of an initial state, a sequence of jump times 0 < t1 < t2 < · · · < tNjump ≤ 1, and the corresponding jump destinations. Between jumps, the chain holds at its current state. The log-likelihood of a path under generator u decomposes as: Z 1 Njump X log u(Xt− → Xtk , tk | s) − λu (Xt , t | s) dt, log Pu (path) = k

k=1

where the ﬁrst term accounts for the jumps that occurred and the second term accounts for the holding (no-jump) periods. The KL divergence follows directly: Puθ (path) KL(Puθ k Puref ) = EPuθ log Puref (path)   Z 1 X uθ (Xt− → Xtk , tk ) k = EPuθ  log λref (Xt , t) − λθ (Xt , t) dt . + uref (Xt− → Xtk , tk ) 0 k:jumps

Monte Carlo estimation. In practice, we estimate this by simulating one trajectory X0 , X1 , . . . , XM from the current generator uθ (with stop-gradient applied to the trajectory) and computing: c KL(s) =

k:jumps

log

uθ (Xt− → Xtk , tk | s) k

uref (Xt− → Xtk , tk | s) k

M−1 X m=0

λref (Xm , tm | s) − λθ (Xm , tm | s) ∆t.

Gradients ﬂow through log uθ at jump times and through λθ along the trajectory, while the trajectory itself is treated as ﬁxed.

C.3

Proof of Proposition 4.1 (Coverage Error)

Proof. Let π̃ be the ideal full-action target policy and let X π̃(a | s). ǫ(s) = a∈A / cand (s)

Assume ǫ(s) < 1. The candidate-restricted policy is π̃cand (a | s) = Therefore,

kπ̃cand − π̃k1 =

a∈Acand

π̃(a | s) 1{a ∈ Acand (s)} . 1 − ǫ(s)

π̃(a) − π̃(a) + 1−ǫ π̃(a) ·

a∈Acand

π̃(a)

a∈A / cand

ǫ ǫ + ǫ = (1 − ǫ) · + ǫ = 2ǫ. 1−ǫ 1−ǫ

For the expectation, action a is excluded only if it is missed by all Nroll independent samples from πref and all Nrand independent uniform samples. Hence P(a ∈ / Acand (s)) = (1 − πref (a | s))Nroll (1 − 1/K)Nrand . Thus, E[ǫ(s)] =

π̃(a | s)(1 − πref (a | s))Nroll (1 − 1/K)Nrand ,

a∈A

and the result follows from kπ̃cand − π̃k1 = 2ǫ(s).

C.4

Proof of Theorem 4.3 (Generator Stability)

Proof. The proof proceeds in three parts. First, we set up the notation and record elementary boundedness facts that will be used throughout. Second, we decompose the generator diﬀerence into three terms, each measuring a distinct source of perturbation. Third, we bound each term separately and combine the results.

Step 1: Setup and elementary bounds. Let d = kπ̃cand − π̃k1 . By construction, both the full and candidate bridges use the same source p0 = Uniform(A) (Eq. 5). Therefore, ṗcand − ṗt = π̃cand − π̃, t

kṗcand − ṗt k1 = d. t

For any ﬁxed source action i and destination action j 6= i, let pt (i) and pcand (i) denote the full and t candidate bridge probabilities of action i at ﬂow time t. The maps x 7→ x+ = max{x, 0} and x 7→ x− = max{−x, 0} are 1-Lipschitz [Rudin, 1976], so X X (ṗcand (a))+ − (ṗt (a))+ ≤ d, (ṗcand (a))− − (ṗt (a))− ≤ d. t t a

The normalizers are Zt =

(ṗt (a))+ ,

Ztcand =

which equal the corresponding total outgoing masses by Ztcand .

(ṗcand (a))+ , t

1 a ṗt (a) = 0. Hence Zt = 2 kṗt k1 and similarly for

To simplify notation, we abbreviate the four quantities entering each rate (the source supply, sink demand, source probability, and total moving mass) using unsubscripted letters for the candidate bridge and a 0subscript for the full bridge: A := (ṗcand (i))− , t

Bj := (ṗcand (j))+ , t

A0 := (ṗt (i))− ,

Bj0 := (ṗt (j))+ ,

and P := pcand (i), t

Z := Ztcand,

P0 := pt (i),

Z0 := Zt .

With this notation,

A0 Bj0 ABj . , u∗t (i → j) = PZ P0 Z0 The lower bounds in the theorem statement give P, P0 ≥ p > 0 and Z, Z0 ≥ Z > 0. Furthermore, since kṗt k1 ≤ kπ̃k1 + kp0 k1 = 2, we have u∗,cand (i → j) = t

A, A0 , Bj , Bj0 ∈ [0, 1],

Z, Z0 ∈ [0, 1],

P, P0 ∈ [0, 1].

These elementary boundedness facts will be used repeatedly below.

Step 2: Three-term decomposition. Inserting and subtracting two intermediate terms, A0 Bj0 A0 (Bj − Bj0 ) ABj (A − A0 )Bj = − + + A0 Bj0 PZ P0 Z0 PZ PZ

1 1 − PZ P0 Z0

Summing over j 6= i and applying the triangle inequality, X ∗,cand ut (i → j) − u∗t (i → j) ≤ (I) + (II) + (III), j6=i

where (I) :=

|A − A0 | X Bj , PZ j6=i

(II) :=

A0 X Bj − Bj0 , PZ

(III) := A0

j6=i

Bj0

j6=i

1 1 . − PZ P0 Z0

Each term measures a distinct perturbation: (I) the change in source supply, (II) the change in sink demand, and (III) the change in the denominator.

Step 3a: Bound on Term (I). Since

j6=i Bj ≤ Z, the factor Z cancels:

|A − A0 | . P Using P ≥ p and the 1-Lipschitz property of x 7→ x− , (I) ≤

|A − A0 | = (ṗcand (i))− − (ṗt (i))− ≤ ṗcand (i) − ṗt (i) = π̃cand (i) − π̃(i) ≤ d, t t where the ﬁnal inequality uses that the pointwise diﬀerence is dominated by the L1 norm. Therefore, d . p

(I) ≤

Step 3b: Bound on Term (II). Since A0 ≤ 1, P ≥ p, and Z ≥ Z, (II) ≤

1 X Bj − Bj0 . pZ j6=i

d . pZ

Step 3c: Bound on Term (III). We expand the denominator diﬀerence as P0 Z0 − P Z (P0 − P )Z0 + P (Z0 − Z) 1 1 = = . − PZ P0 Z0 P Z P0 Z0 P Z P0 Z0 Taking absolute values, the denominator is bounded below by P Z P0 Z0 ≥ p 2 Z 2 . For the numerator, since P ≤ 1 and Z0 ≤ 1, (P0 − P )Z0 + P (Z0 − Z) ≤ |P0 − P | + |Z0 − Z|. Using the shared-source structure of the bridges, P0 − P = pt (i) − pcand (i) = t (π̃(i) − π̃cand (i)), t so |P0 − P | ≤ |π̃(i) − π̃cand (i)| ≤ d (pointwise dominated by L1 ). Likewise, the 1-Lipschitz property applied coordinate-wise and summed gives |Z0 − Z| =

(ṗt (a))+ −

(ṗcand (a))+ ≤ t

(ṗt (a))+ − (ṗcand (a))+ ≤ d. t

Hence the numerator is bounded by 2d. The prefactor satisﬁes A0 ≤ 1 and P A0 j6=i Bj0 ≤ 1. Therefore, 2d (III) ≤ 2 2 . p Z

0 j6=i Bj

≤ Z0 ≤ 1, so

Step 4: Combining the three bounds. Adding the three contributions, X

u∗,cand (i → j) − u∗t (i → j) ≤ t

j6=i

d d 2d + + 2 2. p pZ p Z

Since p, Z ∈ (0, 1], 1 1 ≤ 2 2. pZ p Z

1 1 ≤ 2 2, p p Z

All three terms can therefore be upper bounded by the common denominator, yielding X ∗,cand 4 ut (i → j) − u∗t (i → j) ≤ 2 2 kπ̃cand − π̃k1 . p Z j6=i This establishes the theorem with universal constant C = 4. Remark C.1. The theorem isolates the error caused by restricting the terminal target policy to Acand . When the candidate bridge also changes the source distribution from p0 to pcand , an additional term kpcand − p0 k 1 0 0 cand enters the bound through kṗt − ṗt k1 . Corollary C.2 (Uniform-in-Time Stability). Assume that the full and candidate bridges use the same source distribution p0 . Fix δ ∈ (0, 1) and suppose that, for all t ∈ [0, 1 − δ], pt (i), pcand (i) ≥ p > 0, t

Zt , Ztcand ≥ Z > 0.

Then, for every such t and every i ∈ Acand , X ∗,cand ut (i → j) − u∗t (i → j) ≤ j6=i

C kπ̃cand − π̃k1 , p2 Z 2

where C > 0 is a universal constant. If p and Z are deterministic constants, then combining this bound with Proposition 4.1 gives   X ∗,cand 2C X E ut (i → j) − u∗t (i → j)  ≤ 2 2 π̃(a | s)(1 − πref (a | s))Nroll (1 − 1/K)Nrand . Z p a∈A j6=i

Remark C.3. Using a uniform prior p0 and truncating the ﬂow at t ≤ 1 − δ ensures pt (i) ≥ δ/|A|, providing a practical lower bound for p.

Ablations

All ablation experiments use the Stochastic Goal Gridworld (Toy 5, |A| = 64): a 20×20 grid with four goals whose locations are hidden from the agent. Two “warm” goals G0–G1 (reward +10) appear in the oﬄine data, two “cold” goals G2–G3 (reward +15) are never seen during oﬄine training. A cross-wall layout with 2-cell bottleneck gaps forces indirect navigation. Step penalty −0.2, max 40 steps per episode. This environment isolates oﬄine-to-online transfer and cold-goal discovery. Unless otherwise stated, non-swept hyperparameters are held at α=0.15, M =10, β=0.4, ρ=0.25, averaged over 3 seeds. Figure 2 and Table 4 present the complete results; the top row shows mean reward and the bottom row shows cold-goal visits (G2+G3) across 100 evaluation episodes.

Table 4: Complete ablation results on Toy 5 (mean ± std, 3 seeds). “Cold” = visits to G2+G3 (unseen in oﬄine data, reward +15) out of 100 evaluation episodes.

D.1

Ablation

Conﬁg

Reward

Std

Cold visits

Runtime

KL weight α

α = 0.0 α = 0.05 α = 0.1 α = 0.2 α = 0.5

+2.99 −2.30 −0.55 −3.38 −3.12

1.96 4.64 4.50 3.29 1.66

38 21 24 19 21

169 min 180 min 180 min 182 min 202 min

CTMC steps M

M =2 M =5 M = 10 M = 20

−5.05 −5.11 −2.04 −2.44

0.46 0.65 2.11 1.87

9 10 19 22

59 min 114 min 179 min 330 min

Temperature β

β = 0.2 β = 0.4 β = 0.8 β = 1.5 β = 3.0

−6.26 −0.62 −1.67 −0.55 −1.85

0.34 2.73 0.27 1.41 0.65

5 20 5 4 0

— — — — —

Oﬄine frac. ρ

ρ = 0% ρ = 10% ρ = 25% ρ = 50%

−5.52 −0.51 −0.39 +2.04

2.24 4.55 2.95 1.28

11 28 21 32

— — — —

KL Weight α

α controls how strongly the ﬁne-tuned generator is regularised toward the pre-trained reference. On Toy 5, α=0 achieves the highest reward (+2.99) and the most cold-goal visits (38/100); all non-zero values perform worse. This is expected: Toy 5 has no distribution shift between oﬄine and online phases, so KL regularisation restricts exploration of cold goals without any compensating stability beneﬁt—the pre-trained policy was trained only on warm goals, so penalising drift from it directly penalises cold-goal discovery. The KL penalty’s value surfaces when forgetting is costly, as demonstrated in the Goal-Switch environment (Section E.4), where the oﬄine skills remain useful after the reward change.

D.2

CTMC Sub-Steps M

M controls the resolution of the CTMC Euler discretisation. There is a clear quality threshold at M =10: using M =2 or M =5 causes a ∼3-point reward drop (to ≈−5.1) and reduces cold-goal visits from 19 to 9. Going to M =20 provides no further reward improvement but nearly doubles runtime (330 vs. 179 min). M =10 is therefore the right operating point: suﬃcient ﬂow resolution without the cost of ﬁner discretisation.

(b) CTMC steps M

−2.5 −5.0 0 .05 0.1 0.2 0.5 0

30 20 10 0

−4 −6

Co d goa )isits

−7.5

−2

0 .05 0.1 0.2 0.5 0

2.5

0 −2 −4 −6

0.0 −2.5 −5.0 −7.5

0.2 0.4 0.8 1.5 3.0

M= M= M=1 M=2

0% 10% 25% 50%

20 15 10 5 0

Rewa%d

0.0

Reward

2.5

(d) Off ine f%action ρ

15 10 5 0

2 5 0 0 M= M= =1 =2

M M

Co d goa )isits

(a) KL weight α

Co d goa )isits

5.0

0.2 0.4 0.8 1.5 3.0

30 20 10 0

0% 10% 25% 50%

Figure 2: Full hyperparameter ablation on Toy 5 (|A|=64, 3 seeds). Top: mean evaluation reward (higher is better). Bottom: cold-goal visits out of 100 evaluation episodes (higher = better transfer to unseen goals). Dark bars mark the best value in each sweep. (a) KL weight: α=0 best (no distribution shift in this env). (b) CTMC steps: sharp quality threshold at M =10. (c) Temperature: β=0.4 best balances sharpness and exploration. (d) Oﬄine fraction: monotonic improvement up to ρ=0.5.

Table 5: Cold-start ablation: contribution of each pre-training stage on Toy 5 (mean ± std, 3 seeds). Conﬁguration

Reward

Std

Cold visits

Runtime

Full (Stage 1a + 1b) No Stage 1b (no generator) No Stage 1a (no critic) No pre-training

−2.04 −3.54 +0.20 −3.71

2.11 1.71 6.04 2.50

19 12 24 20

182 min 182 min 178 min 164 min

(a) Pre-training ablation:

mean reward +0.20

Eval reward

4 2 0 −2 −4

-2.04

−6 −8

Full (1a+1b)

-3.54

No 1b (no gen)

-3.71 No 1a (no critic)

Cold goal visits (G2+G3)

(b) Pre-training ablation: cold goal discovery

10 5 0

No re-train

Full (1a+1b)

No 1b (no gen)

No 1a (no critic)

No re-train

Figure 3: Cold-start ablation on Toy 5. (a) Mean reward: removing Stage 1a gives the highest mean but the largest variance. (b) Cold-goal discovery: all conﬁgurations ﬁnd cold goals, but the full method delivers the most reliable balance.

D.3

Temperature β

β controls the sharpness of advantage weighting in π̃(a|s) ∝ πref (a|s) exp(Ā(s, a)/β). Very low β=0.2 makes the target so sharp that the policy locks into oﬄine-biased Q-values and barely reaches any goal (reward −6.26). Very high β≥1.5 spreads probability too uniformly: the policy ﬁnds warm goals but never discovers cold ones (0 cold visits at β=3.0). The sweet spot is β=0.4, achieving the best reward (−0.62) and the most balanced goal coverage (14 warm + 20 cold visits).

D.4

Oﬄine Replay Fraction ρ

ρ controls how much of each training minibatch comes from the oﬄine dataset during ﬁne-tuning. The result is unambiguous: reward increases monotonically from ρ=0 (−5.52) to ρ=0.5 (+2.04), with the largest single gain at the 0→10% jump (+5.01 reward). Cold-goal visits also increase monotonically (11 at 0%, 32 at 50%). Stable Q-values anchored by oﬄine data enable more eﬀective exploration of unseen goals, directly supporting the paper’s claim that mixed replay is essential for oﬄine-to-online transfer.

D.5

Pre-Training Ablation (Cold-Start)

Table 5 and Figure 3 isolate the contribution of each pre-training stage. Removing Stage 1b (generator pre-training via advantage-weighted DFM) costs 1.5 reward points (−3.54 vs. −2.04) and halves cold-goal discovery (12 vs. 19 visits). Removing Stage 1a (critic pre-training) gives the highest mean reward (+0.20) but with very high variance (std = 6.04), indicating unreliable performance across seeds. Removing both stages (−3.71) performs similarly to removing Stage 1b alone. The full method achieves the best balance of performance and consistency. Pre-training costs only ∼0.2 minutes on CPU— negligible against the ∼180-minute online phase—and its beneﬁt will be more pronounced on harder tasks where online steps alone are insuﬃcient for convergence.

Table 6: Path-space KL behaviour across experiments. KL behaviour

Consistent with theory?

Toy 3 (|A| = 16) Toy 4 (|A| = 128) Toy 5 (|A| = 64) MinAtar Asterix

≈ 0 throughout ≈ 0 throughout ≈ 0 throughout Spike to 1–2, settles by 50K

Yes (simple task) Yes (simple task) Yes (no shift) Yes (strongest evidence)

Pa h-space KL es ima e

Environment

Pa h-space KL during fine- uning (MinA ar As erix)

2.0

Ac ive regularisa ion

1.5 1.0 0.5

Bounded (no forge ing)

0.0

100 150 200 250 300 Gradient s ep (×103)

Figure 4: Path-space KL divergence during online ﬁne-tuning on MinAtar Asterix (300,000 gradient steps, smoothed). The KL spikes early and settles to near zero by step 50,000, conﬁrming bounded divergence from the reference policy.

D.6

Path-Space KL Behaviour

Figure 4 shows the path-space KL divergence during online ﬁne-tuning on MinAtar Asterix. The KL displays a characteristic “spike-then-settle” pattern: it rises to ≈1–2 during the ﬁrst 50,000 gradient steps as the policy actively updates, then decays to near zero as the reference generator is refreshed and the two policies converge. This is direct evidence that the regularisation is active—it genuinely constrains the rate of policy change rather than being dormant. Table 6 summarises KL behaviour across all environments, conﬁrming consistency with the theoretical bounded-divergence guarantee.

Toy Experiments

E.1

Environments

We evaluate DRIFT on ﬁve controlled environments, each designed to isolate a speciﬁc property of the method. All toy experiments use two-layer MLPs with 256 hidden units, Adam optimiser (η=3×10−4 ), discount γ=0.95, and soft-update rate τ =0.005. Results are averaged over 3 seeds (42, 123, 2024) unless stated otherwise.

Toy 3: Multimodal Gridworld (|A| = 16). A 6×6 grid with 16 actions (8 directions × 2 step sizes) and three goals at (0, 0), (0, 5), (5, 5) giving rewards +10, +10, +6. The agent starts at centre; a wall barrier at row 2 forces indirect navigation; step penalty −0.1. The oﬄine dataset (400 episodes) is biased 85% toward step-1 actions and 60% toward goal (0, 0). Tests: multimodal policy coverage despite biased oﬄine data.

Table 7: Toy 3 results: Multimodal Gridworld (6×6, |A|=16). Goal visits from 50 evaluation episodes. Both DRIFT and DQN hit the reward ceiling, but DQN mode-collapses to a single goal. Method

Reward

Std

G1 (+10)

G2 (+10)

G3 (+6)

Coverage

DRIFT (ours) DQN PPO

10.00 10.00 7.76

0.00 0.00 3.68

17 0 2

33 50 32

0 0 8

bimodal single mode partial

DRIFT oﬄine-only DQN oﬄine-only

2.92 0.00

4.49 —

9 0

5 0

1 0

trimodal collapsed

Toy 4: Large-Action Gridworld (|A| = 128). A 12×12 grid with 128 actions (8 directions × 16 step sizes). Three goals at (0, 0), (0, 11), (11, 11) give rewards +10, +10, +6. Optimal value from start state (6, 5) is V ∗ (s0 )=9.5 (by value iteration). Many actions correspond to large steps into walls and are eﬀectively “dead.” Oﬄine dataset: 600 biased episodes. Tests: candidate-set approximation in a genuinely large, sparse action space.

Toy 5: Stochastic Goal Gridworld (|A| = 64). Described above (Section D). Tests: oﬄine-toonline transfer and discovery of unseen high-reward goals. Serves as the primary ablation testbed.

Goal-Switch Environment (|A| = 16). A 6×6 grid with 16 actions. The reward structure changes at step 25,000: goal moves from position A (Phase 1) to position B (Phase 2). The oﬄine dataset is collected under Phase 1. Tests: adaptation to reward shift while retaining oﬄine navigation skills; primary demonstration of path-space KL regularisation. Combinatorial Lock Gridworld (|A| = 144). A 16×16 grid with 144 actions (12 directions × 12 step sizes) and four quadrants, each with a cluster of 3 keys and 1 goal chamber. At each episode, one quadrant is randomly designated “active” (hidden from the agent); the agent must collect all 3 keys from that quadrant’s cluster, then enter its goal. Picking up any key commits the agent to that quadrant. Rewards: +1/key, +10 for correct goal with all keys, −3 for wrong goal, −0.1/step, max 80 steps. Oﬄine data covers only quadrants 0 and 1. Tests: multimodal action coverage and candidate-set utility in a large action space where random exploration is ineﬀective.

E.2

Toy 3: Multimodal Coverage

Table 7 and Figure 5 present results on the Multimodal Gridworld. Both DRIFT and DQN reach the reward ceiling of 10.0, while PPO achieves only 7.76 within the 3,000-step budget. The critical diﬀerence is in policy structure. DQN sends all 50 evaluation episodes to goal G2 at (0, 5), completely ignoring the two other goals— classic mode collapse, where the deterministic arg max policy commits to the single goal with the highest Q-value. DRIFT distributes visits across goals (17 to G1 , 33 to G2 ), preserving coverage of multiple modes. The oﬄine-only comparison is equally telling. Oﬄine DQN scores 0.00 despite converging Q-loss, due to greedy action extraction causing action cycles on out-of-distribution states. The DRIFT pre-trained policy achieves 2.92 and visits all goal types, conﬁrming that the CTMC representation transfers more robustly from oﬄine data than deterministic value-based extraction.

E.3

Toy 4: Candidate-Set Phase Transition

Toy 4 validates the candidate-set approximation theory (Proposition 4.1) in a controlled large-action-space setting.

(a) Toy 3: M( timoda Gridwor d (|| = 16) 12

Mean reward

10.00

Goa )isits / 50 episodes

7.76

10.00

8 6 4 2 0

DRIFT (o(rs)

DQN

PPO

(b) Goa )isitation (mode co)erage) DRIFT DQN PPO

50 40 30 20 10 0

G1 (+10) G2 (+10) G3 (+6)

Figure 5: Toy 3 results. (a) Mean reward: DRIFT and DQN both reach the ceiling; PPO falls short. (b) Goal visitation: DQN mode-collapses to G2 ; DRIFT maintains coverage of both +10 goals. Table 8: Toy 4 candidate-set ablation (|A|=128). The phase transition occurs at budget 24 (≈16% coverage). Budget

Nroll

Nrand

Coverage

Mean reward

4 8 12 16 24 32 64 128

1 2 4 5 8 10 21 42

3 6 8 11 16 22 43 86

3.1% 6.0% 9.2% 11.8% 15.9% 19.2% 35.2% 57.4%

8.8 8.1 7.7 8.7 10.0 10.0 10.0 10.0

Candidate-set ablation. Figure 6(a) shows ﬁnal reward as the candidate budget |Acand | increases from 4 to 128. There is a sharp phase transition at budget 24 (≈16% action-space coverage): below this threshold performance is unstable; above it the agent reliably achieves the optimal reward of 10.0 (Table 8). This directly validates the theoretical prediction that generator error decays as candidate coverage increases. Method comparison. Figure 6(b) shows that DQN and PPO achieve a perfect score of 10.0 online, while DRIFT scores 9.36. However, the DRIFT pre-trained policy (before any online interaction) already achieves 9.52, whereas DQN oﬄine scores 0.00. The advantage-weighted DFM pre-training thus produces a usable policy from oﬄine data alone, something greedy DQN extraction cannot. Goal diversity. Figure 6(c) shows that DRIFT distributes episodes across all three goals (13, 29, 8), while DQN sends all 50 to G2 —replicating the Toy 3 mode-collapse ﬁnding at |A|=128.

E.4

Goal-Switch: Adaptation to Distribution Shift

The Goal-Switch environment measures how quickly each method adapts when the reward structure changes at step 25,000. Table 9 and Figure 7 present results. DRIFT reaches the switched goal in 1.42 steps on average—1.9× faster than DQN (2.69) or PPO (2.63). All three methods ultimately solve the task, but DRIFT’s faster adaptation means less wasted exploration

(a) Can i ate-sephase -ransi-ion

(b) Perf)rma(ce (|A| = 128)

10.5

Mean re0ar

Rewar

9.5 9.0 8.5 8.0 7.5

Suboptimal

7.0 0.0

0.1

0.3

0.4

Coverage |Acand |/|A|

0.5

0.6

8 6 4

10.00

Optimal

0.2

10.00 9.36

G)al visits / 50 ep.

10.0

9.52

40 30 20 10

DRIFT DRIFT pre--raine fine--une

DQN

PPO

(+10)

G2 (+10)

G3 (+6)

Figure 6: Toy 4 results (12×12, |A|=128). (a) Candidate-set phase transition: performance saturates above 16% coverage (dashed line); V ∗ =9.5 shown dotted. (b) Method comparison: DRIFT pre-trained achieves 9.52 without any online interaction; DQN oﬄine scores 0.00. (c) Goal diversity: DRIFT visits all three goals; DQN mode-collapses. Table 9: Goal-Switch results (mean ± std, 3 seeds). “Steps to goal” = environment steps to reach the new goal location after the reward switch at step 25,000 (lower is better). Metric

DRIFT (Ours)

DQN

PPO

1.42 ± 0.22 100.0 14.33 ± 0.08

2.69 ± 0.49 99.0 14.07 ± 0.11

2.63 ± 0.43 90.6 14.27 ± 0.21

Steps to goal ↓ Goal rate (%) Mean eval reward

after the switch. This speed advantage comes from the CTMC generator’s ability to rapidly redistribute probability mass over actions when the target policy π̃ changes. DQN must propagate new Q-values through slow Bellman backups before its arg max policy shifts; DRIFT’s ﬂow-matching objective directly reshapes the action distribution in a single generator update.

E.5

Combinatorial Lock Gridworld

The Combinatorial Lock environment (|A|=144) combines a large action space, sequential key collection, and hidden quadrant selection to create genuine diﬃculty. Table 10 presents results. DRIFT achieves the highest mean reward (−2.92) and by far the highest success rate (23.5%), versus DQN (2.0%) and PPO (0.5%). Negative mean rewards reﬂect the step penalty accumulated even in successful episodes. The success rate is the more informative metric: DRIFT’s 12× advantage over DQN demonstrates that the CTMC generator with candidate sets can navigate a complex sequential task that deterministic policies cannot solve through exploration alone.

E.6

Summary of Toy Findings

Table 11 summarises the key ﬁnding from each environment. Across all four properties tested—multimodal coverage, candidate-set scaling, fast adaptation, and sequential task solving—DRIFT’s advantages stem from the same core mechanism: the CTMC generator maintains a distribution over actions rather than committing to a single best action, and can reshape that distribution directly from a ﬂow-matching objective rather than through Bellman backups. The KL penalty’s value is environment-dependent: it helps when forgetting is costly (Goal-Switch) but can hinder exploration when there is no distribution shift (Toy 5).

(a) Po()-(wi)c adap)a)ion (peed 1.9 × faster

3.5 3.0

2.69

105

2.63

2.5 2.0 1.5

1.42

1.0 0.5 0.0

DRIFT (our()

DQN

Goal comple)ion ra)e (%)

S)ep( )o (wi)c ed goal ↓

4.0

PPO

100

(b) G%al c%mple)i%n af)er (wi)c

100.0%

99.0%

90.6%

90 85 80 DRIFT (our()

DQN

PPO

Figure 7: Goal-Switch results. (a) Steps to reach the switched goal (lower is better): DRIFT adapts 1.9× faster than DQN. (b) Goal completion rate: DRIFT achieves 100%; PPO lags at 90.6%. Error bars show ±1 std, 3 seeds. Table 10: Combinatorial Lock results (16×16, |A|=144, mean ± std, 3 seeds). Success = all 3 keys collected and correct goal entered. Mean reward ↑ Success rate (%)

DRIFT (Ours)

DQN

PPO

−2.92 ± 1.51 23.5

−3.97 ± 1.57 2.0

−5.64 ± 0.62 0.5

Benchmark Experiments

F.1

Shared Hyperparameters

Table 12 lists the default DRIFT hyperparameters used across all benchmark experiments. Environmentspeciﬁc overrides are noted in the relevant subsections below.

F.2

MinAtar

F.2.1

Environment and Oﬄine Data

MinAtar [Young and Tian, 2019] implements ﬁve simpliﬁed Atari games on a 10×10 grid with binary feature channels: Breakout (3 actions), Asterix (5), Freeway (3), Seaquest (6), Space Invaders (4). Observations are 10×10×nch binary tensors (nch =4–10 depending on game); rewards are sparse integers. Every method receives the same oﬄine dataset per game: 100,000 transitions from a behavioural DQN with ε linearly annealed from 1.0 to 0.1 over the ﬁrst 50,000 steps. The DQN trains online during collection (η=3×10−4 , target soft-update τ =0.005, 50K-transition replay buﬀer). The dataset is cached to disk and reused across all methods and seeds to guarantee identical oﬄine data.

F.2.2

Network Architecture

All methods share the same encoder: ﬂatten the 10×10×nch observation, then Linear(100·nch , 256)→ReLU→Linear(256, 128)→ReLU, yielding a 128-dim representation. Q-networks add Linear(128, |A|); value networks add Linear(128, 1). The DRIFT generator (CTMCGen) appends a rate head: Linear(128+|A|+1, 128)→ReLU→Linear(128, |A|)→Softplus. The input concatenates the encoded observation, a one-hot current-action vector, and the scalar ﬂow time t. The diagonal is masked to zero (no self-transitions), yielding non-negative oﬀ-diagonal rates.

Table 11: Summary of toy experiments. ↑ = DRIFT wins; ∼ = comparable; ↓ = DRIFT loses on that metric.

Environment

|A|

DRIFT

DQN

PPO

Toy 3 reward Toy 3 diversity Toy 4 score Toy 4 oﬄine Toy 4 diversity Toy 4 ablation Goal-Switch Comb. Lock

16 16 128 128 128 128 16 144

10.00 10.00 7.76 3 goals 1 goal — 9.36 10.00 10.00 9.52 0.00 — 3 goals 1 goal — phase transition at 16% coverage 1.42 2.69 2.63 23.5% 2.0% 0.5%

Key ﬁnding ∼ (both hit ceiling) ↑ multimodal ↓ (score) ↑ oﬄine transfer ↑ multimodal validates Prop. 6.2 ↑ 1.9× faster adaptation ↑ success rate

Table 12: Default DRIFT hyperparameters.

F.2.3

Symbol

Description

Value

γ β α τ c ǫ M Nroll Nrand K ρ ηQ , ηV ηθ B δ

Discount factor Temperature KL weight Soft-update rate Advantage clip Smoothing constant CTMC Euler sub-steps Reference rollouts per state Uniform candidates per state Reference refresh interval Oﬄine mix fraction Critic / value learning rate Actor learning rate Batch size Flow time truncation

0.99 1.0 0.1 0.005 5.0 0.01 20 64 16 1000 0.25 3 × 10−4 1 × 10−4 256 0.05

Training Protocol

Oﬄine phase. Critics are trained for 2,000 gradient steps on the shared oﬄine dataset. Method-speciﬁc oﬄine components run in parallel: IQL runs 2,000 expectile regression steps; AWAC runs 1,000 actor-critic steps; DRIFT trains the CTMC generator for 1,000 epochs of advantage-weighted DFM. Each method is then evaluated over 100 episodes to obtain the oﬄine score. Online phase. Each method interacts with the environment for 300,000 steps. Every training minibatch mixes online and oﬄine data: (1−ρ)·B transitions from the online replay buﬀer and ρ·B from the oﬄine dataset. The ﬁnal policy is evaluated over 100 episodes.

Table 13: DRIFT hyperparameters for MinAtar.

Hyperparameter

Symbol

Value

Discount factor Soft-update rate Learning rate Temperature Path-space KL weight CTMC sub-steps Euler step size Reference refresh Advantage clip Oﬄine replay fraction Critic batch size Actor batch size Critic pre-training Generator pre-training Online steps Replay buﬀer capacity πref smoothing

γ τ η β α M ∆t K c ρ B Bactor Ecritic Egen Nsteps

0.99 0.005 10−4 0.5 0.1 10 1/M 500 steps 3.0 0.25 64 8 2,000 steps 1,000 epochs 300,000 100,000 10−3

F.2.4

DRIFT Hyperparameters

F.2.5

MinAtar Compute

All experiments ran on a university cluster with NVIDIA P4 and V100 GPUs. DRIFT takes ∼8.7 h per run; most baselines take 3.5–6 h (Table 14). The full benchmark (10 methods × 5 games × 3 seeds = 150 runs) consumed approximately 600 GPU-hours. Seeds are set globally for Python, NumPy, and PyTorch; CUDA determinism is enforced via torch.backends.cudnn.deterministic = True and benchmark = False.

Table 14: Mean wall-clock runtime per MinAtar run (one game, one seed). Hours

DRIFT

CQL

Cal-QL

IQL

AWAC

DQN

Rainbow

PPO

PEX

SPA

∼8.7

∼3.5

∼4.2

∼3.5

∼6.0

∼2.5

∼4.5

F.3

Jericho (Interactive Fiction)

F.3.1

Environment and Action Space

Jericho [Hausknecht et al., 2020] provides a Python interface to Z-machine interactive ﬁction games. At each step the agent receives a textual observation and issues a text command from a set of valid actions identiﬁed by the game engine via get_valid_actions(). The number of valid actions per state ranges from 3 to 10 on average across our ten evaluation games (Table 15). This is smaller than the 30–200 range reported in prior work [Yao et al., 2020, Ammanabrolu and Hausknecht, 2020] because those methods augment the action space with language model or template-based generation, while our setup uses only the built-in parser output.

F.3.2

Architecture

All methods share the same text-encoding architecture, adapted from DRRN [He et al., 2016]. Words are tokenised via hash-based indexing (vocabulary size 5003, using MD5 hashing) and embedded into 64 dimensions. A single-layer GRU with hidden size 128 produces ﬁxed-size representations for both observations and actions. Each network (Q1, Q2, V, generator) maintains its own GRU encoder parameters. The Q-network scores each state action pair via Q(s, a) = MLP([hs ; ha ]) with one hidden layer of 128 units and ReLU activation, naturally handling variable-size action sets by scoring each pair independently. The value network computes V (s) = MLP(hs ) with the same hidden layer structure. The CTMC generator uses a bilinear-attention design uθ (i → j, t | s) = softplus c(hs , hai , t) · haj ,

where the context vector c is produced by an MLP from the concatenation of the observation embedding, source action embedding, and scalar ﬂow time t. Rates to each target action are computed via dot products with target action embeddings. Self-transition rates (j = i) are masked to zero. This design handles arbitrary action set sizes without a ﬁxed output head.

F.3.3

Candidate Set Construction

For each actor-update state s, the candidate set Acand (s) is constructed by (1) running Nroll = 8 CTMC rollouts of the frozen reference generator over all valid actions and collecting terminal action counts, (2) sampling Nrand = 32 valid actions uniformly at random, and (3) taking the union after deduplication. The reference distribution uses additive smoothing πref (a | s) ∝ count(a)/Nroll + ε/|Acand | with ε = 10−3 . In practice, the valid action sets in our ten Jericho games average 3 to 9 actions per state, which falls below the candidate budget of 40. In these cases the algorithm uses all valid actions directly without subsampling. The candidate-set mechanism thus operates in its full-coverage regime (ε(s) = 0 in Proposition 4.1), and DRIFT’s advantage on this benchmark stems from the path-space KL regularization and the ﬂow-matching actor update rather than from action-space reduction.

F.3.4

Oﬄine Data Collection

For each game, we collect 200 episodes of oﬄine data using an ε-greedy DQN that trains online during collection. The exploration rate ε is annealed from 1.0 to 0.1 over the ﬁrst half of the collection phase. The resulting transitions are cached to disk and reused across all methods to ensure identical oﬄine data.

F.3.5

Training Protocol

All methods receive the same oﬄine dataset per game and run for 300,000 online interaction steps with a maximum episode length of 200 steps. Critic pre-training runs for 5,000 gradient steps and generator pre-training runs for 2,000 gradient steps, both on the oﬄine dataset. Evaluation uses 20 complete episodes at the end of training. Results are reported for a single seed (42) due to the high computational cost of text-game experiments.

F.3.6

Comparison with Prior Work

The CALM [Yao et al., 2020] and KG-A2C [Ammanabrolu and Hausknecht, 2020] scores in Table 1 are taken from published results under diﬀerent experimental conditions. CALM uses a ﬁne-tuned GPT-2 language model pre-trained on 426 human gameplay transcripts from ClubFloyd to generate action candidates. It does not use the valid-action handicap provided by the game engine. KG-A2C builds a dynamic knowledge graph from game text using Stanford OpenIE, encodes it with a graph neural network, and trains with A2C. It uses the valid-action handicap and template-based action generation. Our methods (DRIFT and all baselines) use the valid-action handicap, GRU text encoders without any pre-trained language model, and 300K online steps. This is a simpler and more self-contained setup that ensures fair comparison among the evaluated methods.

Other recent work on Jericho includes XTX [Tuyls et al., 2022] (ICLR 2022 Spotlight) which achieves state-of-the-art scores via multi-stage episodic control (for example 103 on Zork1) and CSM [Shi et al., 2023] which improves CALM via conﬁdence-based self-imitation. These methods use substantially more compute and game-speciﬁc engineering. We report their existence for context rather than direct comparison.

F.3.7

Per-Game Details

Table 15: Jericho game characteristics. Average valid actions per state measured during DRIFT evaluation (20 episodes). Game Deephome Detective Enchanter Ludicorp Omniquest Pentari Temple Zork1 Zork3 Ztuu

Max score

Avg. valid actions/state

300 360 400 150 50 70 35 350 7 100

7.7 3.7 8.6 2.9 4.3 4.1 8.7 6.4 7.2 6.4

Genre Exploration Mystery Fantasy Puzzle Adventure Adventure Dungeon Dungeon crawl Dungeon crawl Dungeon

We excluded Balances (max score 51) from the main table because all ﬁve methods converge to an identical score of 10.0 (19.6%), providing no discriminative signal within the 300K step budget.

F.3.8

Hyperparameters

F.3.9

Compute

All Jericho experiments ran on a university cluster with NVIDIA Tesla P4 GPUs (12 GB VRAM), 4 CPU cores, and 16 GB RAM per job. Table 17 reports per-game runtimes. The full Jericho benchmark (10 games × 5 methods = 50 runs) consumed approximately 850 GPU-hours. Seeds are set globally for Python, NumPy, and PyTorch with CUDA determinism enforced.

F.4

D4RL

F.4.1

Discretisation Procedure

We discretise the continuous action spaces of Hopper (3-dim), Walker2d (6-dim), and HalfCheetah (6-dim) using MiniBatchKMeans with k=22 clusters ﬁtted on all actions in the oﬄine dataset. Each continuous action is replaced by its nearest cluster index; during online interaction the agent selects a cluster index and executes the corresponding centroid. The median inter-centroid distance ranges from 0.27 (Hopper) to 0.46 (HalfCheetah), conﬁrming reasonable but coarse coverage of the continuous action manifold. We use the Minari versions of the D4RL datasets (v0.5.3) with medium and expert qualities, each containing ≈1M transitions.

F.4.2

Training Protocol

All methods share the same budget: 1M oﬄine gradient steps for critic pre-training, followed by 300k online environment steps. The critic is a 2-layer MLP (256 hidden units, batch size 256, Adam η=3×10−4 ). DRIFT’s generator uses the same architecture with actor batch size 8 and 500K generator pre-training steps. The online phase uses a mixed replay buﬀer with 50% oﬄine data and 200K capacity.

Table 16: DRIFT hyperparameters for Jericho. Parameter

Value

Oﬄine episodes Online steps Max episode steps Evaluation episodes Discount γ Temperature β Path-KL weight α CTMC sub-steps M Reference refresh K Oﬄine mix ratio ρ Candidate budget Advantage clip c Delayed actor updates Kactor Critic batch size Actor batch size Learning rate (all) Soft-update τ Critic pre-training steps Generator pre-training steps GRU hidden size Word embedding dimension Vocabulary size (hash-based)

F.4.3

200 300,000 200 20 0.99 0.1 0.01 5 50 0.75 40 (Nroll =8, Nrand =32) 3.0 4 32 8 10−4 0.005 5,000 2,000 128 64 5,003

Improvement Analysis

Table 18 isolates ﬁne-tuning eﬀectiveness by reporting absolute improvement (online − oﬄine score) for each method. DRIFT’s average improvement (+24.0) exceeds all other methods despite starting from the weakest oﬄine initialisation (0.4 average oﬄine score). The gap in ﬁnal online score relative to PEX and Cal-QL is therefore attributable to the oﬄine phase, not the ﬁne-tuning mechanism. The strongest results appear on Hopper, where the 3-dimensional action space is better covered by 22 cluster centres than the 6-dimensional spaces of Walker2d and HalfCheetah.

F.4.4

Limitations of the Discretised Setting

D4RL was designed for continuous-action algorithms. Discretisation via k-means introduces an approximation ceiling that aﬀects all methods equally, but also obscures each method’s native strengths. With |A|=22, the full action space can be enumerated cheaply, removing DRIFT’s candidate-set advantage entirely. DRIFT’s structural beneﬁts are designed for |A|≫64, as validated by the Jericho results.

F.5

Macro-Action MinAtar (|A| = 216)

F.5.1

Construction

Each macro action m ∈ {0, . . . , 215} maps to a unique 3-tuple (a1 , a2 , a3 ) via base-6 decomposition. Executing m runs the three base actions sequentially in the environment, accumulating reward. The agent observes only the ﬁnal state after all three steps. If the episode terminates mid-sequence the cumulative reward up to termination is returned. This construction expands the action space from 6 to 63 = 216 while preserving the original game dynamics. DRIFT operates on candidate sets of budget 40 (18.5% coverage of the macro-action space). The

Table 17: Mean wall-clock runtime per Jericho run (one game, one method, seed 42). Game

DRIFT

DRRN

DQN

IQL

AWAC

Deephome Detective Enchanter Ludicorp Omniquest Pentari Temple Zork1 Zork3 Ztuu

40.0 h 16.9 h 16.6 h 19.6 h 37.6 h 15.2 h 37.3 h 17.7 h 13.7 h 16.5 h

34.1 h 14.5 h 18.5 h 18.7 h 17.6 h 14.2 h 38.4 h 12.2 h 10.7 h 33.3 h

32.4 h 14.0 h 15.7 h 19.1 h 19.2 h 15.0 h 35.0 h 13.2 h 10.4 h 51.7 h

31.3 h 11.1 h 12.1 h 13.9 h 38.8 h 10.0 h 34.5 h 9.9 h 7.1 h 10.8 h

30.4 h 11.4 h 11.6 h 15.7 h 23.3 h 13.1 h 33.7 h 2.5 h 8.2 h 13.0 h

Table 18: D4RL absolute improvement (online − oﬄine). Red = degradation from oﬄine score.

Task

DRIFT

CQL

IQL

AWAC

Cal-QL

DQN

PPO

PEX

SPA

Hopper-med Hopper-exp Walker-med Walker-exp Cheetah-med Cheetah-exp

+47.8 +46.7 +14.9 +14.7 +10.9 +9.2

+7.0 +20.5 +10.1 +6.7 +12.5 +7.9

+27.5 +17.1 +12.7 +10.3 +7.6 +6.6

−0.7 −0.3 +1.6 −0.4 −0.4 −0.1

+15.3 +16.9 +9.2 +8.4 +12.5 +7.8

+23.3 +22.0 +10.5 +6.3 +14.8 +4.6

— — — — — —

+42.7 +39.4 +12.1 +7.0 +14.4 +7.5

0.0 0.0 0.0 0.0 0.0 0.0

Average

+24.0

+10.8

+13.6

−0.1

+11.7

+13.6

+5.3

+20.5

0.0

CTMC generator uses learned action embeddings (128 dimensions) with bilinear rate computation, replacing the one-hot input used for standard MinAtar. The Q-network output layer is 128 → 216.

Table 19: Macro-Action MinAtar (|A|=216, k=3): mean online score, 5 seeds. DRIFT uses candidate budget 40 (18.5% coverage). Best per game in bold. Game

F.5.2

DRIFT

DRIFT (full)

DQN

DQN-Sub

PPO

Breakout Asterix Freeway Seaquest Space Invaders

8.40 0.78 0.56 1.20 9.45

7.80 0.56 0.21 0.85 6.40

11.98 0.67 0.66 0.52 10.36

7.45 0.53 0.00 0.98 10.77

5.83 0.65 27.90 0.66 14.25

Average

4.08

3.16

4.84

3.95

9.86

Results and Analysis

DQN and PPO substantially outperform DRIFT on this benchmark (Table 19). DRIFT with candidate sets (4.08 average) and DRIFT without candidate sets (3.16 average) perform similarly, conﬁrming that the candidate-set mechanism is not the bottleneck. We attribute the gap to three factors. First, the behavioural DQN that collects the oﬄine dataset explores all 216 actions via ε-greedy, producing broad coverage. The critic learns reasonable Q-values across

this space, but the CTMC generator processes only 8 states per actor update and cannot absorb this breadth, resulting in a weak oﬄine initialisation. Second, with |A| = 216 the action space is large enough to stress the generator but small enough that DQN can still enumerate all actions cheaply via a 216-output network. Third, macro actions compose reward across three sequential steps, creating a credit-assignment problem that PPO’s on-policy gradient handles more naturally than oﬄine Q-learning.

F.5.3

Implications

These results highlight an important boundary condition for DRIFT. The candidate-set mechanism provides theoretical guarantees (Proposition 4.1, Theorem 4.3) and is validated in the controlled Toy 4 ablation (Table 8), but its practical beneﬁt requires that the CTMC generator is well-initialised. When the oﬄine stage produces a weak generator, as occurs here with 216 actions and only 8 states per actor update, the online phase cannot recover regardless of whether candidate sets are used. Improving the oﬄine initialisation is a natural direction for future work.

F.5.4

Compute

Macro-action MinAtar used the same cluster as standard MinAtar (NVIDIA P4 and V100 GPUs). Each run took approximately 3 to 8 hours depending on the method. The full benchmark (5 games × 5 methods × 5 seeds = 125 runs) consumed approximately 500 GPU-hours.

F.6

Baseline Implementation Details

All baselines were originally proposed for continuous-action or Atari-scale settings. We implement faithful discrete-action adaptations that preserve each method’s core mechanism while replacing Gaussian actors with categorical policies and evaluating all discrete actions exactly where needed.

CQL. CQL [Kumar et al., 2020] penalises Q-values on out-of-distribution actions. In discrete action spaces, the penalty is ! X exp Q(s, a) − Q(s, adata ) , LCQL = αCQL log a

with the logsumexp computed exactly over all actions (αCQL =0.1). The penalty applies during both oﬄine pre-training and online ﬁne-tuning.

Cal-QL. Cal-QL [Nakamoto et al., 2023] adds a calibration term to CQL that prevents excessive pessimism:

Lcal = βcal · E 1[Q(s, a) < Q− (s, a) − βcal ] · (Q− (s, a) − βcal − Q(s, a)) ,

with βcal =0.05 on top of the CQL penalty.

IQL. IQL [Kostrikov et al., 2022] avoids out-of-distribution queries via expectile regression on the value function: LV = E L2τ (Q(s, a) − V (s)) , L2τ (u) = |τ − 1[u < 0]| · u2 ,

with τIQL =0.7. At evaluation, we compute advantages A(s, a)=Q(s, a)−V (s) over all actions and sample from π(a|s) ∝ exp(A(s, a)/β). Online ﬁne-tuning continues the expectile update with mixed replay.

AWAC. AWAC [Nair et al., 2020] trains a categorical actor via advantage-weighted regression: Lπ = −E[w(s, a) · log π(a|s)] ,

Ā(s, a) , w(s, a) = exp β

with weights clamped at 20 to prevent gradient explosion (β=0.5). The same AWR update continues online with mixed replay.

SPA. SPA [Li et al., 2026] extends AWAC with a per-state proﬁciency score based on Q-function disagreement: prof(s) = exp

1 −αSPA · |A|

X a

|Q1 (s, a) − Q2 (s, a)| .

Actor weights become w(s, a)= exp(Ā/β) + (1−prof(s)), adding a behaviour-cloning pull when the critic is uncertain (αSPA =0.5). The online phase also adapts the oﬄine replay fraction based on current proﬁciency.

DQN. DQN [Mnih et al., 2015] uses double Q-learning with ε-greedy exploration. The oﬄine phase pretrains two Q-networks and a value network. Online ﬁne-tuning uses ε annealing from 1.0 to 0.05 over the ﬁrst 150,000 steps, target soft-updates (τ =0.005), and mixed replay. Rainbow. Our Rainbow implementation [Hessel et al., 2018] includes all six extensions: (1) C51 distributional learning with 51 atoms on [Vmin , Vmax ]=[−10, 20]; (2) dueling advantage/value streams; (3) factorised noisy layers (σ0 =0.5) replacing ε-greedy; (4) prioritised replay with αPER =0.6 and IS exponent βIS annealing from 0.4 to 1.0; (5) n-step returns (n=3 online, n=1 oﬄine to avoid mixing return horizons); (6) double DQN. The Q-network is warm-started from the oﬄine critic’s encoder. Learning rate is reduced to 6.25×10−5 online. PPO. PPO [Schulman et al., 2017] is a pure online baseline with no oﬄine pre-training. It uses a shared actor-critic with the same encoder. Rollouts of 256 steps, advantages via GAE (λ=0.95, γ=0.99), 4 update epochs per rollout, minibatch size 64, clip ratio 0.2, value coeﬃcient 0.5, entropy bonus 0.01, gradient norm clipping at 0.5.

PEX. PEX [Zhang et al., 2023] maintains a frozen IQL-pretrained oﬄine policy and a learning online Qnetwork (warm-started from the oﬄine critic). At each step, both policies propose an action and the agent selects between them via Boltzmann sampling over the online Q-values: P (use aoﬀ ) =

exp(Qon (s, aoﬀ )/τPEX ) , exp(Qon (s, aoﬀ )/τPEX ) + exp(Qon (s, aon )/τPEX )

with τPEX =0.1. As the online Q-network improves it naturally dominates selection. Online training uses standard double DQN with mixed replay.

F.6.1

Baseline Hyperparameters

Table 20 collects method-speciﬁc hyperparameters. Shared parameters (γ=0.99, τ =0.005, η=10−4 , B=64, buﬀer = 300k) are omitted.

F.7

Total Compute Summary

Table 21 summarises the total computational cost of all experiments reported in this paper, including preliminary runs that informed hyperparameter selection. The full research project including preliminary experiments, debugging runs, and failed conﬁgurations consumed approximately 3,500 GPU-hours total, roughly 50% more than the experiments reported in the paper.

Broader Impact

This paper presents a reinforcement learning algorithm for discrete action spaces. The method is evaluated on game benchmarks and does not target any speciﬁc real-world deployment. We discuss potential positive and negative implications below. On the positive side, improving oﬄine-to-online RL reduces the amount of costly or risky online interaction needed to train eﬀective policies. This is valuable in domains where exploration is expensive or

Table 20: Method-speciﬁc hyperparameters for MinAtar baselines. Method

Parameter

CQL Cal-QL IQL AWAC

Conservative penalty αCQL Calibration threshold βcal Expectile τIQL Weight temperature β Weight clamp Proﬁciency scale αSPA ε range Atoms Natoms Support [Vmin , Vmax ] PER αPER IS βIS start n-step (online / oﬄine) Learning rate Clip ratio GAE λ Entropy bonus Boltzmann temperature τPEX

SPA DQN Rainbow

PPO

PEX

Value 0.1 0.05 0.7 0.5 20 0.5 1.0 → 0.05 51 [−10, 20] 0.6 0.4 3/1 6.25 × 10−5 0.2 0.95 0.01 0.1

Table 21: Total compute across all experiments. Benchmark

Runs

Seeds

GPU type

GPU-hours

Toy environments (5 envs) MinAtar (5 games) Macro-MinAtar (5 games) D4RL discretised (6 tasks) Jericho (10 games) Hyperparameter ablations

150+ 150 125 108 50 138

3 5 5 5 1 3

CPU only P4 / V100 P4 / V100 P4 / V100 P4 CPU only

<10 ∼600 ∼500 ∼400 ∼850 <20

Total

∼2,380

dangerous, such as robotics, healthcare, and autonomous systems. The path-space KL regularization speciﬁcally addresses catastrophic forgetting, which is a safety-relevant failure mode in deployed systems that must be updated without losing previously learned safe behaviors. On the negative side, more capable RL agents could be applied in domains with societal risks. Recommendation systems with large discrete action spaces (a motivating application for candidate sets) can amplify ﬁlter bubbles or manipulate user behavior if reward signals are misaligned with user welfare. Autonomous agents in text-based interfaces could be adapted for automated social engineering or spam generation, though our GRU-based agents are far below the capability threshold for such misuse. More broadly, any improvement in RL sample eﬃciency lowers the barrier to training agents in sensitive domains without adequate safety testing. We believe the beneﬁts of publishing this work outweigh the risks. The method is general-purpose and not designed for any harmful application. The benchmarks we use (MinAtar, Jericho text games, discretised MuJoCo) are standard academic testbeds with no direct path to misuse. We release our code to support reproducibility and enable the research community to build on and scrutinize our results.

Full Training Algorithm

Algorithm 3 provides the complete training procedure with all implementation details.

Algorithm 3 DRIFT: Full Training Procedure Require: Replay source Doﬀ , environment, discount γ, temperature β, KL weight α, soft-update rate τ , advantage clip c, smoothing ǫ, CTMC sub-steps M , candidate sizes Nroll , Nrand , reference refresh interval K, oﬄine mix fraction ρ, learning rates ηQ , ηV , ηθ , ﬂow-time truncation δ ∈ (0, 1) 1: Initialize critics Qφ1 , Qφ2 , value Vψ , generator uθ 2: Initialize target copies Qφ− ← Qφk , Vψ− ← Vψ k 3: Initialize reference generator uref ← uθ 4: Initialize replay buﬀer B with a subset of Doﬀ 5: Observe initial state s 6: for n = 1 to Nsteps do 7: Environment interaction: simulate CTMC under uθ from X0 ∼ Unif(A) for M Euler sub-steps and execute terminal action a = XM 8: Observe (s, a, r, s′ , d), add it to B, and set s ← s′ (reset if d) 9: Mixed minibatch: sample ⌊(1 − ρ)B⌋ transitions from B and ⌊ρB⌋ transitions from Doﬀ 10: Critic update: set y ← r + γVψ− (s′ )(1 − d) and update Qφk toward y for k = 1, 2 11: Value update: compute lagged advantage Ā− and π̃ − (a | s) ∝ exp(Ā− (s, a)/β) from {Qφ− , Vψ− }, then update k

Vψ (s) ≈ sg

" X

π̃ − (a | s) min Qφ− (s, a) k

12: 13: 14: 15:

16:

17:

Actor update: for each sampled state sb in the actor batch do Construct candidate set Acand (sb ) using Nroll rollouts of frozen uref and Nrand uniform actions Deﬁne count(a) ǫ πref (a | sb ) ∝ + , a ∈ Acand (sb ) Nroll |Acand (sb )| Compute clipped advantage Ā(sb , ·) from {Qφk , Vψ− } over Acand (sb ) and set π̃(a | sb ) = P Deﬁne the bridge:

18: 19:

πref (a | sb ) exp(Ā(sb , a)/β) ′ ′ a′ ∈Acand (sb ) πref (a | sb ) exp(Ā(sb , a )/β)

pt (a | sb ) = (1 − t) p0 (a) + t π̃(a | sb ),

Sample t ∼ Unif(0, 1 − δ), i ∼ pt (· | sb ). Construct target generator u∗t (i → · | sb ) by independent coupling on Acand (sb ) Accumulate DFM loss X 2 ℓDFM (sb ) = uθ (i → j, t | sb ) − u∗t (i → j | sb ) j∈Acand (sb ), j6=i

20: 21: 22:

where p0 (a) = Unif(A).

c b ) using frozen uref Accumulate path-space KL estimate KL(s end for Update θ using 1 X 1 Xc ℓDFM (sb ) + α KL(sb ) B b B b

− − 23: Soft-update targets: φ− ← (1 − τ )ψ − + τ ψ k ← (1 − τ )φk + τ φk , ψ 24: Refresh reference generator every K steps: if n mod K = 0, set uref ← uθ 25: end for 26: return Fine-tuned generator uθ

Record · ID 178889 · SHA-256 6200a5d9e91d05a9
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.