Title: Composite Flow Matching for Reinforcement Learning with Shifted-Dynamics Data
Venue: NeurIPS 2025
Accepted: yes
Keywords: Reinforcement Learning

Flow Matching

Dynamics Shift
Abstract: Incorporating pre-collected offline data can substantially improve the sample efficiency of reinforcement learning (RL), but its benefits can break down when the transition dynamics in the offline dataset differ from those encountered online. Existing approaches typically mitigate this issue by penalizing or filtering offline transitions in regions with large dynamics gap. However, their dynamics-gap estimators often rely on KL divergence or mutual information, which can be ill-defined when offline and online dynamics have mismatched support. To address this challenge, we propose CompFlow, a principled framework built on the theoretical connection between flow matching and optimal transport. Specifically, we model the online dynamics as a conditional flow built upon the output distribution of a pretrained offline flow, rather than learning it directly from a Gaussian prior. This composite structure provides two advantages: (1) improved generalization when learning online dynamics under limited interaction data, and (2) a well-defined and stable estimate of the dynamics gap via the Wasserstein distance between offline and online transitions. Building on this dynamics-gap estimator, we further develop an optimistic active data collection strategy that prioritizes exploration in high-gap regions, and show theoretically that it reduces the performance gap to the optimal policy. Empirically, CompFlow consistently outperforms strong baselines across a range of RL benchmarks with shifted-dynamics data.

Composite Flow Matching for Reinforcement Learning with Shifted-Dynamics Data

Abstract

Incorporating pre-collected offline data can substantially improve the sample efficiency of reinforcement learning (RL), but its benefits can break down when the transition dynamics in the offline dataset differ from those encountered online. Existing approaches typically mitigate this issue by penalizing or filtering offline transitions in regions with large dynamics gap. However, their dynamics-gap estimators often rely on KL divergence or mutual information, which can be ill-defined when offline and online dynamics have mismatched support. To address this challenge, we propose COMPFLOW, a principled framework built on the theoretical connection between flow matching and optimal transport. Specifically, we model the online dynamics as a conditional flow built upon the output distribution of a pretrained offline flow, rather than learning it directly from a Gaussian prior. This composite structure provides two advantages: (1) improved generalization when learning online dynamics under limited interaction data, and (2) a well-defined and stable estimate of the dynamics gap via the Wasserstein distance between offline and online transitions. Building on this dynamics-gap estimator, we further develop an optimistic active data collection strategy that prioritizes exploration in high-gap regions, and show theoretically that it reduces the performance gap to the optimal policy. Empirically, COMPFLOW consistently outperforms strong baselines across a range of RL benchmarks with shifted-dynamics data.

1 Introduction

Reinforcement Learning (RL) has demonstrated remarkable performance in complex sequential decision-making tasks such as playing Go and Atari games [53, 55], supported by access to large amounts of online interactions with the environment. However, in many real-world domains such as robotics [26, 60], healthcare [69], and wildlife conservation [29, 30, 66, 67], access to such interactions is often prohibitively expensive, unsafe, or infeasible. The limited availability of interactions presents a major challenge for learning effective and reliable policies. To address this challenge and improve sample efficiency during online training, a promising strategy is to incorporate a pre-collected offline dataset generated by a previous policy [44, 59]. This approach enables the agent to learn from a broader set of experiences, which can help accelerate learning and improve performance [59].

A critical challenge in online RL with offline data arises when the transition dynamics in the offline dataset differ from those in the online environment where the agent actively interacts and learns [49]. This issue, commonly referred to as shifted dynamics , can create severe distribution mismatch, bias policy updates, destabilize learning, and ultimately degrade performance. For instance, in robotics, the transition dynamics specify how the robot's state (e.g., position and velocity) evolves after executing an action. During deployment, changes in physical parameters such as surface friction can alter this state evolution, causing the true next-state distribution to deviate from that implied by the historical data. Similarly, in conservation planning, the transition dynamics describe how the spatial risk of

poaching evolves in response to patrol actions. Data collected in one region may not transfer to another because differences in terrain, accessibility, and human activity patterns lead to different state transitions and behavioral responses under the same patrol strategy [66].

To address these challenges, existing methods either penalize the rewards or value estimates of offline transitions with high dynamics gap [37, 46], or filter out such transitions entirely [64]. However, these approaches face key limitations. Most notably, the estimation of the dynamics gap typically relies on KL divergence or mutual information, both of which can be ill-defined when the offline and online transition dynamics have different supports [2, 48].

In this paper, we propose COMPFLOW, a new method for RL with shifted-dynamics data that leverages the theoretical connection between flow matching and optimal transport. We model the transition dynamics of the online environment using a composite flow architecture, where the online flow is defined on top of the output distribution of a learned offline flow rather than being initialized from a Gaussian prior. This design enables a principled estimation of the dynamics gap using the Wasserstein distance between the offline and online transition dynamics. On the theoretical side, we show that the com-

𝑊2

)

posite flow formulation reduces generalization error compared standard flow matching by reusing structural knowledge embedded in the offline data, particularly when online interactions are limited.

Building on the dynamics gap estimated by composite flow matching using the Wasserstein distance, we go beyond selectively merging offline transitions with low dynamics gap and further propose an active data collection strategy that targets regions in the online environment where the dynamics gap relative to the offline data is high. Such regions are often underrepresented in the replay buffer due to the dominance of low-gap samples. Our theoretical analysis shows that targeted exploration in these high-gap regions can further close the performance gap with respect to the optimal policy.

Our contributions are summarized as follows: (1) We introduce a composite flow model that estimates the dynamics gap by computing the Wasserstein distance between conditional transition distributions. We provide theoretical analysis showing that this approach achieves lower generalization error compared to learning the online dynamics from scratch. (2) Leveraging this principled estimation, we propose a new data collection strategy that encourages the policy to actively explore regions with high dynamics gap. We also provide a theoretical analysis of its performance benefits. (3) We empirically validate our method on various RL benchmarks with shifted dynamics and demonstrate that COMPFLOW outperforms or matches state-of-the-art baselines across these tasks.

2 Problem Statement and Background

2.1 Problem Definition

We consider two infinite-horizon MDPs: M off := ( S , A , p off , r, γ ) and M on := ( S , A , p on , r, γ ) , sharing the same state/action spaces, reward function r : S × A → R , and discount factor γ ∈ (0 , 1) , but differing in transition dynamics:

̸

We assume rewards are bounded, i.e., | r ( s, a ) | ≤ r max for all s, a . For any policy π , let ( s t , a t ) t ≥ 0 be the trajectory generated by M and π . We define the discounted state-action visitation (occupancy) measure as ρ π M ( s, a ) := (1 -γ ) E [ ∑ ∞ t =0 γ t 1 { s t = s, a t = a } ] , and the discounted state visitation as d π M ( s ) := ∑ a ρ π M ( s, a ) . The expected return is η M ( π ) := E ( s,a ) ∼ ρ π M [ r ( s, a )] .

Definition 2.1 (Online Policy Learning with Shifted-Dynamics Offline Data) . Given an offline dataset D off = { ( s i , a i , s ′ i , r i ) } N i =1 from M off and limited online access to M on, the objective is to learn a policy π maximizing η M on ( π ) , ideally approaching η M on ( π ⋆ ) , where π ⋆ := arg max π η M on ( π ) .

Lemma 2.2 (Return Bound between Two Environments [37]) . Let the empirical behavior policy in D off be π D off ( a | s ) . Define C 1 = 2 r max (1 -γ ) 2 . Then for any policy π ,

This bound highlights two key sources of return gap between domains: (1) the mismatch between the learned policy and the behavior policy in the offline dataset, and (2) the shift in environment dynamics. The former can be mitigated through behavior cloning [37, 65], while the latter can be addressed by filtering out source transitions with large dynamics gaps [39, 64]. A central challenge lies in accurately estimating this gap. Existing methods often rely on KL divergence or mutual information [37, 46], which can be ill-defined when the two dynamics have different supports.

2.2 Flow Matching

In this paper, we will adopt flow matching [1, 34, 36] to model the transition dynamics, owing to its ability to capture complex distributions. Flow Matching (FM) offers a simpler alternative to denoising diffusion models [21, 57], which are typically formulated using stochastic differential equations (SDEs). In contrast, FM is based on deterministic ordinary differential equations (ODEs), providing advantages such faster inference, and often improved sample quality.

The goal of FM is to learn a time-dependent velocity field v θ ( x, t ) : R d × [0 , 1] → R d , parameterized by θ , which defines a flow map ψ θ ( x 0 , t ) . This map is the solution to the ODE

and transports samples from a simple source distribution p 0 ( x ) (e.g., an isotropic Gaussian) at time t = 0 to a target distribution p 1 ( x ) at time t = 1 . In practice, generating a sample from the target distribution involves drawing x 0 ∼ p 0 ( x ) and integrating the learned ODE to obtain x 1 = ψ θ ( x 0 , 1) .

A commonly used training objective in flow matching is the linear path matching loss

where x t = (1 -t ) x 0 + tx 1 denotes the linear interpolation between x 0 and x 1 . Using linear interpolation paths encourages the learned flow to follow nearly straight-line trajectories, which reduces discretization error and improves the computational efficiency of ODE solvers during sampling [36].

2.3 Optimal Transport Flow Matching and Wasserstein Distance

Optimal Transport Flow Matching (OT-FM) establishes a direct connection between flow-based modeling and Optimal Transport (OT), providing a principled framework for quantifying the discrepancy between two distributions. This connection is especially valuable in our setting, where a key challenge is to measure the distance between two conditional transition distributions.

Let c : R d × R d → R be a cost function. Optimal transport aims to find a coupling q ⋆ ∈ Π( p 0 , p 1 ) -a joint distribution with marginals p 0 and p 1 -that minimizes the expected transport cost:

This minimum defines the Wasserstein distance W c ( p 0 , p 1 ) between the two distributions under the cost function c . When c ( x 0 , x 1 ) = ∥ x 0 -x 1 ∥ 2 , the resulting distance is known as the squared 2-Wasserstein distance.

In the training objective of Eq. 1, when sample pairs ( x 0 , x 1 ) are drawn from the optimal coupling q ⋆ , the flow model trained with the linear path matching loss learns a vector field that approximates the optimal transport plan. The transport cost of the learned flow approximates the Wasserstein distance

For theoretical justification, see Theorem 4.2 in [47].

Data Filtering

Exploration Bonus

RL Agent

Online Environment

Composite

Flow Dynamics Gap Estimator

Interact

Dynamics Shift

Behavioral Cloning

△(s,a) ≈ W2(Pof(-Is,a),Pon(-Is, a))

(s,a)

Latentxo

Offlines'(x1)

Onlines'(x2)

Doff

V

T

D

on

In practice, the optimal coupling π ⋆ is approximated using mini-batches by solving a discrete OT problem between empirical samples from p 0 and p 1 . The use of mini-batch OT has also been shown to implicitly regularize the transport plan [13, 14], as the stochasticity from independently sampled batches induces behavior similar to entropic regularization [9]. Further details of the OT-FM training procedure are provided in Appendix C.

3 Proposed Method

In this section, we introduce our method, COMPFLOW. We begin by presenting a composite flow matching approach to estimate the dynamics gap via Wasserstein distance. Next, we propose a data collection strategy that actively explores high dynamics-gap regions and provide a theoretical analysis of its benefits. Finally, we describe the practical implementation details of our method.

3.1 Estimating Dynamics Gap via Composite Flow

To estimate the gap between dynamics, we first learn the transition models for both the offline dataset and the online environment. Flow matching provides a flexible framework for modeling complex transition dynamics; however, the limited number of samples available in the online environment presents a significant challenge. Training separate flow models for each environment can result in poor generalization in the online environment, as the model may overfit to the small amount of available data.

To mitigate this, we propose a composite flow formulation. Instead of learning the online transition model p on ( s ′ | s, a ) from scratch, we leverage structural knowledge from a well-trained offline model p off ( s ′ | s, a ) . This enables to incorporate prior knowledge to improve the generalization.

Offline flow. We begin by learning a conditional flow model for the offline data. Let x 0 ∼ N (0 , I ) be the initial latent variable. The offline flow map ψ off θ ( x 0 , t | s, a ) is defined as the solution to the following ODE:

Solving this ODE from t = 0 to 1 produces an intermediate representation: x 1 = ψ off θ ( x 0 , 1 | s, a ) .

Online flow. Instead of learning a online flow directly from a Gaussian prior, we initialize it from the intermediate representation x 1 produced by the offline flow. The target flow map ψ on ϕ ( x, t | s, a ) is defined as:

Solving this ODE from t = 1 to 2 yields the final prediction for the online environment: s ′ . = x 2 = ψ on ϕ ( x 1 , 2 | s, a ) .

Wenow show that when the offline flow induces a distribution ˆ p off ( s ′ | s, a ) that is closer to p on ( s ′ | s, a ) than the standard Gaussian distribution, the proposed composite flow method enjoys a smaller generalization error bound.

Theorem 3.1 (Conditions for Composite Flow Yielding Smaller Errors) . Assume that composite flow and direct flow share the same hypothesis class H of (measurable) vector fields v : R d × [0 , 1] → R d , and that there exists B ∈ (0 , ∞ ) such that sup v ∈H sup x ∈ R d , t ∈ [0 , 1] ∥ v ( x, t ) ∥ 2 ≤ B . Also, assume that max { Tr(Σ on ) , Tr( ̂ Σ off ) } ≤ C TR for some C TR . The composite flow enjoys a strictly tighter high-probability generalization bound than the direct flow if and only if

Here, p G := N (0 , I d ) , Σ on is the covariance of p on , and ̂ Σ off is the covariance of ˆ p off .

Remark 3.2 . In our setting, we assume that the offline data and the online environment share meaningful similarities. Given the abundance of offline data available to accurately learn the offline flow, this assumption is expected to hold well.

To compute the Wasserstein distance between the offline and online transition dynamics, we can use the OT-FM objective to train the online flow, initialized from the offline flow distribution. However, when the state-action pair ( s, a ) lies in a continuous space, it is not feasible to obtain a batch of samples corresponding to a fixed ( s, a ) from the online environment.

To address this issue, we follow prior works [18, 23] and incorporate the conditioning variables directly into the transport cost. Specifically, consider pairs of transition samples ( s off , a off , s ′ off ) and ( s on , a on , s ′ on ) . We define the transport cost as

where η > 0 controls the strength of alignment between conditioning variables.

Let q ∗ denote the optimal coupling between the empirical distributions of ( s off , a off , s ′ off ) and ( s on , a on , s ′ on ) under the above cost. The online flow is trained using the objective

where x t is the linear interpolation between x 1 and x 2 .

Proposition 3.3 (Informal; shared-conditioning coupling is W 2 -optimal) . Assume that, when training the online flow in Eq. (3) , the marginals over the conditioning variables are matched, i.e., p off ( s, a ) = p on ( s, a ) . Then, for any fixed ( s, a ) ∈ S × A , as η →∞ and the minibatch size used to compute the OT coupling tends to infinity, the expected squared displacement induced by the composite flow recovers the squared 2-Wasserstein distance between the offline and online transition distributions:

Monte Carlo Estimator. This result yields a practical Monte Carlo estimator of the dynamics gap ∆( s, a ) :

Remark 3.4 (Enforcing the shared marginal condition) . The marginal condition p off ( s, a ) = p on ( s, a ) can be enforced by construction. In practice, we construct the empirical distribution of offline triples ( s off , a off , s ′ off ) by first sampling ( s off , a off ) from the online dataset D on , and then generating s ′ off ∼ p off ( ·| s off , a off ) using the pretrained offline flow. Separately, we sample online transitions ( s on , a on , s ′ on ) from D on . As a result, the two empirical measures used in the OT coupling have matched ( s, a ) -marginals in expectation, while their conditional transitions differ. The complete training algorithms of the offline flow and the online flow are in Appendix C.

3.2 Data Collection at High Dynamics Gap Region

As discussed in Section 2.1, using behavior cloning and augmenting the replay buffer with low dynamics gap offline data can alleviate performance drop from distribution shift. However, relying solely on such data may limit state-action coverage and hinder policy learning. To address this, we propose a new data collection strategy.

At each training iteration, we construct the replay buffer by selectively incorporating offline transitions with small estimated dynamics gap:

where τ is a predefined threshold. To improve data diversity and encourage better generalization, we actively explore regions with high dynamics gap-areas likely underrepresented in the buffer due to the dominance of low-gap samples. We adopt an optimistic exploration policy that selects actions by

where β is a hyperparameter that trades off return and exploration of underexplored dynamics. To ensure sufficient coverage during training, we can further incorporate stochasticity by adding small perturbations to the selected actions [17] or following a stochastic policy [19] consistent with widely used deep RL frameworks.

Theorem 3.5 (Large Dynamics Gap Exploration Reduces Performance Gap) . Compared to behavior cloning policy π bc on the offline dataset, training a policy ˆ π by replacing all offline samples with a dynamics gap exceeding κ (as estimated by the composite flow) with online environment samples can reduce the performance gap to the optimal online policy π ∗ on with high probability by

Here L r and L P are the Lipschitz constants for the reward and transition function, respectively. γL p < 1 . ∆ W 2 := sup s,a W 2 ( p off ( ·| s, a ) , p on ( ·| s, a ) ) is the largest dynamics gap. C 0 and C 1 are two constants. Γ N on ,δ := R N on ( H ) + √ log(1 /δ ) N on , where H is the same as in Theorem 3.1 and N on is the number of samples used to train the online flow.

From Theorem 3.5, exploration in regions with high dynamics gap can reduce the performance gap relative to the optimal policy. The parameter κ serves as a threshold: the bound holds when the policy is trained without offline samples whose dynamics gap exceeds κ . As β increases, more samples are collected from high-gap regions, which decreases κ . This increases performance gap reduction, resulting in a policy with higher expected return.

3.3 Practical Implementation

Our method can be instantiated using standard actor-critic algorithms with a critic Q ς ( s, a ) and a policy π φ ( a | s ) . To incorporate the dynamics gap, we apply rejection sampling to retain a fixed percentage of offline transitions with the lowest estimated gap in each iteration. The critic is trained by minimizing

where ̂ ∆ ξ % is the ξ -quantile of the estimated dynamics gap in the offline minibatch. The target value is y = r + γQ ς ( s ′ , a ′ ) + β ̂ ∆( s, a ) , with a ′ ∼ π φ ( ·| s ′ ) .

The policy is updated by combining policy improvement with a behavior cloning regularizer, using the objective

where ω > 0 controls the trade-off between maximizing the estimated Q -value and staying close to the offline behavior. The behavior cloning term encourages the policy's sampled action ˜ a to match the offline action a , as motivated by Lemma 2.2.

The pseudocode of COMPFLOW, instantiated with Soft Actor-Critic (SAC) [19], is presented in Appendix D.

4 Related Work

Online RL with offline dataset. Online RL often requires extensive environment interactions [54, 68], which can be costly or impractical in real-world settings. To improve sample efficiency, Offline-to-Online RL leverages pre-collected offline data to bootstrap online learning [43]. A typical two-phase approach trains an initial policy offline, then fine-tunes it online [32, 43, 44]. However, conservative strategies used to mitigate distributional shift, such as pessimistic value estimation [31], can result in suboptimal initial policies and limit effective exploration [40, 44]. To resolve this tension, recent methods propose ensemble-based pessimism [32], value calibration [44], optimistic action selection [32], and policy expansion [70]. Others directly incorporate offline data into the replay buffer of off-policy algorithms [3], improving stability via ensemble distillation and layer normalization. However, these approaches typically assume that the offline dataset is generated under the same transition dynamics as the online environment. In contrast, our work explicitly accounts for dynamics shift between the offline data and the online environment.

RL with dynamics shift. Our work is related to cross-domain RL, where the source and target domains share the same observation and action spaces but differ in transition dynamics. Prior work has addressed such discrepancies via system identification [5, 8, 10], domain randomization [41, 56, 61], imitation learning [20, 24], and meta-RL [42, 50], often assuming shared environment distributions [58] or requiring expert demonstrations. More recent work has relaxed these assumptions. One line of research focuses on reward modification , which adjusts the reward function to penalize source transitions that are unlikely under the target dynamics [12, 35], or down-weights value estimates in regions with high dynamics gap [46]. Another line of work explores data filtering , which selects only source transitions with low estimated dynamics gap [65], using metrics such as transition probability ratios [12], mutual information [64], value inconsistency [65], or representationbased KL divergence [37]. However, these metrics can become unstable or ill-defined when the transition dynamics have different supports, and value-based methods are prone to instability caused by bootstrapping bias.

In contrast, our approach leverages the theoretical connection between optimal transport and flow matching to estimate the dynamics gap in a principled manner. A closely related work by [39] also applies optimal transport to quantify the dynamics gap. However, their setting assumes both the source and target domains are offline, and their method estimates the gap based on the concatenated tuple ( s, a, s ′ ) observed in offline data, rather than comparing conditional transition distributions. As a result, their metric does not accurately capture the gap in transition dynamics and cannot be used to compute exploration bonuses, which require gap estimation conditioned on a given state-action pair.

RL with diffusion and flow models. Diffusion models [21, 28, 57] and flow matching [11, 34] have emerged as powerful generative tools capable of modeling complex, high-dimensional distributions. Their application in RL is consequently expanding. Researchers have employed these generative models for various tasks, including planning and trajectory synthesis [22], representing expressive multimodal policies [7, 51], providing behavior regularization [6], or augmenting training datasets with synthesized experiences. While these works often focus on policy learning or modeling dynamics within a single environment, our approach targets the transfer learning setting. Specifically, we address scenarios characterized by a dynamics gap between the offline data and the online environment. We utilize Flow Matching, leveraging its connection to optimal transport, to estimate this gap.

5 Experiments

In this section 2 , we first evaluate our approach across a range of environments in Gym-MuJoCo that exhibit different types of dynamics shifts. We then conduct ablation and hyperparameter studies to better understand the design choices and behavior of COMPFLOW. Finally, we assess the effectiveness of our method in a real-world inspired wildlife conservation task.

5.1 Gym-MuJoCo

5.1.1 Experimental Setup

Tasks and datasets. We evaluate our algorithm under three types of dynamics shifts, namely morphology, kinematic and friction changes, across three OpenAI Gym locomotion tasks: HalfCheetah,

Hopper, and Walker2d [4]. Each experiment involves an offline environment and an online environment with modified transition dynamics following [38]. Morphology shifts alter the sizes of body parts, kinematic shifts impose constraints on joint angles, and friction shifts modify the static, dynamic, and rolling friction coefficients. For each task, we use three D4RL source datasets: medium, medium replay, and medium expert, which capture varying levels of data quality [16]. Additional environment details are in Appendix K.

Baselines. We compare COMPFLOW against the following baselines: BC-SAC extends SAC by incorporating both offline and online data, with a behavior cloning (BC) term for the offline data. H2O [46] penalizes Q-values for state-action pairs with large dynamics gaps. BC-VGDF [65] selects offline transitions with value targets consistent with the online environment and adds a BC term. BC-PAR [37] applies a reward penalty based on representation mismatch between offline and online transitions, also including a BC term. Implementation details are provided in Appendix K.

5.1.2 Main Results

We evaluate each algorithm by measuring its return in the target environment after 40K environment interactions and 400K gradient updates, reflecting a limited-interaction setting. The results are reported in Table 1. Our key findings are summarized as follows:

(1) COMPFLOW consistently outperforms recent off-dynamics RL baselines across diverse offline dataset qualities and dynamics shifts. It achieves the best return on 19/27 tasks and ties for best on 5 more. On average, COMPFLOW reaches a score of 2193 , compared to 1920 for the strongest baseline (BC-SAC), a 14.2% relative improvement. Additional results under more severe dynamics shifts in Appendix J show the same trend.

(2) Leveraging offline data substantially improves performance: on average, COMPFLOW improves over SAC by 149.8% . Moreover, COMPFLOW outperforms BC-SAC, which directly uses all offline data, on 23/27 tasks and matches it on the remaining 3 , highlighting the benefit of our dynamics-gapaware design.

(3) We also find that many recent baselines perform similarly to BC-SAC, consistent with [38], suggesting limited ability to effectively exploit offline data under dynamics shift. One reason is that methods such as H2O and BC-PAR estimate the dynamics gap using KL divergence or mutual information, which can be ill-defined under large shifts or mismatched support. BC-VGDF instead filters data based on value estimates, which is often harder due to bootstrapping bias and target nonstationarity. In contrast, COMPFLOW models transition dynamics with flow matching and estimates the dynamics gap via the Wasserstein distance through its connection to optimal transport, yielding a more principled and robust criterion.

5.1.3 Ablation and Hyperparameter Analysis

Composite Flow. We first evaluate the effectiveness of the proposed composite flow matching design. For comparison, we train a direct flow model on the target environment initialized from a Gaussian prior. We compute the mean squared error (MSE) on a held-out 10% validation set and report the average MSE across different epochs during RL training. As shown in Figure 3, the composite flow significantly reduces the MSE of the transition dynamics in the environment. This improvement stems from its

0.0

0.2

0.4

0.6

0.8

1.0

1.2

MSE

HalfCheetah Medium (Morph)

0.00

0.02

0.04

0.06

0.08

0.10

0.12

MSE

Hopper Medium Replay (Morph)

Direct Flow

Composite Flow

ability to reuse structural knowledge learned from the abundant offline data. These empirical findings support the theoretical insight presented in Theorem 3.1 that our composite flow can improve the generalization ability.

Impact of data selection ratio ξ % . The data selection ratio ξ decides how many offline data in a sampled batch can be shared for policy training. A larger ξ indicates that more offline data will be admitted. To examine its influence, We sweep ξ across {70, 50, 30, 20}. The results is shown in Figure 4. Although the optimal ξ seems task dependent, moderate values of ξ (e.g., 30 or 50) generally yield good performance across tasks, striking a balance between leveraging useful offline data and avoiding high dynamics mismatch. When ξ is too large (e.g., 70), performance often degrades, particularly in Walker Medium (Morph) and Walker Medium Replay (Morph), likely due to the inclusion of low-quality transitions with large dynamics gap. HalfCheetah consistently achieves higher returns compared to other domains, suggesting that knowledge transfer and policy learning in this environment are easier. Consequently, overly conservative filtering (e.g., ξ = 20 ) may exclude valuable offline data, leading to slower learning. In this case, allowing more offline data (e.g., ξ = 70 ) appears beneficial.

0

10

20

30

40

Env Steps (

×10 4

)

0

1000

2000

Return

HalfCheetah Medium (Morph)

0

10

20

30

40

Env Steps (

×10 4

)

200

400

600

Hopper Medium (Morph)

0

10

20

30

40

Env Steps (

×10 4

)

0

200

400

600

800

Walker2d Medium (Morph)

0

10

20

30

40

Env Steps (

×10 4

)

0

500

1000

1500

Walker2d Medium Replay (Morph)

selection ratio: 70%

selection ratio: 50%

selection ratio: 30%

selection ratio: 20%

Impact of exploration strength β . β controls the strength of exploration toward regions with large dynamics gap. A large β indicates higher incentive to explore such regions. As shown in Figure 5, the effect of β on return is task-dependent, possibly due to differences in the underlying MDPs. In the Friction tasks, we observe that larger exploration bonuses lead to higher asymptotic returns. This suggests that incentivizing exploration helps the algorithm discover high-reward regions more effectively. In contrast, in the Morphology tasks, moderate values of β typically outperform both smaller and larger values. This indicates that excessive exploration may not be effective in some tasks. While the optimal β varies by task, one trend is consistent: across all five experiments, the setting with no exploration ( β = 0 ) consistently ranks as the lowest or the second lowest performers. This

empirical finding confirms the importance of the exploration bonus and provides evidence supporting our algorithmic design.

0

10

20

30

40

Env Steps (

×10 4

)

100

200

300

Return

Hopper Medium (Friction)

0

10

20

30

40

Env Steps (

×10 4

)

100

150

200

250

300

Hopper Medium Replay (Friction)

0

10

20

30

40

Env Steps (

×10 4

)

0

200

400

600

Walker2d Medium Replay (Friction)

0

10

20

30

40

Env Steps (

×10 4

)

0

1000

2000

Return

HalfCheetah Medium (Morph)

0

10

20

30

40

Env Steps (

×10 4

)

200

400

600

Hopper Medium (Morph)

= 0

= 0.001

= 0.01

= 0.1

= 0.2

5.2 Patrol Policy Learning for Wildlife Conservation

We evaluate COMPFLOW in a wildlife conservation setting, where the goal is to learn a patrol policy that allocates limited ranger effort to reduce poaching. We adopt the simulation environment of Xu et al. [67], which represents the park as a 1 × 1 km spatial grid. At each time step, the agent allocates a distribution of patrol effort across cells under a fixed budget. The state, including wildlife density and poaching risk, evolves according to dynamics driven by both poacher behavior and patrol deployment, and the per-step reward reflects the expected number of animals preserved through deterrence and spatial coverage.

We assume access to an offline dataset collected under a previous policy in Murchison Falls National Park, and aim to leverage it to learn an improved patrol policy for Queen Elizabeth National Park with only limited online interactions. Since ecological conditions and poacher behavior differ across parks, the transition dynamics are mismatched. As shown in Figure 6, COMPFLOW achieves the highest reward, outperforming the best baseline, BC-PAR, by 8.8% . It also improves over training from scratch with SAC by 20.8% under the same interaction bud-

SAC

BC-SAC

H2O

BC-VGDF

BC-PAR

Ours

400

450

500

550

600

650

700

750

Return

get. These gains are particularly valuable in large protected areas where ranger capacity is constrained: Murchison Falls spans roughly 3,900 km 2 , while Queen Elizabeth covers about 1,980 km 2 , making data-efficient learning essential for effective protection.

6 Conclusion and Limitations

In this paper, we proposed COMPFLOW, a new method for estimating the dynamics gap in reinforcement learning with shifted-dynamics data. COMPFLOW leverages the theoretical connection between flow matching and optimal transport. To address data scarcity in the online environment, we adopt a composite flow structure that builds the online flow model on top of the output distribution from the offline flow. This composite formulation improves generalization and enables the use of Wasserstein distance between offline and online transitions as a robust measure of the dynamics gap. Using this pricinpled estimation, we further encourage the policy to explore regions with high dynamics gap and provide a theoretical analysis of the benefits. Empirically, we demonstrate that COMPFLOW consistently outperforms or matches state-of-the-art baselines across diverse RL tasks with varying types of dynamics shift.

Due to the space limit, we discuss the limitations of COMPFLOW in Appendix A.

References

Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=li7qeBbCR1t .

Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein gan. In International Conference on Machine Learning , pages 214-223. PMLR, 2017.

Philip J Ball, Laura Smith, Ilya Kostrikov, and Sergey Levine. Efficient online reinforcement learning with offline data. In International Conference on Machine Learning , pages 1577-1594. PMLR, 2023.

Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540 , 2016.

Yevgen Chebotar, Ankur Handa, Viktor Makoviychuk, Miles Macklin, Jan Issac, Nathan Ratliff, and Dieter Fox. Closing the sim-to-real loop: Adapting simulation randomization with real world experience. In 2019 International Conference on Robotics and Automation (ICRA) , pages 8973-8979. IEEE, 2019.

Huayu Chen, Cheng Lu, Zhengyi Wang, Hang Su, and Jun Zhu. Score regularized policy optimization through diffusion behavior. In The Twelfth International Conference on Learning Representations , 2024. URL https://openreview.net/forum?id=xCRr9DrolJ .

Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS) , 2023.

Ignasi Clavera, Anusha Nagabandi, Ronald S Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt: Meta-learning for model-based control. arXiv preprint arXiv:1803.11347 , 3:3, 2018.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems , 26, 2013.

Yuqing Du, Olivia Watkins, Trevor Darrell, Pieter Abbeel, and Deepak Pathak. Auto-tuned sim-to-real transfer. In 2021 IEEE International Conference on Robotics and Automation (ICRA) , pages 1290-1296. IEEE, 2021.

Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first international conference on machine learning , 2024.

Benjamin Eysenbach, Swapnil Asawa, Shreyas Chaudhari, Sergey Levine, and Ruslan Salakhutdinov. Off-dynamics reinforcement learning: Training for transfer with domain classifiers. arXiv preprint arXiv:2006.13916 , 2020.

Kilian Fatras, Younes Zine, Rémi Flamary, Rémi Gribonval, and Nicolas Courty. Learning with minibatch wasserstein: asymptotic and gradient properties. arXiv preprint arXiv:1910.04091 , 2019.

Kilian Fatras, Younes Zine, Szymon Majewski, Rémi Flamary, Rémi Gribonval, and Nicolas Courty. Minibatch optimal transport distances; analysis and applications. arXiv preprint arXiv:2101.01792 , 2021.

Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research , 22(78):1-8, 2021. URL http://jmlr.org/papers/v22/20-451.html .

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 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.

Adam P Generale, Andreas E Robertson, and Surya R Kalidindi. Conditional variable flow matching: Transforming conditional densities with amortized conditional optimal transport. arXiv preprint arXiv:2411.08314 , 2024.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning , pages 1861-1870. Pmlr, 2018.

Donald Hejna, Lerrel Pinto, and Pieter Abbeel. Hierarchically decoupled imitation for morphological transfer. In International Conference on Machine Learning , pages 4159-4171. PMLR, 2020.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020.

Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning , pages 9902-9915. PMLR, 2022.

Gavin Kerrigan, Giosue Migliorini, and Padhraic Smyth. Dynamic conditional optimal transport through simulation-free flows. In The Thirty-eighth Annual Conference on Neural Information Processing Systems , 2024. URL https://openreview.net/forum?id=tk0uaRynhH .

Kuno Kim, Yihong Gu, Jiaming Song, Shengjia Zhao, and Stefano Ermon. Domain adaptive imitation learning. In International Conference on Machine Learning , pages 5286-5295. PMLR, 2020.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Representation Learning , 2015.

Jens Kober, J Andrew Bagnell, and Jan Peters. Reinforcement learning in robotics: A survey. The International Journal of Robotics Research , 32(11):1238-1274, 2013.

Lingkai Kong, Harshavardhan Kamarthi, Peng Chen, B Aditya Prakash, and Chao Zhang. Uncertainty quantification in deep learning. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining , pages 5809-5810, 2023.

Lingkai Kong, Yuanqi Du, Wenhao Mu, Kirill Neklyudov, Valentin De Bortoli, Dongxia Wu, Haorui Wang, Aaron Ferber, Yi-An Ma, Carla P Gomes, et al. Diffusion models as constrained samplers for optimization with unknown constraints. arXiv preprint arXiv:2402.18012 , 2024.

Lingkai Kong, Haichuan Wang, Charles A Emogor, Vincent Börsch-Supan, Lily Xu, and Milind Tambe. Generative ai against poaching: Latent composite flow matching for wildlife conservation. arXiv preprint arXiv:2508.14342 , 2025.

Lingkai Kong, Haichuan Wang, Yuqi Pan, Cheol Woo Kim, Mingxiao Song, Alayna Nguyen, Tonghan Wang, Haifeng Xu, and Milind Tambe. Robust optimization with diffusion models for green security. arXiv preprint arXiv:2503.05730 , 2025.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in neural information processing systems , 33: 1179-1191, 2020.

Seunghyun Lee, Younggyo Seo, Kimin Lee, Pieter Abbeel, and Jinwoo Shin. Offline-to-online reinforcement learning via balanced replay and pessimistic q-ensemble. In Conference on Robot Learning , pages 1702-1712. PMLR, 2022.

Yinghao Li, Lingkai Kong, Yuanqi Du, Yue Yu, Yuchen Zhuang, Wenhao Mu, and Chao Zhang. MUBen: Benchmarking the uncertainty of molecular representation models. Transactions on Machine Learning Research , 2024. ISSN 2835-8856. URL https://openreview.net/ forum?id=qYceFeHgm4 .

Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=PqvMRDCJT9t .

Jinxin Liu, Hongyin Zhang, and Donglin Wang. Dara: Dynamics-aware reward augmentation in offline reinforcement learning. arXiv preprint arXiv:2203.06662 , 2022.

Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003 , 2022.

Jiafei Lyu, Chenjia Bai, Jingwen Yang, Zongqing Lu, and Xiu Li. Cross-domain policy adaptation by capturing representation mismatch. arXiv preprint arXiv:2405.15369 , 2024.

Jiafei Lyu, Kang Xu, Jiacheng Xu, Jing-Wen Yang, Zongzhang Zhang, Chenjia Bai, Zongqing Lu, Xiu Li, et al. Odrl: A benchmark for off-dynamics reinforcement learning. Advances in Neural Information Processing Systems , 37:59859-59911, 2024.

Jiafei Lyu, Mengbei Yan, Zhongjian Qiao, Runze Liu, Xiaoteng Ma, Deheng Ye, Jing-Wen Yang, Zongqing Lu, and Xiu Li. Cross-domain offline policy adaptation with optimal transport and dataset constraint. In The Thirteenth International Conference on Learning Representations , 2025.

Max Sobol Mark, Ali Ghadirzadeh, Xi Chen, and Chelsea Finn. Fine-tuning offline policies with optimistic action selection. In Deep Reinforcement Learning Workshop NeurIPS 2022 , 2022.

Bhairav Mehta, Manfred Diaz, Florian Golemo, Christopher J Pal, and Liam Paull. Active domain randomization. In Conference on Robot Learning , pages 1162-1176. PMLR, 2020.

Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald S Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt in dynamic, real-world environments through metareinforcement learning. arXiv preprint arXiv:1803.11347 , 2018.

Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline 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 offline rl pre-training for efficient online fine-tuning. Advances in Neural Information Processing Systems , 36:62244-62269, 2023.

Ariel Neufeld and Julian Sester. Bounding the difference between the values of robust and non-robust markov decision problems. Journal of Applied Probability , pages 1-14, 2023.

Haoyi Niu, Yiwen Qiu, Ming Li, Guyue Zhou, Jianming Hu, Xianyuan Zhan, et al. When to trust your simulator: Dynamics-aware hybrid offline-and-online reinforcement learning. Advances in Neural Information Processing Systems , 35:36599-36612, 2022.

Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky TQ Chen. Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint arXiv:2304.14772 , 2023.

Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. In International Conference on Machine Learning , pages 5171-5180. PMLR, 2019.

Chengrui Qu, Laixi Shi, Kishan Panaganti, Pengcheng You, and Adam Wierman. Hybrid transfer reinforcement learning: Provable sample efficiency from shifted-dynamics data. In The 28th International Conference on Artificial Intelligence and Statistics , 2025. URL https: //openreview.net/forum?id=ndWOLqVRHC .

Roberta Raileanu, Max Goldstein, Arthur Szlam, and Rob Fergus. Fast adaptation to new environments via policy-dynamics value functions. In Proceedings of the 37th International Conference on Machine Learning , pages 7920-7931, 2020.

Allen Z. Ren, Justin Lidard, Lars Lien Ankile, Anthony Simeonov, Pulkit Agrawal, Anirudha Majumdar, Benjamin Burchfiel, Hongkai Dai, and Max Simchowitz. Diffusion policy policy optimization. In The Thirteenth International Conference on Learning Representations , 2025. URL https://openreview.net/forum?id=mEpqHvbD2h .

Filippo Santambrogio. Optimal transport for applied mathematicians. 2015.

Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature , 588(7839):604-609, 2020.

David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature , 529(7587):484-489, 2016.

David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature , 550(7676):354-359, 2017.

Reda Bahi Slaoui, William R Clements, Jakob N Foerster, and Sébastien Toth. Robust domain randomization for reinforcement learning. 2019.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Yuda Song, Aditi Mavalankar, Wen Sun, and Sicun Gao. Provably efficient model-based policy adaptation. arXiv preprint arXiv:2006.08051 , 2020.

Yuda Song, Yifei Zhou, Ayush Sekhari, J Andrew Bagnell, Akshay Krishnamurthy, and Wen Sun. Hybrid rl: Using both offline and online data can make rl efficient. arXiv preprint arXiv:2210.06718 , 2022.

Chen Tang, Ben Abbatematteo, Jiaheng Hu, Rohan Chandra, Roberto Martín-Martín, and Peter Stone. Deep reinforcement learning for robotics: A survey of real-world successes. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 39, pages 28694-28698, 2025.

Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In 2017 IEEE/RSJ international conference on intelligent robots and systems (IROS) , pages 23-30. IEEE, 2017.

Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems , pages 5026-5033. IEEE, 2012.

Cédric Villani et al. Optimal transport: old and new , volume 338. Springer, 2008.

Xiaoyu Wen, Chenjia Bai, Kang Xu, Xudong Yu, Yang Zhang, Xuelong Li, and Zhen Wang. Contrastive representation for data filtering in cross-domain offline reinforcement learning. arXiv preprint arXiv:2405.06192 , 2024.

Kang Xu, Chenjia Bai, Xiaoteng Ma, Dong Wang, Bin Zhao, Zhen Wang, Xuelong Li, and Wei Li. Cross-domain policy adaptation via value-guided data filtering. Advances in Neural Information Processing Systems , 36:73395-73421, 2023.

Lily Xu, Elizabeth Bondi, Fei Fang, Andrew Perrault, Kai Wang, and Milind Tambe. Dualmandate patrols: Multi-armed bandits for green security. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 14974-14982, 2021.

Lily Xu, Andrew Perrault, Fei Fang, Haipeng Chen, and Milind Tambe. Robust reinforcement learning under minimax regret for green security. In Uncertainty in Artificial Intelligence , pages 257-267. PMLR, 2021.

Deheng Ye, Guibin Chen, Wen Zhang, Sheng Chen, Bo Yuan, Bo Liu, Jia Chen, Zhao Liu, Fuhao Qiu, Hongsheng Yu, et al. Towards playing full moba games with deep reinforcement learning. Advances in Neural Information Processing Systems , 33:621-632, 2020.

Chao Yu, Jiming Liu, Shamim Nemati, and Guosheng Yin. Reinforcement learning in healthcare: A survey. ACM Computing Surveys (CSUR) , 55(1):1-36, 2021.

Haichao Zhang, We Xu, and Haonan Yu. Policy expansion for bridging offline-to-online reinforcement learning. arXiv preprint arXiv:2302.00935 , 2023.

Yuchen Zhuang, Yue Yu, Lingkai Kong, Xiang Chen, and Chao Zhang. Dygen: Learning from noisy labels via dynamics-enhanced generative modeling. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining , pages 3674-3686, 2023.

Appendix for Composite Flow Matching for Reinforcement Learning with Shifted-Dynamics Data

A

Limitations

(1) COMPFLOW is currently limited to settings where the offline data and the online environment share the same state and action spaces. Future work could explore estimating the dynamics gap in a shared latent embedding space to relax this assumption [71]. (2) Our evaluation is conducted entirely in simulated environments; applying COMPFLOW to real-world scenarios remains an important direction for future work. (3) We have not yet incorporated the uncertainty in estimating the dynamics gap. Future work may explore how to quantify this uncertainty and use it to improve both data filtering and exploration strategies [27, 33].

B Broader Impacts

Our work aims to make reinforcement learning more practical in real-world domains such as healthcare, robotics, and conservation, where online interaction is often costly, limited, or unsafe. By addressing the dynamics shift between offline data and the online environment, our method enables more reliable and sample-efficient policy learning. This capability supports safer deployment in high-stakes applications, including clinical decision support and adaptive anti-poaching strategies. Nonetheless, we emphasize that policies trained on historical data should be applied with caution,

as misaligned dynamics or biased datasets may lead to unintended consequences if not carefully validated.

C Algorithms of Training Optimal Transport Flow Matching

The full training algorithm of Optimal Transport Flow Matching is given in Algorithm 1.

Algorithm 1 Training OT Flow Matching (OT-FM)

Require:

Dataset D ; batch size k ; learning rate lr ; vector field v θ

Ensure:

Trained parameters θ

1: while not converged do

2: (Sample) Draw latent codes ( x ( i ) 0 ) k i =1 ∼ N (0 , I )

3: Draw data points ( x ( j ) 1 ) k j =1 ∼ D

4: (OT coupling) Set C ij ←∥ x ( i ) 0 - x ( j ) 1 ∥ 2 2 for all ( i, j ) ∈ [ k ] × [ k ]

5: Compute an optimal transport plan

▷ A is a (scaled) doubly-stochastic coupling

6: Sample index pairs ( i ℓ , j ℓ ) k ℓ =1 i.i.d. from the categorical distribution on [ k ] × [ k ] with probabilities A ij

7: Sample times ( t ℓ ) k ℓ =1 ∼ U [0 , 1]

8: for ℓ = 1 to k do

12: end for

13: L ← 1 k ∑ k ℓ =1 ℓ ℓ

14: θ ← θ -lr ∇ θ L

15: end while

16: return θ

D Algorithms of Training Offline and Online Flows

The full training algorithms of the offline flow and online flow are given in Algorithm 2 and Algorithm 3 respectively.

Algorithm 2 Training the Offline Flow via Flow Matching

Require:

Offline dataset D off ; batch size k ; offline vector field v θ ; learning rate lr

Ensure:

Trained offline flow parameters θ

1: while not converged do

2: Sample transitions ( s ( i ) , a ( i ) , s ′ ( i ) ) k i =1 ∼ D off

3: Set x ( i ) 1 ← s ′ ( i ) for all i ∈ [ k ]

4: Sample latent codes ( x ( i ) 0 ) k i =1 ∼ N (0 , I )

5: Sample times ( t i ) k i =1 ∼ U [0 , 1]

6: for i = 1 to k do

10: end for

11:

L ←

1

k

i

k i =1 12: θ ← θ -lr ∇ θ L

13: end while

14: return θ

∑

ℓ

Algorithm 3 Training the Online Flow via Optimal Transport

Require: Pre-trained offline flow sampler ψ θ ( x, t | s, a ) ; online replay buffer D on ; batch size k ; regularization weight η > 0 ; learning rate lr ; initialized online vector field v ϕ

Ensure: Trained online flow parameters ϕ

1: while not converged do

2: (Offline-synthetic batch) Sample ( s ( i ) off , a ( i ) off ) k i =1 ∼ D on

3: Sample latent codes ( x ( i ) 0 ) k i =1 ∼ N (0 , I )

4: for i = 1 to k do

6: end for

7: (Online batch) Sample ( s ( j ) on , a ( j ) on , s ′ ( j ) on ) k j =1 ∼ D on

8: x ( j ) 2 ← s ′ ( j ) on for all j ∈ [ k ]

9: (OT coupling) For all ( i, j ) ∈ [ k ] × [ k ] , set

10: Compute an optimal transport plan

▷ A is a (scaled) doubly-stochastic coupling

11: Sample index pairs ( i ℓ , j ℓ ) k ℓ =1 i.i.d. from the categorical distribution on [ k ] × [ k ] with probabilities A ij

12: Sample times ( t ℓ ) k ℓ =1 ∼ U [1 , 2]

13: for ℓ = 1 to k do

15: ∆ ( ℓ ) ← x ( j ℓ ) 2 -x ( i ℓ ) 1

17: end for

18: L ← 1 k ∑ k ℓ =1 ℓ ℓ 19: ϕ ← ϕ -lr ∇ ϕ L

20: end while

21: return ϕ

E Algorithm of COMPFLOW built on Soft-Actor-Critic

The full training algorithm of COMPFLOW built on SAC is given in Algorithm 4.

Algorithm 4 COMPFLOW built on Soft Actor-Critic (SAC)

1: Input: Offline dataset D off , online environment M on , max interaction steps T max , update frequency train _ freq , offline selection ratio ξ ∈ (0 , 1] , gap reward scale β , batch size B , behavior cloning weight ω , warmup steps warmup _ steps , gradient steps per interaction K , learning rate lr , target update rate ϖ , discount γ , entropy temperature α .

2: Initialize: Policy π φ , critics { Q ς i } i =1 , 2 , target critics { Q tgt ς i } i =1 , 2 ←{ Q ς i } i =1 , 2 , online replay buffer D on ←∅ .

3: Pretrain offline flow ψ θ ( x, t | s, a ) on D off via Algorithm 2.

4: for t = 1 to T max do

5: Interact with M on using π φ and store transition ( s, a, r, s ′ ) into D on . ▷ a ∼ π φ ( · | s ) , ( r, s ′ ) ∼ M on ( s, a )

6: if t > warmup _ steps and t mod train _ freq = 0 then

7: Train online flow ψ ϕ ( x, t | s, a ) on D on via Algorithm 3.

8: end if

9: for k = 1 to K do

10: Sample offline minibatch b off = { ( s ( i ) off , a ( i ) off , r ( i ) off , s ′ ( i ) off ) } B i =1 ∼ D off .

11: Sample online minibatch b on = { ( s ( i ) on , a ( i ) on , r ( i ) on , s ′ ( i ) on ) } B i =1 ∼ D on .

12: Estimate dynamics gap ̂ ∆( s ( i ) off , a ( i ) off ) for each ( s ( i ) off , a ( i ) off ) via Eq. (4).

13: Select the lowest-gap subset ˜ b off ⊂ b off with | ˜ b off | = ⌈ ξB ⌉ .

14: Shape rewards for selected offline transitions: for each ( s, a, r, s ′ ) ∈ ˜ b off set ˜ r ← r + β ̂ ∆( s, a ) .

15: # Critic update (use online batch + selected offline batch)

16: Define critic batch b Q ← b on ∪ ˜ b off with rewards r for online and ˜ r for selected offline.

18:

19:

Compute Bellman targets for ( s, a, r, s ′ ) ∈ b Q :

Update critic by gradient descent:

20: end for

21: # Soft update target critics

22: for i = 1 , 2 do

24: end for

25: # Actor update (policy improvement + behavior cloning)

26: Define actor state batch b π ←{ s : ( s, · , · , · ) ∈ b on ∪ ˜ b off } .

27: Sample ˜ a ∼ π φ ( · | s ) for each s ∈ b π .

28: Set normalization weight

29: Minimize actor loss:

where ˜ a ∼ π φ ( · | s ) for s ∈ b π , and ¯ a ∼ π φ ( · | s ) for ( s, a ) ∈ b off .

30:

31: end for

32: end for

33: Output: Learned policy π φ .

F Proof of Lemma 2.2

Lemma 2.2 (Return Bound between Two Environments [37]) . Let the empirical behavior policy in D off be π D off ( a | s ) . Define C 1 = 2 r max (1 -γ ) 2 . Then for any policy π ,

Proof. Proof of Lemma 2.2 is already provided as an intermediate step for Theorem A.4 in [37], the offline performance bound case. We include their proof here with slight modifications for completeness.

We first cite the following two necessary lemmas also used in [37].

Lemma F.1 (Extended telescoping lemma) . Denote M 1 = ( S , A , P 1 , r, γ ) and M 2 = ( S , A , P 2 , r, γ ) as two MDPs that only differ in their transition dynamics . Suppose we have two policies π 1 , π 2 , we can reach the following conclusion:

Proof. This is Lemma C.2 in [65].

Lemma F.2. Denote M = ( S , A , P, r, γ ) as the underlying MDP. Suppose we have two policies π 1 , π 2 , then the performance difference of these policies in the MDP gives:

Proof. This is Lemma B.3 in [37].

Proof. We have access to the offline data M off and the empirical behavioral policy π D off , so we bound the performance between η M on ( π ) and η M off ( π D off ) . We have

To bound (a) term, we use Lemma F.1 in the second equality

To bound term (c), we use

where the last inequality comes from the fact that | Q π M on ( s ′ off , a ′ ) | ≤ r max 1 -γ and the definition of TV distance.

where the inequality comes from the fact that Q π M on ( s ′ , a ′ ) weighted by probability is bounded by r max 1 -γ , the upper bound for all Q-values.

Combine the bounds for (c) and (d), we obtain a bound for the (a) term:

Now we try to bound term (b).

where the first equality is a direct application of Lemma F.2. Combine the bound for term (a) and (b), we conclude that

where C 1 = 2 r max (1 -γ ) 2 .

G Proof of Theorem 3.1

Definition G.1. For any distributions p 0 , p 1 on R d with finite second moments, define D 2 ( p 0 , p 1 ) := E || X 1 -X 0 || 2 2 , X 0 ∼ p 0 , X 1 ∼ p 1 independent.

Theorem 3.1 (Conditions for Composite Flow Yielding Smaller Errors) . Assume that composite flow and direct flow share the same hypothesis class H of (measurable) vector fields v : R d × [0 , 1] → R d , and that there exists B ∈ (0 , ∞ ) such that sup v ∈H sup x ∈ R d , t ∈ [0 , 1] ∥ v ( x, t ) ∥ 2 ≤ B . Also, assume that max { Tr(Σ on ) , Tr( ̂ Σ off ) } ≤ C TR for some C TR . The composite flow enjoys a strictly tighter high-probability generalization bound than the direct flow if and only if

Here, p G := N (0 , I d ) , Σ on is the covariance of p on , and ̂ Σ off is the covariance of ˆ p off .

Proof. We consider linear-path flow matching (FM) trained with squared loss. A single training example is generated by: t ∼ Unif[0 , 1] ; X 0 ∼ p 0 (the start distribution), independently of t ; X 1 ∼ p on , independently of ( t, X 0 ) ; The interpolation X t := (1 -t ) X 0 + tX 1 and the label Y := X 1 -X 0 ∈ R d .

Let H be a hypothesis class of (measurable) vector fields v : R d × [0 , 1] → R d and define the population risk

where the expectation is over the sampling mechanism above. Let

denote the Bayes (population) minimizer. Given n i.i.d. samples, let ˆ v be any empirical risk minimizer over H for the squared loss.

Assumptions. A1 ( Uniform boundedness ) There exists B ∈ (0 , ∞ ) such that sup v ∈H sup x ∈ R d , t ∈ [0 , 1] ∥ v ( x, t ) ∥ 2 ≤ B . This ensures integrability and controls the Lipschitz constants used below.

FMas regression; Bayes predictor and excess-risk identity. By definition,

For any measurable v , the Bayes (population) minimizer for the squared loss is the conditional mean

A standard identity for squared loss (we derive it fully) is

Derivation of (13) . Expand and add/subtract v ⋆ ( X t , t ) :

where all functions are evaluated at ( X t , t ) but notationally suppressed for readability. Condition on ( X t , t ) : by definition, E [ Y | X t , t ] = v ⋆ ( X t , t ) , hence E [ v ⋆ -Y | X t , t ] = 0 and the cross term vanishes after taking expectations. Therefore

which rearranges to (13).

In Lemma G.3, we prove that there exists an absolute constant c > 0 such that, for all δ ∈ (0 , 1) , with probability at least 1 -δ (over the sample draw),

where R n ( H ) is the (data-independent) Rademacher complexity of H . In short: symmetrization + vector contraction + Lipschitzness of the squared loss on the B -ball yields (14).)

Subtract R ( v ⋆ ) on both sides of (14) and apply (13) with v = ˆ v :

Identify E ∥ Y ∥ 2 2 as D 2 ( p 0 , p on ) and linearize the square root. By independence of X 0 ∼ p 0 and X 1 ∼ p on ,

By applying Lemma G.4, we have

Therefore, we obtain

This gives the following master bound (19) with C 0 := c ( B + √ Tr(Σ on ) + C TR and C 1 := c .

Master bound. Under A1 , for any start p 0 (with fixed target p on ), with probability at least 1 -δ ,

Composite vs. direct. Consider two trainings sharing the same online p on and class H :

Applying (19) to (20) and (21) yields

When is the composite bound strictly tighter? Applying (19) to p 0 = ˆ p off (composite) and p 0 = p G (direct) gives (22) and (23). Since C 0 and Γ n,δ are identical across the two trainings (they depend on ( B,c, Σ on , C TR ) , n , δ , and H , but not on p 0 ), the only difference in the right-hand sides is the factor W 2 ( · , p on ) .

Therefore the composite bound (22) is strictly smaller than the direct bound (23) if and only if

Remark G.2 (Optional D 2 strengthening) . One can analogously replace W 2 ( p 0 , p on ) by D 2 ( p 0 , p on ) when bounding 15, using the standard inequality trick √ x ≤ 1 2 ( x +1) for all x ≥ 0 . We use W 2 metric because it's more commonly adopted in the literature, but D 2 leads to a tighter bound since we no longer needs the constant C TR . We do not expand the D 2 case here, as the argument mirrors the W 2 case.

Lemma G.3 (Generalization gap for squared loss under uniform boundedness) . Assume A1 . Let ℓ v ( z ) := ∥ v ( x, t ) -y ∥ 2 2 for z = ( x, t, y ) generated as in the theorem. Then there exists a positive constant c > 0 such that, for any δ ∈ (0 , 1) , with probability at least 1 -δ over an i.i.d. sample of size n ,

Proof. We sketch the standard steps and make constants explicit where needed:

(i) Symmetrization. Let ˆ R ( v ) denote the empirical risk. For ERM ˆ v ,

up to negligible terms; this is standard (e.g., by a one-sided symmetrization plus a chaining step). Thus it suffices to bound sup v ∈H ( R ( v ) -ˆ R ( v )) .

(ii) Lipschitz envelope via contraction. For fixed ( x, t, y ) with ∥ v ( x, t ) ∥ ≤ B and any u ∈ R d with ∥ u ∥ ≤ B ,

Thus, as a function of u , the map u ↦→∥ u -y ∥ 2 is (2( B + ∥ y ∥ )) -Lipschitz on the B -ball. Applying the vector contraction inequality to the class { v ( · , · ) ∈ H} gives

By Cauchy-Schwarz and Jensen, E ∥ Y ∥ ≤ √ E ∥ Y ∥ 2 .

(iii) High-probability upgrade. A standard bounded differences (or Bernstein-type) argument for Lipschitz, sub-quadratic losses upgrades the expected sup-gap to a high-probability bound, adding the usual √ log(1 /δ ) /n term. Collecting constants and using E ∥ Y ∥ ≤ √ E ∥ Y ∥ 2 yields the claimed form with some absolute c > 0 :

Since ˆ v is an ERM, ˆ R (ˆ v ) ≤ ˆ R ( v ⋆ ) , and therefore R (ˆ v ) -R ( v ⋆ ) ≤ 2 sup v ( R ( v ) -ˆ R ( v )) in the final step, absorbing the factor 2 into c . □

Lemma G.4. Let µ · and Σ · denote mean and covariance, respectively. We show that if Tr(Σ 0 ) ≤ C TR for some C TR > 0 , we have:

Proof. Let X 0 ∼ p 0 and X 1 ∼ p 1 be independent with means µ 0 , µ 1 and covariances Σ 0 , Σ 1 . We compute

By independence, E ⟨ X 1 , X 0 ⟩ = ⟨ E X 1 , E X 0 ⟩ = ⟨ µ 1 , µ 0 ⟩ . Moreover, for any random vector Z with mean µ and covariance Σ , E ∥ Z ∥ 2 2 = ∥ µ ∥ 2 2 +Tr(Σ) . Thus

Let Γ( p 0 , p 1 ) be all couplings of p 0 , p 1 . For any π ∈ Γ( p 0 , p 1 ) with ( X 0 , X 1 ) ∼ π , we have

by Jensen's inequality since z ↦→∥ z ∥ 2 2 is convex. Taking the infimum over all couplings gives

Therefore, we have ∥ µ 1 -µ 0 ∥ 2 2 ≤ W 2 2 ( p 0 , p 1 ) . Hence,

This completes the proof.

The assumption that Tr(Σ 0 ) ≤ C TR makes sense because we assume bounded state space for all settings considered in our paper.

H Proof of Theorem 3.5

Theorem 3.5 (Large Dynamics Gap Exploration Reduces Performance Gap) . Compared to behavior cloning policy π bc on the offline dataset, training a policy ˆ π by replacing all offline samples with a dynamics gap exceeding κ (as estimated by the composite flow) with online environment samples can reduce the performance gap to the optimal online policy π ∗ on with high probability by

Here L r and L P are the Lipschitz constants for the reward and transition function, respectively. γL p < 1 . ∆ W 2 := sup s,a W 2 ( p off ( ·| s, a ) , p on ( ·| s, a ) ) is the largest dynamics gap. C 0 and C 1 are two constants. Γ N on ,δ := R N on ( H ) + √ log(1 /δ ) N on , where H is the same as in Theorem 3.1 and N on is the number of samples used to train the online flow.

Proof. Let's first list the assumptions:

R1: Lipschitz reward. For every a ∈ A , s ↦→ r ( s, a ) is L r -Lipschitz.

R2: Lipschitz dynamics. For m ∈ { off , on } , W 1 ( p m ( · | s, a ) , p m ( · | s ′ , a ) ) ≤ L p d ( s, s ′ ) for all s, s ′ , a .

R3: Contraction. γL p < 1 .

We proceed in four steps: (A) a policy-dependent model-gap bound; (B) a uniform high-probability W 2 -gap estimation bound for the rectified flow ( α ); (C) learning-gap bounds for π bc and ˆ π ( β ); (D) assembly via the standard three-term decomposition.

Step A: Policy-dependent model-gap bound. For any fixed policy π , we cite Theorem 3.1 of [45] for the following bound. Under R1 -R3 , we have

Averaging over the initial state-distribution µ gives us a bound on what we call the model gap:

Moreover, the bounded state space assumption in Theorem 3.1 implies bounded second moment for both p off and p on , so we have W 1 ( p off , p on ) ≤ W 2 ( p off , p on ) ∀ p off , p on . Hence, we obtain that

Step B: FM-based uniform W 2 estimation bound (construction of α ). Weneed a high-probability uniform control of the error of the rectified-flow W 2 estimator to ensure that thresholding by κ actually caps the on-policy mismatch for ˆ π .

B.1. Flow-to-map error. Let T be the Brenier OT map from ˆ p off ( · | s, a ) to p on ( · | s, a ) (guaranteed by standard assumptions), and let ̂ T be the terminal map obtained by integrating the learned FM velocity field ˆ v from t = 0 to t = 1 . Under the linear path and Lipschitz regularity of v ⋆ and ˆ v (Theorem 3.1 assumptions), a stability argument (Gronwall) yields

for some finite K stab depending on Lipschitz constants of the flow. By the 'master bound' (Eq. (7) in Theorem 3.1), uniformly over t ∈ [0 , 1] , with prob. ≥ 1 -δ ,

So we have

B.2. From map MSE to W 2 error. By the coupling construction (ˆ s, s ) = ( ̂ T ( X 0 ) , T ( X 0 )) with X 0 ∼ ˆ p off and the triangle inequality for W 2 ,

Define the coupling γ ( s,a ) := ( T, ̂ T ) # ˆ p off ( · | s, a ) , i.e., if S := T ( X 0 ) and ̂ S := ̂ T ( X 0 ) then ( S, ̂ S ) ∼ γ ( s,a ) and γ ( s,a ) ∈ Π ( p on ( · | s, a ) , ˆ p on ( · | s, a ) ) . By the definition of the 2 -Wasserstein distance,

Taking square roots gives, for each ( s, a ) ,

Finally, taking the supremum over ( s, a ) preserves the inequality:

Combining (35) and (32) gives

with probability at least 1 -δ .

B.3. Capping the on-policy gap for ˆ π . By construction of the replacement rule, for any ( s, a ) encountered during fine-tuning of ˆ π we have either ̂ ∆ W 2 ( s, a ) ≤ κ (kept offline sample) or we replaced it by an online sample. In either case, on those ( s, a ) we ensure

Therefore the on-policy gap of ˆ π obeys

Step C: Learning-gap bounds for π bc and ̂ π (construction of β ). Recall the three-term decomposition for any policy π :

We now bound the learning term (b) by ERM generalization.

Policies are learned by ERM on a surrogate imitation loss L ( π ) over a policy class Π with Rademacher complexity R N (Π) . There exists a calibration constant C val (depends on concentrability/Lipschitzness; fixed for the class) such that, for any π ,

Moreover, for datasets of sizes N off , N mod , with probability at least 1 -δ ,

for an absolute constant C Π .

Let π bc be the behavior cloning policy trained on the offline dataset of size N off , and let ˆ π be ERM on the modified dataset (size N mod ): for any ( s, a ) encountered, if the FM-estimated gap satisfies ̂ ∆ W 2 ( s, a ) > κ , replace the offline ( s, a, s ′ ) sample by an online ( s, a, s ′ on ) sample (collecting N on such online transitions), leaving all others unchanged. Denote N mod the resulting (modified) dataset size.

By calibration (38) and the generalization inequality (39), with probability at least 1 -δ we have

and similarly

By a union bound (probability 1 -2 δ ) and adding these two upper bounds we obtain a common envelope

with

Step D: Assemble bounds and compare. Incorporate (27) and (40) to (37) with π = π bc , we have with high probability,

Incorporate (27) and (40) to (37) with π = ˆ π , we have with high probability

The difference between these two upper bounds is:

Absorbing K stab into C 0 and C 1 gets the result.

I Proof of Proposition 3.3

We provide a rigorous justification for Proposition 3.3. The key technical point is that the augmented optimal transport (OT) problem on triples ( s, a, s ′ ) with cost

reduces, as η →∞ , to conditional quadratic OT on the next-state variable x = s ′ given each fixed ( s, a ) , provided that the ( s, a ) -marginals match.

I.1 Formal Proposition

Proposition I.1 (Shared-latent coupling recovers conditional W 2 ) . Setup. Let Y := S × A denote the conditioning space and X := S denote the next-state space. Let µ be a probability measure on Y . Let α, β be probability measures on Y × X such that their Y -marginals coincide:

where π sa ( y, x ) = y . Assume α and β admit disintegrations with respect to µ :

where y = ( s, a ) ∈ Y and x ∈ X .

Augmented OT problem. For η > 0 , define the augmented quadratic cost on ( Y × X ) 2 by

where ∥ y -y ′ ∥ 2 2 := ∥ s -s ′ ∥ 2 2 + ∥ a -a ′ ∥ 2 2 . Define the corresponding OT value

where Π( α, β ) is the set of couplings on ( Y × X ) 2 with marginals α and β .

Assumptions (for η → ∞ reduction). Assume α and β have finite second moments, that W 2 ( α y , β y ) < ∞ for µ -a.e. y , and that the conditional squared Wasserstein cost is µ -integrable:

Claim 1 (conditional decomposition). As η →∞ ,

Claim 2 (structure of limiting couplings). Let γ η ∈ arg min γ ∈ Π( α,β ) ∫ c η dγ be any sequence of optimal couplings. Then there exists a subsequence γ η n ⇒ γ ∞ weakly such that:

(i) γ ∞ is supported on the diagonal set { ( y, x ; y, x ′ ) : y = y ′ } ;

(ii) γ ∞ admits a disintegration

where γ y ∈ Π( α y , β y ) is W 2 -optimal for µ -a.e. y .

OT-FM implication (per fixed conditioning y ). Fix any y = ( s, a ) and let p 0 = N (0 , I ) . Assume the offline flow satisfies

For η > 0 and k ∈ N , let ψ on ,k,η ϕ ( · , 2 | y ) denote the online flow obtained by running OT-FM to convergence with minibatch size k , where each iteration computes an optimal transport plan between two empirical measures of size k under the augmented cost (44) .

Assume the standard OT-FM regularity conditions hold for this fixed y (e.g., those in Theorem 4.2 of [47]), including:

(A1) α y and β y have bounded support;

(A2) α y admits a density and the quadratic OT map from α y to β y exists and is continuous;

(A3) the minibatch OT plan is solved exactly at each iteration.

Then, for µ -a.e. y , we have the two-limit convergence

I.2 Proof

Proof. We prove (47) by matching an upper bound and a lower bound, which also implies that any limiting optimal coupling concentrates on the diagonal set { y = y ′ } . Finally, we connect the resulting conditional OT structure to OT-FM trained with the augmented matching cost (44).

Step 1 (Upper bound via a diagonal coupling). Since α and β share the same Y -marginal µ , we may couple them by matching the conditioning variable y exactly. For µ -a.e. y , let q ⋆ y ∈ Π( α y , β y ) be an optimal coupling achieving W 2 2 ( α y , β y ) . A measurable selection y ↦→ q ⋆ y can be chosen (see, e.g., [63, Chapter 5] or [52, Section 1.4]). Define ¯ γ ∈ Π( α, β ) by

where δ y is the Dirac measure at y . Because ¯ γ is supported on { y = y ′ } , for every η > 0 ,

Therefore,

Step 2 ( y -transport vanishes as η →∞ ). Let γ η be an optimal coupling attaining W 2 2 ,η ( α, β ) and define

By Step 1, W 2 2 ,η ( α, β ) ≤ C ⋆ , hence

Since c η (( y, x ) , ( y ′ , x ′ )) ≥ η ∥ y -y ′ ∥ 2 2 , it follows that

Step 3 (Extract a weak limit supported on the diagonal). Because α and β have finite second moments, Π( α, β ) is tight, hence { γ η } η> 0 is tight. Thus there exist η n → ∞ and γ ∞ such that γ η n ⇒ γ ∞ weakly. Let f ( y, y ′ ) := ∥ y -y ′ ∥ 2 2 and define the bounded continuous truncation f M := min { f, M } . Since f M is bounded and continuous,

Moreover, 0 ≤ f M ≤ f , so by (51),

Letting M →∞ and applying monotone convergence yields ∫ f dγ ∞ = 0 , which implies y = y ′ holds γ ∞ -a.s. Hence γ ∞ is supported on { ( y, x ; y, x ′ ) : y = y ′ } .

Step 4 (Lower bound and convergence of values). Since γ ∞ is supported on { y = y ′ } and has marginals α and β , it disintegrates as

where γ y ∈ Π( α y , β y ) for µ -a.e. y . Let g ( x, x ′ ) := ∥ x -x ′ ∥ 2 2 and g M := min { g, M } . Because 0 ≤ g M ≤ c η n ,

By weak convergence, ∫ g M dγ η n → ∫ g M dγ ∞ . Letting M →∞ and using monotone convergence yields

Combining this with (50) proves (47). Moreover, the last inequality can only be tight if

which implies γ y is a W 2 -optimal coupling between α y and β y for µ -a.e. y . This proves Claim 2.

Step 5 (OT-FM with augmented matching cost; then η → ∞ ). Fix any conditioning value y = ( s, a ) and recall that the offline flow satisfies

Let x 1 := ψ off θ ( x 0 , 1 | y ) with x 0 ∼ p 0 , so that x 1 ∼ α y .

(a) OT-FM population limit for a fixed η . Fix η > 0 . Our OT-FM procedure builds minibatch OT matchings between joint samples ( y i , x i ) ∼ α and ( y ′ j , x ′ j ) ∼ β using the augmented cost c η , and then trains a conditional map x ↦→ ψ on ,k,η ϕ ( x, 2 | y ) by regressing onto the matched x ′ values.

Let γ η ∈ Π( α, β ) be an optimal coupling attaining W 2 2 ,η ( α, β ) . Disintegrate γ η with respect to its first marginal α :

where Γ η y,x is a probability kernel on Y × X . Define the induced conditional law of x ′ given ( y, x ) by marginalizing out y ′ :

The associated barycentric projection is

Under assumptions (A1)-(A3) and OT-FM convergence guarantees (e.g., Theorem 4.2 of [47], applied with the quadratic cost in the augmented space), training OT-FM to convergence with minibatch size k yields the population limit

In particular, by conditioning on y , for µ -a.e. fixed y ,

(b) Send η →∞ (reduction to conditional OT). Take a subsequence η n →∞ such that γ η n ⇒ γ ∞ as in Step 3. Let b η n be the barycentric projection induced by γ η n as above. Since β has finite second moment, { b η n } n is uniformly bounded in L 2 ( α ) , so by extracting a further subsequence (not relabeled) we may assume

Because γ ∞ is supported on { y = y ′ } and disintegrates as

the conditional law of x ′ given ( y, x ) under γ ∞ is supported on the conditional coupling γ y .

Under assumption (A2), for µ -a.e. y , the optimal quadratic-cost coupling γ y is induced by a Monge map T y : X → X (i.e., x ′ = T y ( x ) γ y -a.s.). Hence the limiting barycentric projection satisfies

Therefore, for µ -a.e. y ,

Finally, combining (53) with the identification b ∞ ( y, x ) = T y ( x ) and (55), we obtain for µ -a.e. y ,

Substituting x = ψ off θ ( x 0 , 1 | y ) with x 0 ∼ p 0 yields (48). This completes the proof.

Intuition: Why do we need matched ( s, a ) marginals?

When ( s, a ) is continuous, it is unlikely to observe many online transitions s ′ on with exactly the same conditioning pair ( s, a ) . As a result, the OT-FM objective must couple two sets of transitions whose conditioning variables differ across samples. To recover the conditional Wasserstein distance W 2 ( p off ( · | s, a ) , p on ( · | s, a )) for a fixed ( s, a ) , we add ( s, a ) into the OT cost:

As η grows, transporting mass across different ( s, a ) becomes prohibitively expensive, so the optimal coupling is forced to match samples with similar ( s, a ) and transport occurs primarily in the next-state space.

This reduction requires the two empirical measures to share the same marginal distribution over ( s, a ) ; otherwise, the coupling is forced to move mass in ( s, a ) even when η is large. In practice, we ensure this by sampling ( s off , a off ) from the online dataset D on to construct the offline-generated batch (then sampling s ′ off ∼ p off ( · | s, a ) using the offline flow), while sampling online transitions ( s on , a on , s ′ on ) independently from D on .

J Additional Experimental Results

To further increase the task difficulty, we manually modified the benchmark configurations to amplify the dynamic differences beyond the given settings. For example, in the morphology task, we drastically reduced thigh segment size-making the legs much shorter, in the kinematic task, we almost immobilized the back-thigh joint by severely limiting its rotation range. As shown in Table 2, our method clearly outperforms all other baselines in 5 out of 6 settings, while achieving ties in the remaining one.

K Experimental Details of Gym-MuJoCo

In this section, we describe the detailed experimental setup as well as the hyperparameter setup used in this work.

K.1 Environment Setting

K.1.1 Offline Dataset

We use the MuJoCo datasets from D4RL [16] as our offline data. These datasets are collected from continuous control environments in Gym [4], simulated using the MuJoCo physics engine [62]. We focus on three benchmark tasks: HalfCheetah , Hopper , and Walker2d , and evaluate across three dataset types: medium , medium-replay , and medium-expert .

The medium datasets consist of trajectories generated by an SAC policy trained for 1M steps and then early stopped.

The medium-replay datasets capture the replay buffer of a policy trained to the performance level of the medium agent.

The medium-expert datasets are formed by mixing equal proportions of medium and expert data (50-50).

K.1.2 Kinematic Shift Tasks

We use Kinematic Shift Tasks from the benchmark [38]. We select most shift level 'hard' to make the tasks more challenging

HalfCheetah Kinematic Shift: The rotation range of the foot joint is modified to be:

Hopper Kinematic Shift: the rotation range of the foot joint is modified from [ -45 , 45] to [ -9 , 9] :

Walker2D Kinematic Shift: the rotation range of the foot joint is modified from [ -45 , 45] to [ -9 , 9] :

K.1.3 Morphology Shift Tasks

We use Morphology Shift Tasks from the benchmark [38]. We select most shift level 'hard' to make the tasks more challenging

HalfCheetah Morphology Shift: the front thigh size and the back thigh size are modified to be:

Hopper Morphology Shift: the foot size is revised to be 0.4 times of that within the source domain:

Walker2D Morphology Shift: the leg size of the robot is revised to be 0.2 times of that in the source domain.

K.1.4 Friction Shift Tasks

Following [38], the friction shift is implemented by altering the friction attribute in the geom elements. The frictional components are adjusted to 5.0 times the frictional components in the offline environment. The following is an example for the Hopper robot.

K.2 Implementation Details

BC-SAC: This baseline leverages both offline and online transitions for policy learning. Since learning from offline data requires conservatism while online data does not, we incorporate a behavior cloning term into the actor update of the SAC algorithm. Specifically, the critic is updated using

standard Bellman loss on the combined offline and online datasets, and the actor is optimized as:

where λ = ω 1 N ∑ ( s j ,a j ) min i =1 , 2 Q ς i ( s j ,a j ) and ω ∈ R + is a normalization coefficient. We train BC-

SAC for 400K gradient steps, collecting online data every 10 steps. We use the hyperparameters recommended in [38].

H2O: H2O [46] trains domain classifiers to estimate dynamics gaps and uses them as importance sampling weights during critic training. It also incorporates a CQL loss to encourage conservatism. Since the original H2O is designed for the Online-Offline setting, we adapt the objective to the Offline-Online setting. The critic loss is:

where ω ( s, a, s ′ ) is the dynamics-based importance weight, and β CQL is the penalty coefficient. We set β CQL = 10 . 0 , which performs better than the default 0.01. We reproduce H2O using the official codebase, 3 and adopt the suggested hyperparameters. H2O is trained for 40K environment steps, with 10 gradient updates per step.

BC-VGDF: BC-VGDF [65] filters offline transitions whose estimated values align closely with those from the online environment. It trains an ensemble of dynamics models to predict next states from raw state-action pairs under the online dynamics. Each predicted next state is evaluated by the policy to obtain a value ensemble { Q ( s ′ i , a ′ i ) } i M =1 , forming a Gaussian distribution. A fixed percentage ( ξ % ) of offline transitions with the highest likelihood under this distribution are retained. The critic loss is:

where Λ( s, a, s ′ ) denotes the fictitious value proximity (FVP), and Λ ξ % is the ξ -quantile threshold. VGDF also trains an exploration policy. Actor training includes a behavior cloning term, as in BC-SAC. We follow the official implementation, 4 use the recommended hyperparameters, and train for 40K environment steps with 10 gradient updates per step.

BC-PAR: BC-PAR [37] addresses dynamics mismatch through representation mismatch, measured as the deviation between the encoded source state-action pair and its next state. It employs a state encoder f ψ and a state-action encoder g ξ , both trained on the target domain. The encoder loss is:

where SG is the stop-gradient operator. Rewards of the offline data are adjusted as:

where β controls the penalty strength. The actor ( π φ ) and critic ( Q ς i ) are jointly trained using both offline and online data. Actor training includes a behavior cloning term, similar to BC-SAC. We implement BC-PAR using the official codebase, 5 adopt the suggested hyperparameters, and train for 40K environment steps with 10 gradient updates per step.

COMPFLOW. When training the target flow, we use a quadratic cost function and employ the Python Optimal Transport (POT) library [15] to compute the optimal transport plan for each minibatch using the exact solver. Additional hyperparameters are provided in Table 3 and Table 4. Since the exploration bonus term is closely tied to properties of the environment-such as the state space, action space, and reward structure-it is expected that the optimal exploration strength β varies across tasks. We perform a sweep over β ∈ { 0 . 01 , 0 . 1 , 0 . 2 } , and select the offline data selection ratio ξ % from { 30% , 50% } .

L Experimental Details of Wildlife Conservation

We use the green security simulator in [67]. The model is a Markov decision process with state, action, transitions, and a terminal return. We summarize the parts needed to reproduce our experiments.

State and action. At time t , the state is

where w t = ( w 1 t , . . . , w N t ) is wildlife across N cells and a t = ( a 1 t , . . . , a N t ) is patrol effort. The defender chooses a t ∈ [0 , 1] N under the budget ∑ N i =1 a i t ≤ B .

Adversary behavior. At each step, the poacher places a snare in cell i with probability

where z i is the baseline attractiveness of cell i . The parameters β < 0 and η > 0 capture deterrence from prior patrol and displacement from neighboring patrols. The realized attack is

Wildlife transition. After attacks, wildlife in each cell evolves by natural growth and poaching losses:

where ϕ > 1 is the growth rate and α > 0 is the loss per uncovered attack. This defines the transition kernel T z over states given actions:

Return. The episode return is the total wildlife at the horizon,

and the expected return of a policy π under environment parameters z is

We assume access to an offline dataset of 100,000 transitions collected in Murchison Falls National Park using a well-trained SAC policy. For the online environment in Queen Elizabeth National Park, we are allowed a budget of 40,000 interactions.

M Discussion on Computational Cost

The practical cost of data filtering is relatively small. For example, with a batch size of 256 for training the policy network and a Monte Carlo sample size of 30 , the entire filtering process takes just 0 . 03 seconds on an A100 GPU. We will highlight in the paper that this efficiency is due to the simplicity of our flow training objective, which follows a linear path. As a result, solving the corresponding ODE at inference time is very easy-A basic Euler method with just 10 time steps is sufficient.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We propose a new method for reinforcement learning with shifted-dynamics data based on flow matching.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations of our method in Section 6.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provide all proofs for the theorems, propositions, and lemmas in the Appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental result reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide all the experimental details in Appendix K.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We have provided the code in the paper. Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide all the experimental details in Appendix K.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We report the standard deviation across different random seeds for all tasks. Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: All the computing infrastructure information has been provided Appendix K. Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We conform with the NeurIPS code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We have discussed both potential positive societal impacts and negative societal impacts in Appendix B.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have cited all code, data, and models we used in the paper.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: This paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and research with human subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [NA]

Justification: We did not use LLMs.

Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a method that leverages the source offline datasets to improve sample efficiency and generalization in the target environment. Specifically, the method uses optimal transport to calculate the dynamic gap between the source task and the target task, and then selectively expands the replay buffer of the target task for agent training.\n\nSTRENGTHS AND WEAKNESSES:\nAdvantages: The idea presented in this paper is innovative and the theory is solid. \n\nDisadvantages: Some parts of the paper are vaguely expressed, which could cause confusion for readers. Please refer to the issues in the question section and improve the article accordingly.\n\nQUESTIONS:\n1. I suggest that the author specify the range of values for t in the source flow and target flow in Section 3.1.\n2. To address the issue of data scarcity (mentioned in line 136) in the target domain, the author proposes using a two-stage flow to model the target data distribution (line 148), namely, first flowing from the Gaussian prior to the source data distribution, and then from the target data distribution. However, the loss function from the source data distribution to the target data distribution is still based on a small amount of target data, and I believe the issue of data scarcity still persists.\n3. Why is the loss function of the target flow shown in line 174 inconsistent with the loss function in line 16, Algorithm 3?\n4. In Proposition 3.2, suppose both the source flow and target flow can perfectly model the flow map, according to the content introduced by the author in Section 2.3, the corresponding Wasserstein cost for this formula is $(s'_{tar} - s'_{src})^2$. However, the author has mentioned that the Wasserstein cost they intend to use is the cost shown in line 172. Therefore, the author needs to clarify this point.\n5. Essentially, Equation (2) measures how large the gap is between s' that the source dynamics and target dynamics transition to under the same conditions (s, a). Is this understanding correct?\n6. Theorem 3.3 is counterintuitive because, intuitively, selecting samples from the source data that are completely dissimilar to the target data will clearly hinder the model's generalization rather than improve it. Therefore, I am skeptical about the conclusion of Theorem 3.3.\n7. Since Theorem 3.3 mentions that 'exploration in regions with high dynamics gap can reduce the performance gap relative to the optimal policy,' why, in practice, do the authors choose the source data with the lowest gap?\n8. To fully demonstrate the generalization performance of the proposed method, I suggest that the author conduct tests on source and target tasks with larger dynamic differences, such as training on HalfCheetah and testing the generalization performance on Walker2D. Alternatively, the author could consider environments in Meta-World [1], which have the same state and action spaces.\n\n    [1] Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning.\n9. I suggest that the author compare the performance differences between direct flow and composite flow in ablation study, rather than comparing the MSE of the loss, as the loss may be influenced by the values of each dimension in the state and action space.\n10. How is $v^*$ obtained in line 12 of Algorithm 1? Additionally, could the author explain the meaning of $A(i,j)$ and provide a detailed calculation of the optimal transport plan's solution process?\n11. Does the author intend to express that x represents a transition in line 10 of Algorithm 2? However, the v in the loss function provided here is an unconditional velocity field, while the source flow loss in Section 3.1 actually refers to a conditional velocity field.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper tackles the offline-to-online RL setting in which the available offline data come from environments with shifted dynamics. It introduces a novel method for quantifying the dynamics gap that combines a composite-flow model with the Wasserstein distance. The approach shows strong empirical performance on a set of simulated tasks and, in my view, offers a timely contribution to the field.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The idea of using composite flow models and the Wasserstein metric to estimate the dynamics gap is novel.\n- The empirical evaluation is thorough; the proposed method consistently outperforms relevant baselines across most tasks.\n    \n    Weaknesses:\n- Wasserstein distances are approximated by solving discrete OT problems on mini-batches, which adds computational overhead and approximation error. The impact of mini-batch size is not analyzed.\n- **Theorem 3.1**\n    - The bound on the total-variation distance lacks explicit rate information: does it tighten as $N_{tar}$ increases?\n    - The bound can exceed 1, which makes it vacuous; please clarify the conditions under which it is non-trivial.\n    - Because the rest of the paper measures gaps with the Wasserstein distance, presenting this bound in the same metric would improve consistency.\n- **Theorem 3.3**\n    - Is $\\Delta_\\kappa$ always non-negative? If not, under what conditions is it positive?\n\nQUESTIONS:\n- **Assumptions on the source data.** What coverage properties are required? How many source samples are needed for a given state–action pair?\n- **Pure online RL baseline.** To demonstrate that the source data do not hinder learning in the target environment, a comparison with a policy trained from scratch solely on target-environment data will be appreciated.\n\nLIMITATIONS:\n-na-", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces COMPFLOW, a method for knowledge transfer in RL between a source and a target domains with different dynamics. The core contribution is a composite flow matching technique that leverages optimal transport to estimate the dynamics gap between the two domains using Wasserstein distance, where the main idea is to train the flow for the target dynamics on top of the flow for the source dynamics (ie with noise replaced by samples from the latter). COMPFLOW then uses Wasserstein distance to perform both data filtering and guide online exploration to regions with large gap. Empirically, the method is tested on 3 locomotion tasks from D4RL and compares favorably to existing baselines.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The paper tackles a relevant and well-motivated problem, as the setting in which we need to train policies on one domain with scarce data but have additional data under different dynamics frequently arises in practice (eg, in robotics)\n2. The proposed method is novel and sound. I particularly like the idea of \"chaining\" two flow-matching models as a tool to improve predictions and ease the estimation of Wasserstein distance\n3. The proposed method is backed by theoretical results that formally justify its main components: composite flow (Th. 3.1) and exploration bonus (Th. 3.3)\n4. The paper is very well-written and easy to read (apart from some specific details I could not grasp, see questions below)\n5. COMPFLOW performs well empirically. A good amount of ablations is provided\n\nWeaknesses:\n1. Flow-matching, while powerful, may introduce non-negligible computational/implementation overhead (eg estimating the Wasserstein distance requires multiple integrations of the ODE, which may be more inefficient than eg simpler density-based models). There is no discussion about that in the paper\n2. The experiments are carried out only on 3 locomotion domains from D4RL. Adding manipulation and/or navigation tasks (eg from the same suite) would add great value to the experiments\n3. A baseline that learns from scratch in the target domain is missing, which makes it hard to understand the benefits of transferring source samples in these domains\n4. The impact of the exploration parameter \\beta (Figure 5) seems very limited (curves mostly overlap). One wonders if that component is really important: would the performance drop significantly if exploration bonus is removed?\n5. While the paper is generally well written, there are a few theoretical details which are hard to grasp:\n   - Theorem 3.1: the MSE was not formally defined. Also, it is not clear under what conditions we should expect the TV bound assumption to hold\n   - Not clear how the optimal coupling q^* is computed in L_tar(\\phi)\n   - It seems the goal of the FM models is to learn conditional distributions s' | s,a, but many places (eg the paragraph 169 - 178) give the impression that these models should actually capture the joint distributions of triples (s,a,s'). Which one is correct?\n   - Theorem 3.3: what assumptions are made on the composite flow model to require that its estimated dynamics gap is below \\kappa? Is training assumed \"perfect\", so that this model returns the exact dynamics gap and not an estimate?\n6. A bunch of existing works for the problem considered in this paper use instance re-weighting rather than filter samples (see eg [1,2] and references therein). These are not discussed / compared empirically\n\n[1] Tirinzoni, A., Sessa, A., Pirotta, M., & Restelli, M. (2018, July). Importance weighted transfer of samples in reinforcement learning. In International Conference on Machine Learning (pp. 4936-4945). PMLR.\n[2] Xie, A., & Finn, C. (2022, November). Lifelong robotic reinforcement learning by retaining experiences. In Conference on Lifelong Learning Agents (pp. 838-855). PMLR.\n\nQUESTIONS:\n1. Please address limitations above\n2. Following the discussion about the improvement of composite vs direct flow in the scarce-target-data setting after Theorem 3.1, could you provide intuition on why this is the case? If both source and target data are scarce, so the TV bound likely holds even if the two transition kernels may be very far, why the two FM models in composite flow would not \"bias\" each other?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to maximize the sample efficiency and performance of reinforcement learning when there are shifted dynamics between the offline data collected from a source environment and the target environment being currently learned. To this end, it presents a new framework that measures the shifted dynamics between the two environments more accurately and principledly, and utilizes it to perform data filtering and active exploration strategies.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: This paper proposes a novel framework called COMPFLOW for reinforcement learning problems that utilize offline data in dynamically heterogeneous environments. Here, it develops a robust measurement of dynamic gaps based on the Wasserstein distance by utilizing synthetic flow structures and optimal transport theory, and achieves state-of-the-art performance on 15 different RL benchmarks by performing data filtering and active exploration. It is a very interesting topic considering the theoretical approach and addressing dynamics gaps, and it can be very effective in the efficient utilization of replay buffers in RL.\n\nWeaknesses: In this paper, the proposed method is very effective in reducing the dynamic gap; however, it would be beneficial to provide more evidence for the practical costs of data filtering in the replay buffer and the performance improvement due to Large Dynamics Gap Exploration Reducing the Performance Gap (Theorem 3.3), although it is discussed in limitations of this draft. It is experimentally clear that strengthening exploration in areas where the dynamic gap is large helps improve performance, but this approach has the disadvantage of introducing an additional hyperparameter. If there is any way to decide the hyperparameter, it seems that it can help improve the quality of the paper.\n\nQUESTIONS:\nSee the above weaknesses.\n\nLIMITATIONS:\nThere are no negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}